Upload folder using huggingface_hub

graph.py

@@ -18,7 +18,7 @@ from common import get_db
 from config import SheamiConfig
 import logging
 
-from modules import db
+from modules.models import SheamiMilestone, SheamiState, StandardizedReport, TestResultReferenceRange
 from pdf_helper import generate_pdf
 
 logging.basicConfig()
@@ -35,105 +35,6 @@ llm = ChatOpenAI(model=os.getenv("MODEL"), temperature=0.3)
 
 from typing import Optional, List
 from pydantic import BaseModel, Field
-
-
-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] = None
-    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: 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":
-            return "⏳"
-        elif self.status == "completed":
-            return "✅"
-        else:
-            return self.status
-
-
-class SheamiState(TypedDict):
-    user_email: str
-    patient_id: str
-    run_id: str
-    messages: list[str]
-    thread_id: str
-    uploaded_reports: List[HealthReport]
-    standardized_reports: List[StandardizedReport]
-    trends_json: dict
-    interpreted_report: str
-    current_index: int
-    process_desc: str
-    units_processed: int
-    units_total: int
-    overall_units_processed: int
-    overall_units_total: int
-    milestones: list[SheamiMilestone]
-
-
 import re
 
 
@@ -617,6 +518,11 @@ async def fn_interpreter_node(state: SheamiState):
 
     get_db().update_run_stats(run_id=state["run_id"], status="completed")
 
+    ## add parsed reports
+    get_db().add_report_v2(
+        patient_id=state["patient_id"], reports=state["standardized_reports"]
+    )
+
     return state
 
 

modules/db.py

@@ -5,6 +5,7 @@ from datetime import datetime, timezone
 from bson import ObjectId
 from dotenv import load_dotenv
 
+from modules.models import StandardizedReport
 
 class SheamiDB:
     def __init__(self, uri: str, db_name: str = "sheami"):
@@ -68,7 +69,20 @@ class SheamiDB:
     # ---------------------------
     # REPORT FUNCTIONS
     # ---------------------------
-    def add_report(self, patient_id: str, file_name: str, parsed_data: dict) -> str:
+    def add_report_v2(self, patient_id: str, reports: list[StandardizedReport]) -> str:
+        inserted_ids = []
+        for parsed_data in reports:
+            report = {
+                "patient_id": ObjectId(patient_id),
+                "uploaded_at": datetime.now(timezone.utc),
+                "file_name": parsed_data.original_report_file_name,
+                "parsed_data_v2": parsed_data.model_dump(),
+            }
+            result = self.reports.insert_one(report)
+            inserted_ids.append(result.inserted_id)
+        return str(inserted_ids)
+
+    def add_report(self, patient_id: str, file_name: str, parsed_data: any) -> str:
         report = {
             "patient_id": ObjectId(patient_id),
             "uploaded_at": datetime.now(timezone.utc),

modules/models.py

@@ -0,0 +1,100 @@
+from dataclasses import dataclass
+from datetime import 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] = None
+    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: 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":
+            return "⏳"
+        elif self.status == "completed":
+            return "✅"
+        else:
+            return self.status
+
+
+class SheamiState(TypedDict):
+    user_email: str
+    patient_id: str
+    run_id: str
+    messages: list[str]
+    thread_id: str
+    uploaded_reports: List[HealthReport]
+    standardized_reports: List[StandardizedReport]
+    trends_json: dict
+    interpreted_report: str
+    current_index: int
+    process_desc: str
+    units_processed: int
+    units_total: int
+    overall_units_processed: int
+    overall_units_total: int
+    milestones: list[SheamiMilestone]

report_formatter.py

@@ -3,11 +3,11 @@ import gradio as gr
 from typing import List, Optional
 from datetime import datetime
 
-from graph import (
-    CompositeLabResult,
-    LabResult,
+from modules.models import (
     StandardizedReport,
     TestResultReferenceRange,
+    CompositeLabResult,
+    LabResult,
 )
 
 

ui.py

@@ -6,7 +6,8 @@ import os
 import pandas as pd
 from tqdm.auto import tqdm
 from config import SheamiConfig
-from graph import SheamiMilestone, create_graph, SheamiState, HealthReport
+from graph import create_graph
+from modules.models import HealthReport, SheamiMilestone, SheamiState
 from pdf_reader import read_pdf
 from gradio_modal import Modal
 from report_formatter import render_patient_state