models.py

from __future__ import annotations

try:
    from openenv.core.env_server.types import Action, Observation, State
except ImportError:
    try:
        from openenv.core.models import Action, Observation, State
    except ImportError:
        from openenv.core.env_server import Action, Observation, State

from typing import Set, Dict, Any, Optional, List
from pydantic import Field


class AdaptiveAction(Action):
    """
    Agent action for AdaptiveWorld environment.

    action_type options:
      "call_api"       → standard HTTP call to the mock API
      "probe_schema"   → GET /openapi.json, returns current API schema
      "query_history"  → returns last N API responses from this episode
      "declare_belief" → agent states its current world model (no HTTP call)
      "submit_result"  → end episode, agent declares task done + final belief
    """
    action_type: str = Field(default="call_api")
    method: str = Field(default="GET")
    url: str = Field(default="/mock_api/orders")
    headers: dict = Field(default_factory=dict)
    body: dict = Field(default_factory=dict)
    query_params: dict = Field(default_factory=dict)
    belief_state: Optional[Dict[str, Any]] = Field(default=None)
    # belief_state: agent's declared understanding of the world
    # Used when action_type == "declare_belief" or "submit_result"
    history_steps: int = Field(default=3)
    # history_steps: how many past steps to return in query_history


class AdaptiveObservation(Observation):
    """
    Observation for AdaptiveWorld environment.

    NOTE: drift_occurred and drift_type are intentionally ABSENT.
    The agent must infer drift from API responses and query_history.
    """
    task_description: str = Field(default="")
    domain: str = Field(default="e-commerce")
    prior_world_model: Dict[str, Any] = Field(default_factory=dict)
    # prior_world_model: agent's beliefs from previous episodes (v2: cross-episode persistence)
    current_step: int = Field(default=0)
    max_steps: int = Field(default=10)
    last_status_code: int = Field(default=0)
    last_response_body: str = Field(default="")
    last_response_headers: dict = Field(default_factory=dict)
    step_feedback: str = Field(default="")
    episode_history: List[Dict] = Field(default_factory=list)
    # episode_history: populated by query_history action
    task_reward: float = Field(default=0.0)
    belief_accuracy: float = Field(default=0.0)
    difficulty_level: int = Field(default=0)
    # difficulty_level: 0-3, set by DriftDifficultyController
    done: bool = Field(default=False)
    reward: float = Field(default=0.0)


class AdaptiveState(State):
    """Internal environment state for AdaptiveWorld."""
    scenario_id: str = Field(default="easy_field_rename")
    domain: str = Field(default="e-commerce")
    drift_injected: bool = Field(default=False)
    drift_step: int = Field(default=3)
    agent_belief: Dict[str, Any] = Field(default_factory=dict)
    world_truth: Dict[str, Any] = Field(default_factory=dict)
    task_completed: bool = Field(default=False)
    visited_endpoints: Set[str] = Field(default_factory=set)
    step_history: List[Dict] = Field(default_factory=list)
    # step_history: full episode HTTP log for query_history
    curriculum_level: int = Field(default=0)