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)