rag/api.py

#!/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)