models.py

from typing import Dict, List, Optional, Union, Any
from openenv.core.env_server.types import Action, Observation
from pydantic import BaseModel, Field

# Nested State Models Models
class TenantMessage(BaseModel):
    message_id: str
    tenant_id: str
    content: str
    timestamp: str

class ApplianceStatus(BaseModel):
    name: str
    age_years: float
    condition: str

class PropertyStatus(BaseModel):
    property_id: str
    appliances: List[ApplianceStatus]
    lease_terms: Dict[str, str]

class Contractor(BaseModel):
    contractor_id: str
    specialty: str
    hourly_rate: float
    availability: List[str]

class ContractorDirectory(BaseModel):
    contractors: List[Contractor]

# Expected output model for Observation
class TenantDisputeObservation(Observation):
    messages: List[TenantMessage] = Field(default_factory=list)
    property_status: Optional[PropertyStatus] = None
    contractor_directory: Optional[ContractorDirectory] = None
    task_level: int = 1
    reason: str = "Init"

# Expected input model for Action
# Since OpenEnv expects a single Action class, we use a single class with optional fields
class TenantDisputeAction(Action):
    action_type: str = Field(description="Must be RequestClarification, CategorizeIssue, DispatchContractor, EnforceLease, ItemizeDeductions")
    
    # Payload for RequestClarification / EnforceLease
    tenant_id: Optional[str] = None
    message: Optional[str] = None
    clause: Optional[str] = None
    
    # Payload for CategorizeIssue
    issue_type: Optional[str] = None
    urgency: Optional[str] = None
    
    # Payload for DispatchContractor
    contractor_id: Optional[str] = None
    max_budget: Optional[float] = None
    time_window: Optional[str] = None
    
    # Payload for ItemizeDeductions
    deductions: Optional[List[Dict[str, Any]]] = None