evd_agent/llm_client.py

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import TYPE_CHECKING, Any

from openai import OpenAI

try:
    from dotenv import load_dotenv
except ImportError:
    load_dotenv = None

from .models import InterviewState, LLMInterviewPlan

if TYPE_CHECKING:
    from .context_engine import EpidemiologicalContextEngine


class LLMConfigurationError(RuntimeError):
    pass


class ClinicalLLMClient:
    """OpenAI-compatible client for structured clinical interview planning."""

    def __init__(self) -> None:
        self._load_environment_file()
        api_key = os.getenv("EVD_LLM_API_KEY") or os.getenv("OPENAI_API_KEY")
        model = os.getenv("EVD_LLM_MODEL") or os.getenv("OPENAI_MODEL") or "gpt-4.1-mini"
        base_url = os.getenv("EVD_LLM_BASE_URL") or os.getenv("OPENAI_BASE_URL")

        self.model = model
        self.client = OpenAI(api_key=api_key, base_url=base_url) if api_key else None

    @staticmethod
    def _load_environment_file() -> None:
        if load_dotenv is not None:
            load_dotenv()
            return

        env_path = Path(".env")
        if not env_path.exists():
            return

        for raw_line in env_path.read_text(encoding="utf-8").splitlines():
            line = raw_line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue

            key, value = line.split("=", 1)
            key = key.strip()
            value = value.strip().strip('"').strip("'")

            if key and key not in os.environ:
                os.environ[key] = value

    def plan_next_step(self, *, system_prompt: str, user_payload: dict[str, Any]) -> LLMInterviewPlan:
        if self.client is None:
            raise LLMConfigurationError(
                "No LLM credentials configured. Set EVD_LLM_API_KEY or OPENAI_API_KEY, and optionally "
                "EVD_LLM_MODEL and EVD_LLM_BASE_URL. GITHUB_TOKEN is not used."
            )

        completion = self.client.beta.chat.completions.parse(
            model=self.model,
            temperature=0.1,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": json.dumps(user_payload, ensure_ascii=True)},
            ],
            response_format=LLMInterviewPlan,
        )
        message = completion.choices[0].message
        if message.parsed is None:
            raise RuntimeError("The LLM did not return a structured interview plan.")
        return message.parsed

    @staticmethod
    def serialize_state(state: InterviewState, context_engine: "EpidemiologicalContextEngine | None" = None) -> dict[str, Any]:
        facts = state.facts
        compatible_symptoms = [
            name
            for name in (
                "headache",
                "lethargy",
                "loss_of_appetite",
                "muscle_pain",
                "joint_pain",
                "stomach_pain",
                "difficulty_swallowing",
                "vomiting",
                "difficulty_breathing",
                "diarrhea",
                "hiccups",
            )
            if getattr(facts, name) is True
        ]

        payload = {
            "session_id": state.session_id,
            "facts": facts.model_dump(exclude_none=True),
            "decision_summary": {
                "temperature_c": facts.temperature_c,
                "fever_reported": facts.fever_reported,
                "sudden_onset_fever": facts.sudden_onset_fever,
                "compatible_symptom_count": len(compatible_symptoms),
                "compatible_symptoms": compatible_symptoms,
                "unexplained_bleeding": facts.unexplained_bleeding,
                "sudden_unexplained_death": facts.sudden_unexplained_death,
                "exposure_known_case_21d": facts.exposure_known_case_21d,
                "exposure_outbreak_area_21d": facts.exposure_outbreak_area_21d,
                "travel_outbreak_area_21d": facts.travel_outbreak_area_21d,
                "attended_funeral_21d": facts.attended_funeral_21d,
                "healthcare_worker_exposure_21d": facts.healthcare_worker_exposure_21d,
                "epidemiological_link_known_case": facts.epidemiological_link_known_case,
                "lab_confirmation_available": facts.lab_confirmation_available,
                "clinician_assessed_consistent": facts.clinician_assessed_consistent,
                "failed_treatment": facts.failed_treatment,
            },
            "risk_profile": state.risk_profile.model_dump(),
            "context": state.context.model_dump(),
            "asked_questions": sorted(state.asked_questions),
            "history": [turn.model_dump(mode="json") for turn in state.history[-12:]],
            "current_decision": state.decision.model_dump(mode="json"),
        }

        if context_engine is not None:
            payload["kenya_county_context"] = context_engine.get_question_priority_hints(facts)

        return payload