visualization-agent/models.py

"""
Data models for Visualization Agent
Self-contained models for independent deployment
"""
from __future__ import annotations

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


# Input Types
BuildingType = Literal[
    "residential_single_family",
    "residential_multi_family",
    "residential_high_rise",
    "commercial_office",
    "commercial_retail",
    "industrial_warehouse",
    "institutional_school",
    "institutional_hospital",
    "infrastructure_bridge",
    "mixed_use"
]

RiskLevel = Literal["CRITICAL", "HIGH", "MODERATE", "LOW"]


# Base Models
class Coordinates(BaseModel):
    """Geographic coordinates"""
    latitude: float
    longitude: float


class LocationInfo(BaseModel):
    """Location information"""
    name: str
    coordinates: Coordinates
    administrative_area: str


# Risk Assessment Models (needed for input)
class HazardDetail(BaseModel):
    """Detailed information about a specific hazard"""
    status: str
    description: str
    distance: Optional[str] = None
    direction: Optional[str] = None
    severity: Optional[str] = None


class SeismicHazards(BaseModel):
    """Seismic hazard information"""
    active_fault: HazardDetail
    ground_shaking: HazardDetail
    liquefaction: HazardDetail
    tsunami: HazardDetail
    earthquake_induced_landslide: HazardDetail
    fissure: HazardDetail
    ground_rupture: HazardDetail


class VolcanicHazards(BaseModel):
    """Volcanic hazard information"""
    active_volcano: HazardDetail
    potentially_active_volcano: HazardDetail
    inactive_volcano: HazardDetail
    ashfall: HazardDetail
    pyroclastic_flow: HazardDetail
    lahar: HazardDetail
    lava: HazardDetail
    ballistic_projectile: HazardDetail
    base_surge: HazardDetail
    volcanic_tsunami: HazardDetail


class HydroHazards(BaseModel):
    """Hydrometeorological hazard information"""
    flood: HazardDetail
    rain_induced_landslide: HazardDetail
    storm_surge: HazardDetail
    severe_winds: HazardDetail


class HazardData(BaseModel):
    """Complete hazard data from risk assessment"""
    seismic: SeismicHazards
    volcanic: VolcanicHazards
    hydrometeorological: HydroHazards


class RiskSummary(BaseModel):
    """Summary of overall risk assessment"""
    overall_risk_level: RiskLevel
    total_hazards_assessed: int
    high_risk_count: int
    moderate_risk_count: int
    critical_hazards: List[str] = Field(default_factory=list)
    llm_analysis: Optional[str] = None


class FacilityInfo(BaseModel):
    """Critical facilities information from risk assessment"""
    schools: Dict[str, Any] | List[Dict[str, Any]] = Field(default_factory=lambda: {})
    hospitals: Dict[str, Any] | List[Dict[str, Any]] = Field(default_factory=lambda: {})
    road_networks: Dict[str, Any] | List[Dict[str, Any]] = Field(default_factory=lambda: [])


class Metadata(BaseModel):
    """Metadata for data sources"""
    timestamp: str
    source: str
    cache_status: str
    ttl: int


class RiskData(BaseModel):
    """Complete risk assessment data"""
    success: bool
    summary: RiskSummary
    location: LocationInfo
    hazards: HazardData
    facilities: FacilityInfo
    metadata: Metadata


# Construction Recommendations Models (needed for input)
class RecommendationDetail(BaseModel):
    """Detailed construction recommendation"""
    hazard_type: str
    recommendation: str
    rationale: str
    source_url: Optional[str] = None


class BuildingCodeReference(BaseModel):
    """Building code reference"""
    code_name: str
    section: str
    requirement: str


class Recommendations(BaseModel):
    """Construction recommendations"""
    general_guidelines: List[str] = Field(default_factory=list)
    seismic_recommendations: List[RecommendationDetail] = Field(default_factory=list)
    volcanic_recommendations: List[RecommendationDetail] = Field(default_factory=list)
    hydrometeorological_recommendations: List[RecommendationDetail] = Field(default_factory=list)
    priority_actions: List[str] = Field(default_factory=list)
    building_codes: List[BuildingCodeReference] = Field(default_factory=list)


# House Specifications Models
FoundationType = Literal["shallow", "deep_pile", "raft", "combined", "elevated"]
RoofType = Literal["pitched", "flat", "hip", "gable"]
MaterialType = Literal["reinforced_concrete", "steel_frame", "wood", "composite", "masonry"]


class StructuralFeature(BaseModel):
    """Structural feature specification"""
    feature_name: str
    specification: str
    justification: str


class HouseSpecifications(BaseModel):
    """House specifications for construction"""
    number_of_floors: int
    primary_material: MaterialType
    foundation_type: FoundationType
    foundation_depth_meters: float
    roof_type: RoofType
    wall_thickness_mm: int
    reinforcement_details: str
    structural_features: List[StructuralFeature] = Field(default_factory=list)
    compliance_codes: List[str] = Field(default_factory=list)
    decision_rationale: str


# Visualization-specific Models
class VisualizationData(BaseModel):
    """Generated visualization data"""
    image_base64: str  # Base64-encoded PNG image
    prompt_used: str  # The prompt that generated the image
    model_version: str  # Gemini model version used
    generation_timestamp: str  # ISO format timestamp
    image_format: str  # "PNG"
    resolution: str  # "1024x1024"
    features_included: List[str]  # List of disaster-resistant features shown


class VisualizationRequest(BaseModel):
    """Request for visualization generation"""
    risk_data: RiskData
    building_type: BuildingType
    recommendations: Optional[Recommendations] = None


class VisualizationResponse(BaseModel):
    """Agent response format"""
    success: bool
    visualization_data: Optional[VisualizationData] = None
    error: Optional[ErrorDetail] = None


# Error Handling Models
class ErrorDetail(BaseModel):
    """Error detail information"""
    code: str
    message: str
    details: Optional[Dict[str, Any]] = None
    retry_possible: bool = False


class ErrorResponse(BaseModel):
    """Error response structure"""
    success: bool = False
    error: Optional[ErrorDetail] = None
    partial_results: Optional[Dict[str, Any]] = None


# Image Generation Models (Pydantic versions for API validation)
class ImageConfig(BaseModel):
    """Configuration for image generation parameters"""
    model: str = "gemini-2.5-flash-image"  # or "gemini-3-pro-image-preview"
    aspect_ratio: str = "16:9"  # 1:1, 2:3, 3:2, 3:4, 4:3, 4:5, 5:4, 9:16, 16:9, 21:9
    image_size: str = "1K"  # 1K, 2K, 4K (4K only for gemini-3-pro-image-preview)
    use_search: bool = False  # Enable Google Search grounding
    output_format: str = "png"


class VisualizationGenerationRequest(BaseModel):
    """Input data for image generation (API request format)"""
    prompt: str
    construction_data: Optional[Dict[str, Any]] = None
    config: Optional[ImageConfig] = None


class VisualizationGenerationResponse(BaseModel):
    """Output data from image generation (API response format)"""
    success: bool
    image_path: Optional[str] = None
    image_base64: Optional[str] = None
    text_response: Optional[str] = None
    error: Optional[Dict[str, Any]] = None