Spaces:
Sleeping
Sleeping
File size: 7,396 Bytes
d3d0e0e | 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 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 | from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Callable
from data_agent_baseline.benchmark.schema import AnswerTable, PublicTask
from data_agent_baseline.tools.filesystem import (
list_context_tree,
read_csv_preview,
read_doc_preview,
read_json_preview,
resolve_context_path,
)
from data_agent_baseline.tools.python_exec import execute_python_code
from data_agent_baseline.tools.sqlite import execute_read_only_sql, inspect_sqlite_schema
EXECUTE_PYTHON_TIMEOUT_SECONDS = 30
@dataclass(frozen=True, slots=True)
class ToolSpec:
name: str
description: str
input_schema: dict[str, Any]
@dataclass(frozen=True, slots=True)
class ToolExecutionResult:
ok: bool
content: dict[str, Any]
is_terminal: bool = False
answer: AnswerTable | None = None
ToolHandler = Callable[[PublicTask, dict[str, Any]], ToolExecutionResult]
def _list_context(task: PublicTask, action_input: dict[str, Any]) -> ToolExecutionResult:
max_depth = int(action_input.get("max_depth", 4))
return ToolExecutionResult(ok=True, content=list_context_tree(task, max_depth=max_depth))
def _read_csv(task: PublicTask, action_input: dict[str, Any]) -> ToolExecutionResult:
path = str(action_input["path"])
max_rows = int(action_input.get("max_rows", 20))
return ToolExecutionResult(ok=True, content=read_csv_preview(task, path, max_rows=max_rows))
def _read_json(task: PublicTask, action_input: dict[str, Any]) -> ToolExecutionResult:
path = str(action_input["path"])
max_chars = int(action_input.get("max_chars", 10000))
return ToolExecutionResult(ok=True, content=read_json_preview(task, path, max_chars=max_chars))
def _read_doc(task: PublicTask, action_input: dict[str, Any]) -> ToolExecutionResult:
path = str(action_input["path"])
max_chars = int(action_input.get("max_chars", 10000))
return ToolExecutionResult(ok=True, content=read_doc_preview(task, path, max_chars=max_chars))
def _inspect_sqlite_schema(task: PublicTask, action_input: dict[str, Any]) -> ToolExecutionResult:
path = resolve_context_path(task, str(action_input["path"]))
return ToolExecutionResult(ok=True, content=inspect_sqlite_schema(path))
def _execute_context_sql(task: PublicTask, action_input: dict[str, Any]) -> ToolExecutionResult:
path = resolve_context_path(task, str(action_input["path"]))
sql = str(action_input["sql"])
limit = int(action_input.get("limit", 200))
return ToolExecutionResult(ok=True, content=execute_read_only_sql(path, sql, limit=limit))
def _execute_python(task: PublicTask, action_input: dict[str, Any]) -> ToolExecutionResult:
code = str(action_input["code"])
content = execute_python_code(
context_root=task.context_dir,
code=code,
timeout_seconds=EXECUTE_PYTHON_TIMEOUT_SECONDS,
)
return ToolExecutionResult(ok=bool(content.get("success")), content=content)
def _answer(_: PublicTask, action_input: dict[str, Any]) -> ToolExecutionResult:
columns = action_input.get("columns")
rows = action_input.get("rows")
if not isinstance(columns, list) or not columns or not all(isinstance(item, str) for item in columns):
raise ValueError("answer.columns must be a non-empty list of strings.")
if not isinstance(rows, list):
raise ValueError("answer.rows must be a list.")
normalized_rows: list[list[Any]] = []
for row in rows:
if not isinstance(row, list):
raise ValueError("Each answer row must be a list.")
if len(row) != len(columns):
raise ValueError("Each answer row must match the number of columns.")
normalized_rows.append(list(row))
answer = AnswerTable(columns=list(columns), rows=normalized_rows)
return ToolExecutionResult(
ok=True,
content={
"status": "submitted",
"column_count": len(columns),
"row_count": len(normalized_rows),
},
is_terminal=True,
answer=answer,
)
@dataclass(slots=True)
class ToolRegistry:
specs: dict[str, ToolSpec]
handlers: dict[str, ToolHandler]
def describe_for_prompt(self) -> str:
lines = []
for name in sorted(self.specs):
spec = self.specs[name]
lines.append(f"- {spec.name}: {spec.description}")
lines.append(f" input_schema: {spec.input_schema}")
return "\n".join(lines)
def execute(self, task: PublicTask, action: str, action_input: dict[str, Any]) -> ToolExecutionResult:
if action not in self.handlers:
raise KeyError(f"Unknown tool: {action}")
return self.handlers[action](task, action_input)
def create_default_tool_registry() -> ToolRegistry:
specs = {
"answer": ToolSpec(
name="answer",
description="Submit the final answer table. This is the only valid terminating action.",
input_schema={
"columns": ["column_name"],
"rows": [["value_1"]],
},
),
"execute_context_sql": ToolSpec(
name="execute_context_sql",
description="Run a read-only SQL query against a sqlite/db file inside context.",
input_schema={"path": "relative/path/to/file.sqlite", "sql": "SELECT ...", "limit": 200},
),
"execute_python": ToolSpec(
name="execute_python",
description=(
"Execute arbitrary Python code with the task context directory as the "
"working directory. The tool returns the code's captured stdout as `output`. "
f"The execution timeout is fixed at {EXECUTE_PYTHON_TIMEOUT_SECONDS} seconds."
),
input_schema={
"code": "import os\nprint(sorted(os.listdir('.')))",
},
),
"inspect_sqlite_schema": ToolSpec(
name="inspect_sqlite_schema",
description="Inspect tables and columns in a sqlite/db file inside context.",
input_schema={"path": "relative/path/to/file.sqlite"},
),
"list_context": ToolSpec(
name="list_context",
description="List files and directories available under context.",
input_schema={"max_depth": 4},
),
"read_csv": ToolSpec(
name="read_csv",
description="Read a preview of a CSV file inside context.",
input_schema={"path": "relative/path/to/file.csv", "max_rows": 20},
),
"read_doc": ToolSpec(
name="read_doc",
description="Read a text-like document inside context.",
input_schema={"path": "relative/path/to/file.md", "max_chars": 4000},
),
"read_json": ToolSpec(
name="read_json",
description="Read a preview of a JSON file inside context.",
input_schema={"path": "relative/path/to/file.json", "max_chars": 4000},
),
}
handlers = {
"answer": _answer,
"execute_context_sql": _execute_context_sql,
"execute_python": _execute_python,
"inspect_sqlite_schema": _inspect_sqlite_schema,
"list_context": _list_context,
"read_csv": _read_csv,
"read_doc": _read_doc,
"read_json": _read_json,
}
return ToolRegistry(specs=specs, handlers=handlers)
|