From 0f114be5f00c10cda1bdd311e776490dd646e27e Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sun, 11 May 2025 20:44:17 -0700
Subject: [PATCH] --python-tools option on prompt and tools, closes #995

---
 docs/help.md          |  7 +++++--
 llm/cli.py            | 44 +++++++++++++++++++++++++++++++++++++++----
 tests/test_plugins.py | 13 ++++++++++---
 3 files changed, 55 insertions(+), 9 deletions(-)

diff --git a/docs/help.md b/docs/help.md
index b61d551..983a886 100644
--- a/docs/help.md
+++ b/docs/help.md
@@ -127,6 +127,8 @@ Options:
                                   Attachment with explicit mimetype,
                                   --at image.jpg image/jpeg
   -T, --tool TEXT                 Name of a tool to make available to the model
+  --python-tools TEXT             Python code block defining functions to
+                                  register as tools
   -o, --option <TEXT TEXT>...     key/value options for the model
   --schema TEXT                   JSON schema, filepath or ID
   --schema-multi TEXT             JSON schema to use for multiple results
@@ -623,8 +625,9 @@ Usage: llm tools list [OPTIONS]
   List available tools that have been provided by plugins
 
 Options:
-  --json  Output as JSON
-  --help  Show this message and exit.
+  --json               Output as JSON
+  --python-tools TEXT  Python code block defining functions to register as tools
+  --help               Show this message and exit.
 ```
 
 (help-aliases)=
diff --git a/llm/cli.py b/llm/cli.py
index 6b5496a..a5bc9e2 100644
--- a/llm/cli.py
+++ b/llm/cli.py
@@ -15,6 +15,7 @@ from llm import (
     Fragment,
     Response,
     Template,
+    Tool,
     UnknownModelError,
     KeyModel,
     encode,
@@ -340,6 +341,11 @@ def cli():
     multiple=True,
     help="Name of a tool to make available to the model",
 )
+@click.option(
+    "python_tools",
+    "--python-tools",
+    help="Python code block defining functions to register as tools",
+)
 @click.option(
     "options",
     "-o",
@@ -413,6 +419,7 @@ def prompt(
     attachments,
     attachment_types,
     tools,
+    python_tools,
     options,
     schema_input,
     schema_multi,
@@ -669,7 +676,7 @@ def prompt(
     except UnknownModelError as ex:
         raise click.ClickException(ex)
 
-    if conversation is None and tools:
+    if conversation is None and (tools or python_tools):
         conversation = model.conversation()
 
     if conversation:
@@ -731,7 +738,11 @@ def prompt(
     if conversation:
         prompt_method = conversation.prompt
 
-    if tools:
+    extra_tools = []
+    if python_tools:
+        extra_tools = _tools_from_code(python_tools)
+
+    if tools or python_tools:
         prompt_method = lambda *args, **kwargs: conversation.chain(
             *args, **kwargs
         ).details()
@@ -745,7 +756,7 @@ def prompt(
                     ", ".join(bad_tools), ", ".join(registered_tools.keys())
                 )
             )
-        kwargs["tools"] = [registered_tools[tool] for tool in tools]
+        kwargs["tools"] = [registered_tools[tool] for tool in tools] + extra_tools
     try:
         if async_:
 
@@ -2109,9 +2120,17 @@ def tools():
 
 @tools.command(name="list")
 @click.option("json_", "--json", is_flag=True, help="Output as JSON")
-def tools_list(json_):
+@click.option(
+    "python_tools",
+    "--python-tools",
+    help="Python code block defining functions to register as tools",
+)
+def tools_list(json_, python_tools):
     "List available tools that have been provided by plugins"
     tools = get_tools()
+    if python_tools:
+        for tool in _tools_from_code(python_tools):
+            tools[tool.name] = tool
     if json_:
         click.echo(
             json.dumps(
@@ -3380,3 +3399,20 @@ def load_template(name: str) -> Template:
         raise LoadTemplateError(f"Invalid template: {name}")
     content = path.read_text()
     return _parse_yaml_template(name, content)
+
+
+def _tools_from_code(code: str) -> List[Tool]:
+    """
+    Treat all Python functions in the code as tools
+    """
+    globals = {}
+    tools = []
+    try:
+        exec(code, globals)
+    except SyntaxError as ex:
+        raise click.ClickException("Error in --python-tools definition: {}".format(ex))
+    # Register all callables in the locals dict:
+    for name, value in globals.items():
+        if callable(value) and not name.startswith("_"):
+            tools.append(Tool.function(value))
+    return tools
diff --git a/tests/test_plugins.py b/tests/test_plugins.py
index 8be3931..38d9f43 100644
--- a/tests/test_plugins.py
+++ b/tests/test_plugins.py
@@ -247,9 +247,9 @@ def test_register_tools():
             "  Count the number of occurrences of a character in a word.\n"
         )
         # And --json
-        result = runner.invoke(cli.cli, ["tools", "list", "--json"])
-        assert result.exit_code == 0
-        assert json.loads(result.output) == {
+        result2 = runner.invoke(cli.cli, ["tools", "list", "--json"])
+        assert result2.exit_code == 0
+        assert json.loads(result2.output) == {
             "upper": {
                 "description": "Convert text to uppercase.",
                 "arguments": {
@@ -270,6 +270,13 @@ def test_register_tools():
                 },
             },
         }
+        # And test the --python-tools option
+        # llm tools --python-tools 'def reverse(s: str): return s[::-1]'
+        result3 = runner.invoke(
+            cli.cli, ["tools", "--python-tools", "def reverse(s: str): return s[::-1]"]
+        )
+        assert result3.exit_code == 0
+        assert 'reverse(s: str)' in result3.output
     finally:
         plugins.pm.unregister(name="ToolsPlugin")
         assert llm.get_tools() == {}