Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 3 additions & 5 deletions src/any_llm/providers/anthropic/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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}")
Expand Down
3 changes: 0 additions & 3 deletions src/any_llm/providers/gemini/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
9 changes: 4 additions & 5 deletions src/any_llm/providers/openai/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
"""
Comment on lines 82 to 84

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Happy to change this if considered important. let me know

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": {
Expand Down Expand Up @@ -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:
Expand Down
52 changes: 40 additions & 12 deletions tests/unit/providers/test_anthropic_provider.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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:
Expand Down
42 changes: 32 additions & 10 deletions tests/unit/providers/test_gemini_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
66 changes: 56 additions & 10 deletions tests/unit/providers/test_openai_base_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down