File size: 7,772 Bytes
ee1b868 9fe6471 ee1b868 9fe6471 ee1b868 9fe6471 ee1b868 9fe6471 ee1b868 9fe6471 ee1b868 9fe6471 ee1b868 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 | 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}"
)
|