evd_agent/graph.py

from __future__ import annotations

from typing import TypedDict

from langgraph.graph import END, StateGraph

from .context_engine import EpidemiologicalContextEngine
from .llm_client import LLMConfigurationError
from .models import CaseClassification, DecisionOutput, InterviewState, InterviewStatus, LLMInterviewPlan, TurnResult
from .reasoning import ClinicalReasoningAgent


class GraphState(TypedDict):
    interview_state: InterviewState
    clinician_input: str
    llm_plan: LLMInterviewPlan | None
    llm_error: str | None
    state_updates: dict[str, str]


class AgentGraphRunner:
    def __init__(self, context_engine: EpidemiologicalContextEngine) -> None:
        self.reasoner = ClinicalReasoningAgent(context_engine)
        self.graph = self._build_graph()

    def _build_graph(self):
        graph = StateGraph(GraphState)

        graph.add_node("ingest_input", self._ingest_input)
        graph.add_node("llm_reason", self._llm_reason)
        graph.add_node("compose_response", self._compose_response)

        graph.set_entry_point("ingest_input")
        graph.add_edge("ingest_input", "llm_reason")
        graph.add_edge("llm_reason", "compose_response")
        graph.add_edge("compose_response", END)

        return graph.compile()

    def run(self, state: InterviewState, clinician_input: str) -> TurnResult:
        output = self.graph.invoke(
            {
                "interview_state": state,
                "clinician_input": clinician_input,
                "llm_plan": None,
                "llm_error": None,
                "state_updates": {},
            }
        )
        interview_state = output["interview_state"]
        next_key = interview_state.pending_question_key

        if interview_state.decision.should_stop_interview:
            assistant_message = self._alert_message(interview_state, output.get("llm_plan"))
            interview_state.status = InterviewStatus.COMPLETE
        elif output["llm_error"]:
            assistant_message = f"LLM configuration error: {output['llm_error']}"
        else:
            # Surface the LLM's own reasoning summary before the next question
            reasoning_prefix = ""
            llm_plan: LLMInterviewPlan | None = output.get("llm_plan")
            if llm_plan and llm_plan.summary_known:
                reasoning_prefix = f"{llm_plan.summary_known}\n\n"
            next_q = interview_state.pending_question_text or "Please provide the most important missing clinical detail."
            assistant_message = f"{reasoning_prefix}{next_q}"

        return TurnResult(
            assistant_message=assistant_message,
            decision=interview_state.decision,
            risk_profile=interview_state.risk_profile,
            next_question_key=next_key,
            llm_summary=interview_state.llm_summary,
            state_updates=output["state_updates"],
        )

    def _ingest_input(self, graph_state: GraphState) -> GraphState:
        return graph_state

    def _llm_reason(self, graph_state: GraphState) -> GraphState:
        interview_state = graph_state["interview_state"]
        clinician_input = graph_state["clinician_input"]

        try:
            plan = self.reasoner.reason(interview_state, clinician_input)
        except LLMConfigurationError as error:
            graph_state["llm_error"] = str(error)
            interview_state.llm_summary = str(error)
            interview_state.pending_question_key = None
            interview_state.pending_question_text = None
            return graph_state

        graph_state["llm_plan"] = plan
        interview_state.llm_summary = plan.summary_known
        interview_state.missing_evidence = list(plan.missing_evidence)
        interview_state.rationale_log.append(plan.reasoning)
        interview_state.rationale_log.extend(plan.evidence_statements)

        updates = plan.fact_updates.model_dump(exclude_none=True)
        for key, value in updates.items():
            setattr(interview_state.facts, key, value)
            graph_state["state_updates"][key] = str(value)

        interview_state.pending_question_key = None
        interview_state.decision = self._decision_from_llm_plan(plan)

        if plan.should_stop_interview:
            interview_state.pending_question_text = None
            return graph_state

        # Enforce MOH 3-follow-up question hard limit
        MAX_FOLLOWUP_QUESTIONS = 3
        if plan.next_question and interview_state.followup_question_count < MAX_FOLLOWUP_QUESTIONS:
            interview_state.pending_question_text = plan.next_question
            interview_state.asked_questions.add(plan.next_question)
            interview_state.followup_question_count += 1
        elif plan.next_question and interview_state.followup_question_count >= MAX_FOLLOWUP_QUESTIONS:
            # Hard stop: max follow-ups reached without meeting a case definition
            interview_state.pending_question_text = None
            interview_state.decision.should_stop_interview = True
            interview_state.decision.classification = CaseClassification.NOT_SUSPECTED
            interview_state.decision.triggered_rule = "Maximum of 3 follow-up questions reached with insufficient evidence."
            interview_state.decision.recommended_action = (
                "Available information does not meet current MOH Ebola case definitions. "
                "Escalate for clinician review if concern persists."
            )
        else:
            interview_state.pending_question_text = None

        return graph_state

    def _compose_response(self, graph_state: GraphState) -> GraphState:
        return graph_state

    @staticmethod
    def _decision_from_llm_plan(plan: LLMInterviewPlan) -> DecisionOutput:
        classification = plan.classification or CaseClassification.NOT_SUSPECTED
        evidence = plan.criteria_matched or list(plan.evidence_statements)
        triggered_rule = plan.triggered_rule or "No MOH case definition met yet."
        recommended_action = plan.recommended_action or "Continue focused evidence collection."
        return DecisionOutput(
            classification=classification,
            triggered_rule=triggered_rule,
            evidence=evidence,
            recommended_action=recommended_action,
            confidence=plan.confidence,
            should_stop_interview=plan.should_stop_interview,
        )

    @staticmethod
    def _alert_message(interview_state: InterviewState, llm_plan: LLMInterviewPlan | None = None) -> str:
        decision = interview_state.decision
        classification = decision.classification.value

        llm_reasoning = ""
        if llm_plan and llm_plan.reasoning:
            llm_reasoning = f"**Clinical Reasoning:**\n{llm_plan.reasoning}\n\n"

        if decision.classification == CaseClassification.NOT_SUSPECTED:
            missing = "\n".join(f"- {item}" for item in interview_state.missing_evidence) or "- None identified"
            return (
                f"**Classification: No Case**\n\n"
                f"{llm_reasoning}"
                f"The available information does not meet the current Kenya MOH Ebola case definitions.\n\n"
                f"**Reasoning:** {decision.triggered_rule}\n\n"
                f"**Missing or unclear information:**\n{missing}\n\n"
                f"**Recommended Action:** {decision.recommended_action}"
            )

        evidence = "\n".join(f"- {item}" for item in decision.evidence)
        return (
            f"🚨 {classification.upper()} ALERT\n\n"
            f"{llm_reasoning}"
            f"**Criteria Matched:**\n{evidence}\n\n"
            f"**Reasoning:**\n{decision.triggered_rule}\n\n"
            f"**Recommended Action:**\n{decision.recommended_action}"
        )