| from openenv.core.env_server import Action, Observation, State | |
| from pydantic import BaseModel, Field | |
| from typing import List, Dict, Optional, Any | |
| # ============================================================ | |
| # Tool Registry - defines available tools the agent can call | |
| # ============================================================ | |
| class ToolParameter(BaseModel): | |
| """Schema for a single tool parameter.""" | |
| name: str | |
| type: str # string / number / boolean / array | |
| description: str | |
| required: bool = True | |
| enum: Optional[List[str]] = None # allowed values if restricted | |
| class ToolDefinition(BaseModel): | |
| """A tool available to the agent.""" | |
| name: str | |
| description: str | |
| parameters: List[ToolParameter] = Field(default_factory=list) | |
| # ============================================================ | |
| # Scenario - a user query with expected tool call(s) | |
| # ============================================================ | |
| class Scenario(BaseModel): | |
| """A single scenario the agent must handle.""" | |
| id: int | |
| user_query: str # what the user asked | |
| context: str = "" # optional conversation history or extra context | |
| available_tools: List[str] # names of tools available for this scenario | |
| difficulty_tags: List[str] = Field(default_factory=list) # e.g. ["multi_step", "refusal", "param_extraction"] | |
| metadata: Dict[str, str] = Field(default_factory=dict) # extra info: domain, risk_level, etc. | |
| # ============================================================ | |
| # Agent's Action - the tool call it decides to make | |
| # ============================================================ | |
| class ToolCallAction(Action): | |
| """Action taken by the agent - one or more tool calls.""" | |
| scenario_id: int | |
| tool_calls: List[Dict[str, Any]] # [{"tool_name": "...", "parameters": {...}}, ...] | |
| should_refuse: bool = False # agent can signal it should NOT call any tool | |
| reasoning: str = "" # optional chain-of-thought | |
| # ============================================================ | |
| # Observation - what the agent sees | |
| # ============================================================ | |
| class ToolCallObservation(Observation): | |
| """What the agent observes after each step.""" | |
| scenario: Scenario # current scenario to handle | |
| tool_definitions: List[ToolDefinition] # full schema of available tools | |
| queue_size: int # total scenarios in episode | |
| current_step: int # index of current scenario | |
| reward: float # reward from previous step | |
| done: bool # whether episode has ended | |
| # ============================================================ | |
| # Environment State | |
| # ============================================================ | |
| class ToolCallState(State): | |
| """Internal state of the environment.""" | |
| current_index: int | |
| total_scenarios: int | |
| processed_scenario_ids: List[int] | |
| score: float | |
| done: bool | |