evd_agent/conversation.py

from __future__ import annotations

import uuid
from datetime import datetime, timezone

from .config import load_epidemiological_context
from .context_engine import EpidemiologicalContextEngine
from .graph import AgentGraphRunner
from .models import ChatTurn, InterviewState, TurnResult


class ConversationManager:
    def __init__(self, context_path: str | None = None) -> None:
        context = load_epidemiological_context(context_path)
        self.context_engine = EpidemiologicalContextEngine(context)
        self.graph_runner = AgentGraphRunner(self.context_engine)

    def new_state(self) -> InterviewState:
        return InterviewState(session_id=str(uuid.uuid4()), context=self.context_engine.context)

    def start_session(self) -> tuple[InterviewState, TurnResult]:
        state = self.new_state()
        result = self.graph_runner.run(state, "")
        state.history.append(ChatTurn(role="assistant", content=result.assistant_message, timestamp=datetime.now(timezone.utc)))
        return state, result

    def process_turn(self, state: InterviewState, clinician_message: str) -> TurnResult:
        state.history.append(ChatTurn(role="clinician", content=clinician_message, timestamp=datetime.now(timezone.utc)))
        result = self.graph_runner.run(state, clinician_message)
        state.history.append(ChatTurn(role="assistant", content=result.assistant_message, timestamp=datetime.now(timezone.utc)))
        return result