models.py

from pydantic import BaseModel
from typing import List, Optional, Dict, Any

class Patient(BaseModel):
    id: Optional[str] = None
    name: str
    age: int
    gender: str
    medical_history: Optional[str] = None

class AgentMessage(BaseModel):
    role: str  # "user", "assistant", "system"
    content: str

class AgentState(BaseModel):
    patient: Optional[Patient] = None
    messages: List[AgentMessage] = []
    context: Dict[str, Any] = {}
    next_step: Optional[str] = None

class AgentConfig(BaseModel):
    name: str
    type: str  # "intake", "anamnesis", "diagnosis", "planner"
    model: str = "gpt-4o"
    system_prompt: str