File size: 1,977 Bytes
ede2fa4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""
Typed Action and Observation models for the SQL Query Environment.

Action: The agent submits a SQL query string and a task_id.
Observation: The environment returns schema info, query results, feedback, and reward.
"""

from pydantic import Field
from openenv.core.env_server.types import Action, Observation


class SQLAction(Action):
    """Action submitted by the agent: a SQL query to execute."""

    task_id: str = Field(
        ...,
        description="ID of the task being attempted (task_1, task_2, task_3)",
    )
    sql_query: str = Field(
        ...,
        description="The SQL query string to execute against the database",
    )


class SQLObservation(Observation):
    """Observation returned to the agent after each step."""

    # Task information
    task_id: str = Field(default="", description="Current task ID")
    task_description: str = Field(
        default="", description="Natural language question the agent must answer"
    )
    difficulty: str = Field(default="", description="easy, medium, or hard")

    # Database schema
    schema_description: str = Field(
        default="", description="SQL CREATE TABLE statements describing the database"
    )

    # Query result / feedback
    query_result: str = Field(
        default="",
        description="Result of the executed SQL query (rows as text), or error message",
    )
    query_error: bool = Field(
        default=False, description="True if the SQL query caused an error"
    )
    feedback: str = Field(
        default="",
        description="Human-readable feedback on the query result",
    )

    # Scoring
    reward: float = Field(default=0.0, description="Score from 0.0 to 1.0")
    done: bool = Field(default=False, description="True if the episode is complete")

    # Metadata
    step_count: int = Field(default=0, description="Number of steps taken so far")
    max_steps: int = Field(
        default=3, description="Maximum steps allowed per task"
    )