Spaces:
Sleeping
Sleeping
| """Pydantic models for DB-ER actions, observations, and state.""" | |
| from __future__ import annotations | |
| from typing import Any, Dict, List, Literal, Optional, Union | |
| from pydantic import Field | |
| from openenv.core.env_server.types import Action, Observation, State | |
| class DBERAction(Action): | |
| """Actions the agent can take inside the DB-ER environment. | |
| Two action types: | |
| - ``execute_sql`` - run a raw SQL query against the live database. | |
| - ``submit_resolution``- declare the episode finished; triggers the grader. | |
| """ | |
| action_type: Literal["execute_sql", "submit_resolution"] = Field( | |
| ..., description="Type of action to perform" | |
| ) | |
| query: Optional[str] = Field( | |
| None, | |
| description=( | |
| "SQL query to execute. Required when action_type='execute_sql'. " | |
| "Budget cost: -1 for SELECT/PRAGMA, -5 for INSERT/UPDATE/DELETE/ALTER." | |
| ), | |
| ) | |
| notes: Optional[str] = Field( | |
| None, | |
| description=( | |
| "Human-readable resolution notes. Required when action_type='submit_resolution'. " | |
| "Should describe what was diagnosed and fixed." | |
| ), | |
| ) | |
| model_config = {"extra": "forbid"} | |
| class DBERObservation(Observation): | |
| """Full observation returned after every reset() and step(). | |
| All fields are JSON-serialisable so the agent can reason over them directly. | |
| """ | |
| # Incident context | |
| incident_ticket: str = Field( | |
| ..., description="The initial incident alert / pager message" | |
| ) | |
| task_id: int = Field( | |
| 1, description="Active task ID: 1=Easy, 2=Medium, 3=Medium-Hard, 4=Hard, 5=Very Hard" | |
| ) | |
| # Last action result | |
| last_action_type: str = Field( | |
| "NONE", | |
| description="Category of the last executed action: SELECT | MUTATION | META | ERROR | NONE", | |
| ) | |
| last_query_result: Union[List[Dict[str, Any]], str] = Field( | |
| default_factory=list, | |
| description="Rows returned by the last SELECT, or a status string for mutations.", | |
| ) | |
| error_logs: str = Field( | |
| "", description="stderr / exception message from the last action (empty if no error)" | |
| ) | |
| # Environment state | |
| db_summary: Dict[str, Any] = Field( | |
| default_factory=dict, | |
| description='Live DB snapshot: {"tables": [...], "row_counts": {...}}', | |
| ) | |
| violation_count: int = Field( | |
| 0, | |
| description=( | |
| "Current count of broken constraints (FK violations + data integrity issues). " | |
| "Decreasing this gives +0.2 intermediate reward." | |
| ), | |
| ) | |
| budget_remaining: int = Field( | |
| 100, | |
| description="Remaining action budget. SELECT/PRAGMA costs 1; mutations cost 5. Hitting 0 ends the episode.", | |
| ) | |
| step_count: int = Field(0, description="Number of steps taken in the current episode") | |
| model_config = {"extra": "forbid"} | |
| class DBERState(State): | |
| """Internal server-side state (returned by GET /state).""" | |
| task_id: int = Field(1) | |
| budget_remaining: int = Field(100) | |
| violation_count: int = Field(0) | |
| initial_row_counts: Dict[str, int] = Field(default_factory=dict) | |
| query_history: Dict[str, int] = Field( | |
| default_factory=dict, | |
| description="Normalised query -> execution count (for spam detection)", | |
| ) | |