pydantic-ai-slim 1.0.0b1__py3-none-any.whl → 1.0.2__py3-none-any.whl

pydantic_ai/models/groq.py

@@ -7,8 +7,11 @@ from dataclasses import dataclass, field
 from datetime import datetime
 from typing import Any, Literal, cast, overload
 
+from pydantic import BaseModel, Json, ValidationError
 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
@@ -48,7 +51,7 @@ from . import (
 )
 
 try:
-    from groq import NOT_GIVEN, APIStatusError, AsyncGroq, AsyncStream
+    from groq import NOT_GIVEN, APIError, APIStatusError, AsyncGroq, AsyncStream
     from groq.types import chat
     from groq.types.chat.chat_completion_content_part_image_param import ImageURL
 except ImportError as _import_error:
@@ -169,9 +172,24 @@ class GroqModel(Model):
         model_request_parameters: ModelRequestParameters,
     ) -> ModelResponse:
         check_allow_model_requests()
-        response = await self._completions_create(
-            messages, False, cast(GroqModelSettings, model_settings or {}), model_request_parameters
-        )
+        try:
+            response = await self._completions_create(
+                messages, False, cast(GroqModelSettings, model_settings or {}), model_request_parameters
+            )
+        except ModelHTTPError as e:
+            if isinstance(e.body, dict):  # pragma: no branch
+                # The Groq SDK tries to be helpful by raising an exception when generated tool arguments don't match the schema,
+                # but we'd rather handle it ourselves so we can tell the model to retry the tool call.
+                try:
+                    error = _GroqToolUseFailedError.model_validate(e.body)  # pyright: ignore[reportUnknownMemberType]
+                    tool_call_part = ToolCallPart(
+                        tool_name=error.error.failed_generation.name,
+                        args=error.error.failed_generation.arguments,
+                    )
+                    return ModelResponse(parts=[tool_call_part])
+                except ValidationError:
+                    pass
+            raise
         model_response = self._process_response(response)
         return model_response
 
@@ -228,6 +246,18 @@ class GroqModel(Model):
 
         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 +270,7 @@ class GroqModel(Model):
                 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 +416,19 @@ class GroqModel(Model):
             },
         }
 
+    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:
@@ -449,36 +493,52 @@ class GroqStreamedResponse(StreamedResponse):
     _provider_name: str
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
-        async for chunk in self._response:
-            self._usage += _map_usage(chunk)
-
-            try:
-                choice = chunk.choices[0]
-            except IndexError:
-                continue
-
-            # Handle the text part of the response
-            content = choice.delta.content
-            if content is not None:
-                maybe_event = self._parts_manager.handle_text_delta(
-                    vendor_part_id='content',
-                    content=content,
-                    thinking_tags=self._model_profile.thinking_tags,
-                    ignore_leading_whitespace=self._model_profile.ignore_streamed_leading_whitespace,
-                )
-                if maybe_event is not None:  # pragma: no branch
-                    yield maybe_event
-
-            # Handle the tool calls
-            for dtc in choice.delta.tool_calls or []:
-                maybe_event = self._parts_manager.handle_tool_call_delta(
-                    vendor_part_id=dtc.index,
-                    tool_name=dtc.function and dtc.function.name,
-                    args=dtc.function and dtc.function.arguments,
-                    tool_call_id=dtc.id,
-                )
-                if maybe_event is not None:
-                    yield maybe_event
+        try:
+            async for chunk in self._response:
+                self._usage += _map_usage(chunk)
+
+                try:
+                    choice = chunk.choices[0]
+                except IndexError:
+                    continue
+
+                # Handle the text part of the response
+                content = choice.delta.content
+                if content is not None:
+                    maybe_event = self._parts_manager.handle_text_delta(
+                        vendor_part_id='content',
+                        content=content,
+                        thinking_tags=self._model_profile.thinking_tags,
+                        ignore_leading_whitespace=self._model_profile.ignore_streamed_leading_whitespace,
+                    )
+                    if maybe_event is not None:  # pragma: no branch
+                        yield maybe_event
+
+                # Handle the tool calls
+                for dtc in choice.delta.tool_calls or []:
+                    maybe_event = self._parts_manager.handle_tool_call_delta(
+                        vendor_part_id=dtc.index,
+                        tool_name=dtc.function and dtc.function.name,
+                        args=dtc.function and dtc.function.arguments,
+                        tool_call_id=dtc.id,
+                    )
+                    if maybe_event is not None:
+                        yield maybe_event
+        except APIError as e:
+            if isinstance(e.body, dict):  # pragma: no branch
+                # The Groq SDK tries to be helpful by raising an exception when generated tool arguments don't match the schema,
+                # but we'd rather handle it ourselves so we can tell the model to retry the tool call
+                try:
+                    error = _GroqToolUseFailedInnerError.model_validate(e.body)  # pyright: ignore[reportUnknownMemberType]
+                    yield self._parts_manager.handle_tool_call_part(
+                        vendor_part_id='tool_use_failed',
+                        tool_name=error.failed_generation.name,
+                        args=error.failed_generation.arguments,
+                    )
+                    return
+                except ValidationError as e:  # pragma: no cover
+                    pass
+            raise  # pragma: no cover
 
     @property
     def model_name(self) -> GroqModelName:
@@ -510,3 +570,31 @@ def _map_usage(completion: chat.ChatCompletionChunk | chat.ChatCompletion) -> us
         input_tokens=response_usage.prompt_tokens,
         output_tokens=response_usage.completion_tokens,
     )
+
+
+class _GroqToolUseFailedGeneration(BaseModel):
+    name: str
+    arguments: dict[str, Any]
+
+
+class _GroqToolUseFailedInnerError(BaseModel):
+    message: str
+    type: Literal['invalid_request_error']
+    code: Literal['tool_use_failed']
+    failed_generation: Json[_GroqToolUseFailedGeneration]
+
+
+class _GroqToolUseFailedError(BaseModel):
+    # The Groq SDK tries to be helpful by raising an exception when generated tool arguments don't match the schema,
+    # but we'd rather handle it ourselves so we can tell the model to retry the tool call.
+    # Example payload from `exception.body`:
+    # {
+    #     'error': {
+    #         'message': "Tool call validation failed: tool call validation failed: parameters for tool get_something_by_name did not match schema: errors: [missing properties: 'name', additionalProperties 'foo' not allowed]",
+    #         'type': 'invalid_request_error',
+    #         'code': 'tool_use_failed',
+    #         'failed_generation': '{"name": "get_something_by_name", "arguments": {\n  "foo": "bar"\n}}',
+    #     }
+    # }
+
+    error: _GroqToolUseFailedInnerError

pydantic_ai/models/instrumented.py

@@ -221,7 +221,10 @@ class InstrumentationSettings:
                         _otel_messages.ChatMessage(role='system' if is_system else 'user', parts=message_parts)
                     )
             elif isinstance(message, ModelResponse):  # pragma: no branch
-                result.append(_otel_messages.ChatMessage(role='assistant', parts=message.otel_message_parts(self)))
+                otel_message = _otel_messages.OutputMessage(role='assistant', parts=message.otel_message_parts(self))
+                if message.finish_reason is not None:
+                    otel_message['finish_reason'] = message.finish_reason
+                result.append(otel_message)
         return result
 
     def handle_messages(self, input_messages: list[ModelMessage], response: ModelResponse, system: str, span: Span):
@@ -246,12 +249,10 @@ class InstrumentationSettings:
         else:
             output_messages = self.messages_to_otel_messages([response])
             assert len(output_messages) == 1
-            output_message = cast(_otel_messages.OutputMessage, output_messages[0])
-            if response.provider_details and 'finish_reason' in response.provider_details:
-                output_message['finish_reason'] = response.provider_details['finish_reason']
+            output_message = output_messages[0]
             instructions = InstrumentedModel._get_instructions(input_messages)  # pyright: ignore [reportPrivateUsage]
             system_instructions_attributes = self.system_instructions_attributes(instructions)
-            attributes = {
+            attributes: dict[str, AttributeValue] = {
                 'gen_ai.input.messages': json.dumps(self.messages_to_otel_messages(input_messages)),
                 'gen_ai.output.messages': json.dumps([output_message]),
                 **system_instructions_attributes,
@@ -420,12 +421,25 @@ class InstrumentedModel(WrapperModel):
                         return
 
                     self.instrumentation_settings.handle_messages(messages, response, system, span)
-                    span.set_attributes(
-                        {
-                            **response.usage.opentelemetry_attributes(),
-                            'gen_ai.response.model': response_model,
-                        }
-                    )
+
+                    attributes_to_set = {
+                        **response.usage.opentelemetry_attributes(),
+                        'gen_ai.response.model': response_model,
+                    }
+                    try:
+                        attributes_to_set['operation.cost'] = float(response.cost().total_price)
+                    except LookupError:
+                        # The cost of this provider/model is unknown, which is common.
+                        pass
+                    except Exception as e:
+                        warnings.warn(
+                            f'Failed to get cost from response: {type(e).__name__}: {e}', CostCalculationFailedWarning
+                        )
+                    if response.provider_response_id is not None:
+                        attributes_to_set['gen_ai.response.id'] = response.provider_response_id
+                    if response.finish_reason is not None:
+                        attributes_to_set['gen_ai.response.finish_reasons'] = [response.finish_reason]
+                    span.set_attributes(attributes_to_set)
                     span.update_name(f'{operation} {request_model}')
 
                 yield finish
@@ -473,3 +487,7 @@ class InstrumentedModel(WrapperModel):
                 return str(value)
             except Exception as e:
                 return f'Unable to serialize: {e}'
+
+
+class CostCalculationFailedWarning(Warning):
+    """Warning raised when cost calculation fails."""

pydantic_ai/models/openai.py

@@ -24,6 +24,7 @@ from ..messages import (
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
     DocumentUrl,
+    FinishReason,
     ImageUrl,
     ModelMessage,
     ModelRequest,
@@ -72,6 +73,7 @@ try:
     )
     from openai.types.responses import ComputerToolParam, FileSearchToolParam, WebSearchToolParam
     from openai.types.responses.response_input_param import FunctionCallOutput, Message
+    from openai.types.responses.response_status import ResponseStatus
     from openai.types.shared import ReasoningEffort
     from openai.types.shared_params import Reasoning
 except ImportError as _import_error:
@@ -103,6 +105,25 @@ allows this model to be used more easily with other model types (ie, Ollama, Dee
 """
 
 
+_CHAT_FINISH_REASON_MAP: dict[
+    Literal['stop', 'length', 'tool_calls', 'content_filter', 'function_call'], FinishReason
+] = {
+    'stop': 'stop',
+    'length': 'length',
+    'tool_calls': 'tool_call',
+    'content_filter': 'content_filter',
+    'function_call': 'tool_call',
+}
+
+_RESPONSES_FINISH_REASON_MAP: dict[Literal['max_output_tokens', 'content_filter'] | ResponseStatus, FinishReason] = {
+    'max_output_tokens': 'length',
+    'content_filter': 'content_filter',
+    'completed': 'stop',
+    'cancelled': 'error',
+    'failed': 'error',
+}
+
+
 class OpenAIChatModelSettings(ModelSettings, total=False):
     """Settings used for an OpenAI model request."""
 
@@ -225,6 +246,7 @@ class OpenAIChatModel(Model):
             'openrouter',
             'together',
             'vercel',
+            'litellm',
         ]
         | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,
@@ -252,6 +274,7 @@ class OpenAIChatModel(Model):
             'openrouter',
             'together',
             'vercel',
+            'litellm',
         ]
         | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,
@@ -278,6 +301,7 @@ class OpenAIChatModel(Model):
             'openrouter',
             'together',
             'vercel',
+            'litellm',
         ]
         | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,
@@ -471,24 +495,22 @@ class OpenAIChatModel(Model):
         if reasoning_content := getattr(choice.message, 'reasoning_content', None):
             items.append(ThinkingPart(content=reasoning_content))
 
-        vendor_details: dict[str, Any] | None = None
+        vendor_details: dict[str, Any] = {}
 
         # Add logprobs to vendor_details if available
         if choice.logprobs is not None and choice.logprobs.content:
             # Convert logprobs to a serializable format
-            vendor_details = {
-                'logprobs': [
-                    {
-                        'token': lp.token,
-                        'bytes': lp.bytes,
-                        'logprob': lp.logprob,
-                        'top_logprobs': [
-                            {'token': tlp.token, 'bytes': tlp.bytes, 'logprob': tlp.logprob} for tlp in lp.top_logprobs
-                        ],
-                    }
-                    for lp in choice.logprobs.content
-                ],
-            }
+            vendor_details['logprobs'] = [
+                {
+                    'token': lp.token,
+                    'bytes': lp.bytes,
+                    'logprob': lp.logprob,
+                    'top_logprobs': [
+                        {'token': tlp.token, 'bytes': tlp.bytes, 'logprob': tlp.logprob} for tlp in lp.top_logprobs
+                    ],
+                }
+                for lp in choice.logprobs.content
+            ]
 
         if choice.message.content is not None:
             items.extend(split_content_into_text_and_thinking(choice.message.content, self.profile.thinking_tags))
@@ -504,14 +526,21 @@ class OpenAIChatModel(Model):
                     assert_never(c)
                 part.tool_call_id = _guard_tool_call_id(part)
                 items.append(part)
+
+        finish_reason: FinishReason | None = None
+        if raw_finish_reason := choice.finish_reason:  # pragma: no branch
+            vendor_details['finish_reason'] = raw_finish_reason
+            finish_reason = _CHAT_FINISH_REASON_MAP.get(raw_finish_reason)
+
         return ModelResponse(
             parts=items,
             usage=_map_usage(response),
             model_name=response.model,
             timestamp=timestamp,
-            provider_details=vendor_details,
+            provider_details=vendor_details or None,
             provider_response_id=response.id,
             provider_name=self._provider.name,
+            finish_reason=finish_reason,
         )
 
     async def _process_streamed_response(
@@ -606,7 +635,7 @@ class OpenAIChatModel(Model):
     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
@@ -820,6 +849,14 @@ class OpenAIResponsesModel(Model):
                         items.append(TextPart(content.text))
             elif item.type == 'function_call':
                 items.append(ToolCallPart(item.name, item.arguments, tool_call_id=item.call_id))
+
+        finish_reason: FinishReason | None = None
+        provider_details: dict[str, Any] | None = None
+        raw_finish_reason = details.reason if (details := response.incomplete_details) else response.status
+        if raw_finish_reason:
+            provider_details = {'finish_reason': raw_finish_reason}
+            finish_reason = _RESPONSES_FINISH_REASON_MAP.get(raw_finish_reason)
+
         return ModelResponse(
             parts=items,
             usage=_map_usage(response),
@@ -827,6 +864,8 @@ class OpenAIResponsesModel(Model):
             provider_response_id=response.id,
             timestamp=timestamp,
             provider_name=self._provider.name,
+            finish_reason=finish_reason,
+            provider_details=provider_details,
         )
 
     async def _process_streamed_response(
@@ -1166,11 +1205,22 @@ class OpenAIStreamedResponse(StreamedResponse):
         async for chunk in self._response:
             self._usage += _map_usage(chunk)
 
+            if chunk.id and self.provider_response_id is None:
+                self.provider_response_id = chunk.id
+
             try:
                 choice = chunk.choices[0]
             except IndexError:
                 continue
 
+            # When using Azure OpenAI and an async content filter is enabled, the openai SDK can return None deltas.
+            if choice.delta is None:  # pyright: ignore[reportUnnecessaryComparison]
+                continue
+
+            if raw_finish_reason := choice.finish_reason:
+                self.provider_details = {'finish_reason': raw_finish_reason}
+                self.finish_reason = _CHAT_FINISH_REASON_MAP.get(raw_finish_reason)
+
             # Handle the text part of the response
             content = choice.delta.content
             if content is not None:
@@ -1230,6 +1280,13 @@ class OpenAIResponsesStreamedResponse(StreamedResponse):
             if isinstance(chunk, responses.ResponseCompletedEvent):
                 self._usage += _map_usage(chunk.response)
 
+                raw_finish_reason = (
+                    details.reason if (details := chunk.response.incomplete_details) else chunk.response.status
+                )
+                if raw_finish_reason:  # pragma: no branch
+                    self.provider_details = {'finish_reason': raw_finish_reason}
+                    self.finish_reason = _RESPONSES_FINISH_REASON_MAP.get(raw_finish_reason)
+
             elif isinstance(chunk, responses.ResponseContentPartAddedEvent):
                 pass  # there's nothing we need to do here
 
@@ -1237,7 +1294,8 @@ class OpenAIResponsesStreamedResponse(StreamedResponse):
                 pass  # there's nothing we need to do here
 
             elif isinstance(chunk, responses.ResponseCreatedEvent):
-                pass  # there's nothing we need to do here
+                if chunk.response.id:  # pragma: no branch
+                    self.provider_response_id = chunk.response.id
 
             elif isinstance(chunk, responses.ResponseFailedEvent):  # pragma: no cover
                 self._usage += _map_usage(chunk.response)
@@ -1270,12 +1328,7 @@ class OpenAIResponsesStreamedResponse(StreamedResponse):
                         tool_call_id=chunk.item.call_id,
                     )
                 elif isinstance(chunk.item, responses.ResponseReasoningItem):
-                    content = chunk.item.summary[0].text if chunk.item.summary else ''
-                    yield self._parts_manager.handle_thinking_delta(
-                        vendor_part_id=chunk.item.id,
-                        content=content,
-                        signature=chunk.item.id,
-                    )
+                    pass
                 elif isinstance(chunk.item, responses.ResponseOutputMessage):
                     pass
                 elif isinstance(chunk.item, responses.ResponseFunctionWebSearch):
@@ -1291,7 +1344,11 @@ class OpenAIResponsesStreamedResponse(StreamedResponse):
                 pass
 
             elif isinstance(chunk, responses.ResponseReasoningSummaryPartAddedEvent):
-                pass  # there's nothing we need to do here
+                yield self._parts_manager.handle_thinking_delta(
+                    vendor_part_id=f'{chunk.item_id}-{chunk.summary_index}',
+                    content=chunk.part.text,
+                    id=chunk.item_id,
+                )
 
             elif isinstance(chunk, responses.ResponseReasoningSummaryPartDoneEvent):
                 pass  # there's nothing we need to do here
@@ -1301,9 +1358,9 @@ class OpenAIResponsesStreamedResponse(StreamedResponse):
 
             elif isinstance(chunk, responses.ResponseReasoningSummaryTextDeltaEvent):
                 yield self._parts_manager.handle_thinking_delta(
-                    vendor_part_id=chunk.item_id,
+                    vendor_part_id=f'{chunk.item_id}-{chunk.summary_index}',
                     content=chunk.delta,
-                    signature=chunk.item_id,
+                    id=chunk.item_id,
                 )
 
             # TODO(Marcelo): We should support annotations in the future.
@@ -1311,9 +1368,7 @@ class OpenAIResponsesStreamedResponse(StreamedResponse):
                 pass  # there's nothing we need to do here
 
             elif isinstance(chunk, responses.ResponseTextDeltaEvent):
-                maybe_event = self._parts_manager.handle_text_delta(
-                    vendor_part_id=chunk.content_index, content=chunk.delta
-                )
+                maybe_event = self._parts_manager.handle_text_delta(vendor_part_id=chunk.item_id, content=chunk.delta)
                 if maybe_event is not None:  # pragma: no branch
                     yield maybe_event
 

pydantic_ai/providers/__init__.py

@@ -135,6 +135,10 @@ def infer_provider_class(provider: str) -> type[Provider[Any]]:  # noqa: C901
         from .github import GitHubProvider
 
         return GitHubProvider
+    elif provider == 'litellm':
+        from .litellm import LiteLLMProvider
+
+        return LiteLLMProvider
     else:  # pragma: no cover
         raise ValueError(f'Unknown provider: {provider}')
 

pydantic_ai/providers/bedrock.py

@@ -35,11 +35,19 @@ class BedrockModelProfile(ModelProfile):
     ALL FIELDS MUST BE `bedrock_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
     """
 
-    bedrock_supports_tool_choice: bool = True
+    bedrock_supports_tool_choice: bool = False
     bedrock_tool_result_format: Literal['text', 'json'] = 'text'
     bedrock_send_back_thinking_parts: bool = False
 
 
+def bedrock_amazon_model_profile(model_name: str) -> ModelProfile | None:
+    """Get the model profile for an Amazon model used via Bedrock."""
+    profile = amazon_model_profile(model_name)
+    if 'nova' in model_name:
+        return BedrockModelProfile(bedrock_supports_tool_choice=True).update(profile)
+    return profile
+
+
 class BedrockProvider(Provider[BaseClient]):
     """Provider for AWS Bedrock."""
 
@@ -58,13 +66,13 @@ class BedrockProvider(Provider[BaseClient]):
     def model_profile(self, model_name: str) -> ModelProfile | None:
         provider_to_profile: dict[str, Callable[[str], ModelProfile | None]] = {
             'anthropic': lambda model_name: BedrockModelProfile(
-                bedrock_supports_tool_choice=False, bedrock_send_back_thinking_parts=True
+                bedrock_supports_tool_choice=True, bedrock_send_back_thinking_parts=True
             ).update(anthropic_model_profile(model_name)),
             'mistral': lambda model_name: BedrockModelProfile(bedrock_tool_result_format='json').update(
                 mistral_model_profile(model_name)
             ),
             'cohere': cohere_model_profile,
-            'amazon': amazon_model_profile,
+            'amazon': bedrock_amazon_model_profile,
             'meta': meta_model_profile,
             'deepseek': deepseek_model_profile,
         }

pydantic_ai/providers/google_vertex.py

@@ -1,6 +1,7 @@
 from __future__ import annotations as _annotations
 
 import functools
+from asyncio import Lock
 from collections.abc import AsyncGenerator, Mapping
 from pathlib import Path
 from typing import Literal, overload
@@ -118,7 +119,7 @@ class GoogleVertexProvider(Provider[httpx.AsyncClient]):
 class _VertexAIAuth(httpx.Auth):
     """Auth class for Vertex AI API."""
 
-    _refresh_lock: anyio.Lock = anyio.Lock()
+    _refresh_lock: Lock = Lock()
 
     credentials: BaseCredentials | ServiceAccountCredentials | None
 

pydantic_ai/providers/groq.py

@@ -14,6 +14,7 @@ from pydantic_ai.profiles.groq import groq_model_profile
 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 @@ except ImportError as _import_error:  # pragma: no cover
     ) 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 @@ class GroqProvider(Provider[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():