"""Session state management for FDAM AI Pipeline. Provides Pydantic models for session state and localStorage persistence. Images are stored separately (not in localStorage due to size limits). MVP Simplification: Single room assessment (no project-level fields). """ import json import uuid from datetime import datetime from typing import Optional from pydantic import BaseModel, Field from schemas.input import ( ConstructionEra, FacilityClassification, OdorIntensity, CharDensity, ) # --- Form Data Models (for localStorage) --- class RoomFormData(BaseModel): """Form data for Tab 1: Room Assessment. Includes facility_classification and construction_era moved from the removed ProjectFormData (required for calculations). """ id: str = Field(default_factory=lambda: f"room-{uuid.uuid4().hex[:8]}") name: str = "" length_ft: float = 0 width_ft: float = 0 ceiling_height_ft: float = 0 # Moved from ProjectFormData (required for calculations) facility_classification: FacilityClassification = "non-operational" construction_era: ConstructionEra = "post-2000" class ImageFormData(BaseModel): """Form data for a single image (metadata only, not bytes).""" id: str = Field(default_factory=lambda: f"img-{uuid.uuid4().hex[:8]}") filename: str = "" room_id: str = "" description: str = "" # Image bytes stored separately, referenced by id class ObservationsFormData(BaseModel): """Form data for Tab 3: Observations.""" smoke_fire_odor: bool = False odor_intensity: OdorIntensity = "none" visible_soot_deposits: bool = False soot_pattern_description: str = "" large_char_particles: bool = False char_density_estimate: Optional[CharDensity] = None ash_like_residue: bool = False ash_color_texture: str = "" surface_discoloration: bool = False discoloration_description: str = "" dust_loading_interference: bool = False dust_notes: str = "" wildfire_indicators: bool = False wildfire_notes: str = "" additional_notes: str = "" class SessionState(BaseModel): """Complete session state for an assessment. This model is serialized to localStorage for persistence. Images are stored separately and referenced by ID. MVP Simplification: Single room, 2 tabs (Input + Results/Chat). """ # Session metadata session_id: str = Field(default_factory=lambda: uuid.uuid4().hex) created_at: str = Field(default_factory=lambda: datetime.now().isoformat()) updated_at: str = Field(default_factory=lambda: datetime.now().isoformat()) name: str = "" # Display name for history list # Input completion status (single flag replaces 3 tab flags) input_complete: bool = False # Form data - single room (not list) room: RoomFormData = Field(default_factory=RoomFormData) images: list[ImageFormData] = Field(default_factory=list) observations: ObservationsFormData = Field(default_factory=ObservationsFormData) # Results (after generation) has_results: bool = False results_generated_at: Optional[str] = None # Chat history (Gradio 6 messages format) chat_history: list[dict] = Field(default_factory=list) # Serializable subset of PipelineResult (excludes PIL images) pipeline_result_json: Optional[str] = None # Document state for modifications generated_document: Optional[str] = None original_document: Optional[str] = None def update_timestamp(self) -> None: """Update the updated_at timestamp.""" self.updated_at = datetime.now().isoformat() def get_display_name(self) -> str: """Get a display name for the history list.""" if self.name: return self.name if self.room.name: return self.room.name return f"Assessment {self.session_id[:8]}" def validate_tab1(self) -> tuple[bool, list[str]]: """Validate Tab 1 (Room Assessment) is complete.""" errors = [] r = self.room if not r.name: errors.append("Room name is required") if r.length_ft <= 0: errors.append("Length must be greater than 0") if r.width_ft <= 0: errors.append("Width must be greater than 0") if r.ceiling_height_ft <= 0: errors.append("Ceiling height must be greater than 0") return len(errors) == 0, errors def validate_tab2(self) -> tuple[bool, list[str]]: """Validate Tab 2 (Images) is complete.""" errors = [] if not self.images: errors.append("At least one image is required") for img in self.images: if not img.room_id: errors.append(f"Image '{img.filename}': Must be associated with the room") return len(errors) == 0, errors def validate_tab3(self) -> tuple[bool, list[str]]: """Validate Tab 3 (Observations) is complete.""" # Tab 3 has no required fields - all checkboxes default to False return True, [] def can_generate(self) -> tuple[bool, list[str]]: """Check if assessment can be generated.""" all_errors = [] valid1, errors1 = self.validate_tab1() if not valid1: all_errors.extend(errors1) valid2, errors2 = self.validate_tab2() if not valid2: all_errors.extend(errors2) valid3, errors3 = self.validate_tab3() if not valid3: all_errors.extend(errors3) return len(all_errors) == 0, all_errors class AssessmentHistory(BaseModel): """Collection of saved assessments for history list.""" assessments: list[SessionState] = Field(default_factory=list) current_session_id: Optional[str] = None def add_assessment(self, session: SessionState) -> None: """Add or update an assessment in history.""" session.update_timestamp() # Remove existing if present self.assessments = [a for a in self.assessments if a.session_id != session.session_id] # Add to front of list self.assessments.insert(0, session) # Keep only last 20 assessments self.assessments = self.assessments[:20] def get_assessment(self, session_id: str) -> Optional[SessionState]: """Get an assessment by ID.""" for a in self.assessments: if a.session_id == session_id: return a return None def remove_assessment(self, session_id: str) -> None: """Remove an assessment from history.""" self.assessments = [a for a in self.assessments if a.session_id != session_id] def get_history_items(self) -> list[dict]: """Get history items for display in dropdown.""" return [ { "id": a.session_id, "name": a.get_display_name(), "updated": a.updated_at, "has_results": a.has_results, } for a in self.assessments ] # --- Gradio State Helpers --- def create_new_session() -> SessionState: """Create a new empty session.""" return SessionState() def session_to_json(session: SessionState) -> str: """Serialize session to JSON for localStorage.""" return session.model_dump_json() def session_from_json(json_str: str) -> SessionState: """Deserialize session from JSON. Includes migration from old multi-room format to single room. """ try: data = json.loads(json_str) # Migration: Convert old multi-room format to single room if "rooms" in data and isinstance(data["rooms"], list): # Use first room if available, otherwise create empty if data["rooms"]: data["room"] = data["rooms"][0] else: data["room"] = {} del data["rooms"] # Migration: Move facility_classification and construction_era from project to room if "project" in data: project = data["project"] if "room" not in data: data["room"] = {} # Move fields if they exist in project if "facility_classification" in project: data["room"]["facility_classification"] = project["facility_classification"] if "construction_era" in project: data["room"]["construction_era"] = project["construction_era"] del data["project"] # Migration: Remove old tab4_complete (now only 3 tabs) if "tab4_complete" in data: del data["tab4_complete"] # Migration: Convert old tab1/2/3_complete to input_complete if "tab1_complete" in data or "tab2_complete" in data or "tab3_complete" in data: # Input is complete if all three old tabs were complete tab1 = data.pop("tab1_complete", False) tab2 = data.pop("tab2_complete", False) tab3 = data.pop("tab3_complete", False) data["input_complete"] = tab1 and tab2 and tab3 return SessionState.model_validate(data) except Exception: return create_new_session() def history_to_json(history: AssessmentHistory) -> str: """Serialize history to JSON for localStorage.""" return history.model_dump_json() def history_from_json(json_str: str) -> AssessmentHistory: """Deserialize history from JSON.""" try: return AssessmentHistory.model_validate_json(json_str) except Exception: return AssessmentHistory()