Spaces:
Sleeping
Sleeping
File size: 3,308 Bytes
822ca51 5e2a88b 822ca51 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | """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)",
)
|