from sql_env.client import SQLEnvClient
from sql_env.models import SQLAction

URL = "https://hjerpe-sql-env.hf.space"

with SQLEnvClient(base_url=URL) as env:
    r = env.reset()
    print("RESET observation:", r.observation)
    print()

    # Pick a real table name from r.observation.schema_info before running
    # the next two steps — edit this script after reset prints something.
    r = env.step(SQLAction(action_type="DESCRIBE", argument="<TABLE>"))
    print("DESCRIBE:", r.observation)

    r = env.step(SQLAction(action_type="QUERY", argument="SELECT 1"))
    print("QUERY:", r.observation)

    r = env.step(SQLAction(action_type="ANSWER", argument="test"))
    print("ANSWER reward=", r.reward, "done=", r.done)