File size: 4,956 Bytes
9bdf06b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 |
import os
import logging
import sys
from flask import Flask, request, jsonify, Response
# Inicializa o Flask
app = Flask(__name__)
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.core import (
Settings,
SimpleDirectoryReader,
StorageContext,
Document,
)
Settings.llm = OpenAI(model="gpt-3.5-turbo")
Settings.embed_model = OpenAIEmbedding(model_name="text-embedding-3-small")
directory_path = "documentos"
from llama_index.readers.file import PDFReader #concatenar todo o documento já vem nativo no pdfreader
file_extractor = {".pdf": PDFReader(return_full_document = True)}
from drive_downloader import GoogleDriveDownloader
# ID da pasta no Drive e caminho local
folder_id = "1n34bmh9rlbOtCvE_WPZRukQilKeabWsN"
local_path = directory_path
GoogleDriveDownloader().download_from_folder(folder_id, local_path)
documents = SimpleDirectoryReader(
input_dir=directory_path,
file_extractor=file_extractor,
filename_as_id=True,
recursive=True
).load_data()
from document_creator import create_single_document_with_filenames
document = create_single_document_with_filenames(directory_path = directory_path)
documents.append(document)
#from llama_index.core.ingestion import IngestionPipeline
#ingestion pipeline vai entrar em uso quando adicionar o extrator de metadados
from llama_index.core.node_parser import SentenceSplitter
splitter = SentenceSplitter(chunk_size=1024, chunk_overlap=128)
nodes = splitter.get_nodes_from_documents(documents)
from llama_index.core.storage.docstore import SimpleDocumentStore
docstore = SimpleDocumentStore()
docstore.add_documents(nodes)
from llama_index.core import VectorStoreIndex, StorageContext
from llama_index.vector_stores.chroma import ChromaVectorStore
import chromadb
db = chromadb.PersistentClient(path="chroma_db")
chroma_collection = db.get_or_create_collection("dense_vectors")
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
storage_context = StorageContext.from_defaults(
docstore=docstore, vector_store=vector_store
)
index = VectorStoreIndex(nodes = nodes, storage_context=storage_context, show_progress = True)
storage_context.docstore.persist("./docstore.json")
index_retriever = index.as_retriever(similarity_top_k=2)
import nest_asyncio
nest_asyncio.apply()
from llama_index.retrievers.bm25 import BM25Retriever
bm25_retriever = BM25Retriever.from_defaults(
docstore=index.docstore,
similarity_top_k=2,
language = "portuguese",
verbose=True,
)
from llama_index.core.retrievers import QueryFusionRetriever
retriever = QueryFusionRetriever(
[index_retriever, bm25_retriever],
num_queries=1, #desativado = 1
mode="reciprocal_rerank",
use_async=True,
verbose=True,
)
from llama_index.core.memory import ChatMemoryBuffer
from mysqlchatstore import MySQLChatStore
chat_store = MySQLChatStore.from_params(
host=os.getenv("MYSQL_HOST"),
port=os.getenv("MYSQL_PORT"),
user=os.getenv("MYSQL_USER"),
password=os.getenv("MYSQL_PASSWORD"),
database=os.getenv("MYSQL_DATABASE"),
table_name=os.getenv("MYSQL_TABLE")
)
chat_memory = ChatMemoryBuffer.from_defaults(
token_limit=3000,
chat_store=chat_store,
chat_store_key="Sicoob", #Tendo algumas dificuldades ainda pra passar o user
)
from llama_index.core.query_engine import RetrieverQueryEngine
query_engine = RetrieverQueryEngine.from_args(retriever)
from llama_index.core.chat_engine import CondensePlusContextChatEngine
chat_engine = CondensePlusContextChatEngine.from_defaults(
query_engine,
memory=chat_memory,
context_prompt=(
"Você é um assistente virtual capaz de interagir normalmente, além de"
" fornecer informações sobre organogramas e listar funcionários."
" Aqui estão os documentos relevantes para o contexto:\n"
"{context_str}"
"\nInstrução: Use o histórico da conversa anterior, ou o contexto acima, para responder."
"No final da resposta, depois de uma quebra de linha escreva o nome do documento que contém a informação entre dois ||, como ||Documento Nome||"
),
)
@app.route("/chat", methods=["POST"])
def chat():
user_input = request.json.get("message", "")
if not user_input:
return jsonify({"error": "Mensagem vazia"}), 400
def generate_response():
try:
response = chat_engine.stream_chat(user_input)
for token in response.response_gen:
yield token # Envia cada token
except Exception as e:
yield f"Erro: {str(e)}"
return Response(generate_response(), content_type="text/plain")
if __name__ == "__main__":
app.run(port=5001, debug=False)
|