Spaces:
Sleeping
Sleeping
File size: 1,960 Bytes
2194233 e965a47 341aec9 5716a3a e965a47 2194233 5716a3a 2194233 5716a3a e965a47 5716a3a e965a47 5716a3a e965a47 341aec9 2be7532 eecac6e 341aec9 2194233 | 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 | from typing import List, Optional, Any
from openenv.core.env_server.types import Action, Observation, State
from pydantic import Field, BaseModel
class SQLAction(Action):
corrected_query: str
class SQLObservation(Observation):
"""
Observation returned to the agent each step.
Fields:
task_id: Unique identifier for the current task instance.
broken_query: The malformed SQL query the agent must fix.
schema_context: Table/column definitions (hard tasks only).
error_hint: Plain-language hint about the error (easy tasks only).
step_number: Current step within the episode (0 = initial observation).
steps_remaining: How many steps are left before the episode ends.
previous_attempt: The agent's SQL output from the previous step.
feedback: Grader feedback on the previous attempt.
"""
task_id: str
broken_query: str
schema_context: Optional[str] = None
error_hint: Optional[str] = None
step_number: int
steps_remaining: Optional[int] = None
previous_attempt: Optional[str] = None
feedback: Optional[str] = None
class SQLState(State):
task_id: str
difficulty: str
step_count: int
max_steps: int
done: bool
last_reward: float
rewards_history: List[float]
class SQLReward(BaseModel):
# Strictly between 0 and 1 as required by the OpenEnv spec
value: float = Field(gt=0.0, lt=1.0)
reason: str
class SQLTask(BaseModel):
task_id: str
difficulty: str
broken_query: str
canonical_answer: str
schema_context: Optional[str] = None
error_hint: Optional[str] = None
max_steps: int = 5
grader: Optional[Any] = Field(default=None, exclude=True)
model_config = {"arbitrary_types_allowed": True}
class StepResult(BaseModel):
observation: SQLObservation
reward: float
done: bool
info: dict = Field(default_factory=dict)
|