# pip install fastapi uvicorn psycopg2-binary pydantic python-dotenv openaifrom fastapi import FastAPIfrom pydantic import BaseModelimport os, psycopg2, jsonfrom math import sqrtfrom openai import OpenAIclient = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))app = FastAPI()# --- Consulta vetorial simples (pgvector) ---def embed(text: str): # modelo de embedding barato return client.embeddings.create(model="text-embedding-3-small", input=text).data[0].embeddingdef search_chunks(q: str, k: int = 5): v = embed(q) conn = psycopg2.connect(os.getenv("DATABASE_URL")) cur = conn.cursor() cur.execute(""" SELECT doc_title, chapter, page, content, 1 - (embedding <=> %s::vector) AS score FROM chunks ORDER BY embedding <=> %s::vector LIMIT %s; """, (v, v, k)) rows = cur.fetchall() cur.close(); conn.close() return [{"doc": r[0], "chapter": r[1], "page": r[2], "content": r[3], "score": float(r[4])} for r in rows]class Ask(BaseModel): question: str banca: str | None = None@app.post("/ask")def ask(body: Ask): ctx = search_chunks(body.question, k=5) # monta “contexto citável” citations = [] context_str = "" for i, c in enumerate(ctx, 1): context_str += f"[{i}] {c['doc']} | Cap.: {c['chapter']} | pág.: {c['page']}n{c['content']}nn" citations.append(f"{c['doc']}, Cap. {c['chapter']}, p. {c['page']}") system = ( "Você é o Professor IA da Gramática Simplificada. " "Responda somente com base no CONTEXTO a seguir. " "Formato: Explicação objetiva → Regras (tópicos) → Exemplo certo/errado → 2 exercícios com gabarito → Fontes." ) user_msg = f"Pergunta: {body.question}nBanca: {body.banca or 'não informada'}nnCONTEXTO:n{context_str}" completion = client.chat.completions.create( model="gpt-4.1-mini", messages=[{"role":"system","content":system}, {"role":"user","content":user_msg}], temperature=0.2 ) answer = completion.choices[0].message.content return {"answer": answer, "sources": citations}