Support fragments in chat through -f, --sf, and !fragments (#1048)

* add vscode and uv artifacts to gitignore
* !fragment feature
* Update fragment command to accept multiple fragments
* Add support for multiple fragments in !fragment command, factor out fragment processing
* Add support for fragment/system fragment arguments in chat command
* update docs

---------

Co-authored-by: Simon Willison <swillison@gmail.com>

.gitignore

@@ -8,3 +8,5 @@ venv
 *.egg-info
 .DS_Store
 .idea/
+.vscode/
+uv.lock

docs/fragments.md

@@ -39,6 +39,42 @@ This will read the contents of `setup.py` from standard input and use it as a fr
 
 Fragments can also be used as part of your system prompt. Use `--sf value` or `--system-fragment value` instead of `-f`.
 
+## Using fragments in chat
+
+The `chat` command also supports the `-f` and `--sf` arguments to start a chat with fragments.
+
+```bash
+llm chat -f my_doc.txt
+Chatting with gpt-4
+Type 'exit' or 'quit' to exit
+Type '!multi' to enter multiple lines, then '!end' to finish
+Type '!edit' to open your default editor and modify the prompt.
+Type '!fragment <my_fragment> [<another_fragment> ...]' to insert one or more fragments
+> Explain this document to me
+```
+
+Fragments can also be added *during* a chat conversation using the `!fragment <my_fragment>` command.
+
+```bash
+Chatting with gpt-4
+Type 'exit' or 'quit' to exit
+Type '!multi' to enter multiple lines, then '!end' to finish
+Type '!edit' to open your default editor and modify the prompt.
+Type '!fragment <my_fragment> [<another_fragment> ...]' to insert one or more fragments
+> !fragment https://llm.datasette.io/en/stable/fragments.html
+```
+
+This can be combined with `!multi`:
+
+```bash
+> !multi
+Explain the difference between fragments and templates to me
+!fragment https://llm.datasette.io/en/stable/fragments.html https://llm.datasette.io/en/stable/templates.html
+!end
+```
+
+Additionally, any `!fragment` lines found in a prompt created with `!edit` will also be parsed.
+
 (fragments-browsing)=
 ## Browsing fragments
 

docs/help.md

@@ -163,24 +163,27 @@ Usage: llm chat [OPTIONS]
   Hold an ongoing chat with a model.
 
 Options:
-  -s, --system TEXT            System prompt to use
+  -s, --system TEXT             System prompt to use
-  -m, --model TEXT             Model to use
+  -m, --model TEXT              Model to use
-  -c, --continue               Continue the most recent conversation.
+  -c, --continue                Continue the most recent conversation.
-  --cid, --conversation TEXT   Continue the conversation with the given ID.
+  --cid, --conversation TEXT    Continue the conversation with the given ID.
-  -t, --template TEXT          Template to use
+  -f, --fragment TEXT           Fragment (alias, URL, hash or file path) to add
-  -p, --param <TEXT TEXT>...   Parameters for template
+                                to the prompt
-  -o, --option <TEXT TEXT>...  key/value options for the model
+  --sf, --system-fragment TEXT  Fragment to add to system prompt
-  -d, --database FILE          Path to log database
+  -t, --template TEXT           Template to use
-  --no-stream                  Do not stream output
+  -p, --param <TEXT TEXT>...    Parameters for template
-  --key TEXT                   API key to use
+  -o, --option <TEXT TEXT>...   key/value options for the model
-  -T, --tool TEXT              Name of a tool to make available to the model
+  -d, --database FILE           Path to log database
-  --functions TEXT             Python code block or file path defining functions
+  --no-stream                   Do not stream output
-                               to register as tools
+  --key TEXT                    API key to use
-  --td, --tools-debug          Show full details of tool executions
+  -T, --tool TEXT               Name of a tool to make available to the model
-  --ta, --tools-approve        Manually approve every tool execution
+  --functions TEXT              Python code block or file path defining
-  --cl, --chain-limit INTEGER  How many chained tool responses to allow, default
+                                functions to register as tools
-                               5, set 0 for unlimited
+  --td, --tools-debug           Show full details of tool executions
-  --help                       Show this message and exit.
+  --ta, --tools-approve         Manually approve every tool execution
+  --cl, --chain-limit INTEGER   How many chained tool responses to allow,
+                                default 5, set 0 for unlimited
+  --help                        Show this message and exit.
 ```
 
 (help-keys)=

docs/index.md

@@ -76,6 +76,7 @@ Chatting with gpt-4o
 Type 'exit' or 'quit' to exit
 Type '!multi' to enter multiple lines, then '!end' to finish
 Type '!edit' to open your default editor and modify the prompt.
+Type '!fragment <my_fragment> [<another_fragment> ...]' to insert one or more fragments
 > Tell me a joke about a pelican
 Why don't pelicans like to tip waiters?
 

docs/usage.md

@@ -405,6 +405,7 @@ Chatting with gpt-4
 Type 'exit' or 'quit' to exit
 Type '!multi' to enter multiple lines, then '!end' to finish
 Type '!edit' to open your default editor and modify the prompt
+Type '!fragment <my_fragment> [<another_fragment> ...]' to insert one or more fragments
 > who are you?
 I am a sentient cheesecake, meaning I am an artificial
 intelligence embodied in a dessert form, specifically a
@@ -426,6 +427,7 @@ Chatting with gpt-4
 Type 'exit' or 'quit' to exit
 Type '!multi' to enter multiple lines, then '!end' to finish
 Type '!edit' to open your default editor and modify the prompt.
+Type '!fragment <my_fragment> [<another_fragment> ...]' to insert one or more fragments
 > !multi custom-end
  Explain this error:
 
@@ -445,6 +447,7 @@ Chatting with gpt-4
 Type 'exit' or 'quit' to exit
 Type '!multi' to enter multiple lines, then '!end' to finish
 Type '!edit' to open your default editor and modify the prompt.
+Type '!fragment <my_fragment> [<another_fragment> ...]' to insert one or more fragments
 > !edit
 ```
 

llm/cli.py

@@ -163,6 +163,39 @@ def resolve_fragments(
     return resolved
 
 
+def process_fragments_in_chat(
+    db: sqlite_utils.Database, prompt: str
+) -> tuple[str, list[Fragment], list[Attachment]]:
+    """
+    Process any !fragment commands in a chat prompt and return the modified prompt plus resolved fragments and attachments.
+    """
+    prompt_lines = []
+    fragments = []
+    attachments = []
+    for line in prompt.splitlines():
+        if line.startswith("!fragment "):
+            try:
+                fragment_strs = line.strip().removeprefix("!fragment ").split()
+                fragments_and_attachments = resolve_fragments(
+                    db, fragments=fragment_strs, allow_attachments=True
+                )
+                fragments += [
+                    fragment
+                    for fragment in fragments_and_attachments
+                    if isinstance(fragment, Fragment)
+                ]
+                attachments += [
+                    attachment
+                    for attachment in fragments_and_attachments
+                    if isinstance(attachment, Attachment)
+                ]
+            except FragmentNotFound as ex:
+                raise click.ClickException(str(ex))
+        else:
+            prompt_lines.append(line)
+    return "\n".join(prompt_lines), fragments, attachments
+
+
 class AttachmentError(Exception):
     """Exception raised for errors in attachment resolution."""
 
@@ -888,6 +921,20 @@ def prompt(
     "--conversation",
     help="Continue the conversation with the given ID.",
 )
+@click.option(
+    "fragments",
+    "-f",
+    "--fragment",
+    multiple=True,
+    help="Fragment (alias, URL, hash or file path) to add to the prompt",
+)
+@click.option(
+    "system_fragments",
+    "--sf",
+    "--system-fragment",
+    multiple=True,
+    help="Fragment to add to system prompt",
+)
 @click.option("-t", "--template", help="Template to use")
 @click.option(
     "-p",
@@ -953,6 +1000,8 @@ def chat(
     model_id,
     _continue,
     conversation_id,
+    fragments,
+    system_fragments,
     template,
     param,
     options,
@@ -1054,15 +1103,48 @@ def chat(
     if key and isinstance(model, KeyModel):
         kwargs["key"] = key
 
+    try:
+        fragments_and_attachments = resolve_fragments(
+            db, fragments, allow_attachments=True
+        )
+        argument_fragments = [
+            fragment
+            for fragment in fragments_and_attachments
+            if isinstance(fragment, Fragment)
+        ]
+        argument_attachments = [
+            attachment
+            for attachment in fragments_and_attachments
+            if isinstance(attachment, Attachment)
+        ]
+        argument_system_fragments = resolve_fragments(db, system_fragments)
+    except FragmentNotFound as ex:
+        raise click.ClickException(str(ex))
+
     click.echo("Chatting with {}".format(model.model_id))
     click.echo("Type 'exit' or 'quit' to exit")
     click.echo("Type '!multi' to enter multiple lines, then '!end' to finish")
     click.echo("Type '!edit' to open your default editor and modify the prompt")
+    click.echo(
+        "Type '!fragment <my_fragment> [<another_fragment> ...]' to insert one or more fragments"
+    )
     in_multi = False
+
     accumulated = []
+    accumulated_fragments = []
+    accumulated_attachments = []
     end_token = "!end"
     while True:
         prompt = click.prompt("", prompt_suffix="> " if not in_multi else "")
+        fragments = []
+        attachments = []
+        if argument_fragments:
+            fragments += argument_fragments
+            # fragments from --fragments will get added to the first message only
+            argument_fragments = []
+        if argument_attachments:
+            attachments = argument_attachments
+            argument_attachments = []
         if prompt.strip().startswith("!multi"):
             in_multi = True
             bits = prompt.strip().split()
@@ -1074,17 +1156,28 @@ def chat(
             if edited_prompt is None:
                 click.echo("Editor closed without saving.", err=True)
                 continue
-            prompt = edited_prompt.strip()
+            prompt, fragments, attachments = process_fragments_in_chat(
-            if not prompt:
+                db, edited_prompt.strip()
+            )
+            if not prompt and not fragments and not attachments:
                 continue
-            click.echo(prompt)
+        if prompt.strip().startswith("!fragment "):
+            prompt, fragments, attachments = process_fragments_in_chat(db, prompt)
+
         if in_multi:
             if prompt.strip() == end_token:
                 prompt = "\n".join(accumulated)
+                fragments = accumulated_fragments
+                attachments = accumulated_attachments
                 in_multi = False
                 accumulated = []
+                accumulated_fragments = []
+                accumulated_attachments = []
             else:
-                accumulated.append(prompt)
+                if prompt:
+                    accumulated.append(prompt)
+                accumulated_fragments += fragments
+                accumulated_attachments += attachments
                 continue
         if template_obj:
             try:
@@ -1100,9 +1193,21 @@ def chat(
                 prompt = new_prompt
         if prompt.strip() in ("exit", "quit"):
             break
-        response = conversation.chain(prompt, system=system, **kwargs)
+
+        response = conversation.chain(
+            prompt,
+            fragments=[str(fragment) for fragment in fragments],
+            system_fragments=[
+                str(system_fragment) for system_fragment in argument_system_fragments
+            ],
+            attachments=attachments,
+            system=system,
+            **kwargs,
+        )
+
         # System prompt only sent for the first message:
         system = None
+        system_fragments = []
         for chunk in response:
             print(chunk, end="")
             sys.stdout.flush()

tests/test_chat.py

@@ -25,6 +25,7 @@ def test_chat_basic(mock_model, logs_db):
         "\nType 'exit' or 'quit' to exit"
         "\nType '!multi' to enter multiple lines, then '!end' to finish"
         "\nType '!edit' to open your default editor and modify the prompt"
+        "\nType '!fragment <my_fragment> [<another_fragment> ...]' to insert one or more fragments"
         "\n> Hi"
         "\none world"
         "\n> Hi two"
@@ -91,6 +92,7 @@ def test_chat_basic(mock_model, logs_db):
         "\nType 'exit' or 'quit' to exit"
         "\nType '!multi' to enter multiple lines, then '!end' to finish"
         "\nType '!edit' to open your default editor and modify the prompt"
+        "\nType '!fragment <my_fragment> [<another_fragment> ...]' to insert one or more fragments"
         "\n> Continue"
         "\ncontinued"
         "\n> quit"
@@ -140,6 +142,7 @@ def test_chat_system(mock_model, logs_db):
         "\nType 'exit' or 'quit' to exit"
         "\nType '!multi' to enter multiple lines, then '!end' to finish"
         "\nType '!edit' to open your default editor and modify the prompt"
+        "\nType '!fragment <my_fragment> [<another_fragment> ...]' to insert one or more fragments"
         "\n> Hi"
         "\nI am mean"
         "\n> quit"
@@ -303,6 +306,7 @@ def test_chat_tools(logs_db):
         "Type 'exit' or 'quit' to exit\n"
         "Type '!multi' to enter multiple lines, then '!end' to finish\n"
         "Type '!edit' to open your default editor and modify the prompt\n"
+        "Type '!fragment <my_fragment> [<another_fragment> ...]' to insert one or more fragments\n"
         '> {"prompt": "Convert hello to uppercase", "tool_calls": [{"name": "upper", '
         '"arguments": {"text": "hello"}}]}\n'
         "{\n"