Enable streaming when response_format is set for gemini, openai, and anthropic completion endpoints.

Author: liukidar

src/any_llm/providers/anthropic/utils.py

@@ -335,13 +335,14 @@ def _convert_tool_spec(openai_tools: list[dict[str, Any]]) -> list[dict[str, Any
 
     anthropic_tools = []
     for tool in generic_tools:
+        params: dict[str, Any] = tool["parameters"] or {}
         anthropic_tool = {
             "name": tool["name"],
             "description": tool["description"],
             "input_schema": {
                 "type": "object",
-                "properties": tool["parameters"].get("properties") or {},
-                "required": tool["parameters"].get("required", []),
+                "properties": params.get("properties") or {},
+                "required": params.get("required", []),
             },
         }
         anthropic_tools.append(anthropic_tool)
@@ -395,9 +396,6 @@ def _convert_params(params: CompletionParams, **kwargs: Any) -> dict[str, Any]:
     result_kwargs: dict[str, Any] = kwargs.copy()
 
     if params.response_format:
-        if params.stream:
-            msg = "stream and response_format"
-            raise UnsupportedParameterError(msg, provider_name)
         result_kwargs["output_config"] = _convert_response_format(params.response_format, provider_name)
     if params.max_tokens is None:
         logger.warning(f"max_tokens is required for Anthropic, setting to {DEFAULT_MAX_TOKENS}")

src/any_llm/providers/gemini/base.py

@@ -124,9 +124,6 @@ def _convert_completion_params(params: CompletionParams, **kwargs: Any) -> dict[
         if params.parallel_tool_calls is not None:
             error_message = "parallel_tool_calls"
             raise UnsupportedParameterError(error_message, provider_name)
-        if params.stream and params.response_format is not None:
-            error_message = "stream and response_format"
-            raise UnsupportedParameterError(error_message, provider_name)
 
         if params.frequency_penalty is not None:
             kwargs["frequency_penalty"] = params.frequency_penalty

src/any_llm/providers/openai/base.py

@@ -82,7 +82,9 @@ def _convert_completion_params(params: CompletionParams, **kwargs: Any) -> dict[
         Plain dataclasses are converted to JSON schema dicts since the
         OpenAI SDK's ``.parse()`` only supports Pydantic BaseModel types.
         """
-        if is_structured_output_type(params.response_format) and not issubclass(params.response_format, BaseModel):
+        if is_structured_output_type(params.response_format) and (
+            params.stream or not issubclass(params.response_format, BaseModel)
+        ):
             params.response_format = {
                 "type": "json_schema",
                 "json_schema": {
@@ -202,10 +204,7 @@ async def _acompletion(
         response_format = completion_kwargs.get("response_format")
         use_parse = is_structured_output_type(response_format)
 
-        if response_format:
-            if params.stream:
-                msg = "stream is not supported for response_format"
-                raise ValueError(msg)
+        if response_format and not params.stream:
             completion_kwargs.pop("stream", None)
 
         if use_parse:

tests/unit/providers/test_anthropic_provider.py

@@ -1,7 +1,8 @@
 import dataclasses
-from contextlib import contextmanager
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager, contextmanager
 from datetime import UTC
-from typing import Any, get_args
+from typing import Any, cast, get_args
 from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
@@ -626,23 +627,50 @@ async def test_completion_with_response_format_dict_json_object_raises() -> None
 
 
 @pytest.mark.asyncio
-async def test_stream_with_response_format_raises() -> None:
+async def test_stream_with_response_format_passes_output_config() -> None:
     api_key = "test-api-key"
     model = "model-id"
     messages = [{"role": "user", "content": "Hello"}]
+    schema = {"type": "object", "properties": {"city_name": {"type": "string"}}, "required": ["city_name"]}
+    response_format: dict[str, Any] = {
+        "type": "json_schema",
+        "json_schema": {"name": "Foo", "schema": schema},
+    }
 
-    provider = AnthropicProvider(api_key=api_key)
+    async def empty_events() -> AsyncIterator[Any]:
+        if False:
+            yield None
 
-    with pytest.raises(UnsupportedParameterError, match="stream and response_format"):
-        await provider._acompletion(
-            CompletionParams(
-                model_id=model,
-                messages=messages,
-                response_format={"type": "json_schema", "json_schema": {"name": "Foo", "schema": {}}},
-                stream=True,
-            )
+    @asynccontextmanager
+    async def empty_stream() -> AsyncIterator[AsyncIterator[Any]]:
+        yield empty_events()
+
+    with mock_anthropic_provider() as mock_anthropic:
+        mock_anthropic.return_value.messages.stream = Mock(return_value=empty_stream())
+        provider = AnthropicProvider(api_key=api_key)
+        stream = cast(
+            "AsyncIterator[Any]",
+            await provider._acompletion(
+                CompletionParams(
+                    model_id=model,
+                    messages=messages,
+                    response_format=response_format,
+                    stream=True,
+                )
+            ),
         )
 
+        async for _ in stream:
+            pass
+
+        expected_output_config = {"format": {"type": "json_schema", "schema": transform_schema(schema)}}
+        mock_anthropic.return_value.messages.stream.assert_called_once()
+        call_kwargs = mock_anthropic.return_value.messages.stream.call_args.kwargs
+        assert call_kwargs["model"] == model
+        assert call_kwargs["messages"] == messages
+        assert call_kwargs["max_tokens"] == DEFAULT_MAX_TOKENS
+        assert call_kwargs["output_config"] == expected_output_config
+
 
 @pytest.mark.asyncio
 async def test_completion_with_response_format_dict_unknown_type_raises() -> None:

tests/unit/providers/test_gemini_provider.py

@@ -392,22 +392,44 @@ async def test_completion_with_unsupported_dict_response_format_raises() -> None
 
 
 @pytest.mark.asyncio
-async def test_completion_with_stream_and_response_format_raises() -> None:
+async def test_stream_with_response_format_passes_generation_config() -> None:
     api_key = "test-api-key"
     model = "gemini-pro"
     messages = [{"role": "user", "content": "Hello"}]
+    expected_schema: dict[str, Any] = {
+        "type": "object",
+        "properties": {
+            "name": {"type": "string"},
+            "value": {"type": "integer"},
+        },
+        "required": ["name", "value"],
+    }
+    response_format: dict[str, Any] = {
+        "type": "json_schema",
+        "json_schema": {
+            "name": "TestOutput",
+            "schema": expected_schema,
+        },
+    }
 
-    with mock_gemini_provider():
+    with mock_gemini_provider() as mock_genai:
+        mock_genai.return_value.aio.models.generate_content_stream = AsyncMock()
         provider = GeminiProvider(api_key=api_key)
-        with pytest.raises(UnsupportedParameterError):
-            await provider._acompletion(
-                CompletionParams(
-                    model_id=model,
-                    messages=messages,
-                    stream=True,
-                    response_format={"type": "json_object"},
-                )
+        await provider._acompletion(
+            CompletionParams(
+                model_id=model,
+                messages=messages,
+                stream=True,
+                response_format=response_format,
             )
+        )
+
+        _, call_kwargs = mock_genai.return_value.aio.models.generate_content_stream.call_args
+        generation_config = call_kwargs["config"]
+
+    assert call_kwargs["model"] == model
+    assert generation_config.response_mime_type == "application/json"
+    assert generation_config.response_schema == expected_schema
 
 
 @pytest.mark.asyncio

tests/unit/providers/test_openai_base_provider.py

@@ -224,24 +224,70 @@ class ListModelsProvider(BaseOpenAIProvider):
 
 
 @pytest.mark.asyncio
-async def test_stream_with_response_format_raises() -> None:
+async def test_stream_with_dict_response_format_uses_create() -> None:
     class TestProvider(BaseOpenAIProvider):
         PROVIDER_NAME = "TestProvider"
         ENV_API_KEY_NAME = "TEST_API_KEY"
         PROVIDER_DOCUMENTATION_URL = "https://example.com"
 
-    with patch("any_llm.providers.openai.base.AsyncOpenAI"):
+    response_format = {"type": "json_object"}
+
+    with patch("any_llm.providers.openai.base.AsyncOpenAI") as mock_openai_class:
+        mock_client = AsyncMock()
+        mock_openai_class.return_value = mock_client
+        mock_client.chat.completions.create = AsyncMock(return_value=MagicMock())
+
         provider = TestProvider(api_key="test-key")
 
-        with pytest.raises(ValueError, match="stream is not supported for response_format"):
-            await provider._acompletion(
-                CompletionParams(
-                    model_id="test-model",
-                    messages=[{"role": "user", "content": "Hello"}],
-                    stream=True,
-                    response_format={"type": "json_object"},
-                )
+        await provider._acompletion(
+            CompletionParams(
+                model_id="test-model",
+                messages=[{"role": "user", "content": "Hello"}],
+                stream=True,
+                response_format=response_format,
             )
+        )
+
+        mock_client.chat.completions.create.assert_awaited_once_with(
+            model="test-model",
+            messages=[{"role": "user", "content": "Hello"}],
+            stream=True,
+            response_format=response_format,
+        )
+        mock_client.chat.completions.parse.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_stream_with_basemodel_response_format_uses_create_json_schema() -> None:
+    class TestProvider(BaseOpenAIProvider):
+        PROVIDER_NAME = "TestProvider"
+        ENV_API_KEY_NAME = "TEST_API_KEY"
+        PROVIDER_DOCUMENTATION_URL = "https://example.com"
+
+    with patch("any_llm.providers.openai.base.AsyncOpenAI") as mock_openai_class:
+        mock_client = AsyncMock()
+        mock_openai_class.return_value = mock_client
+        mock_client.chat.completions.create = AsyncMock(return_value=MagicMock())
+
+        provider = TestProvider(api_key="test-key")
+
+        await provider._acompletion(
+            CompletionParams(
+                model_id="test-model",
+                messages=[{"role": "user", "content": "Hello"}],
+                stream=True,
+                response_format=_City,
+            )
+        )
+
+        create_kwargs = mock_client.chat.completions.create.call_args.kwargs
+        assert create_kwargs["stream"] is True
+        assert create_kwargs["response_format"]["type"] == "json_schema"
+        assert create_kwargs["response_format"]["json_schema"]["name"] == "_City"
+        assert create_kwargs["response_format"]["json_schema"]["schema"]["properties"] == {
+            "city_name": {"title": "City Name", "type": "string"}
+        }
+        mock_client.chat.completions.parse.assert_not_called()
 
 
 def test_base_provider_maps_max_tokens_to_max_completion_tokens() -> None: