Fix for async usage of schemas, refs #22

Author: simonw

llm_anthropic.py

@@ -416,8 +416,13 @@ async def execute(self, prompt, stream, response, conversation, key):
             async with messages_client.stream(**kwargs) as stream_obj:
                 if prefill_text:
                     yield prefill_text
-                async for text in stream_obj.text_stream:
-                    yield text
+                async for chunk in stream_obj:
+                    if hasattr(chunk, "delta"):
+                        delta = chunk.delta
+                        if hasattr(delta, "text"):
+                            yield delta.text
+                        elif hasattr(delta, "partial_json"):
+                            yield delta.partial_json
             response.response_json = (await stream_obj.get_final_message()).model_dump()
         else:
             completion = await messages_client.create(**kwargs)

tests/test_anthropic.py

@@ -98,16 +98,17 @@ def test_image_prompt():
     assert response.token_details is None
 
 
+class Dog(BaseModel):
+    name: str
+    age: int
+    bio: str
+
+
 @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 == {
@@ -117,6 +118,30 @@ class Dog(BaseModel):
     }
 
 
+@pytest.mark.vcr
+@pytest.mark.asyncio
+async def test_schema_prompt_async():
+    model = llm.get_async_model("claude-3.7-sonnet")
+    model.key = model.key or ANTHROPIC_API_KEY  # don't override existing key
+    response = await model.prompt("Invent a terrific dog", schema=Dog)
+    dog_json = await response.text()
+    dog = json.loads(dog_json)
+    assert dog == {
+        "name": "Maximus Thunder",
+        "age": 3,
+        "bio": (
+            "Maximus Thunder is an extraordinary golden retriever with a natural "
+            "talent for search and rescue operations. His keen sense of smell "
+            "can detect people trapped under debris from over a mile away. When "
+            "he's not saving lives, Maximus enjoys surfing at the beach and has "
+            "won three local dog surfing competitions. He's also incredibly "
+            "gentle with children and regularly visits hospitals as a therapy "
+            "dog. His favorite treat is peanut butter, and he has a unique howl "
+            'that sounds remarkably like he\'s saying "hello."'
+        ),
+    }
+
+
 @pytest.mark.vcr
 def test_prompt_with_prefill_and_stop_sequences():
     model = llm.get_model("claude-3.5-haiku")