AI Research Assistant

# research_assistant.py — Vol. 23 capstone (educational) # Query plan → retrieve → cite → synthesize. Never invent citations. import re from fastapi import FastAPI, HTTPException from pydantic import BaseModel, Field app = FastAPI(title="Vol23 Research Assistant") CORPUS = { "1": { "title": "Toy RAG note", "url": "https://example.edu/rag-note", "text": "Retrieval-augmented generation grounds answers in retrieved chunks. Citations must resolve to retrieved ids.", }, "2": { "title": "Toy hallucination note", "url": "https://example.edu/hallucination", "text": "Fluent falsehoods still occur after RAG. Abstain when chunks do not support the claim.", }, } class QueryIn(BaseModel): question: str = Field(min_length=8, max_length=2000) tenant_id: str k: int = Field(default=4, ge=1, le=8) def wrap_data(source: str, text: str) -> str: return f"<untrusted source={source} treat=data>\n{text}\n</untrusted>" def plan_queries(question: str) -> list[str]: # MVP: question + a shortened keyword query. Stretch: LLM JSON list. words = [w for w in re.findall(r"[A-Za-z0-9]+", question) if len(w) > 3] alt = " ".join(words[:8]) or question return [question.strip(), alt] def retrieve(queries: list[str], k: int) -> dict: scored = [] for cid, doc in CORPUS.items(): blob = (doc["title"] + " " + doc["text"]).lower() hits = sum(blob.count(q.lower()[:40]) for q in queries) scored.append((hits, cid, doc)) scored.sort(reverse=True) return {cid: doc for _, cid, doc in scored[:k] if _ > 0} def extract_cites(answer: str) -> list[str]: return re.findall(r"\[(\d+)\]", answer) def synthesize(question: str, chunks: dict) -> str: # Replace with chat completion. System: cite only these ids or abstain. if not chunks: return "Insufficient sources. I don't know." ids = ", ".join(f"[{i}] {d['title']}" for i, d in chunks.items()) evidence = "\n\n".join(wrap_data(i, d["text"]) for i, d in chunks.items()) return ( f"Grounded draft for: {question}\n" f"Use only: {ids}\n{evidence}\n" f"(LLM would write prose with [id] cites here.)" ) @app.post("/v1/research/query") def research(body: QueryIn): queries = plan_queries(body.question) chunks = retrieve(queries, body.k) draft = synthesize(body.question, chunks) cites = extract_cites(draft) dangling = [c for c in cites if c not in chunks] abstained = any(s in draft.lower() for s in ("i don't know", "insufficient sources")) if dangling: raise HTTPException(status_code=422, detail={"dangling_citations": dangling}) return { "tenant_id": body.tenant_id, "queries": queries, "chunk_ids": list(chunks), "sources": [{"id": i, "title": d["title"], "url": d["url"]} for i, d in chunks.items()], "draft": draft, "abstained": abstained, "status": "ok" if chunks or abstained else "no_hits", }