Fragment plugins can now optionally return attachments (#974)

Closes #972

docs/fragments.md

@@ -143,7 +143,9 @@ This plugin turns a single call to `-f github:simonw/s3-credentials` into multip
 
 Running `llm logs -c` will show that this prompt incorporated 26 fragments, one for each file.
 
-Running `llm logs -c --usage --expand` includes token usage information and turns each fragment ID into a full copy of that file. [Here's the output of that command](https://gist.github.com/simonw/c9bbbc5f6560b01f4b7882ac0194fb25).
+Running `llm logs -c --usage --expand` (shortcut: `llm logs -cue`) includes token usage information and turns each fragment ID into a full copy of that file. [Here's the output of that command](https://gist.github.com/simonw/c9bbbc5f6560b01f4b7882ac0194fb25).
+
+Fragment plugins can return {ref}`attachments <usage-attachments>` (such as images) as well.
 
 See the {ref}`register_fragment_loaders() plugin hook <plugin-hooks-register-fragment-loaders>` documentation for details on writing your own custom fragment plugin.
 

docs/plugins/plugin-hooks.md

@@ -113,12 +113,20 @@ The loader function should raise a `ValueError` if the template cannot be found
 
 Plugins can register new fragment loaders using the `register_template_loaders` hook. These can then be used with the `llm -f prefix:argument` syntax.
 
+Fragment loader plugins differ from template loader plugins in that you can stack more than one fragment loader call together in the same prompt.
+
+A fragment loader can return one or more string fragments or attachments, or a mixture of the two. The fragments will be concatenated together into the prompt string, while any attachments will be added to the list of attachments to be sent to the model.
+
 The `prefix` specifies the loader. The `argument` will be passed to that registered callback..
 
-The callback works in a very similar way to template loaders, but returns either a single `llm.Fragment` or a list of `llm.Fragment` objects.
+The callback works in a very similar way to template loaders, but returns either a single `llm.Fragment`, a list of `llm.Fragment` objects, a single `llm.Attachment`, or a list that can mix `llm.Attachment` and `llm.Fragment` objects.
 
 The `llm.Fragment` constructor takes a required string argument (the content of the fragment) and an optional second `source` argument, which is a string that may be displayed as debug information. For files this is a path and for URLs it is a URL. Your plugin can use anything you like for the `source` value.
 
+See {ref}`the Python API documentation for attachments <python-api-attachments>` for details of the `llm.Attachment` class.
+
+Here is some example code:
+
 ```python
 import llm
 
@@ -138,11 +146,12 @@ def my_fragment_loader(argument: str) -> llm.Fragment:
             f"Fragment 'my-fragments:{argument}' could not be loaded: {str(ex)}"
         )
 
-# Or for the case where you want to return multiple fragments:
+# Or for the case where you want to return multiple fragments and attachments:
 def my_fragment_loader(argument: str) -> list[llm.Fragment]:
     return [
         llm.Fragment("Fragment 1 content", "my-fragments:{argument}"),
         llm.Fragment("Fragment 2 content", "my-fragments:{argument}"),
+        llm.Attachment(path="/path/to/image.png"),
     ]
 ```
 A plugin like this one can be called like so:
@@ -151,4 +160,4 @@ llm -f my-fragments:argument
 ```
 If multiple fragments are returned they will be used as if the user passed multiple `-f X` arguments to the command.
 
-Multiple fragments are useful for things like plugins that return every file in a directory. By giving each file its own fragment we can avoid having multiple copies of the full collection stored if only a single file has changed.
+Multiple fragments are particularly useful for things like plugins that return every file in a directory. If these were concatenated together by the plugin, a change to a single file would invalidate the de-duplicatino cache for that whole fragment. Giving each file its own fragment means we can avoid storing multiple copies of that full collection if only a single file has changed.

llm/__init__.py

@@ -120,9 +120,10 @@ def get_template_loaders() -> Dict[str, Callable[[str], Template]]:
     return _get_loaders(pm.hook.register_template_loaders)
 
 
-def get_fragment_loaders() -> (
-    Dict[str, Callable[[str], Union[Fragment, List[Fragment]]]]
-):
+def get_fragment_loaders() -> Dict[
+    str,
+    Callable[[str], Union[Fragment, Attachment, List[Union[Fragment, Attachment]]]],
+]:
     """Get fragment loaders registered by plugins."""
     return _get_loaders(pm.hook.register_fragment_loaders)
 

llm/cli.py

@@ -89,13 +89,13 @@ def validate_fragment_alias(ctx, param, value):
 
 
 def resolve_fragments(
-    db: sqlite_utils.Database, fragments: Iterable[str]
-) -> List[Fragment]:
+    db: sqlite_utils.Database, fragments: Iterable[str], allow_attachments: bool = False
+) -> List[Union[Fragment, Attachment]]:
     """
-    Resolve fragments into a list of (content, source) tuples
+    Resolve fragment strings into a mixed of llm.Fragment() and llm.Attachment() objects.
     """
 
-    def _load_by_alias(fragment):
+    def _load_by_alias(fragment: str) -> Tuple[Optional[str], Optional[str]]:
         rows = list(
             db.query(
                 """
@@ -111,8 +111,8 @@ def resolve_fragments(
             return row["content"], row["source"]
         return None, None
 
-    # These can be URLs or paths or plugin references
-    resolved = []
+    # The fragment strings could be URLs or paths or plugin references
+    resolved: List[Union[Fragment, Attachment]] = []
     for fragment in fragments:
         if fragment.startswith("http://") or fragment.startswith("https://"):
             client = httpx.Client(follow_redirects=True, max_redirects=3)
@@ -131,6 +131,14 @@ def resolve_fragments(
                 result = loader(rest)
                 if not isinstance(result, list):
                     result = [result]
+                if not allow_attachments and any(
+                    isinstance(r, Attachment) for r in result
+                ):
+                    raise FragmentNotFound(
+                        "Fragment loader {} returned a disallowed attachment".format(
+                            prefix
+                        )
+                    )
                 resolved.extend(result)
             except Exception as ex:
                 raise FragmentNotFound(
@@ -687,8 +695,20 @@ def prompt(
     response = None
 
     try:
-        fragments = resolve_fragments(db, fragments)
-        system_fragments = resolve_fragments(db, system_fragments)
+        fragments_and_attachments = resolve_fragments(
+            db, fragments, allow_attachments=True
+        )
+        resolved_fragments = [
+            fragment
+            for fragment in fragments_and_attachments
+            if isinstance(fragment, Fragment)
+        ]
+        resolved_attachments.extend(
+            attachment
+            for attachment in fragments_and_attachments
+            if isinstance(attachment, Attachment)
+        )
+        resolved_system_fragments = resolve_fragments(db, system_fragments)
     except FragmentNotFound as ex:
         raise click.ClickException(str(ex))
 
@@ -706,8 +726,8 @@ def prompt(
                         attachments=resolved_attachments,
                         system=system,
                         schema=schema,
-                        fragments=fragments,
-                        system_fragments=system_fragments,
+                        fragments=resolved_fragments,
+                        system_fragments=resolved_system_fragments,
                         **kwargs,
                     )
                     async for chunk in response:
@@ -717,11 +737,11 @@ def prompt(
                 else:
                     response = prompt_method(
                         prompt,
-                        fragments=fragments,
+                        fragments=resolved_fragments,
                         attachments=resolved_attachments,
                         schema=schema,
                         system=system,
-                        system_fragments=system_fragments,
+                        system_fragments=resolved_system_fragments,
                         **kwargs,
                     )
                     text = await response.text()
@@ -736,11 +756,11 @@ def prompt(
         else:
             response = prompt_method(
                 prompt,
-                fragments=fragments,
+                fragments=resolved_fragments,
                 attachments=resolved_attachments,
                 system=system,
                 schema=schema,
-                system_fragments=system_fragments,
+                system_fragments=resolved_system_fragments,
                 **kwargs,
             )
             if should_stream:

tests/conftest.py

@@ -85,6 +85,7 @@ class MockModel(llm.Model):
 class EchoModel(llm.Model):
     model_id = "echo"
     can_stream = True
+    attachment_types = {"image/png"}
 
     class Options(llm.Options):
         example_int: Optional[int] = Field(
@@ -100,6 +101,10 @@ class EchoModel(llm.Model):
         }
         if non_null_options:
             yield "\n\noptions: {}".format(json.dumps(non_null_options))
+        if prompt.attachments:
+            yield "\n\nattachments:\n"
+            for attachment in prompt.attachments:
+                yield f"  - {attachment.url}\n"
 
 
 class MockKeyModel(llm.KeyModel):

tests/test_plugins.py

@@ -3,6 +3,7 @@ import click
 import importlib
 import llm
 from llm import cli, hookimpl, plugins, get_template_loaders, get_fragment_loaders
+import textwrap
 
 
 def test_register_commands():
@@ -90,7 +91,15 @@ def test_register_template_loaders():
         assert get_template_loaders() == {}
 
 
-def test_register_fragment_loaders(logs_db):
+def test_register_fragment_loaders(logs_db, httpx_mock):
+    httpx_mock.add_response(
+        method="HEAD",
+        url="https://example.com/attachment.png",
+        content=b"attachment",
+        headers={"Content-Type": "image/png"},
+        is_reusable=True,
+    )
+
     assert get_fragment_loaders() == {}
 
     def single_fragment(argument):
@@ -104,6 +113,12 @@ def test_register_fragment_loaders(logs_db):
             llm.Fragment(f"three:{argument}", "three"),
         ]
 
+    def fragment_and_attachment(argument):
+        return [
+            llm.Fragment(f"one:{argument}", "one"),
+            llm.Attachment(url="https://example.com/attachment.png"),
+        ]
+
     class FragmentLoadersPlugin:
         __name__ = "FragmentLoadersPlugin"
 
@@ -111,6 +126,7 @@ def test_register_fragment_loaders(logs_db):
         def register_fragment_loaders(self, register):
             register("single", single_fragment)
             register("three", three_fragments)
+            register("mixed", fragment_and_attachment)
 
     try:
         plugins.pm.register(FragmentLoadersPlugin(), name="FragmentLoadersPlugin")
@@ -118,6 +134,7 @@ def test_register_fragment_loaders(logs_db):
         assert loaders == {
             "single": single_fragment,
             "three": three_fragments,
+            "mixed": fragment_and_attachment,
         }
 
         # Test the CLI command
@@ -137,9 +154,30 @@ def test_register_fragment_loaders(logs_db):
             "\n"
             "three:\n"
             "  Undocumented\n"
+            "\n"
+            "mixed:\n"
+            "  Undocumented\n"
         )
         assert result2.output == expected2
 
+        # Test the one that includes an attachment
+        result3 = runner.invoke(
+            cli.cli, ["-m", "echo", "-f", "mixed:x"], catch_exceptions=False
+        )
+        assert result3.exit_code == 0
+        result3.output.strip == textwrap.dedent(
+            """\
+            system:
+
+
+            prompt:
+            one:x
+
+            attachments:
+            - https://example.com/attachment.png
+            """
+        ).strip()
+
     finally:
         plugins.pm.unregister(name="FragmentLoadersPlugin")
         assert get_fragment_loaders() == {}