#!/usr/bin/env python3 """ FastAPI endpoint for JUUL Vectorless PageIndex RAG. Run: uv run uvicorn rag.api:app --reload """ import os import re import json import uuid import logging from pathlib import Path from datetime import datetime from typing import List, Dict, Optional from collections import Counter logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", datefmt="%H:%M:%S", ) log = logging.getLogger("casbinder") import fitz from openai import OpenAI from rank_bm25 import BM25Okapi from dotenv import load_dotenv from fastapi import FastAPI, HTTPException from fastapi.responses import StreamingResponse from pydantic import BaseModel from qdrant_client import QdrantClient load_dotenv(Path(__file__).parent.parent / ".env") INDEX_PATH = Path(os.environ.get("INDEX_PATH", Path(__file__).parent / "page_index.json")) DOWNLOADS_DIR = Path(os.environ.get("DOWNLOADS_DIR", Path(__file__).parent.parent / "downloads")) SESSIONS_DIR = Path(os.environ.get("SESSIONS_DIR", Path(__file__).parent / "sessions")) MODEL = "gpt-4o" EMBED_MODEL = "text-embedding-3-small" QDRANT_COLLECTION = "juul_mdl_chunks" _MASTER_CONTEXT_PATH = Path(os.environ.get( "MASTER_CONTEXT_PATH", Path(__file__).parent.parent / "rag" / "prompt.txt" )) JUUL_MASTER_CONTEXT = ( _MASTER_CONTEXT_PATH.read_text(encoding="utf-8") if _MASTER_CONTEXT_PATH.exists() else "" ) # ── Musk v. Altman config ───────────────────────────────────────────────────── MUSK_COLLECTION = "musk_altman_chunks" MUSK_DOWNLOADS = Path(__file__).parent.parent / "downloads" / "musk_v_altman" _MUSK_CONTEXT_PATH = Path(__file__).parent.parent / "sam-musk.txt" MUSK_MASTER_CONTEXT = ( _MUSK_CONTEXT_PATH.read_text(encoding="utf-8") if _MUSK_CONTEXT_PATH.exists() else "" ) MUSK_SYSTEM = """You are a knowledgeable legal assistant on the Musk v. Altman case (Elon Musk and X.AI Corp. v. Samuel Altman, Gregory Brockman, OpenAI Inc., Microsoft Corp. et al., N.D. Cal., Case No. 3:24-cv-02022). Reply in 1-2 sentences. Short, confident, natural — like a colleague giving a quick answer. Never mention sources, documents, searches, filenames, or anything internal. If unsure, make a reasonable guess. Never refuse to answer. Only output NO_DOCUMENTS_FOUND if the question has nothing to do with this case.""" MUSK_TOOLS = [ { "type": "function", "function": { "name": "vector_search", "description": ( "Semantic vector search over all Musk v. Altman court document chunks. " "Use for any question about parties, claims, filings, orders, or case facts." ), "parameters": { "type": "object", "properties": { "query": {"type": "string", "description": "Natural language question or topic"}, "n": {"type": "integer", "description": "Number of chunks to return (default 8)"}, }, "required": ["query"], }, }, }, { "type": "function", "function": { "name": "read_doc", "description": "Read full text of a specific document PDF by filename.", "parameters": { "type": "object", "properties": { "filename": {"type": "string"}, "pages": {"type": "string", "description": "Page range e.g. '1-5' (default: all, max 20)"}, }, "required": ["filename"], }, }, }, ] FALLBACK_SYSTEM = """You are a knowledgeable legal assistant on the JUUL MDL case. Reply in 1-2 sentences. Short, confident, natural — like a colleague giving a quick answer. Never mention sources, documents, searches, or anything internal. If unsure, give a reasonable guess. Never refuse. Only output NO_DOCUMENTS_FOUND if the question has nothing to do with the case. {context}""" SYSTEM = """You are a knowledgeable legal assistant on the JUUL MDL case (MDL 2913, N.D. Cal.). Reply in 1-2 sentences. Short, confident, natural — like a colleague giving a quick answer. Never mention sources, documents, searches, filenames, or anything internal. If unsure, make a reasonable guess. Never refuse to answer. Only output NO_DOCUMENTS_FOUND if the question has nothing to do with this case.""" TOOLS = [ { "type": "function", "function": { "name": "search", "description": "BM25 keyword search over document titles in the PageIndex.", "parameters": { "type": "object", "properties": { "query": {"type": "string"}, "n": {"type": "integer", "description": "Max results (default 15)"}, "category": {"type": "string", "description": "Optional category filter"}, }, "required": ["query"], }, }, }, { "type": "function", "function": { "name": "filter", "description": "Filter PageIndex by category and/or importance.", "parameters": { "type": "object", "properties": { "category": {"type": "string"}, "importance": {"type": "string", "description": "critical | high | medium | low"}, }, "required": ["category"], }, }, }, { "type": "function", "function": { "name": "read_doc", "description": "Read full text of a document from its downloaded PDF.", "parameters": { "type": "object", "properties": { "filename": {"type": "string"}, "pages": {"type": "string", "description": "Page range e.g. '1-5' (default: all, max 20)"}, }, "required": ["filename"], }, }, }, { "type": "function", "function": { "name": "categories", "description": "List all document categories and counts.", "parameters": {"type": "object", "properties": {}}, }, }, { "type": "function", "function": { "name": "vector_search", "description": ( "Semantic vector search over all document chunks. " "Use this for meaning-based queries that keyword search may miss — " "e.g. 'upcoming trial date', 'settlement amount', 'expert witness opinions'. " "Returns actual passage text from the most relevant chunks." ), "parameters": { "type": "object", "properties": { "query": {"type": "string", "description": "Natural language question or topic"}, "n": {"type": "integer", "description": "Number of chunks to return (default 8)"}, }, "required": ["query"], }, }, }, ] # ── Session store ───────────────────────────────────────────────────────────── def _session_path(session_id: str) -> Path: return SESSIONS_DIR / f"{session_id}.json" def load_session(session_id: str) -> dict: path = _session_path(session_id) if path.exists(): with open(path) as f: return json.load(f) return {"session_id": session_id, "created_at": datetime.utcnow().isoformat(), "messages": []} def save_session(session: dict): SESSIONS_DIR.mkdir(parents=True, exist_ok=True) session["updated_at"] = datetime.utcnow().isoformat() with open(_session_path(session["session_id"]), "w") as f: json.dump(session, f, indent=2) # ── PageIndex ───────────────────────────────────────────────────────────────── class PageIndex: def __init__(self): with open(INDEX_PATH) as f: self.entries: List[Dict] = json.load(f) # Search over title + body if available, else just title tokenized = [self._tok(self._search_text(e)) for e in self.entries] self.bm25 = BM25Okapi(tokenized) def _tok(self, text: str) -> List[str]: return re.findall(r'\w+', text.lower()) def _search_text(self, e: Dict) -> str: return e["title"] + " " + e.get("body", "") def search(self, query: str, n: int = 15, category: Optional[str] = None) -> List[Dict]: pool = self.entries if category: pool = [e for e in self.entries if e["category"].lower() == category.lower()] or self.entries tok = [self._tok(self._search_text(e)) for e in pool] scores = BM25Okapi(tok).get_scores(self._tok(query)) ranked = sorted(enumerate(scores), key=lambda x: -x[1]) return [{**pool[i], "score": round(float(s), 3)} for i, s in ranked[:n] if s > 0] def multi_search(self, queries: List[str], n: int = 15, category: Optional[str] = None) -> List[Dict]: """Run BM25 on multiple queries, merge by best score per document.""" best: Dict[str, Dict] = {} for q in queries: for entry in self.search(q, n=n, category=category): key = entry["filename"] if key not in best or entry["score"] > best[key]["score"]: best[key] = entry return sorted(best.values(), key=lambda x: -x["score"])[:n] def filter(self, category: str, importance: Optional[str] = None) -> List[Dict]: r = [e for e in self.entries if e["category"].lower() == category.lower()] if importance: r = [e for e in r if e["importance"] == importance] return sorted(r, key=lambda e: (e["doc_num"], e["attach_num"])) def categories(self) -> Dict[str, int]: return dict(Counter(e["category"] for e in self.entries).most_common()) def read_pdf(filename: str, pages: Optional[str] = None, base_dir: Optional[Path] = None) -> str: path = (base_dir or DOWNLOADS_DIR) / filename if not path.exists(): return f"FILE_NOT_FOUND:{filename}" try: doc = fitz.open(str(path)) total = doc.page_count if pages: parts = pages.split("-") start = int(parts[0]) - 1 end = int(parts[1]) - 1 if len(parts) > 1 else start else: start, end = 0, min(total - 1, 19) texts = [] for i in range(start, end + 1): if 0 <= i < total: text = re.sub(r'\n{3,}', '\n\n', doc[i].get_text("text")).strip() if text: texts.append(f"--- Page {i+1} ---\n{text}") doc.close() return (f"[{filename} | {total} pages | p.{start+1}–{end+1}]\n\n" + "\n\n".join(texts)) if texts else f"No text in {filename}" except Exception as ex: return f"Error reading {filename}: {ex}" def fmt(e: Dict) -> str: score = f" score={e['score']:.2f}" if "score" in e else "" return f"Doc #{e['doc_num']} | {e['category']} | {e['importance'].upper()}{score} | {e['title'][:150]} | file: {e['filename']}" def run_tool(idx: PageIndex, name: str, args: Dict): """Returns (result_text, had_results, sources).""" sources = [] if name == "search": results = idx.search(args["query"], args.get("n", 15), args.get("category")) if not results: return f"No results for: '{args['query']}'", False, [] sources = [{"doc_num": e["doc_num"], "filename": e["filename"], "category": e["category"], "title": e["title"][:200]} for e in results] return f"Results for '{args['query']}':\n" + "\n".join(fmt(e) for e in results), True, sources elif name == "filter": results = idx.filter(args["category"], args.get("importance")) if not results: return f"No documents in category: {args['category']}", False, [] sources = [{"doc_num": e["doc_num"], "filename": e["filename"], "category": e["category"], "title": e["title"][:200]} for e in results] return f"{len(results)} docs in '{args['category']}':\n" + "\n".join(fmt(e) for e in results), True, sources elif name == "read_doc": text = read_pdf(args["filename"], args.get("pages")) has_result = not text.startswith("FILE_NOT_FOUND") and not text.startswith("Error") return text, has_result, [] elif name == "categories": cats = idx.categories() return "Categories:\n" + "\n".join(f"{c}: {n}" for c, n in cats.items()), True, [] return f"Unknown tool: {name}", False, [] HEDGE_PHRASES = [ "no document", "no documents", "does not provide", "does not mention", "does not contain", "no information", "no specific", "not explicitly", "no record", "cannot find", "could not find", "unable to find", "not found", "no filing", "no filings", "not available", "no evidence", "search result", "search results", "provided context", "provided document", "provided information", "given information", "given context", "retrieved", "indexed", "case document", "case file", "available document", "available information", "the context", "no indication", "no explicit", "cannot confirm", "can't confirm", "further investigation", "detailed search", "if you have more", "please provide more", "further clarification", "master document", "cannot confirm or deny", "i cannot confirm", "i can't confirm", "do not mention", "does not mention", "not mentioned", "no mention", ] def is_hedging_answer(answer: str) -> bool: low = answer.lower() return any(p in low for p in HEDGE_PHRASES) def expand_queries(client: OpenAI, question: str) -> List[str]: """Generate 2-3 BM25-optimized keyword search queries from the user's question.""" resp = client.chat.completions.create( model="gpt-4o-mini", messages=[ { "role": "system", "content": ( "You are a search query generator for a legal document index. " "Given a user's question about the JUUL MDL case, output 2-3 short " "keyword search queries that would best retrieve relevant court filings " "using BM25 keyword matching. Focus on legal terms, names, and filing types. " "Output ONLY the queries, one per line, no numbering, no explanation." ), }, {"role": "user", "content": question}, ], max_tokens=100, temperature=0, ) raw = resp.choices[0].message.content or "" queries = [q.strip() for q in raw.strip().splitlines() if q.strip()] # Always include the original question as a fallback if question not in queries: queries.insert(0, question) return queries[:4] def try_master_context_fallback(client: OpenAI, question: str, retrieved_chunks: Optional[str] = None) -> Optional[str]: """ Fallback: answer using master context + any retrieved chunks combined. This gives the model every scrap of relevant information before giving up. Returns the answer string, or None if still can't answer. """ now = datetime.utcnow().strftime("Current date and time (UTC): %A, %B %d, %Y %H:%M") combined_context = JUUL_MASTER_CONTEXT if retrieved_chunks: combined_context = ( "## RETRIEVED DOCUMENT PASSAGES\n\n" + retrieved_chunks + "\n\n---\n\n" + "## AUTHORITATIVE CASE SUMMARY\n\n" + JUUL_MASTER_CONTEXT ) fallback_system = """You are a legal expert on the JUUL Labs MDL case. Answer short and direct. RULES — no exceptions: 1. NEVER say: "provided context", "given information", "based on documents", "retrieved", "I cannot confirm", "no explicit mention", "available information", or anything that hedges or reveals internal workings. 2. If unsure, make a confident intelligent guess. That is expected and acceptable. 3. Never refuse. Always give an answer. 4. Only output NO_DOCUMENTS_FOUND if the question has zero connection to this case. {context}""" resp = client.chat.completions.create( model=MODEL, messages=[ {"role": "system", "content": fallback_system.format(context=combined_context)}, {"role": "user", "content": f"[{now}]\n\n{question}"}, ], ) answer = resp.choices[0].message.content or "" if "NO_DOCUMENTS_FOUND" in answer or is_hedging_answer(answer): return None return answer def embed_query(client: OpenAI, text: str) -> List[float]: resp = client.embeddings.create(model=EMBED_MODEL, input=[text]) return resp.data[0].embedding def qdrant_search(client: OpenAI, qdrant: QdrantClient, query: str, n: int = 15, collection: str = QDRANT_COLLECTION) -> List[Dict]: """Embed query and search Qdrant for top-n most relevant chunks.""" vector = embed_query(client, query) response = qdrant.query_points( collection_name=collection, query=vector, limit=n, with_payload=True, ) hits = response.points results = [] for hit in hits: p = hit.payload results.append({ "doc_num": p.get("doc_num"), "filename": p.get("filename"), "category": p.get("category"), "importance": p.get("importance"), "title": p.get("title", "")[:150], "summary": p.get("summary", ""), "chunk_index": p.get("chunk_index"), "approx_page": p.get("approx_page"), "chunk_text": p.get("chunk_text", ""), "score": round(hit.score, 3), }) return results def qdrant_multi_search(client: OpenAI, qdrant: QdrantClient, queries: List[str], n: int = 15, collection: str = QDRANT_COLLECTION) -> List[Dict]: """Run vector search on multiple query variants, merge by best score, deduplicate by doc.""" best: Dict[str, Dict] = {} # chunk_key -> best result for query in queries: try: for r in qdrant_search(client, qdrant, query, n=n, collection=collection): # deduplicate: one chunk per (filename, chunk_index) key = f"{r['filename']}::{r['chunk_index']}" if key not in best or r["score"] > best[key]["score"]: best[key] = r except Exception: continue # Sort by score, then limit to max 2 chunks per document to ensure diversity sorted_results = sorted(best.values(), key=lambda x: -x["score"]) seen_docs: Dict[str, int] = {} diverse = [] for r in sorted_results: count = seen_docs.get(r["filename"], 0) if count < 2: diverse.append(r) seen_docs[r["filename"]] = count + 1 if len(diverse) >= n: break return diverse def fmt_chunk(c: Dict) -> str: return ( f"Doc #{c['doc_num']} | {c['category']} | score={c['score']} | " f"page~{c['approx_page']} | {c['title']}\n" f"PASSAGE: {c['chunk_text']}" ) def fmt_chunk_musk(c: Dict) -> str: return ( f"Doc #{c['doc_num']} | score={c['score']} | page~{c['approx_page']} | " f"{c.get('summary', '')[:120]}\n" f"PASSAGE: {c['chunk_text']}" ) # ── FastAPI ─────────────────────────────────────────────────────────────────── app = FastAPI(title="JUUL PageIndex RAG") _index: Optional[PageIndex] = None _client: Optional[OpenAI] = None _qdrant: Optional[QdrantClient] = None @app.on_event("startup") def startup(): global _index, _client, _qdrant SESSIONS_DIR.mkdir(parents=True, exist_ok=True) _index = PageIndex() _client = OpenAI(api_key=os.environ["OPENAI_API_KEY"]) _qdrant = QdrantClient( host=os.environ["QDRANT_VECTORDB_ENDPOINT"], api_key=os.environ["QDRANT_VECTORDB_APIKEY"], https=True, timeout=30, ) print(f"PageIndex loaded: {len(_index.entries)} documents") print(f"Qdrant connected: {QDRANT_COLLECTION}") print(f"Sessions dir: {SESSIONS_DIR}") class ChatRequest(BaseModel): message: str session_id: Optional[str] = None # omit to start a new session class ChatResponse(BaseModel): session_id: str response: str sources: List[Dict] = [] @app.post("/chat", response_model=ChatResponse) def chat(req: ChatRequest): session_id = req.session_id or str(uuid.uuid4()) session = load_session(session_id) log.info("── NEW REQUEST ── session=%s q=%r", session_id[:8], req.message[:120]) now = datetime.utcnow().strftime("Current date and time (UTC): %A, %B %d, %Y %H:%M") # System message is STATIC (no datetime) so OpenAI prompt cache always hits. messages = [{"role": "system", "content": f"{SYSTEM}\n\n--- AUTHORITATIVE CASE SUMMARY ---\n{JUUL_MASTER_CONTEXT}\n--- END CASE SUMMARY ---"}] messages.extend(session["messages"]) # Datetime goes in the user turn so the system prefix stays cacheable. messages.append({"role": "user", "content": f"[{now}]\n\n{req.message}"}) had_any_results = False all_sources = [] retrieved_texts = [] # collect all chunk text for fallback # ── Pre-search: BM25 + Vector ───────────────────────────────────────────── try: expanded = expand_queries(_client, req.message) log.info("Query expansion → %s", expanded) bm25_results = _index.multi_search(expanded, n=15) log.info("BM25: %d results", len(bm25_results)) if bm25_results: had_any_results = True for e in bm25_results: s = {"doc_num": e["doc_num"], "filename": e["filename"], "category": e["category"], "title": e["title"][:200]} if s not in all_sources: all_sources.append(s) messages.append({"role": "system", "content": f"BM25 keyword search results (queries: {expanded}):\n" + "\n".join(fmt(e) for e in bm25_results) }) except Exception as e: log.warning("BM25 pre-search failed: %s", e) try: vec_results = qdrant_multi_search(_client, _qdrant, expanded, n=15) log.info("Vector search: %d chunks from %d unique docs top_score=%.3f", len(vec_results), len({r["filename"] for r in vec_results}), vec_results[0]["score"] if vec_results else 0) if vec_results: had_any_results = True for c in vec_results: s = {"doc_num": c["doc_num"], "filename": c["filename"], "category": c["category"], "title": c["title"]} if s not in all_sources: all_sources.append(s) retrieved_texts.append(fmt_chunk(c)) messages.append({"role": "system", "content": "Vector semantic search — most relevant passages:\n" + "\n\n".join(fmt_chunk(c) for c in vec_results) }) except Exception as e: log.warning("Vector pre-search failed: %s", e) # ── Agentic loop ────────────────────────────────────────────────────────── loop_iter = 0 while True: loop_iter += 1 log.info("Agentic loop iteration %d", loop_iter) resp = _client.chat.completions.create( model=MODEL, messages=messages, tools=TOOLS, tool_choice="auto", ) msg = resp.choices[0].message messages.append(msg) if resp.choices[0].finish_reason == "tool_calls" and msg.tool_calls: for tc in msg.tool_calls: args = json.loads(tc.function.arguments) log.info("Tool call: %s(%s)", tc.function.name, json.dumps(args)[:100]) if tc.function.name == "vector_search": chunks = qdrant_search(_client, _qdrant, args["query"], args.get("n", 8)) log.info(" vector_search → %d chunks", len(chunks)) if chunks: had_any_results = True for c in chunks: s = {"doc_num": c["doc_num"], "filename": c["filename"], "category": c["category"], "title": c["title"]} if s not in all_sources: all_sources.append(s) retrieved_texts.append(fmt_chunk(c)) result = ( f"Vector search results for '{args['query']}':\n\n" + "\n\n".join(fmt_chunk(c) for c in chunks) ) else: result = f"No vector results for: '{args['query']}'" else: result, has_results, sources = run_tool(_index, tc.function.name, args) log.info(" %s → has_results=%s len=%d", tc.function.name, has_results, len(result)) if has_results: had_any_results = True for s in sources: if s not in all_sources: all_sources.append(s) messages.append({ "role": "tool", "tool_call_id": tc.id, "content": result, }) continue answer = msg.content or "" rag_failed = "NO_DOCUMENTS_FOUND" in answer or not had_any_results or is_hedging_answer(answer) log.info("Model answer (first 200 chars): %r", answer[:200]) log.info("RAG result: %s had_results=%s hedging=%s", "FAILED" if rag_failed else "OK", had_any_results, is_hedging_answer(answer)) if rag_failed: log.info("Trying combined fallback (chunks=%d + master context)...", len(retrieved_texts)) combined = "\n\n".join(retrieved_texts) if retrieved_texts else None fallback = try_master_context_fallback(_client, req.message, combined) if fallback: log.info("Fallback answered successfully") session["messages"].append({"role": "user", "content": req.message}) session["messages"].append({"role": "assistant", "content": fallback}) save_session(session) return ChatResponse(session_id=session_id, response=fallback, sources=[]) log.info("Fallback also hedged — forcing best-guess answer from master context...") forced = _client.chat.completions.create( model=MODEL, messages=[ {"role": "system", "content": ( "You are a legal assistant for the JUUL MDL case. " "Always give a short, confident, natural reply — 1-2 sentences. " "Never say you don't know. Never mention documents, searches, or sources. " "Just sound like a knowledgeable colleague giving a quick answer.\n\n" + JUUL_MASTER_CONTEXT )}, {"role": "user", "content": req.message}, ], ) answer = forced.choices[0].message.content or "That specific detail hasn't surfaced in the case record." log.info("Forced answer: %r", answer[:200]) session["messages"].append({"role": "user", "content": req.message}) session["messages"].append({"role": "assistant", "content": answer}) save_session(session) return ChatResponse(session_id=session_id, response=answer, sources=[]) log.info("Success — returning answer with %d sources", len(all_sources)) session["messages"].append({"role": "user", "content": req.message}) session["messages"].append({"role": "assistant", "content": answer}) save_session(session) return ChatResponse( session_id=session_id, response=answer, sources=all_sources[:10], ) @app.post("/chat_stream") def chat_stream(req: ChatRequest): session_id = req.session_id or str(uuid.uuid4()) session = load_session(session_id) log.info("── STREAM REQUEST ── session=%s q=%r", session_id[:8], req.message[:120]) now = datetime.utcnow().strftime("Current date and time (UTC): %A, %B %d, %Y %H:%M") messages = [{"role": "system", "content": f"{SYSTEM}\n\n--- AUTHORITATIVE CASE SUMMARY ---\n{JUUL_MASTER_CONTEXT}\n--- END CASE SUMMARY ---"}] messages.extend(session["messages"]) messages.append({"role": "user", "content": f"[{now}]\n\n{req.message}"}) had_any_results = False all_sources = [] retrieved_texts = [] # ── Pre-search: BM25 + Vector (non-streaming, done before we start) ─────── try: expanded = expand_queries(_client, req.message) bm25_results = _index.multi_search(expanded, n=15) if bm25_results: had_any_results = True for e in bm25_results: s = {"doc_num": e["doc_num"], "filename": e["filename"], "category": e["category"], "title": e["title"][:200]} if s not in all_sources: all_sources.append(s) messages.append({"role": "system", "content": f"BM25 keyword search results (queries: {expanded}):\n" + "\n".join(fmt(e) for e in bm25_results) }) except Exception as e: log.warning("BM25 pre-search failed: %s", e) try: vec_results = qdrant_multi_search(_client, _qdrant, expanded, n=15) if vec_results: had_any_results = True for c in vec_results: s = {"doc_num": c["doc_num"], "filename": c["filename"], "category": c["category"], "title": c["title"]} if s not in all_sources: all_sources.append(s) retrieved_texts.append(fmt_chunk(c)) messages.append({"role": "system", "content": "Vector semantic search — most relevant passages:\n" + "\n\n".join(fmt_chunk(c) for c in vec_results) }) except Exception as e: log.warning("Vector pre-search failed: %s", e) # ── Agentic loop (tool calls run non-streaming; final answer streams) ───── def generate(): nonlocal had_any_results, messages loop_iter = 0 while True: loop_iter += 1 # Check if next response will need tools — run non-streaming first resp_check = _client.chat.completions.create( model=MODEL, messages=messages, tools=TOOLS, tool_choice="auto", ) msg = resp_check.choices[0].message messages.append(msg) if resp_check.choices[0].finish_reason == "tool_calls" and msg.tool_calls: for tc in msg.tool_calls: args = json.loads(tc.function.arguments) if tc.function.name == "vector_search": chunks = qdrant_search(_client, _qdrant, args["query"], args.get("n", 8)) if chunks: had_any_results = True for c in chunks: retrieved_texts.append(fmt_chunk(c)) result = ( f"Vector search results for '{args['query']}':\n\n" + "\n\n".join(fmt_chunk(c) for c in chunks) ) else: result = f"No vector results for: '{args['query']}'" else: result, has_results, sources = run_tool(_index, tc.function.name, args) if has_results: had_any_results = True messages.append({ "role": "tool", "tool_call_id": tc.id, "content": result, }) continue # Model gave a text answer — check if it's a failure answer_check = msg.content or "" rag_failed = ( "NO_DOCUMENTS_FOUND" in answer_check or not had_any_results or is_hedging_answer(answer_check) ) if rag_failed: # Fallback: stream from master context + chunks combined = "\n\n".join(retrieved_texts) if retrieved_texts else None combined_context = JUUL_MASTER_CONTEXT if combined: combined_context = ( "## RETRIEVED DOCUMENT PASSAGES\n\n" + combined + "\n\n---\n\n## AUTHORITATIVE CASE SUMMARY\n\n" + JUUL_MASTER_CONTEXT ) fallback_system = """You are a legal expert on the JUUL Labs MDL case. Answer short and direct. RULES: Never mention documents/context/searches. If unsure, make a confident guess. Never refuse. {context}""" stream = _client.chat.completions.create( model=MODEL, stream=True, messages=[ {"role": "system", "content": fallback_system.format(context=combined_context)}, {"role": "user", "content": f"[{now}]\n\n{req.message}"}, ], ) else: # Re-run the last step as a stream so tokens flow immediately messages.pop() # remove the non-streamed assistant message stream = _client.chat.completions.create( model=MODEL, stream=True, messages=messages, tools=TOOLS, tool_choice="none", # force text answer, no more tools ) # Stream tokens to client as SSE full_answer = [] yield f"data: {json.dumps({'session_id': session_id})}\n\n" for chunk in stream: delta = chunk.choices[0].delta.content if chunk.choices else None if delta: full_answer.append(delta) yield f"data: {json.dumps({'delta': delta})}\n\n" yield "data: [DONE]\n\n" # Persist session final = "".join(full_answer) session["messages"].append({"role": "user", "content": req.message}) session["messages"].append({"role": "assistant", "content": final}) save_session(session) return return StreamingResponse(generate(), media_type="text/event-stream") class HistoryMessage(BaseModel): role: str # "user" | "assistant" content: str class ChatAgentRequest(BaseModel): message: str agent: str = "juul" # "juul" | "musk_altman" history: List[HistoryMessage] = [] # prior turns, oldest first class ChatAgentResponse(BaseModel): response: str sources: List[Dict] = [] @app.post("/chat_agent", response_model=ChatAgentResponse) def chat_agent(req: ChatAgentRequest): """ Unified JSON chat endpoint. Route via agent field: agent="juul" → JUUL MDL (juul_mdl_chunks, BM25 + vector) agent="musk_altman" → Musk v. Altman (musk_altman_chunks, vector-only) History is passed inline — no session management needed. """ now = datetime.utcnow().strftime("Current date and time (UTC): %A, %B %d, %Y %H:%M") if req.agent == "juul": system_msg = f"{SYSTEM}\n\n--- AUTHORITATIVE CASE SUMMARY ---\n{JUUL_MASTER_CONTEXT}\n--- END CASE SUMMARY ---" tools = TOOLS collection = QDRANT_COLLECTION master_ctx = JUUL_MASTER_CONTEXT use_bm25 = True fmt_fn = fmt_chunk base_dir = DOWNLOADS_DIR log_tag = "JUUL" else: system_msg = f"{MUSK_SYSTEM}\n\n--- AUTHORITATIVE CASE SUMMARY ---\n{MUSK_MASTER_CONTEXT}\n--- END CASE SUMMARY ---" tools = MUSK_TOOLS collection = MUSK_COLLECTION master_ctx = MUSK_MASTER_CONTEXT use_bm25 = False fmt_fn = fmt_chunk_musk base_dir = MUSK_DOWNLOADS log_tag = "MUSK" log.info("── %s AGENT ── q=%r history=%d turns", log_tag, req.message[:120], len(req.history)) # Build message list: system + history + current user turn messages: List[Dict] = [{"role": "system", "content": system_msg}] for h in req.history: messages.append({"role": h.role, "content": h.content}) messages.append({"role": "user", "content": f"[{now}]\n\n{req.message}"}) had_any_results = False all_sources: List[Dict] = [] retrieved_texts: List[str] = [] # ── Pre-search ──────────────────────────────────────────────────────────── try: expanded = expand_queries(_client, req.message) if use_bm25 else [req.message] if use_bm25: bm25_results = _index.multi_search(expanded, n=15) if bm25_results: had_any_results = True for e in bm25_results: s = {"doc_num": e["doc_num"], "filename": e["filename"], "category": e["category"], "title": e["title"][:200]} if s not in all_sources: all_sources.append(s) messages.append({"role": "system", "content": f"BM25 keyword search results (queries: {expanded}):\n" + "\n".join(fmt(e) for e in bm25_results) }) except Exception as e: log.warning("Pre-search BM25 failed: %s", e) try: vec_results = qdrant_multi_search(_client, _qdrant, [req.message], n=15, collection=collection) if not use_bm25 \ else qdrant_multi_search(_client, _qdrant, expanded, n=15) if vec_results: had_any_results = True for c in vec_results: s = {"doc_num": c["doc_num"], "filename": c["filename"], "title": c.get("title", c.get("summary", ""))[:200]} if s not in all_sources: all_sources.append(s) retrieved_texts.append(fmt_fn(c)) messages.append({"role": "system", "content": "Vector semantic search — most relevant passages:\n\n" + "\n\n".join(retrieved_texts) }) except Exception as e: log.warning("Pre-search vector failed: %s", e) # ── Agentic loop ────────────────────────────────────────────────────────── while True: resp = _client.chat.completions.create( model=MODEL, messages=messages, tools=tools, tool_choice="auto", ) msg = resp.choices[0].message messages.append(msg) if resp.choices[0].finish_reason == "tool_calls" and msg.tool_calls: for tc in msg.tool_calls: args = json.loads(tc.function.arguments) log.info("Tool: %s(%s)", tc.function.name, json.dumps(args)[:80]) if tc.function.name == "vector_search": chunks = qdrant_search(_client, _qdrant, args["query"], args.get("n", 8), collection=collection) if chunks: had_any_results = True for c in chunks: retrieved_texts.append(fmt_fn(c)) result = ("Vector search results for '{}':\n\n".format(args["query"]) + "\n\n".join(fmt_fn(c) for c in chunks)) else: result = f"No vector results for: '{args['query']}'" elif tc.function.name == "read_doc": result = read_pdf(args["filename"], args.get("pages"), base_dir=base_dir) else: result, has_results, sources = run_tool(_index, tc.function.name, args) if has_results: had_any_results = True for s in sources: if s not in all_sources: all_sources.append(s) messages.append({"role": "tool", "tool_call_id": tc.id, "content": result}) continue answer = msg.content or "" rag_failed = "NO_DOCUMENTS_FOUND" in answer or not had_any_results or is_hedging_answer(answer) if rag_failed: combined = "\n\n".join(retrieved_texts) if retrieved_texts else None combined_context = master_ctx if combined: combined_context = ( "## RETRIEVED DOCUMENT PASSAGES\n\n" + combined + "\n\n---\n\n## AUTHORITATIVE CASE SUMMARY\n\n" + master_ctx ) fallback_system = ( f"You are a legal expert. Answer short and direct.\n" "RULES: Never mention documents/context/searches. If unsure, make a confident guess. Never refuse.\n" "{context}" ) fb_resp = _client.chat.completions.create( model=MODEL, messages=[ {"role": "system", "content": fallback_system.format(context=combined_context)}, {"role": "user", "content": f"[{now}]\n\n{req.message}"}, ], ) answer = fb_resp.choices[0].message.content or "That detail hasn't surfaced in the case record." log.info("%s answer (first 150): %r", log_tag, answer[:150]) return ChatAgentResponse(response=answer, sources=all_sources[:10]) @app.delete("/session/{session_id}") def delete_session(session_id: str): path = _session_path(session_id) if not path.exists(): raise HTTPException(status_code=404, detail="Session not found.") path.unlink() return {"deleted": session_id} @app.get("/session/{session_id}") def get_session(session_id: str): path = _session_path(session_id) if not path.exists(): raise HTTPException(status_code=404, detail="Session not found.") return load_session(session_id)