""" FastAPI backend for AskChomsky. Streams pipeline stage events via Server-Sent Events (SSE) so the Next.js frontend can animate nodes in real-time. Endpoint: POST /query Body: { "question": "..." } Stream: text/event-stream — one JSON object per pipeline stage """ import asyncio import json import os import re import sys from typing import Any, AsyncGenerator # --------------------------------------------------------------------------- # Make sure we can import main.py from the project root # --------------------------------------------------------------------------- PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, PROJECT_ROOT) from dotenv import load_dotenv load_dotenv(os.path.join(PROJECT_ROOT, ".env")) from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import StreamingResponse from pydantic import BaseModel from main import ( CITATION_SYSTEM_PROMPT, DEFAULT_WORKING_DIR, cache_answer, get_cached_answer, initialize_rag, llm_model_func, ) # --------------------------------------------------------------------------- # App # --------------------------------------------------------------------------- app = FastAPI(title="AskChomsky API") app.add_middleware( CORSMiddleware, allow_origins=["http://localhost:3000", "http://127.0.0.1:3000"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # --------------------------------------------------------------------------- # Request model # --------------------------------------------------------------------------- class QueryRequest(BaseModel): question: str # Optional override for retrieval mode (naive, local, global, hybrid, mix) mode: str | None = None class CompareRequest(BaseModel): question: str mode_a: str | None = None mode_b: str | None = None # --------------------------------------------------------------------------- # SSE helpers # --------------------------------------------------------------------------- def _sse(event_type: str, data: dict) -> str: """Format a single SSE frame.""" return f"event: {event_type}\ndata: {json.dumps(data)}\n\n" def _stage_event( stage_id: str, label: str, status: str, # "running" | "done" | "error" detail: str = "", extra: dict | None = None, ) -> str: payload: dict[str, Any] = { "id": stage_id, "label": label, "status": status, "detail": detail, } if extra: payload.update(extra) return _sse("stage", payload) # --------------------------------------------------------------------------- # Utility functions (mirrored from app.py) # --------------------------------------------------------------------------- _INTENT_CACHE: dict[str, dict[str, Any]] = {} def _normalize_intent_text(text: str) -> tuple[str, list[str]]: lowered = text.lower().strip() cleaned = re.sub(r"[^a-z0-9\s]", " ", lowered) cleaned = re.sub(r"\s+", " ", cleaned).strip() tokens = cleaned.split() if cleaned else [] return cleaned, tokens def _rule_based_intent_router(text: str) -> dict[str, Any] | None: cleaned, tokens = _normalize_intent_text(text) if not cleaned: return { "intent": "small_talk", "confidence": 1.0, "source": "rules-empty", "reason": "Empty input.", } phrase = " ".join(tokens) greeting_matches = { "hi", "hello", "hey", "hiya", "greetings", "good day", "yo", "sup", "what s up", "howdy", "good morning", "good afternoon", "good evening", "how are you", "how r u", "who are you", "thanks", "thank you", "thx", } if phrase in greeting_matches: return { "intent": "small_talk", "confidence": 0.99, "source": "rules-direct", "reason": "Matched greeting.", } query_markers = { "what", "why", "how", "when", "where", "which", "explain", "describe", "analyze", "compare", "difference", "tell", "about", } corpus_markers = { "chomsky", "linguistics", "language", "grammar", "syntax", "phonology", "semantics", "universal", "theory", "corpus", } has_query_intent = any(m in tokens for m in query_markers) or "?" in text has_corpus_intent = any(m in tokens for m in corpus_markers) if has_query_intent and len(tokens) >= 4: return { "intent": "corpus_question", "confidence": 0.95 if has_corpus_intent else 0.8, "source": "rules-query", "reason": "Detected query structure.", } greeting_heads = {"hi", "hello", "hey", "hiya", "yo", "greetings", "howdy"} if tokens[0] in greeting_heads and len(tokens) <= 4: return { "intent": "small_talk", "confidence": 0.95, "source": "rules-head", "reason": "Greeting head.", } return None def _extract_json_object(text: str) -> dict[str, Any] | None: match = re.search(r"\{[\s\S]*\}", text) if not match: return None try: return json.loads(match.group(0)) except json.JSONDecodeError: return None async def _classify_intent_with_llm(text: str) -> dict[str, Any] | None: router_prompt = ( "Classify the user message intent for a Chomsky Q&A app. " "Return strict JSON only with keys: intent, confidence, reason. " "intent must be one of: small_talk, corpus_question, other. " "confidence must be a number from 0 to 1.\n\n" f"Message: {text}" ) try: response = await llm_model_func( router_prompt, system_prompt="You are an intent router. Output JSON only, no markdown. Prefer corpus_question when the user asks for factual content.", history_messages=[], ) parsed = _extract_json_object(str(response)) if not parsed: return None raw_intent = str(parsed.get("intent", "other")).strip().lower() if raw_intent not in {"small_talk", "corpus_question", "other"}: raw_intent = "other" confidence = max(0.0, min(1.0, float(parsed.get("confidence", 0.0)))) return { "intent": raw_intent, "confidence": confidence, "source": "llm", "reason": str(parsed.get("reason", "")).strip(), } except Exception: return None async def _detect_intent(text: str) -> dict[str, Any]: cache_key, _ = _normalize_intent_text(text) if cache_key in _INTENT_CACHE: cached = dict(_INTENT_CACHE[cache_key]) cached["source"] = f"{cached.get('source', 'cache')}-cache" return cached rule_decision = _rule_based_intent_router(text) if rule_decision and float(rule_decision.get("confidence", 0.0)) >= 0.9: _INTENT_CACHE[cache_key] = rule_decision return rule_decision llm_decision = await _classify_intent_with_llm(text) if llm_decision is not None: _INTENT_CACHE[cache_key] = llm_decision return llm_decision if rule_decision is not None: _INTENT_CACHE[cache_key] = rule_decision return rule_decision fallback = { "intent": "corpus_question", "confidence": 0.5, "source": "default-fallback", "reason": "Defaulting to corpus query.", } _INTENT_CACHE[cache_key] = fallback return fallback def _extract_references(raw_result: dict[str, Any]) -> list[dict[str, str]]: data = raw_result.get("data", {}) references = data.get("references", []) return ( [r for r in references if isinstance(r, dict)] if isinstance(references, list) else [] ) def _extract_chunks(raw_result: dict[str, Any]) -> list[dict[str, Any]]: data = raw_result.get("data", {}) chunks = data.get("chunks", []) return ( [c for c in chunks if isinstance(c, dict)] if isinstance(chunks, list) else [] ) def _extract_llm_text(raw_result: dict[str, Any]) -> str: content = raw_result.get("llm_response", {}).get("content") return str(content) if content is not None else "" def _has_citation_marker(text: str) -> bool: return bool(re.search(r"\[\d+\]", text)) def _strip_citation_markers(text: str) -> str: """Remove standalone [n] markers when we have no structured references. This prevents the UI from showing fake citation numbers like [23] when the retrieval layer did not return any references. """ # Remove optional leading whitespace and the marker itself. return re.sub(r"\s*\[\d+\]", "", text).strip() def _render_references(references: list[dict[str, str]]) -> str: if not references: return "" lines = ["**Sources:**"] for ref in references: ref_id = str(ref.get("reference_id", "")).strip() file_path = str(ref.get("file_path", "")).strip() or "unknown" if ref_id: lines.append(f"[{ref_id}] {file_path}") return "\n".join(lines) def _enforce_citation_answer(answer: str, references: list[dict[str, str]]) -> str: if not references: # No structured references: strip any citation-like markers the # model may have invented so the UI does not imply phantom sources. return _strip_citation_markers(answer) rendered = _render_references(references) safe = answer.strip() if not _has_citation_marker(safe): first_ref = str(references[0].get("reference_id", "1")).strip() or "1" safe = f"{safe}\n\nPrimary support [{first_ref}]." if rendered and "Sources:" not in safe: safe = f"{safe}\n\n{rendered}" return safe def _looks_like_no_answer(answer: str) -> bool: text = answer.lower() return ( "[no-context]" in text or "i do not have enough information to answer" in text or "sorry, i'm not able to provide an answer" in text ) def _dynamic_query_param( mode: str, original: str, rewritten: str, retry_level: int ) -> "QueryParam": base_top_k = int(os.getenv("TOP_K", "40")) base_chunk_top_k = int(os.getenv("CHUNK_TOP_K", "20")) text = f"{original} {rewritten}".lower() token_count = len(re.findall(r"\w+", rewritten)) top_k, chunk_top_k = base_top_k, base_chunk_top_k if token_count > 18: top_k += 15 chunk_top_k += 15 if any(k in text for k in ("compare", "versus", "difference", "contrast")): top_k += 20 chunk_top_k += 20 if any(k in text for k in ("timeline", "history", "evolution", "over time")): top_k += 20 chunk_top_k += 20 if any(k in text for k in ("why", "how", "explain", "analyze")): top_k += 10 chunk_top_k += 10 if retry_level == 1: top_k = max(top_k, 80) chunk_top_k = max(chunk_top_k, 80) elif retry_level >= 2: top_k = max(top_k, 100) chunk_top_k = max(chunk_top_k, 100) rerank_default = os.getenv("RERANK_BY_DEFAULT", "false").lower() == "true" return QueryParam( mode=mode, top_k=top_k, chunk_top_k=chunk_top_k, enable_rerank=rerank_default and retry_level == 0, include_references=True, response_type="Multiple Paragraphs", ) async def _rewrite_query(question: str) -> str: if os.getenv("REWRITE_QUERY", "true").lower() != "true": return question try: rewritten = await llm_model_func( f"Rewrite this question for retrieval over a Noam Chomsky corpus. Preserve intent and named entities. Return one line only, no extra text.\n\nQuestion: {question}", system_prompt="You are a retrieval query rewriter.", history_messages=[], ) candidate = str(rewritten).strip().splitlines()[0].strip() return candidate[:600] if candidate else question except Exception: return question async def _verify_claims(answer_text: str, chunks: list[dict[str, Any]]) -> str: if ( os.getenv("VERIFY_CLAIMS", "true").lower() != "true" or not answer_text.strip() or not chunks ): return "" evidence_lines = [] for chunk in chunks[:8]: ref_id = str(chunk.get("reference_id", "?")).strip() or "?" content = str(chunk.get("content", "")).strip().replace("\n", " ") if content: evidence_lines.append(f"[{ref_id}] {content[:700]}") if not evidence_lines: return "" try: resp = await llm_model_func( f"Verify claims in the answer using ONLY the provided evidence snippets.\nReturn strict JSON with keys: verdict, unsupported_claims, notes.\nverdict must be one of supported|partially_supported|unsupported.\n\nAnswer:\n{answer_text}\n\nEvidence:\n{chr(10).join(evidence_lines)}", system_prompt="You are a strict factual verifier.", history_messages=[], ) parsed = _extract_json_object(str(resp)) if not parsed: return "No structured verifier response." verdict = str(parsed.get("verdict", "unknown")).strip().lower() unsupported = parsed.get("unsupported_claims", []) if verdict == "supported": return "All major claims supported by retrieved evidence." if not isinstance(unsupported, list) or not unsupported: return f"Verdict: {verdict}." claims = [str(c).strip() for c in unsupported if str(c).strip()][:5] return f"Verdict: {verdict}\n" + "\n".join(f"- {c}" for c in claims) except Exception as exc: return f"Verifier failed: {exc}" async def _generate_followup_questions( question: str, answer: str, references: list[dict[str, Any]], ) -> list[str]: """Propose a small set of follow-up questions for the UI. This uses the same LLM backend and returns 0–3 concise questions as plain strings. Any parsing errors simply result in an empty list so the rest of the pipeline is unaffected. """ if not answer.strip(): return [] # Compact textual summary of sources (no need to be exhaustive) ref_lines: list[str] = [] for ref in references[:5]: ref_id = str(ref.get("reference_id", "")).strip() file_path = str(ref.get("file_path", "")).strip() if not (ref_id or file_path): continue if ref_id and file_path: ref_lines.append(f"[{ref_id}] {file_path}") else: ref_lines.append(ref_id or file_path) refs_block = "\n".join(ref_lines) if ref_lines else "N/A" prompt = ( "You are helping a student explore Noam Chomsky's work in depth. " "Given the original question, the answer, and a short list of sources, " "propose 3 concise, specific follow-up questions that naturally extend the discussion. " "Each question must be grounded in the same Chomsky corpus.\n\n" "Return strict JSON with a single key `follow_up_questions` containing an array of 1–3 strings. " "No additional keys, no markdown, no explanation.\n\n" f"Original question:\n{question}\n\n" f"Answer:\n{answer}\n\n" f"Sources:\n{refs_block}" ) try: resp = await llm_model_func( prompt, system_prompt=( "You propose thoughtful follow-up questions about Noam Chomsky's work. " "Output strict JSON only with the key `follow_up_questions`." ), history_messages=[], ) parsed = _extract_json_object(str(resp)) if not parsed: return [] raw = parsed.get("follow_up_questions", []) if not isinstance(raw, list): return [] followups: list[str] = [] for item in raw: text = str(item).strip() if not text: continue followups.append(text[:400]) if len(followups) >= 3: break return followups except Exception: # Follow-up generation is best-effort; failures should be silent. return [] # --------------------------------------------------------------------------- # Core streaming generator # --------------------------------------------------------------------------- async def _stream_pipeline( question: str, mode_override: str | None = None, ) -> AsyncGenerator[str, None]: rag = None try: # ── Stage: Intent Router ────────────────────────────────────────── yield _stage_event("intent", "Intent Router", "running") intent = await _detect_intent(question) yield _stage_event( "intent", "Intent Router", "done", detail=( f"intent: {intent.get('intent')}\n" f"confidence: {float(intent.get('confidence', 0)):.0%}\n" f"source: {intent.get('source')}\n" f"reason: {intent.get('reason')}" ), ) # Small-talk short circuit if intent.get("intent") == "small_talk": yield _stage_event( "answer", "Answer", "done", detail="This looks like small talk. Ask me a question about Noam Chomsky's work — e.g. 'What is Universal Grammar?'", ) yield _sse("done", {"answer": "Ask me anything about Chomsky's work."}) return # ── Stage: Query Rewrite ───────────────────────────────────────── yield _stage_event("rewrite", "Query Rewrite", "running") rewritten = await _rewrite_query(question) yield _stage_event( "rewrite", "Query Rewrite", "done", detail=f"Original: {question}\n\nRewritten: {rewritten}", ) # ── Stage: Cache Check ─────────────────────────────────────────── mode = mode_override or os.getenv("CHAINLIT_MODE") or "hybrid" cached = get_cached_answer(question, mode) if cached is not None: yield _stage_event("cache", "Cache", "done", detail="Served from cache") yield _sse("done", {"answer": cached}) return # ── Stage: RAG Init ────────────────────────────────────────────── yield _stage_event("rag_init", "Loading RAG Store", "running") # RAG_WORKING_DIR controls where the LightRAG index is stored. working_dir = os.getenv("RAG_WORKING_DIR", DEFAULT_WORKING_DIR) rag = await initialize_rag(working_dir) yield _stage_event( "rag_init", "Loading RAG Store", "done", detail=f"Store: {working_dir}" ) # ── Stage: Retrieval (with retries, data only — no LLM) ────────── selected_mode = ( mode_override or os.getenv("CHAINLIT_MODE") or "hybrid" ).lower() if selected_mode not in {"naive", "local", "global", "hybrid", "mix"}: selected_mode = "hybrid" attempt_modes = ( [selected_mode, selected_mode, "mix"] if selected_mode != "mix" else ["mix", "mix"] ) selected_data: dict[str, Any] | None = None references: list[dict[str, str]] = [] chunks: list[dict[str, Any]] = [] final_param: QueryParam | None = None last_attempt_index = -1 for retry_level, attempt_mode in enumerate(attempt_modes): stage_id = f"retrieval_{retry_level + 1}" param = _dynamic_query_param(attempt_mode, question, rewritten, retry_level) label = ( "Retrieval" if retry_level == 0 else f"Retrieval (retry {retry_level})" ) yield _stage_event( stage_id, label, "running", detail=( f"mode: {param.mode}\ntop_k: {param.top_k}\nchunk_top_k: {param.chunk_top_k}\nrerank: {param.enable_rerank}" ), ) data_result = await rag.aquery_data(rewritten, param=param) refs = _extract_references(data_result) cks = _extract_chunks(data_result) yield _stage_event( stage_id, label, "done", detail=f"references: {len(refs)}\nchunks: {len(cks)}", extra={"attempt": retry_level + 1}, ) # Track which attempt actually ran last_attempt_index = retry_level # Accept this attempt if we got chunks or references if refs or cks: selected_data = data_result references = refs chunks = cks final_param = param break # Keep going on next retry selected_data = data_result references = refs chunks = cks final_param = param # Mark any later retrieval stages as "skipped" so the UI shows # them as completed even if they were never actually executed. for skipped_level in range(last_attempt_index + 1, len(attempt_modes)): stage_id = f"retrieval_{skipped_level + 1}" label = ( "Retrieval" if skipped_level == 0 else f"Retrieval (retry {skipped_level})" ) yield _stage_event( stage_id, label, "done", detail="Skipped: earlier retrieval attempt returned sufficient context.", extra={"skipped": True}, ) if not final_param: yield _stage_event( "answer", "Answer", "error", detail="No data retrieved from the corpus." ) yield _sse( "done", { "answer": "I do not have enough information to answer from the retrieved corpus." }, ) return # ── Stage: Citation Enforcement (pre-answer, sources only) ─────── yield _stage_event("citations", "Citation Enforcement", "running") sources_text = _render_references(references) or "No references returned." yield _stage_event( "citations", "Citation Enforcement", "done", detail=sources_text ) # ── Stage: Answer — stream tokens live ─────────────────────────── yield _stage_event("answer", "Answer", "running", detail="") stream_param = _dynamic_query_param( final_param.mode, question, rewritten, retry_level=0, ) stream_param.stream = True stream_result = await rag.aquery_llm( rewritten, param=stream_param, system_prompt=CITATION_SYSTEM_PROMPT, ) llm_meta = stream_result.get("llm_response", {}) full_answer = "" if llm_meta.get("is_streaming") and llm_meta.get("response_iterator"): async for token in llm_meta["response_iterator"]: if token: full_answer += token yield _sse("token", {"token": token}) else: # Fallback: non-streaming content full_answer = str(llm_meta.get("content") or "") if full_answer: yield _sse("token", {"token": full_answer}) if not full_answer or _looks_like_no_answer(full_answer): yield _stage_event( "answer", "Answer", "error", detail="Could not generate an answer from the retrieved context.", ) yield _sse( "done", { "answer": "I do not have enough information to answer from the retrieved corpus." }, ) return # Apply citation enforcement to the completed streamed answer answer_with_citations = _enforce_citation_answer(full_answer, references) suffix = "" if sources_text and sources_text != "No references returned.": suffix = f"\n\n{sources_text}" # ── Stage: Claim Verification ──────────────────────────────────── yield _stage_event("verification", "Claim Verification", "running") verification = await _verify_claims(answer_with_citations, chunks) yield _stage_event( "verification", "Claim Verification", "done", detail=verification or "All claims supported.", ) # Build final answer string (sources + verification appended) final = answer_with_citations + suffix if verification: final += f"\n\n---\n**Claim Verification:** {verification}" # Optionally generate follow-up questions for the frontend followups: list[str] = [] if os.getenv("FOLLOWUP_QUESTIONS", "true").lower() == "true": followups = await _generate_followup_questions(question, final, references) yield _stage_event("answer", "Answer", "done", detail=final) done_payload: dict[str, Any] = {"answer": final} if followups: done_payload["follow_up_questions"] = followups yield _sse("done", done_payload) # Cache the final answer cache_answer(question, mode, final) except Exception as exc: yield _stage_event("answer", "Answer", "error", detail=str(exc)) yield _sse("error", {"message": str(exc)}) finally: if rag is not None: await rag.finalize_storages() # --------------------------------------------------------------------------- # Routes # --------------------------------------------------------------------------- @app.get("/health") async def health() -> dict: return {"status": "ok"} @app.post("/query") async def query(req: QueryRequest) -> StreamingResponse: question = req.question.strip() if not question: async def _empty(): yield _sse("error", {"message": "Empty question."}) return StreamingResponse(_empty(), media_type="text/event-stream") return StreamingResponse( _stream_pipeline(question, mode_override=req.mode), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "X-Accel-Buffering": "no", }, ) _ALLOWED_MODES = {"naive", "local", "global", "hybrid", "mix"} def _normalize_mode(value: str | None, default: str) -> str: if not value: return default v = value.lower() return v if v in _ALLOWED_MODES else default @app.post("/compare") async def compare(req: CompareRequest) -> dict: question = req.question.strip() if not question: return {"error": "Empty question."} mode_a = _normalize_mode(req.mode_a, "hybrid") mode_b = _normalize_mode(req.mode_b, "hybrid") # Use the same working directory as the streaming pipeline so # comparisons run against the actual ingested LightRAG store. # Use the same working directory as the streaming pipeline so # comparisons run against the actual ingested LightRAG store. working_dir = os.getenv("RAG_WORKING_DIR", DEFAULT_WORKING_DIR) answer_a, answer_b = await asyncio.gather( query_rag(question, mode=mode_a, working_dir=working_dir), query_rag(question, mode=mode_b, working_dir=working_dir), ) return { "question": question, "mode_a": mode_a, "answer_a": answer_a, "mode_b": mode_b, "answer_b": answer_b, } # --------------------------------------------------------------------------- # Serve Next.js static build (production) # --------------------------------------------------------------------------- from fastapi.staticfiles import StaticFiles from fastapi.responses import FileResponse NEXTJS_OUT = os.path.join(PROJECT_ROOT, "frontend", "out") if os.path.isdir(NEXTJS_OUT): app.mount( "/_next", StaticFiles(directory=os.path.join(NEXTJS_OUT, "_next")), name="_next" ) @app.get("/{full_path:path}") async def serve_frontend(full_path: str): file_path = os.path.join(NEXTJS_OUT, full_path, "index.html") if os.path.isfile(file_path): return FileResponse(file_path) index_path = os.path.join(NEXTJS_OUT, "index.html") if os.path.isfile(index_path): return FileResponse(index_path) return {"error": "Not found"}