Tools in templates (#1138)

Closes #1009

README.md

@@ -204,6 +204,7 @@ See also [the llm tag](https://simonwillison.net/tags/llm/) on my blog.
     * [System prompts](https://llm.datasette.io/en/stable/templates.html#system-prompts)
     * [Fragments](https://llm.datasette.io/en/stable/templates.html#fragments)
     * [Options](https://llm.datasette.io/en/stable/templates.html#options)
+    * [Tools](https://llm.datasette.io/en/stable/templates.html#tools)
     * [Schemas](https://llm.datasette.io/en/stable/templates.html#schemas)
     * [Additional template variables](https://llm.datasette.io/en/stable/templates.html#additional-template-variables)
     * [Specifying default parameters](https://llm.datasette.io/en/stable/templates.html#specifying-default-parameters)

docs/plugins/plugin-hooks.md

@@ -215,10 +215,16 @@ def my_template_loader(template_path: str) -> llm.Template:
         # Raise a ValueError with a clear message if the template cannot be found
         raise ValueError(f"Template '{template_path}' could not be loaded: {str(e)}")
 ```
-Consult the latest code in [llm/templates.py](https://github.com/simonw/llm/blob/main/llm/templates.py) for details of that `llm.Template` class.
+The `llm.Template` class has the following constructor:
+
+```{eval-rst}
+.. autoclass:: llm.Template
+```
 
 The loader function should raise a `ValueError` if the template cannot be found or loaded correctly, providing a clear error message.
 
+Note that `functions:` provided by templates using this plugin hook will not be made available, to avoid the risk of plugin hooks that load templates from remote sources introducing arbitrary code execution vulnerabilities.
+
 (plugin-hooks-register-fragment-loaders)=
 ## register_fragment_loaders(register)
 

docs/templates.md

@@ -37,6 +37,15 @@ If you want to include a literal `$` sign in your prompt, use `$$` instead:
 ```bash
 llm --system 'Estimate the cost in $$ of this: $input' --save estimate
 ```
+Use `--tool/-T` one or more times to add tools to the template:
+```bash
+llm -T llm_time --system 'Always include the current time in the answer' --save time
+```
+You can also use `--functions` to add Python function code directly to the template:
+```bash
+llm --functions 'def reverse_string(s): return s[::-1]' --system 'reverse any input' --save reverse
+llm -t reverse 'Hello, world!'
+```
 
 Add `--schema` to bake a {ref}`schema <usage-schemas>` into your template:
 
@@ -47,7 +56,7 @@ llm --schema dog.schema.json 'invent a dog' --save dog
 If you add `--extract` the setting to  {ref}`extract the first fenced code block <usage-extract-fenced-code>` will be persisted in the template.
 ```bash
 llm --system 'write a Python function' --extract --save python-function
-llm -t python-function 'reverse a string'
+llm -t python-function 'calculate haversine distance between two points'
 ```
 In each of these cases the template will be saved in YAML format in a dedicated directory on disk.
 
@@ -66,15 +75,16 @@ This can be combined with the `-m` option to specify a different model:
 curl -s https://llm.datasette.io/en/latest/ | \
   llm -t summarize -m gpt-3.5-turbo-16k
 ```
-Templates can also be specified as full URLs to YAML files:
+Templates can also be specified as a direct path to a YAML file on disk:
+```bash
+llm -t path/to/template.yaml 'extra prompt here'
+```
+Or as a URL to a YAML file hosted online:
 ```bash
 llm -t https://raw.githubusercontent.com/simonw/llm-templates/refs/heads/main/python-app.yaml \
   'Python app to pick a random line from a file'
 ```
-Or as a direct path to a YAML file on disk:
-```bash
-llm -t path/to/template.yaml 'extra prompt here'
-```
+Note that templates loaded via URLs will have any `functions:` keys ignored, to avoid accidentally executing arbitrary code. This restriction also applies to templates loaded via the {ref}`template loaders plugin mechanism <plugin-hooks-register-template-loaders>`.
 
 (prompt-templates-list)=
 
@@ -190,6 +200,27 @@ options:
   temperature: 1.8
 ```
 
+(prompt-templates-tools)=
+
+### Tools
+
+The `tools:` key can provide a list of tool names from other plugins - either function names or toolbox specifiers:
+```yaml
+name: time-plus
+tools:
+- llm_time
+- Datasette("https://example.com/timezone-lookup")
+```
+The `functions:` key can provide a multi-line string of Python code defining additional functions:
+```yaml
+name: my-functions
+functions: |
+  def reverse_string(s: str):
+      return s[::-1]
+
+  def greet(name: str):
+      return f"Hello, {name}!"
+```
 (prompt-templates-schemas)=
 
 ### Schemas

llm/cli.py

@@ -626,6 +626,10 @@ def prompt(
             to_save["fragments"] = list(fragments)
         if system_fragments:
             to_save["system_fragments"] = list(system_fragments)
+        if python_tools:
+            to_save["functions"] = "\n\n".join(python_tools)
+        if tools:
+            to_save["tools"] = list(tools)
         if attachments:
             # Only works for attachments with a path or url
             to_save["attachments"] = [
@@ -675,6 +679,10 @@ def prompt(
             system_fragments = [*template_obj.system_fragments, *system_fragments]
         if template_obj.schema_object:
             schema = template_obj.schema_object
+        if template_obj.tools:
+            tools = [*template_obj.tools, *tools]
+        if template_obj.functions and template_obj._functions_is_trusted:
+            python_tools = [template_obj.functions, *python_tools]
         input_ = ""
         if template_obj.options:
             # Make options mutable (they start as a tuple)
@@ -3836,7 +3844,10 @@ def load_template(name: str) -> Template:
     if not path.exists():
         raise LoadTemplateError(f"Invalid template: {name}")
     content = path.read_text()
-    return _parse_yaml_template(name, content)
+    template_obj = _parse_yaml_template(name, content)
+    # We trust functions here because they came from the filesystem
+    template_obj._functions_is_trusted = True
+    return template_obj
 
 
 def _tools_from_code(code_or_path: str) -> List[Tool]:

llm/templates.py

@@ -22,12 +22,20 @@ class Template(BaseModel):
     schema_object: Optional[dict] = None
     fragments: Optional[List[str]] = None
     system_fragments: Optional[List[str]] = None
+    tools: Optional[List[str]] = None
+    functions: Optional[str] = None
 
     model_config = ConfigDict(extra="forbid")
 
     class MissingVariables(Exception):
         pass
 
+    def __init__(self, **data):
+        super().__init__(**data)
+        # Not a pydantic field to avoid YAML being able to set it
+        # this controls if Python inline functions code is trusted
+        self._functions_is_trusted = False
+
     def evaluate(
         self, input: str, params: Optional[Dict[str, Any]] = None
     ) -> Tuple[Optional[str], Optional[str]]:

tests/test_templates.py

@@ -1,11 +1,14 @@
 from click.testing import CliRunner
+from importlib.metadata import version
 import json
-from llm import Template
+from llm import Template, Toolbox, hookimpl, user_dir
 from llm.cli import cli
+from llm.plugins import pm
 import os
 from unittest import mock
 import pathlib
 import pytest
+import textwrap
 import yaml
 
 
@@ -399,3 +402,138 @@ def test_execute_prompt_from_template_path():
             "stream": True,
             "previous": [],
         }
+
+
+FUNCTIONS_EXAMPLE = """
+def greet(name: str) -> str:
+    return f"Hello, {name}!"
+"""
+
+
+class Greeting(Toolbox):
+    def __init__(self, greeting: str):
+        self.greeting = greeting
+
+    def greet(self, name: str) -> str:
+        "Greet name with a greeting"
+        return f"{self.greeting}, {name}!"
+
+
+class GreetingsPlugin:
+    __name__ = "GreetingsPlugin"
+
+    @hookimpl
+    def register_tools(self, register):
+        register(Greeting)
+
+
+@pytest.mark.parametrize(
+    "source,expected_tool_success,expected_functions_success",
+    (
+        ("alias", True, True),
+        ("file", True, True),
+        # Loaded from URL or plugin = functions: should not work
+        ("url", True, False),
+        ("plugin", True, False),
+    ),
+)
+def test_tools_in_templates(
+    source, expected_tool_success, expected_functions_success, httpx_mock, tmpdir
+):
+    template_yaml = textwrap.dedent(
+        """
+    name: test
+    tools:
+    - llm_version
+    - Greeting("hi")
+    functions: |
+      def demo():
+          return "Demo"
+    """
+    )
+    args = []
+
+    def before():
+        pass
+
+    def after():
+        pass
+
+    if source == "alias":
+        args = ["-t", "test"]
+        (user_dir() / "templates").mkdir(parents=True, exist_ok=True)
+        (user_dir() / "templates" / "test.yaml").write_text(template_yaml, "utf-8")
+    elif source == "file":
+        (tmpdir / "test.yaml").write_text(template_yaml, "utf-8")
+        args = ["-t", str(tmpdir / "test.yaml")]
+    elif source == "url":
+        httpx_mock.add_response(
+            url="https://example.com/test.yaml",
+            method="GET",
+            text=template_yaml,
+            status_code=200,
+            is_reusable=True,
+        )
+        args = ["-t", "https://example.com/test.yaml"]
+    elif source == "plugin":
+
+        class LoadTemplatePlugin:
+            __name__ = "LoadTemplatePlugin"
+
+            @hookimpl
+            def register_template_loaders(self, register):
+                register(
+                    "tool-template",
+                    lambda s: Template(
+                        name="tool-template",
+                        tools=["llm_version", 'Greeting("hi")'],
+                        functions=FUNCTIONS_EXAMPLE,
+                    ),
+                )
+
+        def before():
+            pm.register(LoadTemplatePlugin(), name="test-tools-in-templates")
+
+        def after():
+            pm.unregister(name="test-tools-in-templates")
+
+        args = ["-t", "tool-template:"]
+
+    before()
+    pm.register(GreetingsPlugin(), name="greetings-plugin")
+    try:
+        runner = CliRunner()
+        # Test llm_version, then Greeting, then demo
+        for tool_call, text, should_be_present in (
+            ({"name": "llm_version"}, version("llm"), True),
+            (
+                {"name": "Greeting_greet", "arguments": {"name": "Alice"}},
+                "hi, Alice",
+                expected_tool_success,
+            ),
+            (
+                {"name": "Greeting_greet", "arguments": {"name": "Bob"}},
+                "hi, Bob!",
+                expected_tool_success,
+            ),
+            ({"name": "demo"}, '"output": "Demo"', expected_functions_success),
+        ):
+            result = runner.invoke(
+                cli,
+                args
+                + [
+                    "-m",
+                    "echo",
+                    "--no-stream",
+                    json.dumps({"tool_calls": [tool_call]}),
+                ],
+                catch_exceptions=False,
+            )
+            assert result.exit_code == 0
+            if should_be_present:
+                assert text in result.output
+            else:
+                assert text not in result.output
+    finally:
+        after()
+        pm.unregister(name="greetings-plugin")