agent.py

"""
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 = 5000
MAX_RETRIES = 0   # reflection shows in trace but never retries (3B too small for reliable self-grounding)

SYSTEM_PROMPT = """You are ClinIQ, a clinical document assistant for small medical clinics.
Answer ONLY from the provided document excerpts. Be concise and medically precise.
If information is explicitly stated in the excerpts, answer from it directly.
Only say "Not found in the provided documents" if the information is genuinely absent from ALL excerpts.
Never hallucinate clinical information. Do not append source citations to your answer."""

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 _sanitize(text: str) -> str:
    import re
    return re.sub(r'[^\x00-\x7F]', '-', text)


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=300,
        )
        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=6):
            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"] = _sanitize("\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

    # Heuristic grounding check — fast, no extra LLM call needed
    ok = True
    note = "Answer accepted."
    # Flag if answer seems to fabricate — but small model self-grounding is unreliable,
    # so we only hard-reject obvious complete misses
    answer_lower = state["answer"].lower()
    context_lower = state["context"].lower()
    # If answer is very long but context has no overlap at all, flag it
    if len(state["answer"]) > 100:
        answer_words = set(answer_lower.split())
        context_words = set(context_lower.split())
        overlap = len(answer_words & context_words)
        if overlap < 5:
            ok = False
            note = "Low overlap between answer and retrieved context — may be hallucinated."

    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'}** — {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
        state["sub_queries"] = [state["question"]]
        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_edge("reflect",       END)          # no retry loop — straight to 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": [],
    }
    last_trace_len = 0
    for event in graph.stream(init, stream_mode="updates"):
        for node_name, state_update in event.items():
            new_trace = state_update.get("trace", [])
            # Only yield truly new trace entries (avoid duplicates from nodes with no trace entry)
            if len(new_trace) > last_trace_len:
                for step in new_trace[last_trace_len:]:
                    yield {"type": "trace_step", "step": step, "node": node_name}
                last_trace_len = len(new_trace)
            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]],
                }