Automated tests for attachments, refs #587

tests/conftest.py

@@ -49,6 +49,7 @@ def env_setup(monkeypatch, user_path):
 
 class MockModel(llm.Model):
     model_id = "mock"
+    attachment_types = {"image/png"}
 
     class Options(llm.Options):
         max_tokens: Optional[int] = Field(

tests/test_attachments.py

@@ -0,0 +1,48 @@
+from click.testing import CliRunner
+from unittest.mock import ANY
+import llm
+
+TINY_PNG = (
+    b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\xa6\x00\x00\x01\x1a"
+    b"\x02\x03\x00\x00\x00\xe6\x99\xc4^\x00\x00\x00\tPLTE\xff\xff\xff"
+    b"\x00\xff\x00\xfe\x01\x00\x12t\x01J\x00\x00\x00GIDATx\xda\xed\xd81\x11"
+    b"\x000\x08\xc0\xc0.]\xea\xaf&Q\x89\x04V\xe0>\xf3+\xc8\x91Z\xf4\xa2\x08EQ\x14E"
+    b"Q\x14EQ\x14EQ\xd4B\x91$I3\xbb\xbf\x08EQ\x14EQ\x14EQ\x14E\xd1\xa5"
+    b"\xd4\x17\x91\xc6\x95\x05\x15\x0f\x9f\xc5\t\x9f\xa4\x00\x00\x00\x00IEND\xaeB`"
+    b"\x82"
+)
+
+
+def test_prompt_image(mock_model, logs_db):
+    runner = CliRunner()
+    mock_model.enqueue(["two boxes"])
+    result = runner.invoke(
+        llm.cli.cli,
+        ["prompt", "-m", "mock", "describe image", "-a", "-"],
+        input=TINY_PNG,
+        catch_exceptions=False,
+    )
+    assert result.exit_code == 0
+    assert result.output == "two boxes\n"
+    assert mock_model.history[0][0].attachments[0] == llm.Attachment(
+        type="image/png", path=None, url=None, content=TINY_PNG, _id=ANY
+    )
+
+    # Check it was logged correctly
+    conversations = list(logs_db["conversations"].rows)
+    assert len(conversations) == 1
+    conversation = conversations[0]
+    assert conversation["model"] == "mock"
+    assert conversation["name"] == "describe image"
+    response = list(logs_db["responses"].rows)[0]
+    attachment = list(logs_db["attachments"].rows)[0]
+    assert attachment == {
+        "id": ANY,
+        "type": "image/png",
+        "path": None,
+        "url": None,
+        "content": TINY_PNG,
+    }
+    prompt_attachment = list(logs_db["prompt_attachments"].rows)[0]
+    assert prompt_attachment["attachment_id"] == attachment["id"]
+    assert prompt_attachment["response_id"] == response["id"]