diff --git a/pydantic_ai_slim/pydantic_ai/models/groq.py b/pydantic_ai_slim/pydantic_ai/models/groq.py index 730d0486f4..881364dcad 100644 --- a/pydantic_ai_slim/pydantic_ai/models/groq.py +++ b/pydantic_ai_slim/pydantic_ai/models/groq.py @@ -9,6 +9,8 @@ from typing_extensions import assert_never +from pydantic_ai._output import DEFAULT_OUTPUT_TOOL_NAME, OutputObjectDefinition + from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage from .._run_context import RunContext from .._thinking_part import split_content_into_text_and_thinking @@ -228,6 +230,18 @@ async def _completions_create( groq_messages = self._map_messages(messages) + response_format: chat.completion_create_params.ResponseFormat | None = None + if model_request_parameters.output_mode == 'native': + output_object = model_request_parameters.output_object + assert output_object is not None + response_format = self._map_json_schema(output_object) + elif ( + model_request_parameters.output_mode == 'prompted' + and not tools + and self.profile.supports_json_object_output + ): # pragma: no branch + response_format = {'type': 'json_object'} + try: extra_headers = model_settings.get('extra_headers', {}) extra_headers.setdefault('User-Agent', get_user_agent()) @@ -240,6 +254,7 @@ async def _completions_create( tool_choice=tool_choice or NOT_GIVEN, stop=model_settings.get('stop_sequences', NOT_GIVEN), stream=stream, + response_format=response_format or NOT_GIVEN, max_tokens=model_settings.get('max_tokens', NOT_GIVEN), temperature=model_settings.get('temperature', NOT_GIVEN), top_p=model_settings.get('top_p', NOT_GIVEN), @@ -385,6 +400,19 @@ def _map_tool_definition(f: ToolDefinition) -> chat.ChatCompletionToolParam: }, } + def _map_json_schema(self, o: OutputObjectDefinition) -> chat.completion_create_params.ResponseFormat: + response_format_param: chat.completion_create_params.ResponseFormatResponseFormatJsonSchema = { + 'type': 'json_schema', + 'json_schema': { + 'name': o.name or DEFAULT_OUTPUT_TOOL_NAME, + 'schema': o.json_schema, + 'strict': o.strict, + }, + } + if o.description: # pragma: no branch + response_format_param['json_schema']['description'] = o.description + return response_format_param + @classmethod def _map_user_message(cls, message: ModelRequest) -> Iterable[chat.ChatCompletionMessageParam]: for part in message.parts: diff --git a/pydantic_ai_slim/pydantic_ai/models/openai.py b/pydantic_ai_slim/pydantic_ai/models/openai.py index 7ff5b2f7f9..d712ca6b6e 100644 --- a/pydantic_ai_slim/pydantic_ai/models/openai.py +++ b/pydantic_ai_slim/pydantic_ai/models/openai.py @@ -606,7 +606,7 @@ def _map_tool_call(t: ToolCallPart) -> ChatCompletionMessageFunctionToolCallPara def _map_json_schema(self, o: OutputObjectDefinition) -> chat.completion_create_params.ResponseFormat: response_format_param: chat.completion_create_params.ResponseFormatJSONSchema = { # pyright: ignore[reportPrivateImportUsage] 'type': 'json_schema', - 'json_schema': {'name': o.name or DEFAULT_OUTPUT_TOOL_NAME, 'schema': o.json_schema, 'strict': True}, + 'json_schema': {'name': o.name or DEFAULT_OUTPUT_TOOL_NAME, 'schema': o.json_schema}, } if o.description: response_format_param['json_schema']['description'] = o.description diff --git a/pydantic_ai_slim/pydantic_ai/providers/groq.py b/pydantic_ai_slim/pydantic_ai/providers/groq.py index 95e1913deb..ed75db024d 100644 --- a/pydantic_ai_slim/pydantic_ai/providers/groq.py +++ b/pydantic_ai_slim/pydantic_ai/providers/groq.py @@ -14,6 +14,7 @@ from pydantic_ai.profiles.meta import meta_model_profile from pydantic_ai.profiles.mistral import mistral_model_profile from pydantic_ai.profiles.moonshotai import moonshotai_model_profile +from pydantic_ai.profiles.openai import openai_model_profile from pydantic_ai.profiles.qwen import qwen_model_profile from pydantic_ai.providers import Provider @@ -26,6 +27,23 @@ ) from _import_error +def groq_moonshotai_model_profile(model_name: str) -> ModelProfile | None: + """Get the model profile for an MoonshotAI model used with the Groq provider.""" + return ModelProfile(supports_json_object_output=True, supports_json_schema_output=True).update( + moonshotai_model_profile(model_name) + ) + + +def meta_groq_model_profile(model_name: str) -> ModelProfile | None: + """Get the model profile for a Meta model used with the Groq provider.""" + if model_name in {'llama-4-maverick-17b-128e-instruct', 'llama-4-scout-17b-16e-instruct'}: + return ModelProfile(supports_json_object_output=True, supports_json_schema_output=True).update( + meta_model_profile(model_name) + ) + else: + return meta_model_profile(model_name) + + class GroqProvider(Provider[AsyncGroq]): """Provider for Groq API.""" @@ -44,13 +62,14 @@ def client(self) -> AsyncGroq: def model_profile(self, model_name: str) -> ModelProfile | None: prefix_to_profile = { 'llama': meta_model_profile, - 'meta-llama/': meta_model_profile, + 'meta-llama/': meta_groq_model_profile, 'gemma': google_model_profile, 'qwen': qwen_model_profile, 'deepseek': deepseek_model_profile, 'mistral': mistral_model_profile, - 'moonshotai/': moonshotai_model_profile, + 'moonshotai/': groq_moonshotai_model_profile, 'compound-': groq_model_profile, + 'openai/': openai_model_profile, } for prefix, profile_func in prefix_to_profile.items(): diff --git a/tests/models/cassettes/test_groq/test_groq_native_output.yaml b/tests/models/cassettes/test_groq/test_groq_native_output.yaml new file mode 100644 index 0000000000..be3cb35eee --- /dev/null +++ b/tests/models/cassettes/test_groq/test_groq_native_output.yaml @@ -0,0 +1,89 @@ +interactions: +- request: + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '416' + content-type: + - application/json + host: + - api.groq.com + method: POST + parsed_body: + messages: + - content: What is the largest city in Mexico? + role: user + model: openai/gpt-oss-120b + n: 1 + response_format: + json_schema: + description: A city and its country. + name: CityLocation + schema: + additionalProperties: false + properties: + city: + type: string + country: + type: string + required: + - city + - country + type: object + strict: true + type: json_schema + stream: false + uri: https://api.groq.com/openai/v1/chat/completions + response: + headers: + alt-svc: + - h3=":443"; ma=86400 + cache-control: + - private, max-age=0, no-store, no-cache, must-revalidate + connection: + - keep-alive + content-length: + - '947' + content-type: + - application/json + transfer-encoding: + - chunked + vary: + - Origin + parsed_body: + choices: + - finish_reason: stop + index: 0 + logprobs: null + message: + content: '{"city":"Mexico City","country":"Mexico"}' + reasoning: 'The user asks: "What is the largest city in Mexico?" The system expects a JSON object conforming to + CityLocation schema: properties city (string) and country (string), required both. Provide largest city in Mexico: + Mexico City. So output JSON: {"city":"Mexico City","country":"Mexico"} in compact format, no extra text.' + role: assistant + created: 1756843265 + id: chatcmpl-92437948-262c-49fe-87d1-774e54201105 + model: openai/gpt-oss-120b + object: chat.completion + service_tier: on_demand + system_fingerprint: fp_213abb2467 + usage: + completion_time: 0.186978247 + completion_tokens: 94 + prompt_time: 0.008149307 + prompt_tokens: 178 + queue_time: 1.095749809 + total_time: 0.195127554 + total_tokens: 272 + usage_breakdown: null + x_groq: + id: req_01k4609j8zemga5d9erfnwbfma + status: + code: 200 + message: OK +version: 1 diff --git a/tests/models/cassettes/test_groq/test_groq_prompted_output.yaml b/tests/models/cassettes/test_groq/test_groq_prompted_output.yaml new file mode 100644 index 0000000000..b52172b762 --- /dev/null +++ b/tests/models/cassettes/test_groq/test_groq_prompted_output.yaml @@ -0,0 +1,81 @@ +interactions: +- request: + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '516' + content-type: + - application/json + host: + - api.groq.com + method: POST + parsed_body: + messages: + - content: |- + Always respond with a JSON object that's compatible with this schema: + + {"properties": {"city": {"type": "string"}, "country": {"type": "string"}}, "required": ["city", "country"], "title": "CityLocation", "type": "object"} + + Don't include any text or Markdown fencing before or after. + role: system + - content: What is the largest city in Mexico? + role: user + model: openai/gpt-oss-120b + n: 1 + response_format: + type: json_object + stream: false + uri: https://api.groq.com/openai/v1/chat/completions + response: + headers: + alt-svc: + - h3=":443"; ma=86400 + cache-control: + - private, max-age=0, no-store, no-cache, must-revalidate + connection: + - keep-alive + content-length: + - '926' + content-type: + - application/json + transfer-encoding: + - chunked + vary: + - Origin + parsed_body: + choices: + - finish_reason: stop + index: 0 + logprobs: null + message: + content: '{"city":"Mexico City","country":"Mexico"}' + reasoning: 'We need to respond with JSON object with properties city and country. The question: "What is the largest + city in Mexico?" The answer: City is Mexico City, country is Mexico. Must output compact JSON without any extra + text or markdown. So {"city":"Mexico City","country":"Mexico"} Ensure valid JSON.' + role: assistant + created: 1756843266 + id: chatcmpl-d7085def-1e9f-45d7-b90b-65633ef23489 + model: openai/gpt-oss-120b + object: chat.completion + service_tier: on_demand + system_fingerprint: fp_ed9190d8b7 + usage: + completion_time: 0.173182068 + completion_tokens: 87 + prompt_time: 0.006958709 + prompt_tokens: 177 + queue_time: 0.212268627 + total_time: 0.180140777 + total_tokens: 264 + usage_breakdown: null + x_groq: + id: req_01k4609m4zf9qsbf0wwgvw8398 + status: + code: 200 + message: OK +version: 1 diff --git a/tests/models/test_groq.py b/tests/models/test_groq.py index d6d8e81ec3..714f82d903 100644 --- a/tests/models/test_groq.py +++ b/tests/models/test_groq.py @@ -13,6 +13,7 @@ import pytest from dirty_equals import IsListOrTuple from inline_snapshot import snapshot +from pydantic import BaseModel from typing_extensions import TypedDict from pydantic_ai import Agent, ModelHTTPError, ModelRetry, UnexpectedModelBehavior @@ -37,6 +38,7 @@ ToolReturnPart, UserPromptPart, ) +from pydantic_ai.output import NativeOutput, PromptedOutput from pydantic_ai.usage import RequestUsage, RunUsage from ..conftest import IsDatetime, IsInstance, IsNow, IsStr, raise_if_exception, try_import @@ -1004,3 +1006,90 @@ async def test_groq_model_thinking_part_iter(allow_model_requests: None, groq_ap length=996, ) ) + + +async def test_groq_native_output(allow_model_requests: None, groq_api_key: str): + m = GroqModel('openai/gpt-oss-120b', provider=GroqProvider(api_key=groq_api_key)) + + class CityLocation(BaseModel): + """A city and its country.""" + + city: str + country: str + + agent = Agent(m, output_type=NativeOutput(CityLocation)) + + result = await agent.run('What is the largest city in Mexico?') + assert result.output == snapshot(CityLocation(city='Mexico City', country='Mexico')) + + assert result.all_messages() == snapshot( + [ + ModelRequest( + parts=[ + UserPromptPart( + content='What is the largest city in Mexico?', + timestamp=IsDatetime(), + ) + ] + ), + ModelResponse( + parts=[ + ThinkingPart( + content='The user asks: "What is the largest city in Mexico?" The system expects a JSON object conforming to CityLocation schema: properties city (string) and country (string), required both. Provide largest city in Mexico: Mexico City. So output JSON: {"city":"Mexico City","country":"Mexico"} in compact format, no extra text.' + ), + TextPart(content='{"city":"Mexico City","country":"Mexico"}'), + ], + usage=RequestUsage(input_tokens=178, output_tokens=94), + model_name='openai/gpt-oss-120b', + timestamp=IsDatetime(), + provider_name='groq', + provider_response_id=IsStr(), + ), + ] + ) + + +async def test_groq_prompted_output(allow_model_requests: None, groq_api_key: str): + m = GroqModel('openai/gpt-oss-120b', provider=GroqProvider(api_key=groq_api_key)) + + class CityLocation(BaseModel): + city: str + country: str + + agent = Agent(m, output_type=PromptedOutput(CityLocation)) + + result = await agent.run('What is the largest city in Mexico?') + assert result.output == snapshot(CityLocation(city='Mexico City', country='Mexico')) + + assert result.all_messages() == snapshot( + [ + ModelRequest( + parts=[ + UserPromptPart( + content='What is the largest city in Mexico?', + timestamp=IsDatetime(), + ) + ], + instructions="""\ +Always respond with a JSON object that's compatible with this schema: + +{"properties": {"city": {"type": "string"}, "country": {"type": "string"}}, "required": ["city", "country"], "title": "CityLocation", "type": "object"} + +Don't include any text or Markdown fencing before or after.\ +""", + ), + ModelResponse( + parts=[ + ThinkingPart( + content='We need to respond with JSON object with properties city and country. The question: "What is the largest city in Mexico?" The answer: City is Mexico City, country is Mexico. Must output compact JSON without any extra text or markdown. So {"city":"Mexico City","country":"Mexico"} Ensure valid JSON.' + ), + TextPart(content='{"city":"Mexico City","country":"Mexico"}'), + ], + usage=RequestUsage(input_tokens=177, output_tokens=87), + model_name='openai/gpt-oss-120b', + timestamp=IsDatetime(), + provider_name='groq', + provider_response_id=IsStr(), + ), + ] + ) diff --git a/tests/providers/test_groq.py b/tests/providers/test_groq.py index dc29e8536f..1ac38312b6 100644 --- a/tests/providers/test_groq.py +++ b/tests/providers/test_groq.py @@ -14,6 +14,7 @@ from pydantic_ai.profiles.meta import meta_model_profile from pydantic_ai.profiles.mistral import mistral_model_profile from pydantic_ai.profiles.moonshotai import moonshotai_model_profile +from pydantic_ai.profiles.openai import openai_model_profile from pydantic_ai.profiles.qwen import qwen_model_profile from ..conftest import TestEnv, try_import @@ -78,10 +79,27 @@ def test_groq_provider_model_profile(mocker: MockerFixture): qwen_model_profile_mock = mocker.patch(f'{ns}.qwen_model_profile', wraps=qwen_model_profile) moonshotai_model_profile_mock = mocker.patch(f'{ns}.moonshotai_model_profile', wraps=moonshotai_model_profile) groq_model_profile_mock = mocker.patch(f'{ns}.groq_model_profile', wraps=groq_model_profile) + openai_model_profile_mock = mocker.patch(f'{ns}.openai_model_profile', wraps=openai_model_profile) meta_profile = provider.model_profile('meta-llama/Llama-Guard-4-12B') meta_model_profile_mock.assert_called_with('llama-guard-4-12b') assert meta_profile is not None + assert meta_profile.supports_json_object_output is False + assert meta_profile.supports_json_schema_output is False + assert meta_profile.json_schema_transformer == InlineDefsJsonSchemaTransformer + + meta_profile = provider.model_profile('meta-llama/llama-4-maverick-17b-128e-instruct') + meta_model_profile_mock.assert_called_with('llama-4-maverick-17b-128e-instruct') + assert meta_profile is not None + assert meta_profile.supports_json_object_output is True + assert meta_profile.supports_json_schema_output is True + assert meta_profile.json_schema_transformer == InlineDefsJsonSchemaTransformer + + meta_profile = provider.model_profile('meta-llama/llama-4-scout-17b-16e-instruct') + meta_model_profile_mock.assert_called_with('llama-4-scout-17b-16e-instruct') + assert meta_profile is not None + assert meta_profile.supports_json_object_output is True + assert meta_profile.supports_json_schema_output is True assert meta_profile.json_schema_transformer == InlineDefsJsonSchemaTransformer meta_profile = provider.model_profile('llama-3.3-70b-versatile') @@ -117,7 +135,21 @@ def test_groq_provider_model_profile(mocker: MockerFixture): moonshotai_profile = provider.model_profile('moonshotai/kimi-k2-instruct') moonshotai_model_profile_mock.assert_called_with('kimi-k2-instruct') assert moonshotai_profile is not None + assert moonshotai_profile.supports_json_object_output is True + assert moonshotai_profile.supports_json_schema_output is True assert moonshotai_profile.ignore_streamed_leading_whitespace is True + openai_profile = provider.model_profile('openai/gpt-oss-20b') + openai_model_profile_mock.assert_called_with('gpt-oss-20b') + assert openai_profile is not None + assert openai_profile.supports_json_object_output is True + assert openai_profile.supports_json_schema_output is True + + openai_profile = provider.model_profile('openai/gpt-oss-120b') + openai_model_profile_mock.assert_called_with('gpt-oss-120b') + assert openai_profile is not None + assert openai_profile.supports_json_object_output is True + assert openai_profile.supports_json_schema_output is True + unknown_profile = provider.model_profile('unknown-model') assert unknown_profile is None