mulebot 0.1.2__tar.gz → 0.2.0__tar.gz

{mulebot-0.1.2 → mulebot-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mulebot
-Version: 0.1.2
+Version: 0.2.0
 Summary: A simple chatbot for Datamule's APIs
 Home-page: https://github.com/john-friedman/mulebot
 Author: John Friedman

mulebot-0.2.0/mulebot/__init__.py

@@ -0,0 +1,4 @@
+from .agent import Agent
+from .tool import Tool, tool
+
+__all__ = ["Agent", "Tool", "tool"]

{mulebot-0.1.2 → mulebot-0.2.0}/mulebot/agent.py

@@ -5,7 +5,8 @@ import urllib.request
 from datetime import date
 from time import perf_counter, time
 
-from .tools import TOOL_DEFINITIONS, Tools, documentation_context
+from .tools import Tools, documentation_context
+from .tool import validate_tools
 
 
 DEFAULT_MODEL = "deepseek/deepseek-v4.1-flash-1017:nitro"
@@ -68,15 +69,67 @@ class Agent:
         api_key = api_key or os.environ.get("DATAMULE_API_KEY")
         if not api_key or not self.model_api_key:
             raise ValueError("Set DATAMULE_API_KEY and OPENROUTER_API_KEY.")
-        self.tools = Tools(api_key, allowed_tables=allowed_tables)
-        context = f"Today is {date.today().isoformat()}.\n\n<documentation>\n"
-        context += documentation_context() + "\n</documentation>\n\n" + INSTRUCTIONS
-        self.messages = [{"role": "system", "content": context}]
+        self._tool_state = Tools(api_key, allowed_tables=allowed_tables)
+        self.instructions = INSTRUCTIONS
+        self.doc = documentation_context()
+        self.messages = []
+        self._update_system_message()
+
+    @property
+    def tools(self):
+        """All active tools, including built-ins. Use set_tools/append_tools to change them."""
+        return validate_tools(self._tool_state.registry)
+
+    @property
+    def results(self):
+        return self._tool_state.results
+
+    @property
+    def artifacts(self):
+        return self._tool_state.artifacts
+
+    def set_instructions(self, text):
+        if not isinstance(text, str):
+            raise TypeError("Instructions must be a string.")
+        self.instructions = text
+        self._update_system_message()
+        return self
+
+    def append_instructions(self, text):
+        if not isinstance(text, str):
+            raise TypeError("Instructions must be a string.")
+        return self.set_instructions("\n\n".join(part for part in (self.instructions, text) if part))
+
+    def set_doc(self, text):
+        if not isinstance(text, str):
+            raise TypeError("Documentation must be a string.")
+        self.doc = text
+        self._update_system_message()
+        return self
+
+    def append_doc(self, text):
+        if not isinstance(text, str):
+            raise TypeError("Documentation must be a string.")
+        return self.set_doc("\n\n".join(part for part in (self.doc, text) if part))
+
+    def set_tools(self, tools):
+        self._tool_state.registry = validate_tools(tools)
+        return self
+
+    def append_tools(self, tools):
+        if not isinstance(tools, (list, tuple)):
+            raise TypeError("Tools must be a list or tuple.")
+        return self.set_tools([*self.tools, *tools])
+
+    def _update_system_message(self):
+        content = (f"Today is {date.today().isoformat()}.\n\n<documentation>\n"
+                   + self.doc + "\n</documentation>\n\n" + self.instructions)
+        self.messages[:1] = [{"role": "system", "content": content}]
 
     def export_state(self):
         """Return browser-held chat data, without credentials or system instructions."""
-        state = {"messages": self.messages[1:], "results": self.tools.results,
-                 "artifacts": self.tools.artifacts}
+        state = {"messages": self.messages[1:], "results": self._tool_state.results,
+                 "artifacts": self._tool_state.artifacts}
         return json.loads(self._redact(json.dumps(state, ensure_ascii=False)))
 
     def load_state(self, state):
@@ -129,29 +182,30 @@ class Agent:
         if not all(isinstance(artifact, dict) for artifact in artifacts):
             raise ValueError("Invalid saved artifacts.")
         self.messages = self.messages[:1] + restored
-        self.tools.results = results
-        self.tools.artifacts = artifacts
+        self._tool_state.results = results
+        self._tool_state.artifacts = artifacts
 
     def _redact(self, text):
-        for key in [self.model_api_key, self.tools.api_key]:
+        for key in [self.model_api_key, self._tool_state.api_key]:
             text = text.replace(key, "[redacted]")
         return text
 
     def _complete(self, tool_choice):
         messages = list(self.messages)
-        if self.tools.results:
+        if self._tool_state.results:
             saved = {key: {"columns": value["columns"], "row_count": len(value["rows"]), "sql": value["sql"]}
-                     for key, value in self.tools.results.items()}
+                     for key, value in self._tool_state.results.items()}
             messages.append({"role": "system", "content":
                 "These results are already saved and remain available even if older chat was trimmed. "
                 "Use query(result_id=...) or plot.sql to read them locally instead of paying to fetch "
                 "the same data again. Local SQL reads the table data using SQLite.\n" + json.dumps(saved)})
         payload = {
             "model": self.model, "messages": messages,
-            "tools": TOOL_DEFINITIONS, "tool_choice": tool_choice,
-            "max_tokens": 8192, "parallel_tool_calls": False,
+            "max_tokens": 8192,
             "reasoning": {"effort": os.environ.get("MULEBOT_REASONING_EFFORT", "medium"), "exclude": False},
         }
+        if self._tool_state.registry:
+            payload.update(tools=self._tool_state.definitions, tool_choice=tool_choice, parallel_tool_calls=False)
         request = urllib.request.Request(
             "https://openrouter.ai/api/v1/chat/completions",
             data=json.dumps(payload).encode("utf-8"),
@@ -233,8 +287,8 @@ class Agent:
                            "llm_cost": (completion.get("usage") or {}).get("cost"),
                            "llm_cost_shared": len(calls) > 1,
                            "reasoning": self._redact(reasoning_text(response))}
-                    result_ids = set(self.tools.results)
-                    first_artifact = len(self.tools.artifacts)
+                    result_ids = set(self._tool_state.results)
+                    first_artifact = len(self._tool_state.artifacts)
                     started = perf_counter()
                     try:
                         if remaining <= 0:
@@ -243,7 +297,8 @@ class Agent:
                         arguments = json.loads(call["function"]["arguments"])
                         if not isinstance(arguments, dict):
                             raise ValueError("Tool arguments must be an object.")
-                        result = self.tools.call(name, arguments)
+                        result = self._tool_state.call(name, arguments)
+                        result = json.loads(json.dumps(result, ensure_ascii=False, allow_nan=False))
                     except Exception as exc:
                         result = {"error": self._redact(str(exc))[:2000]}
                         if getattr(exc, "metadata", None):
@@ -253,10 +308,10 @@ class Agent:
                                           "content": json.dumps(result, ensure_ascii=False)})
                     yield {"type": "tool_result", "id": call["id"], "name": name,
                            "result": result, "duration_ms": duration_ms}
-                    for result_id in self.tools.results.keys() - result_ids:
-                        yield {"type": "query_result", "result_id": result_id, **self.tools.results[result_id]}
-                    for i in range(first_artifact, len(self.tools.artifacts)):
-                        yield {"type": "artifact", "artifact_id": i + 1, "artifact": self.tools.artifacts[i]}
+                    for result_id in self._tool_state.results.keys() - result_ids:
+                        yield {"type": "query_result", "result_id": result_id, **self._tool_state.results[result_id]}
+                    for i in range(first_artifact, len(self._tool_state.artifacts)):
+                        yield {"type": "artifact", "artifact_id": i + 1, "artifact": self._tool_state.artifacts[i]}
         except Exception as exc:
             detail = self._redact(str(exc))[:2000]
             interruption = f"Turn interrupted: {detail}"

mulebot-0.2.0/mulebot/tool.py

@@ -0,0 +1,139 @@
+import inspect
+import math
+import re
+import types
+from copy import deepcopy
+from dataclasses import dataclass, replace
+from typing import Annotated, Any, Literal, Union, get_args, get_origin, get_type_hints
+
+
+@dataclass(frozen=True)
+class Tool:
+    name: str
+    description: str
+    parameters: dict
+    execute: Any
+
+    @property
+    def definition(self):
+        return {"type": "function", "function": {
+            "name": self.name, "description": self.description,
+            "parameters": deepcopy(self.parameters),
+        }}
+
+    def __call__(self, *args, **kwargs):
+        return self.execute(*args, **kwargs)
+
+
+def _schema(annotation):
+    if annotation is Any:
+        return {}
+    if annotation is None or annotation is type(None):
+        return {"type": "null"}
+    primitives = {str: "string", int: "integer", float: "number", bool: "boolean"}
+    if annotation in (str, int, float, bool):
+        return {"type": primitives[annotation]}
+    origin, args = get_origin(annotation), get_args(annotation)
+    if origin is Annotated:
+        return _schema(args[0])
+    if origin in (Union, types.UnionType):
+        return {"anyOf": [_schema(arg) for arg in args]}
+    if origin is Literal:
+        return {"enum": list(args)}
+    if annotation is list or origin is list:
+        return {"type": "array", "items": _schema(args[0]) if args else {}}
+    if annotation is dict or origin is dict:
+        if args and args[0] is not str:
+            raise TypeError("Tool dictionary keys must be strings.")
+        return {"type": "object", "additionalProperties": _schema(args[1]) if args else True}
+    raise TypeError(f"Unsupported tool annotation: {annotation!r}. Supply an explicit parameters schema.")
+
+
+def tool(function=None, *, name=None, description=None, parameters=None):
+    """Declare a synchronous tool using type hints, or an explicit JSON Schema."""
+    def decorate(function):
+        if inspect.iscoroutinefunction(function):
+            raise TypeError("Python tools must be synchronous functions.")
+        schema = parameters
+        if schema is None:
+            hints = get_type_hints(function, include_extras=True)
+            properties, required = {}, []
+            for key, parameter in inspect.signature(function).parameters.items():
+                if parameter.kind not in (parameter.POSITIONAL_OR_KEYWORD, parameter.KEYWORD_ONLY):
+                    raise TypeError("Tool parameters must accept keyword arguments.")
+                if key not in hints:
+                    raise TypeError(f"Tool parameter {key} needs a type annotation.")
+                properties[key] = _schema(hints[key])
+                if parameter.default is parameter.empty:
+                    required.append(key)
+            schema = {"type": "object", "properties": properties,
+                      "required": required, "additionalProperties": False}
+        return validate_tools([Tool(name or function.__name__,
+                                    description if description is not None else inspect.getdoc(function) or "",
+                                    schema, function)])[0]
+    return decorate if function is None else decorate(function)
+
+
+def validate_tools(tools):
+    if not isinstance(tools, (list, tuple)):
+        raise TypeError("Tools must be a list or tuple.")
+    names, result = set(), []
+    for entry in tools:
+        if (not isinstance(entry, Tool) or not isinstance(entry.name, str)
+                or not re.fullmatch(r"[a-zA-Z0-9_-]{1,64}", entry.name)
+                or not isinstance(entry.description, str) or not isinstance(entry.parameters, dict)
+                or entry.parameters.get("type") != "object"
+                or not isinstance(entry.parameters.get("properties"), dict) or not callable(entry.execute)):
+            raise ValueError("Each tool needs a name, description, object parameter schema, and execute function.")
+        if inspect.iscoroutinefunction(entry.execute):
+            raise TypeError("Python tools must be synchronous functions.")
+        if entry.name in names:
+            raise ValueError(f"Duplicate tool name: {entry.name}")
+        names.add(entry.name)
+        result.append(replace(entry, parameters=deepcopy(entry.parameters)))
+    return result
+
+
+def validate_arguments(value, schema, path):
+    if "anyOf" in schema:
+        for option in schema["anyOf"]:
+            try:
+                validate_arguments(value, option, path)
+                return
+            except ValueError:
+                pass
+        raise ValueError(f"{path} does not match any allowed type.")
+    checks = {
+        "string": lambda: isinstance(value, str),
+        "integer": lambda: type(value) is int,
+        "number": lambda: type(value) is int or (type(value) is float and math.isfinite(value)),
+        "boolean": lambda: isinstance(value, bool),
+        "null": lambda: value is None,
+        "object": lambda: isinstance(value, dict),
+        "array": lambda: isinstance(value, list),
+    }
+    kind = schema.get("type")
+    if kind in checks and not checks[kind]():
+        raise ValueError(f"{path} must be a {kind}.")
+    if "enum" in schema and not any(type(value) is type(entry) and value == entry for entry in schema["enum"]):
+        raise ValueError(f"{path} must be one of: {schema['enum']}.")
+    if isinstance(value, dict):
+        properties = schema.get("properties", {})
+        missing = [key for key in schema.get("required", []) if key not in value]
+        if missing:
+            raise ValueError(f"{path} is missing required field(s): {', '.join(missing)}.")
+        unexpected = [key for key in value if key not in properties]
+        additional = schema.get("additionalProperties", True)
+        if additional is False and unexpected:
+            raise ValueError(f"{path} has unknown field(s): {', '.join(unexpected)}. "
+                             f"Allowed fields: {', '.join(properties)}.")
+        for key, entry in value.items():
+            child = properties.get(key, additional if isinstance(additional, dict) else None)
+            if child is not None:
+                validate_arguments(entry, child, f"{path}.{key}")
+    if isinstance(value, list):
+        if len(value) < schema.get("minItems", 0):
+            raise ValueError(f"{path} needs at least {schema['minItems']} item(s).")
+        if "items" in schema:
+            for index, entry in enumerate(value):
+                validate_arguments(entry, schema["items"], f"{path}[{index}]")

{mulebot-0.1.2 → mulebot-0.2.0}/mulebot/tools.py

@@ -12,6 +12,7 @@ from sqlglot.optimizer.annotate_types import annotate_types
 from sqlglot.optimizer.scope import traverse_scope
 
 from .assets import resource
+from .tool import Tool, validate_arguments, validate_tools
 
 
 SITE = "https://datamule.xyz"
@@ -73,7 +74,7 @@ def documentation_context():
     )
 
 
-def tool(name, description, properties, required=()):
+def _tool_definition(name, description, properties, required=()):
     return {
         "type": "function",
         "function": {
@@ -89,7 +90,7 @@ def tool(name, description, properties, required=()):
 
 STRING = {"type": "string"}
 TOOL_DEFINITIONS = [
-    tool("query", "Fetch the user's complete requested dataset with one read-only Athena SELECT. "
+    _tool_definition("query", "Fetch the user's complete requested dataset with one read-only Athena SELECT. "
          "Resolve identifiers in subqueries/CTEs within this SQL instead of separate lookup calls. "
          "Filter comparable records and deduplicate/aggregate before returning rows; do not fetch "
          "raw samples just to inspect them. Athena date differences use date_diff('day', start, end), "
@@ -100,9 +101,9 @@ TOOL_DEFINITIONS = [
              "sql": {"type": "string", "description": "Athena SQL for a new Datamule query, SQLite SQL for saved results. SQLite durations use julianday(end) - julianday(start), not CAST AS DATE or year subtraction."},
              "result_id": {"type": "string", "description": "Optional saved result to query locally using SQLite; omit only when fetching new data from Datamule."},
          }, ["sql"]),
-    tool("write_markdown", "Create a text Markdown report for the caller to display or save. HTML and SVG are displayed as text, not rendered. Use plot for charts; fix a failed plot call instead of embedding SVG in Markdown.",
+    _tool_definition("write_markdown", "Create a text Markdown report for the caller to display or save. HTML and SVG are displayed as text, not rendered. Use plot for charts; fix a failed plot call instead of embedding SVG in Markdown.",
          {"title": STRING, "markdown": STRING}, ["title", "markdown"]),
-    tool("plot", "Create one chart artifact from a required saved result_id returned by query. This tool never fetches new data. For separate charts, call plot separately with the same result_id and different local sql filters. Chart the final analysis rows used in your answer without querying Datamule again. y is a numeric column or a list of numeric columns for multiple lines; group is a categorical column, not another y. Optional sql filters or reshapes the saved rows using a read-only SQLite SELECT FROM data (SQLite dates use julianday; arrays/objects are JSON). Use this to reuse data on follow-ups. Line charts sort x ascending. Supply human-readable axis labels including units.", {
+    _tool_definition("plot", "Create one chart artifact from a required saved result_id returned by query. This tool never fetches new data. For separate charts, call plot separately with the same result_id and different local sql filters. Chart the final analysis rows used in your answer without querying Datamule again. y is a numeric column or a list of numeric columns for multiple lines; group is a categorical column, not another y. Optional sql filters or reshapes the saved rows using a read-only SQLite SELECT FROM data (SQLite dates use julianday; arrays/objects are JSON). Use this to reuse data on follow-ups. Line charts sort x ascending. Supply human-readable axis labels including units.", {
              "title": STRING, "result_id": {"type": "string", "description": "Required ID of a saved query result, such as query_5. Use its column names for x, y, and group. Fetch missing data with query first; never put a fresh Athena query into plot.sql."},
              "kind": {"type": "string", "enum": ["line", "bar", "scatter"]},
              "x": STRING, "y": {"anyOf": [STRING, {"type": "array", "items": STRING, "minItems": 1}]},
@@ -120,6 +121,11 @@ class Tools:
         self.catalog = None
         self.results = {}
         self.artifacts = []
+        self.registry = validate_tools([
+            Tool(definition["function"]["name"], definition["function"]["description"],
+                 definition["function"]["parameters"], getattr(self, definition["function"]["name"]))
+            for definition in TOOL_DEFINITIONS
+        ])
 
     def _saved_results_help(self):
         saved = [f"{result_id} ({len(result['rows'])} rows; columns: {', '.join(result['columns'])})"
@@ -335,24 +341,21 @@ class Tools:
             response.update(columns=columns, rows=rows, row_count=len(rows))
         return response
 
+    @property
+    def definitions(self):
+        return [entry.definition for entry in self.registry]
+
     def call(self, name, arguments):
-        definition = next((tool["function"] for tool in TOOL_DEFINITIONS
-                           if tool["function"]["name"] == name), None)
-        if definition is None:
+        entry = next((entry for entry in self.registry if entry.name == name), None)
+        if entry is None:
             raise ValueError(f"Unknown tool: {name}")
         if not isinstance(arguments, dict):
             raise ValueError("Tool arguments must be an object.")
-        properties = definition["parameters"]["properties"]
-        missing = [key for key in definition["parameters"]["required"] if key not in arguments]
-        if missing:
-            hint = (" Every plot requires a saved result_id. Optional sql is SQLite SELECT FROM data, "
-                    f"not a new Athena query. {self._saved_results_help()}" if name == "plot" else "")
-            raise ValueError(f"{name} is missing required field(s): {', '.join(missing)}.{hint}")
-        unexpected = [key for key in arguments if key not in properties]
-        if unexpected:
-            raise ValueError(f"{name} has unknown field(s): {', '.join(unexpected)}. "
-                             f"Allowed fields: {', '.join(properties)}.")
-        for key, value in arguments.items():
-            if properties[key].get("type") == "string" and not isinstance(value, str):
-                raise ValueError(f"{name}.{key} must be a string.")
-        return getattr(self, name)(**arguments)
+        try:
+            validate_arguments(arguments, entry.parameters, name)
+        except ValueError as exc:
+            if name == "plot" and "missing required field(s)" in str(exc):
+                raise ValueError(f"{exc} Every plot requires a saved result_id. Optional sql is SQLite "
+                                 f"SELECT FROM data, not a new Athena query. {self._saved_results_help()}") from None
+            raise
+        return entry.execute(**arguments)

{mulebot-0.1.2 → mulebot-0.2.0}/mulebot.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mulebot
-Version: 0.1.2
+Version: 0.2.0
 Summary: A simple chatbot for Datamule's APIs
 Home-page: https://github.com/john-friedman/mulebot
 Author: John Friedman

{mulebot-0.1.2 → mulebot-0.2.0}/mulebot.egg-info/SOURCES.txt

@@ -8,11 +8,11 @@ mulebot/agent.py
 mulebot/assets.py
 mulebot/cli.py
 mulebot/server.py
+mulebot/tool.py
 mulebot/tools.py
 mulebot.egg-info/PKG-INFO
 mulebot.egg-info/SOURCES.txt
 mulebot.egg-info/dependency_links.txt
 mulebot.egg-info/entry_points.txt
 mulebot.egg-info/requires.txt
-mulebot.egg-info/top_level.txt
-tests/test_plot.py
+mulebot.egg-info/top_level.txt

{mulebot-0.1.2 → mulebot-0.2.0}/setup.py

@@ -31,7 +31,7 @@ class Sdist(sdist):
 setup(
     name="mulebot",
     author="John Friedman",
-    version="0.1.2",
+    version="0.2.0",
     description="A simple chatbot for Datamule's APIs",
     url="https://github.com/john-friedman/mulebot",
     license="MIT",

mulebot-0.1.2/mulebot/__init__.py

@@ -1,3 +0,0 @@
-from .agent import Agent
-
-__all__ = ["Agent"]

mulebot-0.1.2/tests/test_plot.py

@@ -1,133 +0,0 @@
-import json
-import unittest
-from pathlib import Path
-from unittest.mock import patch
-
-from mulebot.agent import Agent, INSTRUCTIONS
-from mulebot.tools import TOOL_DEFINITIONS, Tools
-
-
-PLOT = {"title": "RPO", "kind": "line", "x": "filingdate", "y": "rpo_b"}
-COMPANIES = ["Oracle", "Microsoft", "Alphabet"]
-
-
-def saved_tools():
-    tools = Tools("test-key")
-    tools.results["query_5"] = {
-        "sql": "SELECT company, filingdate, rpo_b FROM simple_xbrl",
-        "columns": ["company", "filingdate", "rpo_b"],
-        "rows": [{"company": company, "filingdate": date, "rpo_b": amount + index}
-                 for index, company in enumerate(COMPANIES)
-                 for date, amount in [("2025-06-01", 100), ("2025-09-01", 400)]],
-        "local_types": {"company": "TEXT", "filingdate": "TEXT", "rpo_b": "REAL"},
-        "metadata": {},
-    }
-    return tools
-
-
-class PlotTests(unittest.TestCase):
-    def setUp(self):
-        network = patch("urllib.request.urlopen", side_effect=AssertionError("Unexpected network request"))
-        network.start()
-        self.addCleanup(network.stop)
-        self.tools = saved_tools()
-
-    def test_malformed_plot_has_actionable_error_and_preserves_state(self):
-        with self.assertRaisesRegex(ValueError, r"missing required field\(s\): result_id") as caught:
-            self.tools.call("plot", {**PLOT, "sql": "SELECT * FROM simple_xbrl"})
-        self.assertIn("query_5", str(caught.exception))
-        self.assertIn("filingdate, rpo_b", str(caught.exception))
-        self.assertIn("SQLite SELECT FROM data", str(caught.exception))
-        self.assertEqual(self.tools.artifacts, [])
-        self.assertEqual(list(self.tools.results), ["query_5"])
-
-    def test_separate_charts_reuse_saved_data(self):
-        for index, company in enumerate(COMPANIES):
-            result = self.tools.call("plot", {**PLOT, "title": company, "result_id": "query_5",
-                "sql": f"SELECT * FROM data WHERE company = '{company}' ORDER BY filingdate"})
-            self.assertEqual(result["artifact_id"], index + 1)
-            self.assertEqual(result["row_count"], 2)
-            saved = self.tools.results[result["result_id"]]
-            self.assertEqual([row["company"] for row in saved["rows"]], [company, company])
-            self.assertEqual(saved["source_result_id"], "query_5")
-        self.assertEqual([artifact["title"] for artifact in self.tools.artifacts], COMPANIES)
-        self.assertTrue(all(artifact["kind"] == "line" for artifact in self.tools.artifacts))
-        self.assertEqual(len(self.tools.results["query_5"]["rows"]), 6)
-
-    def test_sql_errors_can_be_corrected_without_fetching(self):
-        for name in ["query", "plot"]:
-            args = PLOT if name == "plot" else {}
-            with self.assertRaisesRegex(ValueError, "Local SQLite query failed") as caught:
-                self.tools.call(name, {**args, "result_id": "query_5", "sql": "SELECT * FROM simple_xbrl"})
-            self.assertIn("columns: company, filingdate, rpo_b", str(caught.exception))
-        filtered = self.tools.call("query", {"result_id": "query_5",
-            "sql": "SELECT * FROM data WHERE company = 'Oracle'"})
-        self.assertEqual(filtered["row_count"], 2)
-        result = self.tools.call("plot", {**PLOT, "result_id": filtered["result_id"]})
-        self.assertEqual(result["artifact_id"], 1)
-
-    def test_invalid_arguments_are_specific(self):
-        cases = [
-            ({**PLOT, "result_id": "query_99"}, "existing saved result_id.*query_5"),
-            ({**PLOT, "result_id": "query_5", "y": "missing"}, "Available columns: company, filingdate, rpo_b"),
-            ({**PLOT, "result_id": "query_5", "extra": True}, r"unknown field\(s\): extra"),
-            ({**PLOT, "result_id": None}, "plot.result_id must be a string"),
-        ]
-        for arguments, message in cases:
-            with self.subTest(arguments=arguments), self.assertRaisesRegex(ValueError, message):
-                self.tools.call("plot", arguments)
-        with self.assertRaisesRegex(ValueError, "No results are saved yet"):
-            Tools("test-key").call("plot", PLOT)
-        self.assertEqual(self.tools.artifacts, [])
-
-    def test_markdown_rejects_svg_chart_but_accepts_reports_and_code(self):
-        with self.assertRaisesRegex(ValueError, "cannot render SVG charts.*plot"):
-            self.tools.call("write_markdown", {"title": "Chart", "markdown": "# RPO\n<svg><path/></svg>"})
-        self.assertEqual(self.tools.artifacts, [])
-        self.tools.call("write_markdown", {"title": "Report", "markdown": "# RPO\nA text report."})
-        self.tools.call("write_markdown", {"title": "Example", "markdown": "```html\n<svg></svg>\n```"})
-        self.assertEqual(len(self.tools.artifacts), 2)
-
-    def test_agent_recovers_from_rejected_plot_and_emits_separate_artifacts(self):
-        agent = Agent(api_key="test-data-key", model_api_key="test-model-key")
-        agent.tools.results = self.tools.results
-        completions = 0
-
-        def complete(tool_choice):
-            nonlocal completions
-            step = completions
-            completions += 1
-            if step == 1:
-                rejection = json.loads(next(message["content"] for message in agent.messages
-                                            if message["role"] == "tool"))
-                self.assertIn("missing required field(s): result_id", rejection["error"])
-                self.assertIn("query_5", rejection["error"])
-            if step == 4:
-                return {"choices": [{"message": {"content": "Created three charts."}}]}
-            arguments = ({**PLOT, "sql": "SELECT * FROM simple_xbrl"} if step == 0 else
-                         {**PLOT, "title": COMPANIES[step - 1], "result_id": "query_5",
-                          "sql": f"SELECT * FROM data WHERE company = '{COMPANIES[step - 1]}'"})
-            return {"choices": [{"message": {"tool_calls": [{"id": f"call_{step}", "type": "function",
-                    "function": {"name": "plot", "arguments": json.dumps(arguments)}}]}}]}
-
-        with patch.object(agent, "_complete", side_effect=complete):
-            events = list(agent.run("Plot each company separately."))
-        outcomes = [event for event in events if event["type"] == "tool_result"]
-        self.assertEqual(len(outcomes), 4)
-        self.assertIn("result_id", outcomes[0]["result"]["error"])
-        self.assertTrue(all(event["result"].get("artifact_id") for event in outcomes[1:]))
-        self.assertEqual([event["artifact"]["title"] for event in events if event["type"] == "artifact"], COMPANIES)
-        self.assertFalse(any(event["type"] == "error" for event in events))
-        self.assertEqual(completions, 5)
-
-    def test_tool_contract_and_instructions_match_typescript(self):
-        config = (Path(__file__).resolve().parents[2] / "typescript/src/config.ts").read_text(encoding="utf-8")
-        prefix = "export const INSTRUCTIONS = "
-        line = next(line for line in config.splitlines() if line.startswith(prefix))
-        self.assertEqual(INSTRUCTIONS, json.loads(line[len(prefix):-1]))
-        definitions = json.loads(config.split("export const TOOL_DEFINITIONS = ", 1)[1].strip().removesuffix(";"))
-        self.assertEqual(TOOL_DEFINITIONS, definitions)
-
-
-if __name__ == "__main__":
-    unittest.main()