Update backend/main.py

backend/main.py

@@ -13,8 +13,8 @@ from fastapi.responses import StreamingResponse
 
 # RAG Imports
 from langchain_community.document_loaders import PyPDFLoader
-from langchain_community.embeddings import HuggingFaceEmbeddings # MUDANÇA: Novo import
-from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_text_splitters import RecursiveCharacterTextSplitter # CORRIGIDO: Nova importação
 from langchain_community.vectorstores import FAISS
 from langchain_core.runnables import RunnablePassthrough, RunnableLambda
 from langchain_core.output_parsers import StrOutputParser
@@ -40,11 +40,10 @@ HF_EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
 
 # Inicializa o modelo Groq e o modelo de embedding
 model = ChatGroq(model=os.getenv("GROQ_MODEL", "mixtral-8x7b-32768"))
-# MUDANÇA: Inicializa o HuggingFaceEmbeddings
+# MUDANÇA: Inicializa o HuggingFaceEmbeddings na CPU
 embeddings = HuggingFaceEmbeddings(
     model_name=HF_EMBEDDING_MODEL,
-    # O device="cpu" garante que ele será executado na CPU, o que é ideal em ambientes sem GPU.
-    model_kwargs={'device': 'cpu'}
+    model_kwargs={'device': 'cpu'} 
 ) 
 
 # Prompt RAG modificado para receber contexto
@@ -72,7 +71,6 @@ class ChatRequest(BaseModel):
 
 def format_docs(docs):
     """Formata os documentos recuperados em uma string única para injeção no prompt."""
-    # Adicionar metadados (se existirem) pode ajudar o LLM a "citar" o documento
     return "\n\n---\n\n".join(
         f"Conteúdo: {doc.page_content}\n(Fonte: Página {doc.metadata.get('page', 'N/A')})" 
         for doc in docs
@@ -92,10 +90,8 @@ async def upload_document(file: UploadFile = File(...)):
         raise HTTPException(status_code=400, detail="Apenas arquivos PDF são suportados neste exemplo.")
 
     # 1. Salvar o arquivo temporariamente para o Loader poder ler
-    # Mantenha essa lógica pois o PyPDFLoader precisa de um caminho de arquivo
     try:
         with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
-            # Garante que o arquivo temporário está no disco para o PyPDFLoader
             content = await file.read()
             tmp_file.write(content)
             temp_path = tmp_file.name
@@ -109,9 +105,8 @@ async def upload_document(file: UploadFile = File(...)):
         splits = text_splitter.split_documents(docs)
 
         # 4. Vetorizar e Armazenar (Vector Store)
-        # Atenção: O download do modelo HF (se for a primeira vez) pode demorar!
         vectorstore = FAISS.from_documents(documents=splits, embedding=embeddings)
-        retriever = vectorstore.as_retriever(search_kwargs={"k": 4}) # Recupera os 4 melhores chunks
+        retriever = vectorstore.as_retriever(search_kwargs={"k": 4}) 
         
         # 5. Criar a nova Chain RAG
         rag_chain = (
@@ -145,7 +140,6 @@ async def chat(request: ChatRequest):
         base_prompt = ChatPromptTemplate.from_messages(
             [("system", "Você é um assistente prestativo. Nenhuma informação de documento foi fornecida."), ("human", "{input}")]
         )
-        # Aqui, mantemos o output parser para consistência com a chain RAG
         current_chain = base_prompt | model | StrOutputParser()
 
     # Função geradora que produz os pedaços (chunks) da resposta
@@ -154,12 +148,9 @@ async def chat(request: ChatRequest):
             # 'astream' é o método de streaming assíncrono do LangChain
             async for chunk in current_chain.astream({"input": request.content}):
                 if chunk:
-                    # print(f"Enviando chunk: {chunk}") # Log para depuração
                     yield chunk
         except Exception as e:
-             # Caso a chamada Groq falhe ou outro erro ocorra
              print(f"Erro no streaming: {e}")
-             # Emite o erro para o cliente
              yield f"Erro no serviço de IA: {e}"
 
     # Retorna uma resposta de streaming