145 lines
3.8 KiB
Python
145 lines
3.8 KiB
Python
from fastapi import FastAPI, HTTPException
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from pydantic import BaseModel
|
|
import requests
|
|
import pandas as pd
|
|
import numpy as np
|
|
import faiss
|
|
from sentence_transformers import SentenceTransformer
|
|
|
|
app = FastAPI()
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
INDEX_PATH = "/data/indice_documentos.faiss"
|
|
CHUNKS_PATH = "/data/dataset_chunks.json"
|
|
|
|
LLAMA_URL = "http://192.168.1.39:8089/v1/chat/completions"
|
|
MODEL_NAME = "Qwen3-1.7B-Instruct-Q4_K_M.gguf"
|
|
|
|
modelo_embeddings = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
|
|
index = faiss.read_index(INDEX_PATH)
|
|
dataset_chunks = pd.read_json(CHUNKS_PATH)
|
|
|
|
class PerguntaRequest(BaseModel):
|
|
pergunta: str
|
|
|
|
@app.get("/")
|
|
def root():
|
|
return {"status": "ok"}
|
|
|
|
def pesquisar_chunks(query, modelo_embeddings, index, dataset_chunks, top_k=5):
|
|
query_embedding = modelo_embeddings.encode([query], convert_to_numpy=True)
|
|
query_embedding = np.array(query_embedding, dtype="float32")
|
|
|
|
distancias, indices = index.search(query_embedding, top_k)
|
|
|
|
resultados = dataset_chunks.iloc[indices[0]].copy()
|
|
resultados["score_distancia"] = distancias[0]
|
|
return resultados
|
|
|
|
def construir_contexto_rag(query, resultados, max_chars_por_chunk=500):
|
|
partes = []
|
|
|
|
for _, row in resultados.iterrows():
|
|
chunk_curto = row["texto_chunk"][:max_chars_por_chunk]
|
|
partes.append(
|
|
f"[Documento: {row['nome_pdf']} | Categoria: {row['categoria']}]\n{chunk_curto}"
|
|
)
|
|
|
|
contexto = "\n\n---\n\n".join(partes)
|
|
|
|
return f"""
|
|
És um assistente especializado em documentos do Exército Português.
|
|
|
|
Responde apenas com base no contexto fornecido.
|
|
|
|
Pergunta:
|
|
{query}
|
|
|
|
Contexto:
|
|
{contexto}
|
|
|
|
Regras:
|
|
- Responde em português.
|
|
- Sê objetivo.
|
|
- Dá uma resposta curta.
|
|
- Se a resposta não estiver no contexto, diz isso explicitamente.
|
|
- Se possível, menciona o nome do documento.
|
|
"""
|
|
|
|
def perguntar_qwen(prompt):
|
|
payload = {
|
|
"model": MODEL_NAME,
|
|
"messages": [
|
|
{"role": "user", "content": prompt}
|
|
],
|
|
"temperature": 0.2,
|
|
"max_tokens": 256,
|
|
"stream": False
|
|
}
|
|
|
|
try:
|
|
response = requests.post(
|
|
LLAMA_URL,
|
|
json=payload,
|
|
timeout=18000
|
|
)
|
|
|
|
print("STATUS LLAMA:", response.status_code)
|
|
print("BODY LLAMA:", response.text[:1000])
|
|
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
|
|
return data["choices"][0]["message"]["content"]
|
|
|
|
except Exception as e:
|
|
print("ERRO AO CHAMAR LLAMA:", repr(e))
|
|
raise HTTPException(status_code=500, detail=f"Erro llama.cpp: {str(e)}")
|
|
|
|
@app.post("/perguntar")
|
|
def perguntar(request: PerguntaRequest):
|
|
try:
|
|
query = request.pergunta
|
|
print("PERGUNTA:", query)
|
|
|
|
resultados = pesquisar_chunks(
|
|
query=query,
|
|
modelo_embeddings=modelo_embeddings,
|
|
index=index,
|
|
dataset_chunks=dataset_chunks,
|
|
top_k=1
|
|
)
|
|
|
|
print("RESULTADOS ENCONTRADOS:", len(resultados))
|
|
|
|
prompt = construir_contexto_rag(query, resultados)
|
|
resposta = perguntar_qwen(prompt)
|
|
|
|
fontes = []
|
|
for _, row in resultados.iterrows():
|
|
fontes.append({
|
|
"nome_pdf": row["nome_pdf"],
|
|
"categoria": row["categoria"],
|
|
"chunk_id": int(row["chunk_id"]),
|
|
"score_distancia": float(row["score_distancia"])
|
|
})
|
|
|
|
return {
|
|
"pergunta": query,
|
|
"resposta": resposta,
|
|
"fontes": fontes
|
|
}
|
|
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
print("ERRO NO ENDPOINT:", repr(e))
|
|
raise HTTPException(status_code=500, detail=str(e)) |