Anthropic streaming support

instructor/client.py

@@ -128,8 +128,6 @@ def create_partial(
         strict: bool = True,
         **kwargs: Any,
     ) -> Generator[T, None, None] | AsyncGenerator[T, None]:
-        assert self.provider != Provider.ANTHROPIC, "Anthropic doesn't support partial"
-
         kwargs["stream"] = True
 
         kwargs = self.handle_kwargs(kwargs)
@@ -311,8 +309,6 @@ async def create_iterable(
         strict: bool = True,
         **kwargs: Any,
     ) -> AsyncGenerator[T, None]:
-        assert self.provider != Provider.ANTHROPIC, "Anthropic doesn't support iterable"
-
         kwargs = self.handle_kwargs(kwargs)
         kwargs["stream"] = True
         async for item in await self.create_fn(

instructor/dsl/iterable.py

@@ -79,7 +79,12 @@ def extract_json(
     ) -> Generator[str, None, None]:
         for chunk in completion:
             try:
-                if chunk.choices:
+                if mode == Mode.ANTHROPIC_JSON:
+                    if json_chunk := chunk.delta.text:
+                        yield json_chunk
+                if mode == Mode.ANTHROPIC_TOOLS:
+                    yield chunk.model_extra.get("delta", "").get("partial_json", "")
+                elif chunk.choices:
                     if mode == Mode.FUNCTIONS:
                         if json_chunk := chunk.choices[0].delta.function_call.arguments:
                             yield json_chunk
@@ -102,7 +107,12 @@ async def extract_json_async(
     ) -> AsyncGenerator[str, None]:
         async for chunk in completion:
             try:
-                if chunk.choices:
+                if mode == Mode.ANTHROPIC_JSON:
+                    if json_chunk := chunk.delta.text:
+                        yield json_chunk
+                if mode == Mode.ANTHROPIC_TOOLS:
+                    yield chunk.model_extra.get("delta", "").get("partial_json", "")
+                elif chunk.choices:
                     if mode == Mode.FUNCTIONS:
                         if json_chunk := chunk.choices[0].delta.function_call.arguments:
                             yield json_chunk

instructor/dsl/partial.py

@@ -150,7 +150,12 @@ def extract_json(
     ) -> Generator[str, None, None]:
         for chunk in completion:
             try:
-                if chunk.choices:
+                if mode == Mode.ANTHROPIC_JSON:
+                    if json_chunk := chunk.delta.text:
+                        yield json_chunk
+                if mode == Mode.ANTHROPIC_TOOLS:
+                    yield chunk.model_extra.get("delta", "").get("partial_json", "")
+                elif chunk.choices:
                     if mode == Mode.FUNCTIONS:
                         if json_chunk := chunk.choices[0].delta.function_call.arguments:
                             yield json_chunk
@@ -173,7 +178,12 @@ async def extract_json_async(
     ) -> AsyncGenerator[str, None]:
         async for chunk in completion:
             try:
-                if chunk.choices:
+                if mode == Mode.ANTHROPIC_JSON:
+                    if json_chunk := chunk.delta.text:
+                        yield json_chunk
+                if mode == Mode.ANTHROPIC_TOOLS:
+                    yield chunk.model_extra.get("delta", "").get("partial_json", "")
+                elif chunk.choices:
                     if mode == Mode.FUNCTIONS:
                         if json_chunk := chunk.choices[0].delta.function_call.arguments:
                             yield json_chunk

tests/llm/test_anthropic/conftest.py

@@ -0,0 +1,39 @@
+# conftest.py
+from anthropic import AsyncAnthropic, Anthropic
+import pytest
+import os
+
+try:
+    import braintrust
+
+    wrap_anthropic = braintrust.wrap_anthropic
+except ImportError:
+
+    def wrap_anthropic(x):
+        return x
+
+
+@pytest.fixture(scope="session")
+def client():
+    if os.environ.get("BRAINTRUST_API_KEY"):
+        yield wrap_anthropic(
+            Anthropic(
+                api_key=os.environ["BRAINTRUST_API_KEY"],
+                base_url="https://braintrustproxy.com/v1",
+            )
+        )
+    else:
+        yield Anthropic()
+
+
+@pytest.fixture(scope="session")
+def aclient():
+    if os.environ.get("BRAINTRUST_API_KEY"):
+        yield wrap_anthropic(
+            AsyncAnthropic(
+                api_key=os.environ["BRAINTRUST_API_KEY"],
+                base_url="https://braintrustproxy.com/v1",
+            )
+        )
+    else:
+        yield AsyncAnthropic()

tests/llm/test_anthropic/test_stream.py

@@ -0,0 +1,87 @@
+from itertools import product
+from collections.abc import Iterable
+from pydantic import BaseModel
+import pytest
+import instructor
+from instructor.dsl.partial import Partial
+
+from .util import models, modes
+
+
+class UserExtract(BaseModel):
+    name: str
+    age: int
+
+
+@pytest.mark.parametrize("model, mode, stream", product(models, modes, [True, False]))
+def test_iterable_model(model, mode, stream, client):
+    client = instructor.from_anthropic(client, mode=mode)
+    model = client.messages.create(
+        model=model,
+        response_model=Iterable[UserExtract],
+        max_retries=2,
+        stream=stream,
+        max_tokens=1024,
+        messages=[
+            {"role": "user", "content": "Make two up people"},
+        ],
+    )
+    for m in model:
+        assert isinstance(m, UserExtract)
+
+
+@pytest.mark.parametrize("model, mode, stream", product(models, modes, [True, False]))
+@pytest.mark.asyncio
+async def test_iterable_model_async(model, mode, stream, aclient):
+    aclient = instructor.from_anthropic(aclient, mode=mode)
+    model = await aclient.messages.create(
+        model=model,
+        response_model=Iterable[UserExtract],
+        max_retries=2,
+        stream=stream,
+        max_tokens=1024,
+        messages=[
+            {"role": "user", "content": "Make two up people"},
+        ],
+    )
+    if stream:
+        async for m in model:
+            assert isinstance(m, UserExtract)
+    else:
+        for m in model:
+            assert isinstance(m, UserExtract)
+
+
+@pytest.mark.parametrize("model,mode", product(models, modes))
+def test_partial_model(model, mode, client):
+    client = instructor.from_anthropic(client, mode=mode)
+    model = client.messages.create(
+        model=model,
+        response_model=Partial[UserExtract],
+        max_retries=2,
+        max_tokens=1024,
+        stream=True,
+        messages=[
+            {"role": "user", "content": "Jason Liu is 12 years old"},
+        ],
+    )
+    for m in model:
+        assert isinstance(m, UserExtract)
+
+
+@pytest.mark.parametrize("model,mode", product(models, modes))
+@pytest.mark.asyncio
+async def test_partial_model_async(model, mode, aclient):
+    aclient = instructor.from_anthropic(aclient, mode=mode)
+    model = await aclient.messages.create(
+        model=model,
+        response_model=Partial[UserExtract],
+        max_retries=2,
+        stream=True,
+        max_tokens=1024,
+        messages=[
+            {"role": "user", "content": "Jason Liu is 12 years old"},
+        ],
+    )
+    async for m in model:
+        assert isinstance(m, UserExtract)

tests/llm/test_anthropic/util.py

@@ -0,0 +1,6 @@
+import instructor
+
+models = ["claude-3-haiku-20240307"]
+modes = [
+    instructor.Mode.ANTHROPIC_TOOLS,
+]