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
10 changes: 9 additions & 1 deletion python/packages/ollama/agent_framework_ollama/_chat_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,8 @@ class OllamaChatOptions(ChatOptions[ResponseModelT], Generic[ResponseModelT], to
presence_penalty: Presence penalty, translates to ``options.presence_penalty``.
tools: List of function tools.
response_format: Output format, translates to ``format``.
Use 'json' for JSON mode or a JSON schema dict for structured output.
Use 'json' for JSON mode, a JSON schema dict, or a Pydantic model class
(converted to its JSON schema) for structured output.
# Options not supported in Ollama:
tool_choice: Ollama only supports auto tool choice.
Expand Down Expand Up @@ -415,6 +416,13 @@ def _prepare_options(self, messages: Sequence[Message], options: Mapping[str, An
else:
# Apply top-level translations (e.g., response_format -> format)
translated_key = OLLAMA_OPTION_TRANSLATIONS.get(key, key)
if translated_key == "format" and isinstance(value, type) and issubclass(value, BaseModel):
# Ollama's `format` accepts '', 'json', or a JSON-schema dict, not a
# Pydantic model class. Convert the class to its JSON schema, matching
# OpenAIChatClient/FoundryChatClient and Ollama's documented usage
# (https://ollama.com/blog/structured-outputs). The original class is
# kept in `options` for typed parsing of the response.
value = value.model_json_schema()
run_options[translated_key] = value

# Add model options to run_options if any
Expand Down
38 changes: 38 additions & 0 deletions python/packages/ollama/tests/test_ollama_chat_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
from ollama._types import ChatResponse as OllamaChatResponse
from ollama._types import Message as OllamaMessage
from openai import AsyncStream
from pydantic import BaseModel
from pytest import fixture

from agent_framework_ollama import OllamaChatClient
Expand Down Expand Up @@ -281,6 +282,43 @@ async def test_cmc_response_format_dict(
assert result.value["answer"] == "test"


@patch.object(AsyncClient, "chat", new_callable=AsyncMock)
async def test_cmc_response_format_pydantic_model(
mock_chat: AsyncMock,
ollama_unit_test_env: dict[str, str],
chat_history: list[Message],
) -> None:
"""A Pydantic model class is converted to a JSON schema dict for Ollama's ``format``.
Ollama only accepts ``''``, ``'json'``, or a JSON-schema dict for ``format``; a model
class would fail request construction. The class is still kept for typed parsing of
the response, matching OpenAI/Foundry behavior.
"""

class Answer(BaseModel):
answer: str

mock_chat.return_value = OllamaChatResponse(
message=OllamaMessage(content='{"answer": "test"}', role="assistant"),
model="test",
eval_count=1,
prompt_eval_count=1,
created_at="2024-01-01T00:00:00Z",
)
chat_history.append(Message(contents=["hello world"], role="user"))

ollama_client = OllamaChatClient()
result = await ollama_client.get_response(messages=chat_history, options={"response_format": Answer})

# Outgoing ``format`` must be the JSON schema dict, not the model class.
assert mock_chat.await_args is not None
assert mock_chat.await_args.kwargs["format"] == Answer.model_json_schema()

# Typed parsing still works because the original model class is preserved.
assert isinstance(result.value, Answer)
assert result.value.answer == "test"


@patch.object(AsyncClient, "chat", new_callable=AsyncMock)
async def test_cmc_reasoning(
mock_chat: AsyncMock,
Expand Down