models.py

"""
Banking Support Agent Environment — Type Definitions

Defines the typed contracts for agent-environment interaction:
- BankingSupportAction: What the agent can do (use tools + send messages)
- BankingSupportObservation: What the agent sees after each step
- BankingSupportState: Episode metadata and tracking
"""

from typing import Dict, List, Optional, Any
from dataclasses import dataclass, field

# Determine which base classes to use
_USE_OPENENV = False
_USE_PYDANTIC = False

try:
    from openenv.core.env_server import Action, Observation, State
    _USE_OPENENV = True
except ImportError:
    try:
        from pydantic import BaseModel
        _USE_PYDANTIC = True
    except ImportError:
        pass


if _USE_OPENENV:
    # Production: use OpenEnv's Pydantic-based models
    class BankingSupportAction(Action):
        tool_name: Optional[str] = None
        tool_args: Dict[str, Any] = {}
        message: Optional[str] = None
        resolve: bool = False

    class BankingSupportObservation(Observation):
        customer_message: str = ""
        tool_result: Optional[str] = None
        available_tools: List[str] = []
        task_description: str = ""
        conversation_history: List[Dict[str, str]] = []
        error: Optional[str] = None
        last_action_error: Optional[str] = None

    class BankingSupportState(State):
        task_name: str = ""
        max_steps: int = 15
        tools_used: List[str] = []
        resolution_status: str = "open"
        score: float = 0.0

elif _USE_PYDANTIC:
    # Pydantic available but not openenv
    class BankingSupportAction(BaseModel):
        metadata: Dict[str, Any] = {}
        tool_name: Optional[str] = None
        tool_args: Dict[str, Any] = {}
        message: Optional[str] = None
        resolve: bool = False

    class BankingSupportObservation(BaseModel):
        done: bool = False
        reward: Optional[float] = None
        metadata: Dict[str, Any] = {}
        customer_message: str = ""
        tool_result: Optional[str] = None
        available_tools: List[str] = []
        task_description: str = ""
        conversation_history: List[Dict[str, str]] = []
        error: Optional[str] = None
        last_action_error: Optional[str] = None

    class BankingSupportState(BaseModel):
        episode_id: Optional[str] = None
        step_count: int = 0
        task_name: str = ""
        max_steps: int = 15
        tools_used: List[str] = []
        resolution_status: str = "open"
        score: float = 0.0

else:
    # Pure dataclass fallback for testing
    @dataclass
    class BankingSupportAction:
        metadata: Dict[str, Any] = field(default_factory=dict)
        tool_name: Optional[str] = None
        tool_args: Dict[str, Any] = field(default_factory=dict)
        message: Optional[str] = None
        resolve: bool = False

    @dataclass
    class BankingSupportObservation:
        done: bool = False
        reward: Optional[float] = None
        metadata: Dict[str, Any] = field(default_factory=dict)
        customer_message: str = ""
        tool_result: Optional[str] = None
        available_tools: List[str] = field(default_factory=list)
        task_description: str = ""
        conversation_history: List[Dict[str, str]] = field(default_factory=list)
        error: Optional[str] = None
        last_action_error: Optional[str] = None

    @dataclass
    class BankingSupportState:
        episode_id: Optional[str] = None
        step_count: int = 0
        task_name: str = ""
        max_steps: int = 15
        tools_used: List[str] = field(default_factory=list)
        resolution_status: str = "open"
        score: float = 0.0