Upload agent.py

agent.py

@@ -0,0 +1,385 @@
+"""
+ClinIQ — Multi-step LangGraph agent.
+
+Graph:  classify_query → [decompose] → retrieve → generate → reflect ──→ END
+                                                                  ↑__________| (retry if not grounded, max 2x)
+
+Query types:
+  simple      — single-hop fact lookup  ("What is the diagnosis?")
+  structured  — extract as JSON list    ("List all medications")
+  complex     — multi-hop reasoning     ("Do any meds interact with the allergy?")
+  comparison  — across multiple docs    ("How did medications change between visits?")
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import time
+from typing import Any, Dict, Iterator, List, Literal, Optional, TypedDict
+
+import httpx
+from langgraph.graph import END, StateGraph
+
+from retriever import Chunk, HybridRetriever
+
+MODAL_ENDPOINT = os.getenv("MODAL_ENDPOINT", "")
+MAX_CONTEXT_CHARS = 4000
+MAX_RETRIES = 2
+
+SYSTEM_PROMPT = """You are ClinIQ, a clinical document assistant for small medical clinics.
+Answer ONLY from the provided document excerpts. Cite the source document for every claim.
+If information is not in the excerpts, say clearly: "Not found in the provided documents."
+Be concise and medically precise. Never guess or hallucinate clinical information."""
+
+STRUCTURED_SCHEMAS = {
+    "medications":  ("List every medication with dose and frequency as JSON array: "
+                     '[{"name":"...","dose":"...","frequency":"...","route":"..."}]'),
+    "allergies":    ('List every allergy with reaction as JSON array: '
+                     '[{"substance":"...","reaction":"...","severity":"..."}]'),
+    "diagnoses":    ('List all diagnoses/conditions as JSON array: '
+                     '[{"diagnosis":"...","type":"primary|secondary"}]'),
+    "followup":     ('List all follow-up appointments/instructions as JSON array: '
+                     '[{"action":"...","date":"...","provider":"..."}]'),
+    "vitals":       ('Extract all vital signs as JSON object: '
+                     '{"bp":"...","hr":"...","rr":"...","temp":"...","spo2":"...","weight":"..."}'),
+}
+
+
+# ── State ──────────────────────────────────────────────────────────────────────
+
+class AgentState(TypedDict):
+    question: str
+    query_type: str                    # simple | structured | complex | comparison
+    struct_key: Optional[str]          # which STRUCTURED_SCHEMAS key if structured
+    sub_queries: List[str]             # decomposed queries for complex/comparison
+    chunks: List[Chunk]
+    context: str
+    answer: str
+    structured_data: Optional[Any]     # parsed JSON for structured answers
+    reflection_ok: bool
+    reflection_note: str
+    retry_count: int
+    trace: List[Dict[str, Any]]
+
+
+# ── Helpers ────────────────────────────────────────────────────────────────────
+
+def _call_model(prompt: str, max_tokens: int = 600, json_mode: bool = False) -> str:
+    """Call Modal llama.cpp endpoint (or local fallback)."""
+    if MODAL_ENDPOINT:
+        resp = httpx.post(
+            MODAL_ENDPOINT,
+            json={"prompt": prompt, "max_tokens": max_tokens, "json_mode": json_mode},
+            timeout=120,
+        )
+        resp.raise_for_status()
+        return resp.json()["text"].strip()
+    return _local_fallback(prompt, max_tokens)
+
+
+def _local_fallback(prompt: str, max_tokens: int) -> str:
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline as hf_pipeline
+
+    MODEL_ID = os.getenv("MODEL_ID", "Qwen/Qwen2.5-3B-Instruct")
+    if not hasattr(_local_fallback, "_pipe"):
+        tok = AutoTokenizer.from_pretrained(MODEL_ID)
+        model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32)
+        _local_fallback._pipe = hf_pipeline("text-generation", model=model, tokenizer=tok)
+    out = _local_fallback._pipe(prompt, max_new_tokens=max_tokens, do_sample=False,
+                                temperature=None, top_p=None)
+    raw = out[0]["generated_text"]
+    return raw[len(prompt):].strip()
+
+
+def _build_prompt(system: str, user: str) -> str:
+    return (f"<|im_start|>system\n{system}<|im_end|>\n"
+            f"<|im_start|>user\n{user}<|im_end|>\n"
+            f"<|im_start|>assistant\n")
+
+
+def _detect_struct_key(question: str) -> Optional[str]:
+    q = question.lower()
+    if any(w in q for w in ["medication", "medicine", "drug", "prescribed", "taking", "rx"]):
+        return "medications"
+    if any(w in q for w in ["allerg"]):
+        return "allergies"
+    if any(w in q for w in ["diagnos", "condition", "problem list"]):
+        return "diagnoses"
+    if any(w in q for w in ["follow", "appointment", "next visit", "schedule"]):
+        return "followup"
+    if any(w in q for w in ["vital", "blood pressure", "heart rate", "bp ", "hr ", "temp"]):
+        return "vitals"
+    return None
+
+
+# ── Nodes ──────────────────────────────────────────────────────────────────────
+
+def node_classify(state: AgentState) -> AgentState:
+    t0 = time.time()
+    q = state["question"].lower()
+
+    # Rule-based classification (fast, no LLM call needed)
+    is_list_question = any(w in q for w in ["list", "all ", "what are", "enumerate", "summarize"])
+    struct_key = _detect_struct_key(state["question"])
+    is_comparison = any(w in q for w in ["changed", "different", "compare", "between", "versus", "vs", "previous"])
+    is_complex = any(w in q for w in ["interact", "relate", "given", "safe", "why", "how does", "affect"])
+
+    if struct_key and is_list_question:
+        qtype = "structured"
+    elif is_comparison:
+        qtype = "comparison"
+    elif is_complex:
+        qtype = "complex"
+    else:
+        qtype = "simple"
+
+    state["query_type"] = qtype
+    state["struct_key"] = struct_key if qtype == "structured" else None
+    state["trace"].append({
+        "step": "classify",
+        "icon": "🔍",
+        "detail": f"Query type: **{qtype}**" + (f" (extracting: {struct_key})" if struct_key else ""),
+        "ms": int((time.time() - t0) * 1000),
+    })
+    return state
+
+
+def node_decompose(state: AgentState) -> AgentState:
+    """For complex/comparison queries, break into focused sub-queries."""
+    t0 = time.time()
+    if state["query_type"] in ("simple", "structured"):
+        state["sub_queries"] = [state["question"]]
+        return state
+
+    prompt = _build_prompt(
+        "You are a query decomposition assistant. Break the user's question into 2-3 focused sub-queries "
+        "that can each be answered independently from a medical document. Output ONLY a JSON array of strings.",
+        f"Question: {state['question']}\nOutput format: [\"sub-query 1\", \"sub-query 2\"]"
+    )
+    try:
+        raw = _call_model(prompt, max_tokens=200)
+        match = re.search(r'\[.*?\]', raw, re.DOTALL)
+        sub_queries = json.loads(match.group()) if match else [state["question"]]
+    except Exception:
+        sub_queries = [state["question"]]
+
+    state["sub_queries"] = sub_queries[:3]
+    state["trace"].append({
+        "step": "decompose",
+        "icon": "✂️",
+        "detail": f"Split into **{len(sub_queries)}** sub-queries",
+        "sub_queries": sub_queries,
+        "ms": int((time.time() - t0) * 1000),
+    })
+    return state
+
+
+def node_retrieve(state: AgentState, retriever: HybridRetriever) -> AgentState:
+    t0 = time.time()
+    seen_ids: set = set()
+    all_chunks: List[Chunk] = []
+
+    for sq in state["sub_queries"]:
+        for c in retriever.retrieve(sq, top_k=4):
+            uid = (c.source, c.index)
+            if uid not in seen_ids:
+                seen_ids.add(uid)
+                all_chunks.append(c)
+
+    # Limit total chunks
+    all_chunks = all_chunks[:8]
+    state["chunks"] = all_chunks
+
+    sources = list(dict.fromkeys(c.source for c in all_chunks))
+    state["trace"].append({
+        "step": "retrieve",
+        "icon": "📄",
+        "detail": f"Retrieved **{len(all_chunks)}** chunks from {len(sources)} document(s)",
+        "sources": sources,
+        "ms": int((time.time() - t0) * 1000),
+    })
+    return state
+
+
+def node_build_context(state: AgentState) -> AgentState:
+    parts: List[str] = []
+    total = 0
+    for c in state["chunks"]:
+        snippet = f"[Document: {c.source}]\n{c.text}"
+        if total + len(snippet) > MAX_CONTEXT_CHARS:
+            break
+        parts.append(snippet)
+        total += len(snippet)
+    state["context"] = "\n\n---\n\n".join(parts)
+    return state
+
+
+def node_generate(state: AgentState) -> AgentState:
+    t0 = time.time()
+    qtype = state["query_type"]
+
+    if qtype == "structured" and state["struct_key"]:
+        schema_hint = STRUCTURED_SCHEMAS[state["struct_key"]]
+        user_msg = (
+            f"Document excerpts:\n{state['context']}\n\n"
+            f"Task: {schema_hint}\n"
+            f"Question: {state['question']}\n"
+            f"Output ONLY valid JSON, no explanation."
+        )
+        raw = _call_model(_build_prompt(SYSTEM_PROMPT, user_msg), max_tokens=800, json_mode=True)
+        # Try to parse JSON
+        try:
+            match = re.search(r'(\[.*?\]|\{.*?\})', raw, re.DOTALL)
+            parsed = json.loads(match.group()) if match else None
+        except Exception:
+            parsed = None
+        state["structured_data"] = parsed
+        state["answer"] = raw if not parsed else json.dumps(parsed, indent=2)
+    else:
+        user_msg = (
+            f"Document excerpts:\n{state['context']}\n\n"
+            f"Question: {state['question']}"
+        )
+        state["answer"] = _call_model(_build_prompt(SYSTEM_PROMPT, user_msg), max_tokens=600)
+        state["structured_data"] = None
+
+    state["trace"].append({
+        "step": "generate",
+        "icon": "🧠",
+        "detail": f"Generated answer via Qwen2.5-3B-Instruct ({'structured JSON' if qtype == 'structured' else 'free text'})",
+        "ms": int((time.time() - t0) * 1000),
+    })
+    return state
+
+
+def node_reflect(state: AgentState) -> AgentState:
+    """Check if the answer is grounded in the context. Flag if hallucinated."""
+    t0 = time.time()
+
+    # Skip reflection for structured data (it's JSON, harder to verify this way)
+    if state["query_type"] == "structured" and state["structured_data"]:
+        state["reflection_ok"] = True
+        state["reflection_note"] = "Structured extraction — skipped."
+        return state
+
+    # Quick heuristic: if answer says "not found" it's honest, skip
+    answer_lower = state["answer"].lower()
+    if "not found" in answer_lower or "not mentioned" in answer_lower or "not in" in answer_lower:
+        state["reflection_ok"] = True
+        state["reflection_note"] = "Model self-reported missing information — answer is honest."
+        state["trace"].append({
+            "step": "reflect",
+            "icon": "✅",
+            "detail": "Answer honest (model flagged missing info)",
+            "ms": int((time.time() - t0) * 1000),
+        })
+        return state
+
+    prompt = _build_prompt(
+        "You are a clinical fact-checker. Given an answer and source excerpts, "
+        "decide if the answer is fully supported by the excerpts. "
+        "Reply with JSON: {\"grounded\": true|false, \"note\": \"reason\"}",
+        f"Answer:\n{state['answer'][:800]}\n\nSource excerpts:\n{state['context'][:1500]}"
+    )
+    try:
+        raw = _call_model(prompt, max_tokens=150)
+        match = re.search(r'\{.*?\}', raw, re.DOTALL)
+        data = json.loads(match.group()) if match else {"grounded": True, "note": "parse error"}
+        ok = bool(data.get("grounded", True))
+        note = data.get("note", "")
+    except Exception:
+        ok, note = True, "Reflection parse error — defaulting to accept."
+
+    state["reflection_ok"] = ok
+    state["reflection_note"] = note
+    state["trace"].append({
+        "step": "reflect",
+        "icon": "✅" if ok else "⚠️",
+        "detail": f"Grounded: **{'yes' if ok else 'no — retrying'}** — {note}",
+        "ms": int((time.time() - t0) * 1000),
+    })
+    return state
+
+
+def _should_retry(state: AgentState) -> str:
+    if not state["reflection_ok"] and state["retry_count"] < MAX_RETRIES:
+        state["retry_count"] += 1
+        # Broaden query on retry
+        state["sub_queries"] = [state["question"] + " (provide only information explicitly stated)"]
+        return "retrieve"
+    return END
+
+
+# ── Graph ──────────────────────────────────────────────────────────────────────
+
+def build_graph(retriever: HybridRetriever):
+    g = StateGraph(AgentState)
+    g.add_node("classify",      node_classify)
+    g.add_node("decompose",     node_decompose)
+    g.add_node("retrieve",      lambda s: node_retrieve(s, retriever))
+    g.add_node("build_context", node_build_context)
+    g.add_node("generate",      node_generate)
+    g.add_node("reflect",       node_reflect)
+
+    g.set_entry_point("classify")
+    g.add_edge("classify",      "decompose")
+    g.add_edge("decompose",     "retrieve")
+    g.add_edge("retrieve",      "build_context")
+    g.add_edge("build_context", "generate")
+    g.add_edge("generate",      "reflect")
+    g.add_conditional_edges("reflect", _should_retry, {"retrieve": "retrieve", END: END})
+    return g.compile()
+
+
+def run_query(graph, question: str) -> AgentState:
+    return graph.invoke({
+        "question": question,
+        "query_type": "simple",
+        "struct_key": None,
+        "sub_queries": [question],
+        "chunks": [],
+        "context": "",
+        "answer": "",
+        "structured_data": None,
+        "reflection_ok": True,
+        "reflection_note": "",
+        "retry_count": 0,
+        "trace": [],
+    })
+
+
+def stream_query(graph, question: str) -> Iterator[Dict[str, Any]]:
+    """
+    Yields trace dicts as each node completes, then a final 'done' event.
+    Uses LangGraph streaming.
+    """
+    init: AgentState = {
+        "question": question,
+        "query_type": "simple",
+        "struct_key": None,
+        "sub_queries": [question],
+        "chunks": [],
+        "context": "",
+        "answer": "",
+        "structured_data": None,
+        "reflection_ok": True,
+        "reflection_note": "",
+        "retry_count": 0,
+        "trace": [],
+    }
+    for event in graph.stream(init, stream_mode="updates"):
+        for node_name, state_update in event.items():
+            new_trace = state_update.get("trace", [])
+            if new_trace:
+                yield {"type": "trace_step", "step": new_trace[-1], "node": node_name}
+            if "answer" in state_update and state_update["answer"]:
+                yield {
+                    "type": "answer",
+                    "answer": state_update["answer"],
+                    "structured_data": state_update.get("structured_data"),
+                    "chunks": [{"source": c.source, "excerpt": c.text[:350]}
+                               for c in state_update.get("chunks", [])[:4]],
+                }