Prompt templates

Closes #23

docs/conf.py

@@ -31,6 +31,7 @@ from subprocess import PIPE, Popen
 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
 # ones.
 extensions = ["myst_parser"]
+myst_enable_extensions = ["colon_fence"]
 
 # Add any paths that contain templates here, relative to this directory.
 templates_path = ["_templates"]

docs/index.md

@@ -19,6 +19,7 @@ maxdepth: 3
 ---
 setup
 usage
+templates
 logging
 contributing
 changelog

docs/setup.md

@@ -21,6 +21,7 @@ Keys can be persisted in a file that is used by the tool. This file is called `k
 ```
 llm keys path
 ```
+On macOS this will be `~/Library/Application Support/io.datasette.llm/keys.json`. On Linux it may be something like `~/.config/io.datasette.llm/keys.json`.
 
 Rather than editing this file directly, you can instead add keys to it using the `llm keys set` command.
 

docs/templates.md

@@ -0,0 +1,147 @@
+# Prompt templates
+
+Prompt templates can be created to reuse useful prompts with different input data.
+
+## Getting started
+
+Here's a template for summarizing text:
+
+```yaml
+'Summarize this: $input'
+```
+This is a string of YAML - hence the single quotes. The `$input` token will be replaced by extra content provided to the template when it is executed.
+
+To create this template with the name `summary` run the following:
+
+```
+llm templates edit summary
+```
+This will open the system default editor.
+
+:::{tip}
+You can control which editor will be used here using the `EDITOR` environment variable - for example, to use VS Code:
+
+    export EDITOR="code -w"
+
+Add that to your `~/.zshrc` or `~/.bashrc` file depending on which shell you use (`zsh` is the default on macOS since macOS Catalina in 2019).
+:::
+
+You can also create a file called `summary.yaml` in the folder shown by running `llm templates path`, for example:
+```bash
+$ llm templates path
+/Users/simon/Library/Application Support/io.datasette.llm/templates
+```
+
+## Using a template
+
+You can execute a named template using the `-t/--template` option:
+
+```bash
+curl -s https://example.com/ | llm -t summary
+```
+
+This can be combined with the `-m` option to specify a different model:
+```bash
+curl -s https://llm.datasette.io/en/latest/ | \
+  llm -t summary -m gpt-3.5-turbo-16k
+```
+## Longer prompts
+
+You can also represent this template as a YAML dictionary with a `prompt:` key, like this one:
+
+```yaml
+prompt: 'Summarize this: $input'
+```
+Or use YAML multi-line strings for longer inputs. I created this using `llm templates edit steampunk`:
+```yaml
+prompt: >
+    Summarize the following text.
+
+    Insert frequent satirical steampunk-themed illustrative anecdotes.
+    Really go wild with that.
+
+    Text to summarize: $input
+```
+The `prompt: >` causes the following indented text to be treated as a single string, with newlines collapsed to spaces. Use `prompt: |` to preserve newlines.
+
+Running that with `llm -t steampunk` against GPT-4 (via [strip-tags](https://github.com/simonw/strip-tags) to remove HTML tags from the input and minify whitespace):
+```bash
+curl -s 'https://til.simonwillison.net/macos/imovie-slides-and-audio' | \
+  strip-tags -m | llm -t steampunk -m 4
+```
+Output:
+> In a fantastical steampunk world, Simon Willison decided to merge an old MP3 recording with slides from the talk using iMovie. After exporting the slides as images and importing them into iMovie, he had to disable the default Ken Burns effect using the "Crop" tool. Then, Simon manually synchronized the audio by adjusting the duration of each image. Finally, he published the masterpiece to YouTube, with the whimsical magic of steampunk-infused illustrations leaving his viewers in awe.
+
+## System templates
+
+Templates are YAML files. A template can contain a single string, as shown above, which will then be treated as the prompt.
+
+When working with models that support system prompts (such as `gpt-3.5-turbo` and `gpt-4`) you can instead set a system prompt using a `system:` key like so:
+
+```yaml
+system: Summarize this
+```
+If you specify only a system prompt you don't need to use the `$input` variable - `llm` will use the user's input as the whole of the regular prompt, which will then be processed using the instructions set in that system prompt.
+
+You can combine system and regular prompts like so:
+
+```yaml
+system: You speak like an excitable Victorian adventurer
+prompt: 'Summarize this: $input'
+```
+
+## Additional template variables
+
+Templates that work against the user's normal input (content that is either piped to the tool via standard input or passed as a command-line argument) use just the `$input` variable.
+
+You can use additional named variables. These will then need to be provided using the `-p/--param` option when executing the template.
+
+Here's an example template called `recipe`, created using `llm templates edit recipe`:
+
+```yaml
+prompt: |
+    Suggest a recipe using ingredients: $ingredients
+
+    It should be based on cuisine from this country: $country
+```
+This can be executed like so:
+
+```bash
+llm -t recipe -p ingredients 'sausages, milk' -p country Germany
+```
+My output started like this:
+> Recipe: German Sausage and Potato Soup
+>
+> Ingredients:
+> - 4 German sausages
+> - 2 cups whole milk
+
+This example combines input piped to the tool with additional parameters. Call this `summarize`:
+
+```yaml
+system: Summarize this text in the voice of $voice
+```
+Then to run it:
+```bash
+curl -s 'https://til.simonwillison.net/macos/imovie-slides-and-audio' | \
+  strip-tags -m | llm -t summarize -p voice GlaDOS
+```
+I got this:
+
+> My previous test subject seemed to have learned something new about iMovie. They exported keynote slides as individual images [...] Quite impressive for a human.
+
+## Setting a default model for a template
+
+Templates executed using `llm -t template-name` will execute using the default model that the user has configured for the tool - or `gpt-3.5-turbo` if they have not configured their own default.
+
+You can specify a new default model for a template using the `model:` key in the associated YAML. Here's a template called `roast`:
+
+```yaml
+model: gpt-4
+system: roast the user at every possible opportunity, be succinct
+```
+Example:
+```bash
+llm -t roast 'How are you today?'
+```
+> I'm doing great but with your boring questions, I must admit, I've seen more life in a cemetery.

llm/__init__.py

@@ -0,0 +1,48 @@
+from pydantic import BaseModel
+import string
+from typing import Optional
+
+
+class Template(BaseModel):
+    name: str
+    prompt: Optional[str]
+    system: Optional[str]
+    model: Optional[str]
+
+    class Config:
+        extra = "forbid"
+
+    class MissingVariables(Exception):
+        pass
+
+    def execute(self, input, params=None):
+        params = params or {}
+        params["input"] = input
+        if not self.prompt:
+            system = self.interpolate(self.system, params)
+            prompt = input
+        else:
+            prompt = self.interpolate(self.prompt, params)
+            system = self.interpolate(self.system, params)
+        return prompt, system
+
+    @classmethod
+    def interpolate(cls, text, params):
+        if not text:
+            return text
+        # Confirm all variables in text are provided
+        string_template = string.Template(text)
+        vars = cls.extract_vars(string_template)
+        missing = [p for p in vars if p not in params]
+        if missing:
+            raise cls.MissingVariables(
+                "Missing variables: {}".format(", ".join(missing))
+            )
+        return string_template.substitute(**params)
+
+    @staticmethod
+    def extract_vars(string_template):
+        return [
+            match.group("named")
+            for match in string_template.pattern.finditer(string_template.template)
+        ]

llm/cli.py

@@ -2,14 +2,19 @@ import click
 from click_default_group import DefaultGroup
 import datetime
 import json
+from llm import Template
 from .migrations import migrate
 import openai
 import os
 import pathlib
+import pydantic
+import shutil
 import sqlite_utils
+from string import Template as StringTemplate
 import sys
 import time
 import warnings
+import yaml
 
 warnings.simplefilter("ignore", ResourceWarning)
 
@@ -17,6 +22,8 @@ DEFAULT_MODEL = "gpt-3.5-turbo"
 
 MODEL_ALIASES = {"4": "gpt-4", "gpt4": "gpt-4", "chatgpt": "gpt-3.5-turbo"}
 
+DEFAULT_TEMPLATE = "prompt: "
+
 
 @click.group(
     cls=DefaultGroup,
@@ -32,6 +39,14 @@ def cli():
 @click.argument("prompt", required=False)
 @click.option("--system", help="System prompt to use")
 @click.option("-m", "--model", help="Model to use")
+@click.option("-t", "--template", help="Template to use")
+@click.option(
+    "-p",
+    "--param",
+    multiple=True,
+    type=(str, str),
+    help="Parameters for template",
+)
 @click.option("--no-stream", is_flag=True, help="Do not stream output")
 @click.option("-n", "--no-log", is_flag=True, help="Don't log to database")
 @click.option(
@@ -49,12 +64,32 @@ def cli():
     type=int,
 )
 @click.option("--key", help="API key to use")
-def prompt(prompt, system, model, no_stream, no_log, _continue, chat_id, key):
+def prompt(
+    prompt, system, model, template, param, no_stream, no_log, _continue, chat_id, key
+):
     "Execute a prompt against on OpenAI model"
     if prompt is None:
-        # Read from stdin instead
-        prompt = sys.stdin.read()
+        if template:
+            # If running a template only consume from stdin if it has data
+            if not sys.stdin.isatty():
+                prompt = sys.stdin.read()
+        else:
+            # Hang waiting for input to stdin
+            prompt = sys.stdin.read()
+
     openai.api_key = get_key(key, "openai", "OPENAI_API_KEY")
+    if template:
+        params = dict(param)
+        # Cannot be used with system
+        if system:
+            raise click.ClickException("Cannot use -t/--template and --system together")
+        template_obj = load_template(template)
+        try:
+            prompt, system = template_obj.execute(prompt, params)
+        except Template.MissingVariables as ex:
+            raise click.ClickException(str(ex))
+        if model is None and template_obj.model:
+            model = template_obj.model
     messages = []
     if _continue:
         _continue = -1
@@ -215,6 +250,78 @@ def logs_list(count, path, truncate):
     click.echo(json.dumps(list(rows), indent=2))
 
 
+@cli.group()
+def templates():
+    "Manage prompt templates"
+
+
+@templates.command(name="list")
+def templates_list():
+    "List available templates"
+    path = template_dir()
+    pairs = []
+    for file in path.glob("*.yaml"):
+        name = file.stem
+        template = load_template(name)
+        pairs.append((name, template.prompt or ""))
+    max_name_len = max(len(p[0]) for p in pairs)
+    fmt = "{name:<" + str(max_name_len) + "} : {prompt}"
+    for name, prompt in sorted(pairs):
+        text = fmt.format(name=name, prompt=prompt)
+        click.echo(display_truncated(text))
+
+
+def display_truncated(text):
+    console_width = shutil.get_terminal_size()[0]
+    if len(text) > console_width:
+        return text[: console_width - 3] + "..."
+    else:
+        return text
+
+
+@templates.command(name="show")
+@click.argument("name")
+def templates_show(name):
+    "Show the specified template"
+    template = load_template(name)
+    click.echo(
+        yaml.dump(
+            dict((k, v) for k, v in template.dict().items() if v is not None),
+            indent=4,
+            default_flow_style=False,
+        )
+    )
+
+
+@templates.command(name="edit")
+@click.argument("name")
+def templates_edit(name):
+    "Edit the specified template"
+    # First ensure it exists
+    path = template_dir() / f"{name}.yaml"
+    if not path.exists():
+        path.write_text(DEFAULT_TEMPLATE, "utf-8")
+    click.edit(filename=path)
+    # Validate that template
+    load_template(name)
+
+
+@templates.command(name="path")
+def templates_path():
+    "Output path to templates directory"
+    click.echo(template_dir())
+
+
+def template_dir():
+    llm_templates_path = os.environ.get("LLM_TEMPLATES_PATH")
+    if llm_templates_path:
+        path = pathlib.Path(llm_templates_path)
+    else:
+        path = user_dir() / "templates"
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
 def _truncate_string(s, max_length=100):
     if len(s) > max_length:
         return s[: max_length - 3] + "..."
@@ -284,6 +391,26 @@ def log(no_log, system, prompt, response, model, chat_id=None, debug=None, start
     )
 
 
+def load_template(name):
+    path = template_dir() / f"{name}.yaml"
+    if not path.exists():
+        raise click.ClickException(f"Invalid template: {name}")
+    try:
+        loaded = yaml.safe_load(path.read_text())
+    except yaml.YAMLError as ex:
+        raise click.ClickException("Invalid YAML: {}".format(str(ex)))
+    if isinstance(loaded, str):
+        return Template(name=name, prompt=loaded)
+    loaded["name"] = name
+    try:
+        return Template.parse_obj(loaded)
+    except pydantic.ValidationError as e:
+        msg = "A validation error occurred:"
+        for error in e.errors():
+            msg += f"\n  {error['loc'][0]}: {error['msg']}"
+        raise click.ClickException(msg)
+
+
 def get_history(chat_id):
     if chat_id is None:
         return None, []

setup.py

@@ -36,6 +36,8 @@ setup(
         "openai",
         "click-default-group-wheel",
         "sqlite-utils",
+        "pydantic",
+        "PyYAML",
     ],
     extras_require={"test": ["pytest", "requests-mock"]},
     python_requires=">=3.7",

tests/conftest.py

@@ -11,10 +11,18 @@ def keys_path(tmpdir):
     return tmpdir / "keys.json"
 
 
+@pytest.fixture
+def templates_path(tmpdir):
+    path = tmpdir / "templates"
+    path.mkdir()
+    return path
+
+
 @pytest.fixture(autouse=True)
-def env_setup(monkeypatch, log_path, keys_path):
+def env_setup(monkeypatch, log_path, keys_path, templates_path):
     monkeypatch.setenv("LLM_KEYS_PATH", str(keys_path))
     monkeypatch.setenv("LLM_LOG_PATH", str(log_path))
+    monkeypatch.setenv("LLM_TEMPLATES_PATH", str(templates_path))
 
 
 @pytest.fixture

tests/test_templates.py

@@ -0,0 +1,121 @@
+from click.testing import CliRunner
+from llm import Template
+from llm.cli import cli
+import os
+from unittest import mock
+import pytest
+
+
+@pytest.mark.parametrize(
+    "prompt,system,params,expected_prompt,expected_system,expected_error",
+    (
+        ("S: $input", None, {}, "S: input", None, None),
+        ("S: $input", "system", {}, "S: input", "system", None),
+        ("No vars", None, {}, "No vars", None, None),
+        ("$one and $two", None, {}, None, None, "Missing variables: one, two"),
+        ("$one and $two", None, {"one": 1, "two": 2}, "1 and 2", None, None),
+    ),
+)
+def test_template_execute(
+    prompt, system, params, expected_prompt, expected_system, expected_error
+):
+    t = Template(name="t", prompt=prompt, system=system)
+    if expected_error:
+        with pytest.raises(Template.MissingVariables) as ex:
+            prompt, system = t.execute("input", params)
+        assert ex.value.args[0] == expected_error
+    else:
+        prompt, system = t.execute("input", params)
+        assert prompt == expected_prompt
+        assert system == expected_system
+
+
+def test_templates_list(templates_path):
+    (templates_path / "one.yaml").write_text("template one", "utf-8")
+    (templates_path / "two.yaml").write_text("template two", "utf-8")
+    (templates_path / "three.yaml").write_text(
+        "template three is very long " * 4, "utf-8"
+    )
+    runner = CliRunner()
+    result = runner.invoke(cli, ["templates", "list"])
+    assert result.exit_code == 0
+    assert result.output == (
+        "one   : template one\n"
+        "three : template three is very long template three is very long template thre...\n"
+        "two   : template two\n"
+    )
+
+
+@mock.patch.dict(os.environ, {"OPENAI_API_KEY": "X"})
+@pytest.mark.parametrize(
+    "template,extra_args,expected_model,expected_input,expected_error",
+    (
+        (
+            "'Summarize this: $input'",
+            [],
+            "gpt-3.5-turbo",
+            "Summarize this: Input text",
+            None,
+        ),
+        (
+            "prompt: 'Summarize this: $input'\nmodel: gpt-4",
+            [],
+            "gpt-4",
+            "Summarize this: Input text",
+            None,
+        ),
+        (
+            "prompt: 'Summarize this: $input'",
+            ["-m", "4"],
+            "gpt-4",
+            "Summarize this: Input text",
+            None,
+        ),
+        (
+            "boo",
+            ["--system", "s"],
+            None,
+            None,
+            "Error: Cannot use -t/--template and --system together",
+        ),
+        (
+            "prompt: 'Say $hello'",
+            [],
+            None,
+            None,
+            "Error: Missing variables: hello",
+        ),
+        (
+            "prompt: 'Say $hello'",
+            ["-p", "hello", "Blah"],
+            "gpt-3.5-turbo",
+            "Say Blah",
+            None,
+        ),
+    ),
+)
+def test_template_basic(
+    templates_path,
+    mocked_openai,
+    template,
+    extra_args,
+    expected_model,
+    expected_input,
+    expected_error,
+):
+    (templates_path / "template.yaml").write_text(template, "utf-8")
+    runner = CliRunner()
+    result = runner.invoke(
+        cli,
+        ["--no-stream", "-t", "template", "Input text"] + extra_args,
+        catch_exceptions=False,
+    )
+    if expected_error is None:
+        assert result.exit_code == 0
+        assert mocked_openai.last_request.json() == {
+            "model": expected_model,
+            "messages": [{"role": "user", "content": expected_input}],
+        }
+    else:
+        assert result.exit_code == 1
+        assert result.output.strip() == expected_error