from dataclasses import dataclass from datetime import date, datetime from pydantic import BaseModel, Field from typing import List, Literal, Optional, TypedDict, Union class PatientInfo(BaseModel): name: Optional[str] = Field(None, description="Patient's full name") age: Optional[int] = Field(None, description="Patient's age in years") sex: Optional[str] = Field(None, description="Male/Female/Other") medical_record_number: Optional[str] = None class Config: extra = "forbid" # 🚨 ensures schema matches OpenAI’s strict rules class TestResultReferenceRange(BaseModel): min: Optional[float] = None max: Optional[float] = None class LabResult(BaseModel): test_name: str result_value: str test_unit: str test_reference_range: Optional[TestResultReferenceRange] = None test_date: Optional[str] = Field( None, description="Date of test in ISO format (YYYY-MM-DD). Convert all dates into this format.", ) inferred_range: Literal["low", "normal", "high"] class Config: extra = "forbid" class CompositeLabResult(BaseModel): section_name: str # e.g., "CUE - COMPLETE URINE ANALYSIS" sub_results: List[LabResult] # each individual parameter under the section class Config: extra = "forbid" class StandardizedReport(BaseModel): original_report_file_name: str patient_info: PatientInfo lab_results: List[ Union[LabResult, CompositeLabResult] ] # ✅ supports both flat + grouped results diagnosis: List[str] recommendations: List[str] class Config: extra = "forbid" @dataclass class HealthReport: report_file_name_with_path: str report_file_name: str report_contents: str @dataclass class SheamiMilestone: step_name: str status: str # e.g., "started", "in_progress", "completed", "failed" start_time: Optional[datetime] = None end_time: Optional[datetime] = None @property def time_taken(self) -> Optional[float]: """Return time taken in seconds if available.""" if self.start_time and self.end_time: return (self.end_time - self.start_time).total_seconds() return None @property def status_icon(self) -> str: if self.status == "started" or self.status == "inprogress": return "⏳" elif self.status == "completed": return "✅" elif self.status == "failed": return "❌" else: return self.status class SheamiState(TypedDict): user_email: str patient_id: str report_id_list: str run_id: str messages: list[str] thread_id: str uploaded_reports: List[HealthReport] standardized_reports: List[StandardizedReport] trends_json: dict pdf_path: str current_index: int process_desc: str units_processed: int units_total: int overall_units_processed: int overall_units_total: int milestones: list[SheamiMilestone] interpretation_html: str @dataclass class SheamiUser: email : str name : str picture_url : str