PDF attachment support for Claude 3.5 Sonnet, refs #22

simonw · Nov 1, 2024 · ec3aa12 · ec3aa12
1 parent e065c6d
commit ec3aa12
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 16 deletions.
diff --git a/llm_claude_3.py b/llm_claude_3.py
@@ -13,9 +13,9 @@ def register_models(register):
     register(ClaudeMessages("claude-3-haiku-20240307"), aliases=("claude-3-haiku",))
     # 3.5 models
     register(ClaudeMessagesLong("claude-3-5-sonnet-20240620"))
-    register(ClaudeMessagesLong("claude-3-5-sonnet-20241022"))
+    register(ClaudeMessagesLong("claude-3-5-sonnet-20241022", supports_pdf=True)),
     register(
-        ClaudeMessagesLong("claude-3-5-sonnet-latest"),
+        ClaudeMessagesLong("claude-3-5-sonnet-latest", supports_pdf=True),
         aliases=("claude-3.5-sonnet", "claude-3.5-sonnet-latest"),
     )
     # register(
@@ -89,19 +89,25 @@ 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): ...
 
-    def __init__(self, model_id, claude_model_id=None, extra_headers=None):
+    def __init__(
+        self, model_id, claude_model_id=None, extra_headers=None, supports_pdf=False
+    ):
         self.model_id = model_id
         self.claude_model_id = claude_model_id or model_id
-        self.extra_headers = extra_headers
+        self.extra_headers = extra_headers or {}
+        if supports_pdf:
+            self.extra_headers["anthropic-beta"] = "pdfs-2024-09-25"
+        self.attachment_types = {
+            "image/png",
+            "image/jpeg",
+            "image/webp",
+            "image/gif",
+        }
+        if supports_pdf:
+            self.attachment_types.add("application/pdf")
 
     def build_messages(self, prompt, conversation) -> List[dict]:
         messages = []
@@ -110,7 +116,11 @@ def build_messages(self, prompt, conversation) -> List[dict]:
                 if response.attachments:
                     content = [
                         {
-                            "type": "image",
+                            "type": (
+                                "document"
+                                if attachment.resolve_type() == "application/pdf"
+                                else "image"
+                            ),
                             "source": {
                                 "data": attachment.base64_content(),
                                 "media_type": attachment.resolve_type(),
@@ -136,7 +146,11 @@ def build_messages(self, prompt, conversation) -> List[dict]:
                     "role": "user",
                     "content": [
                         {
-                            "type": "image",
+                            "type": (
+                                "document"
+                                if attachment.resolve_type() == "application/pdf"
+                                else "image"
+                            ),
                             "source": {
                                 "data": attachment.base64_content(),
                                 "media_type": attachment.resolve_type(),

diff --git a/tests/test_claude_3.py b/tests/test_claude_3.py
@@ -43,10 +43,7 @@ def test_prompt():
 @pytest.mark.vcr
 def test_image_prompt():
     model = llm.get_model("claude-3.5-sonnet")
-    model.key = (
-        model.key
-        or "sk-..."
-    )
+    model.key = model.key or "sk-..."
     response = model.prompt(
         "Describe image in three words",
         attachments=[llm.Attachment(content=TINY_PNG)],