Implement schema support, refs #22

Refs simonw/llm#776

llm_anthropic.py

@@ -201,6 +201,7 @@ class _Shared:
     can_stream = True
 
     supports_thinking = False
+    supports_schema = True
     default_max_tokens = 4096
 
     class Options(ClaudeOptions): ...
@@ -348,6 +349,15 @@ def build_kwargs(self, prompt, conversation):
             if "thinking" in kwargs:
                 kwargs["extra_body"] = {"thinking": kwargs.pop("thinking")}
 
+        if prompt.schema:
+            kwargs["tools"] = [
+                {
+                    "name": "output_structured_data",
+                    "input_schema": prompt.schema,
+                }
+            ]
+            kwargs["tool_choice"] = {"type": "tool", "name": "output_structured_data"}
+
         return kwargs
 
     def set_usage(self, response):
@@ -374,8 +384,13 @@ def execute(self, prompt, stream, response, conversation, key):
             with messages_client.stream(**kwargs) as stream:
                 if prefill_text:
                     yield prefill_text
-                for text in stream.text_stream:
-                    yield text
+                for chunk in stream:
+                    if hasattr(chunk, "delta"):
+                        delta = chunk.delta
+                        if hasattr(delta, "text"):
+                            yield delta.text
+                        elif hasattr(delta, "partial_json"):
+                            yield delta.partial_json
                 # This records usage and other data:
                 response.response_json = stream.get_final_message().model_dump()
         else:

pyproject.toml

@@ -9,7 +9,7 @@ classifiers = [
     "License :: OSI Approved :: Apache Software License"
 ]
 dependencies = [
-    "llm>=0.22",
+    "llm>=0.23a0",
     "anthropic>=0.47.2",
 ]
 

tests/cassettes/test_anthropic/test_schema_prompt.yaml

@@ -0,0 +1,239 @@
+interactions:
+- request:
+    body: '{"max_tokens":8192,"messages":[{"role":"user","content":"Invent a good
+      dog"}],"model":"claude-3-7-sonnet-latest","temperature":1.0,"tools":[{"name":"output_structured_data","input_schema":{"properties":{"name":{"title":"Name","type":"string"},"age":{"title":"Age","type":"integer"},"bio":{"title":"Bio","type":"string"}},"required":["name","age","bio"],"title":"Dog","type":"object"}}],"tool_choice":{"type":"tool","name":"output_structured_data"},"stream":true}'
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      anthropic-version:
+      - '2023-06-01'
+      connection:
+      - keep-alive
+      content-length:
+      - '462'
+      content-type:
+      - application/json
+      host:
+      - api.anthropic.com
+      user-agent:
+      - Anthropic/Python 0.47.2
+      x-stainless-arch:
+      - arm64
+      x-stainless-async:
+      - 'false'
+      x-stainless-lang:
+      - python
+      x-stainless-os:
+      - MacOS
+      x-stainless-package-version:
+      - 0.47.2
+      x-stainless-read-timeout:
+      - '600'
+      x-stainless-retry-count:
+      - '0'
+      x-stainless-runtime:
+      - CPython
+      x-stainless-runtime-version:
+      - 3.12.0
+      x-stainless-stream-helper:
+      - messages
+      x-stainless-timeout:
+      - NOT_GIVEN
+    method: POST
+    uri: https://api.anthropic.com/v1/messages
+  response:
+    body:
+      string: 'event: message_start
+
+        data: {"type":"message_start","message":{"id":"msg_011RyssDpYDzSyU3zK8Bhvbf","type":"message","role":"assistant","model":"claude-3-7-sonnet-20250219","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":429,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":25}}}
+
+
+        event: content_block_start
+
+        data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu_01BEUcgSZ295kTDkxpovXdpA","name":"output_structured_data","input":{}}            }
+
+
+        event: ping
+
+        data: {"type": "ping"}
+
+
+        event: content_block_delta
+
+        data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":""}   }
+
+
+        event: content_block_delta
+
+        data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"n"}      }
+
+
+        event: content_block_delta
+
+        data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"am"}    }
+
+
+        event: content_block_delta
+
+        data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"e\":
+        \""}             }
+
+
+        event: content_block_delta
+
+        data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"Buddy\""}
+        }
+
+
+        event: content_block_delta
+
+        data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":",
+        \"age\":"}  }
+
+
+        event: content_block_delta
+
+        data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"
+        3"}         }
+
+
+        event: content_block_delta
+
+        data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":",
+        \"bio\":"}    }
+
+
+        event: content_block_delta
+
+        data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"
+        \"Buddy is "}            }
+
+
+        event: content_block_delta
+
+        data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"a
+        loy"}}
+
+
+        event: content_block_delta
+
+        data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"al
+        and "}             }
+
+
+        event: content_block_delta
+
+        data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"ene"}             }
+
+
+        event: content_block_delta
+
+        data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"rgetic
+        Gol"}           }
+
+
+        event: content_block_delta
+
+        data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"den
+        Ret"}       }
+
+
+        event: content_block_delta
+
+        data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"ri"}        }
+
+
+        event: content_block_delta
+
+        data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"ever
+        w"}              }
+
+
+        event: content_block_delta
+
+        data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"ho
+        love"}               }
+
+
+        event: content_block_delta
+
+        data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"s
+        long wal"}     }
+
+
+        event: content_block_delta
+
+        data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"ks.\"}"}          }
+
+
+        event: content_block_stop
+
+        data: {"type":"content_block_stop","index":0      }
+
+
+        event: message_delta
+
+        data: {"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":140}         }
+
+
+        event: message_stop
+
+        data: {"type":"message_stop"      }
+
+
+        '
+    headers:
+      CF-RAY:
+      - 9184543f7bb208cb-LAX
+      Cache-Control:
+      - no-cache
+      Connection:
+      - keep-alive
+      Content-Type:
+      - text/event-stream; charset=utf-8
+      Date:
+      - Thu, 27 Feb 2025 01:16:23 GMT
+      Server:
+      - cloudflare
+      Transfer-Encoding:
+      - chunked
+      X-Robots-Tag:
+      - none
+      anthropic-organization-id:
+      - b320df16-c763-4455-b901-3f7d54e9a012
+      anthropic-ratelimit-input-tokens-limit:
+      - '200000'
+      anthropic-ratelimit-input-tokens-remaining:
+      - '200000'
+      anthropic-ratelimit-input-tokens-reset:
+      - '2025-02-27T01:16:22Z'
+      anthropic-ratelimit-output-tokens-limit:
+      - '80000'
+      anthropic-ratelimit-output-tokens-remaining:
+      - '76000'
+      anthropic-ratelimit-output-tokens-reset:
+      - '2025-02-27T01:16:25Z'
+      anthropic-ratelimit-requests-limit:
+      - '4000'
+      anthropic-ratelimit-requests-remaining:
+      - '3999'
+      anthropic-ratelimit-requests-reset:
+      - '2025-02-27T01:16:22Z'
+      anthropic-ratelimit-tokens-limit:
+      - '280000'
+      anthropic-ratelimit-tokens-remaining:
+      - '276000'
+      anthropic-ratelimit-tokens-reset:
+      - '2025-02-27T01:16:22Z'
+      cf-cache-status:
+      - DYNAMIC
+      request-id:
+      - req_01R68QTrbL1hYTBF3ozzrsxj
+      via:
+      - 1.1 google
+    status:
+      code: 200
+      message: OK
+version: 1

tests/test_anthropic.py

@@ -1,6 +1,8 @@
+import json
 import llm
 import os
 import pytest
+from pydantic import BaseModel
 
 TINY_PNG = (
     b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\xa6\x00\x00\x01\x1a"
@@ -96,6 +98,25 @@ def test_image_prompt():
     assert response.token_details is None
 
 
+@pytest.mark.vcr
+def test_schema_prompt():
+    model = llm.get_model("claude-3.7-sonnet")
+    model.key = model.key or ANTHROPIC_API_KEY
+
+    class Dog(BaseModel):
+        name: str
+        age: int
+        bio: str
+
+    response = model.prompt("Invent a good dog", schema=Dog)
+    dog = json.loads(response.text())
+    assert dog == {
+        "name": "Buddy",
+        "age": 3,
+        "bio": "Buddy is a loyal and energetic Golden Retriever who loves long walks.",
+    }
+
+
 @pytest.mark.vcr
 def test_prompt_with_prefill_and_stop_sequences():
     model = llm.get_model("claude-3.5-haiku")