Image attachment support

Refs #19 Refs simonw/llm#587
simonw · Oct 29, 2024 · 97f2aee · 97f2aee
1 parent 1aa490c
commit 97f2aee
Show file tree

Hide file tree

Showing 4 changed files with 441 additions and 6 deletions.
diff --git a/llm_claude_3.py b/llm_claude_3.py
@@ -15,7 +15,8 @@ def register_models(register):
     register(ClaudeMessagesLong("claude-3-5-sonnet-20240620"))
     register(ClaudeMessagesLong("claude-3-5-sonnet-20241022"))
     register(
-        ClaudeMessagesLong("claude-3-5-sonnet-latest"), aliases=("claude-3.5-sonnet", "claude-3.5-sonnet-latest")
+        ClaudeMessagesLong("claude-3-5-sonnet-latest"),
+        aliases=("claude-3.5-sonnet", "claude-3.5-sonnet-latest"),
     )
     # register(
     #     ClaudeMessagesLong("claude-3-5-haiku-latest"), aliases=("claude-3.5-haiku",)
@@ -88,6 +89,12 @@ class ClaudeMessages(llm.Model):
     needs_key = "claude"
     key_env_var = "ANTHROPIC_API_KEY"
     can_stream = True
+    attachment_types = {
+        "image/png",
+        "image/jpeg",
+        "image/webp",
+        "image/gif",
+    }
 
     class Options(ClaudeOptions): ...
 
@@ -100,16 +107,48 @@ def build_messages(self, prompt, conversation) -> List[dict]:
         messages = []
         if conversation:
             for response in conversation.responses:
+                if response.attachments:
+                    content = [
+                        {
+                            "type": "image",
+                            "source": {
+                                "data": attachment.base64_content(),
+                                "media_type": attachment.resolve_type(),
+                                "type": "base64",
+                            },
+                        }
+                        for attachment in response.attachments
+                    ]
+                else:
+                    content = response.prompt.prompt
                 messages.extend(
                     [
                         {
                             "role": "user",
-                            "content": response.prompt.prompt,
+                            "content": content,
                         },
                         {"role": "assistant", "content": response.text()},
                     ]
                 )
-        messages.append({"role": "user", "content": prompt.prompt})
+        if prompt.attachments:
+            messages.append(
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "image",
+                            "source": {
+                                "data": attachment.base64_content(),
+                                "media_type": attachment.resolve_type(),
+                                "type": "base64",
+                            },
+                        }
+                        for attachment in prompt.attachments
+                    ],
+                }
+            )
+        else:
+            messages.append({"role": "user", "content": prompt.prompt})
         return messages
 
     def execute(self, prompt, stream, response, conversation):

diff --git a/pyproject.toml b/pyproject.toml
@@ -9,8 +9,8 @@ classifiers = [
     "License :: OSI Approved :: Apache Software License"
 ]
 dependencies = [
-    "llm",
-    "anthropic>=0.17.0",
+    "llm>=0.17a0",
+    "anthropic>=0.37.1",
 ]
 
 [project.urls]