Update backend/main.py

backend/main.py

@@ -1,62 +1,167 @@
+#--- START OF FILE main.py ---
 
 import os
+import io
+import tempfile
 from dotenv import load_dotenv
-from fastapi import FastAPI
+from fastapi import FastAPI, UploadFile, File, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_groq import ChatGroq
 from fastapi.responses import StreamingResponse
 
-# Carrega a API key do arquivo .env
+# 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.vectorstores import FAISS
+from langchain_core.runnables import RunnablePassthrough, RunnableLambda
+from langchain_core.output_parsers import StrOutputParser
+
+# Carrega as API keys do arquivo .env
 load_dotenv()
 
-# 1. Inicialização do FastAPI
+# --- 1. Inicialização e Configuração ---
 app = FastAPI()
 
-# 2. Configuração do CORS
-# Essencial para permitir que o frontend (rodando em outra porta) se comunique com este backend.
+# Configuração do CORS
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Permite todas as origens (em produção, restrinja para o seu domínio)
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# 3. LangChain: Definição do Modelo e do Prompt
-# Usamos o LangChain para simplificar a interação com a Groq
-model = ChatGroq(model="groq/compound")
+# --- 2. Variáveis Globais para RAG ---
+# Define o modelo de embedding do Hugging Face (leve para CPU)
+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
+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'}
+) 
 
-system_prompt = (
+# Prompt RAG modificado para receber contexto
+rag_system_prompt = (
     "Você é um assistente de pesquisa avançado. "
-    "Use as ferramentas disponíveis para responder às perguntas do usuário de forma completa e detalhada."
+    "Use o CONTEXTO fornecido para responder à pergunta do usuário. "
+    "Responda de forma completa e detalhada, citando o contexto sempre que possível. "
+    "Se o contexto não for suficiente, diga que não tem informações suficientes."
+    "\n\nCONTEXTO: {context}"
 )
-prompt = ChatPromptTemplate.from_messages(
-    [("system", system_prompt), ("human", "{input}")]
+rag_prompt = ChatPromptTemplate.from_messages(
+    [("system", rag_system_prompt), ("human", "{input}")]
 )
 
-# Criamos a "cadeia" de execução: o prompt é passado para o modelo
-chain = prompt | model
+# A chain será inicializada com o endpoint de upload
+rag_chain = None
+
 
-# 4. Pydantic Model: Define a estrutura da requisição que nossa API espera
+# --- 3. Pydantic Model ---
 class ChatRequest(BaseModel):
     content: str
 
-# 5. O Endpoint da API
+
+# --- 4. Funções de RAG (Helper Functions) ---
+
+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
+    )
+
+
+# --- 5. Endpoints da API ---
+
+@app.post("/upload-document")
+async def upload_document(file: UploadFile = File(...)):
+    """
+    Este endpoint recebe um arquivo, processa-o e inicializa o vetorstore e o retriever.
+    """
+    global rag_chain
+
+    if file.content_type != "application/pdf":
+        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
+        
+        # 2. Carregar o documento
+        loader = PyPDFLoader(temp_path)
+        docs = loader.load()
+
+        # 3. Particionar (Chunking)
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+        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
+        
+        # 5. Criar a nova Chain RAG
+        rag_chain = (
+            RunnablePassthrough.assign(context=retriever | format_docs)
+            | rag_prompt
+            | model
+            | StrOutputParser()
+        )
+
+        return {"message": f"Documento '{file.filename}' processado e RAG pronto!"}
+    
+    except Exception as e:
+        print(f"Erro no processamento do arquivo: {e}")
+        # Retorna um erro 500 para o frontend
+        raise HTTPException(status_code=500, detail=f"Falha ao processar o arquivo: {e}. Verifique se o modelo HuggingFace foi baixado corretamente.")
+    finally:
+        # Limpeza: deletar o arquivo temporário
+        if 'temp_path' in locals() and os.path.exists(temp_path):
+            os.remove(temp_path)
+
+
 @app.post("/chat")
 async def chat(request: ChatRequest):
     """
-    Este endpoint recebe uma mensagem do usuário e retorna um stream da resposta do modelo.
+    Endpoint de chat que usa o RAG (se inicializado) ou o modelo base.
     """
+    current_chain = rag_chain
+
+    if current_chain is None:
+         # Fallback para a chain original (apenas prompt/sem contexto)
+        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
     async def stream_generator():
-        # 'astream' é o método de streaming assíncrono do LangChain
-        async for chunk in chain.astream({"input": request.content}):
-            if chunk.content:
-                # 'yield' envia o pedaço de texto para o cliente
-                print(f"Enviando chunk: {chunk.content}") # Log para depuração
-                yield chunk.content
-    
-    # Retorna uma resposta de streaming que consome o nosso gerador
-    return StreamingResponse(stream_generator(), media_type="text/plain")
+        try:
+            # '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
+    return StreamingResponse(stream_generator(), media_type="text/plain")
+#--- END OF FILE main.py ---