from __future__ import annotations as _annotations import warnings from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Generator, Mapping from contextlib import asynccontextmanager, contextmanager from dataclasses import dataclass, field from datetime import datetime from typing import Any, Literal, cast, overload from pydantic import BaseModel, ValidationError from pydantic_core import from_json from typing_extensions import assert_never from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage from .._output import DEFAULT_OUTPUT_TOOL_NAME from .._run_context import RunContext from .._thinking_part import split_content_into_text_and_thinking from .._utils import generate_tool_call_id, guard_tool_call_id as _guard_tool_call_id, number_to_datetime from ..exceptions import ModelAPIError, UserError from ..messages import ( AudioUrl, BinaryContent, CachePoint, CompactionPart, DocumentUrl, FilePart, FinishReason, ImageUrl, ModelMessage, ModelRequest, ModelResponse, ModelResponsePart, ModelResponseStreamEvent, NativeToolCallPart, NativeToolReturnPart, RetryPromptPart, SpeechPart, SystemPromptPart, TextContent, TextPart, ThinkingPart, ToolCallPart, ToolReturnPart, UploadedFile, UserContent, UserPromptPart, VideoUrl, ) from ..native_tools import AbstractNativeTool, WebSearchTool from ..output import OutputObjectDefinition from ..profiles import DEFAULT_THINKING_TAGS, ModelProfile, ModelProfileSpec from ..profiles.groq import GROQ_GPT_OSS_REASONING_EFFORT_MAP from ..providers import Provider, infer_provider from ..settings import ModelSettings from ..tools import ToolDefinition from . import ( Model, ModelRequestParameters, StreamedResponse, _suggest_known_model_id_from_provider_error, # pyright: ignore[reportPrivateUsage] _unconverted_speech_part_error, # pyright: ignore[reportPrivateUsage] check_allow_model_requests, download_item, get_user_agent, ) from ._tool_choice import resolve_tool_choice try: from groq import NOT_GIVEN, APIConnectionError, APIError, APIStatusError, AsyncGroq, AsyncStream, NotGiven from groq.types import chat from groq.types.chat.chat_completion_content_part_image_param import ImageURL from groq.types.chat.chat_completion_message import ExecutedTool from groq.types.chat.chat_completion_named_tool_choice_param import ChatCompletionNamedToolChoiceParam from groq.types.chat.chat_completion_tool_choice_option_param import ChatCompletionToolChoiceOptionParam from groq.types.chat.completion_create_params import SearchSettings except ImportError as _import_error: raise ImportError( 'Please install `groq` to use the Groq model, ' 'you can use the `groq` optional group — `pip install "pydantic-ai-slim[groq]"`' ) from _import_error @contextmanager def _map_api_errors(model_name: str, model_id_namespace: str = 'groq') -> Generator[None]: try: yield except APIStatusError as e: if (status_code := e.status_code) >= 400: body: object | None = e.body suggested_model_id = None if _utils.is_str_dict(body) and _utils.is_str_dict(error := body.get('error')): if error.get('code') == 'model_not_found': suggested_model_id = _suggest_known_model_id_from_provider_error(model_id_namespace, model_name) raise ModelHTTPError( status_code=status_code, model_name=model_name, body=body, headers=dict(e.response.headers), suggested_model_id=suggested_model_id, ) from e raise ModelAPIError(model_name=model_name, message=e.message) from e # pragma: lax no cover except APIConnectionError as e: raise ModelAPIError(model_name=model_name, message=e.message) from e ProductionGroqModelNames = Literal[ 'llama-3.1-8b-instant', 'llama-3.3-70b-versatile', 'meta-llama/llama-guard-4-12b', 'openai/gpt-oss-120b', 'openai/gpt-oss-20b', 'whisper-large-v3', 'whisper-large-v3-turbo', ] """Production Groq models from .""" PreviewGroqModelNames = Literal[ 'meta-llama/llama-4-maverick-17b-128e-instruct', 'meta-llama/llama-prompt-guard-2-22m', 'meta-llama/llama-prompt-guard-2-86m', 'openai/gpt-oss-safeguard-20b', 'playai-tts', 'playai-tts-arabic', ] """Preview Groq models from .""" GroqModelName = str | ProductionGroqModelNames | PreviewGroqModelNames """Possible Groq model names. Since Groq supports a variety of models and the list changes frequently, we explicitly list the named models as of 2025-03-31 but allow any name in the type hints. See for an up to date list of models and more details. """ _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', } class GroqModelSettings(ModelSettings, total=False): """Settings used for a Groq model request.""" # ALL FIELDS MUST BE `groq_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS. groq_reasoning_format: Literal['hidden', 'raw', 'parsed'] """The format of the reasoning output. See [the Groq docs](https://console.groq.com/docs/reasoning#reasoning-format) for more details. """ groq_reasoning_effort: Literal['none', 'default', 'low', 'medium', 'high'] """The reasoning effort level. See [the Groq docs](https://console.groq.com/docs/reasoning#reasoning-effort) for more details. """ @dataclass(init=False) class GroqModel(Model[AsyncGroq]): """A model that uses the Groq API. Internally, this uses the [Groq Python client](https://github.com/groq/groq-python) to interact with the API. Apart from `__init__`, all methods are private or match those of the base class. """ _model_name: GroqModelName = field(repr=False) _provider: Provider[AsyncGroq] = field(repr=False) def __init__( self, model_name: GroqModelName, *, provider: Literal['groq', 'gateway'] | Provider[AsyncGroq] = 'groq', profile: ModelProfileSpec | None = None, settings: ModelSettings | None = None, ): """Initialize a Groq model. Args: model_name: The name of the Groq model to use. List of model names available [here](https://console.groq.com/docs/models). provider: The provider to use for authentication and API access. Can be either the string 'groq' or an instance of `Provider[AsyncGroq]`. If not provided, a new provider will be created using the other parameters. profile: The model profile to use. Defaults to a profile picked by the provider based on the model name. settings: Model-specific settings that will be used as defaults for this model. """ self._model_name = model_name if isinstance(provider, str): provider = infer_provider('gateway/groq' if provider == 'gateway' else provider) self._provider = provider super().__init__(settings=settings, profile=profile) @property def client(self) -> AsyncGroq: return self._provider.client @property def base_url(self) -> str: return str(self.client.base_url) @property def model_name(self) -> GroqModelName: """The model name.""" return self._model_name @property def system(self) -> str: """The model provider.""" return self._provider.name @classmethod def supported_native_tools(cls) -> frozenset[type[AbstractNativeTool]]: """Return the set of builtin tool types this model can handle.""" return frozenset({WebSearchTool}) async def request( self, messages: list[ModelMessage], model_settings: ModelSettings | None, model_request_parameters: ModelRequestParameters, ) -> ModelResponse: check_allow_model_requests() model_settings, model_request_parameters = self.prepare_request( model_settings, model_request_parameters, ) try: response = await self._completions_create( messages, False, cast(GroqModelSettings, model_settings or {}), model_request_parameters ) except ModelHTTPError as e: # 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. if (failed_generation := _parse_tool_use_failed_error(e.body)) is not None: if isinstance(failed_generation, _GroqToolUseFailedGeneration): part = ToolCallPart( tool_name=failed_generation.name, args=failed_generation.arguments, ) elif failed_generation: part = TextPart(content=failed_generation) else: # pragma: no cover part = None return ModelResponse( parts=[part] if part else [], model_name=e.model_name, provider_name=self._provider.name, provider_url=self.base_url, finish_reason='error', ) raise model_response = self._process_response(response) return model_response @asynccontextmanager async def request_stream( self, messages: list[ModelMessage], model_settings: ModelSettings | None, model_request_parameters: ModelRequestParameters, run_context: RunContext[Any] | None = None, ) -> AsyncGenerator[StreamedResponse]: check_allow_model_requests() model_settings, model_request_parameters = self.prepare_request( model_settings, model_request_parameters, ) response = await self._completions_create( messages, True, cast(GroqModelSettings, model_settings or {}), model_request_parameters ) async with response: yield await self._process_streamed_response(response, model_request_parameters) def _translate_thinking( self, model_settings: GroqModelSettings, model_request_parameters: ModelRequestParameters, disable_via_effort: bool, ) -> Literal['hidden', 'raw', 'parsed'] | NotGiven: """Get reasoning format, falling back to unified thinking when provider-specific setting is not set.""" if fmt := model_settings.get('groq_reasoning_format'): return fmt thinking = model_request_parameters.thinking if thinking is False: if disable_via_effort: # qwen3 truly disables reasoning via `reasoning_effort='none'` (set in `extra_body`), # so no reasoning format is needed. return NOT_GIVEN # Other reasoning models have no true disable; 'hidden' only suppresses reasoning output. return 'hidden' if thinking is not None: return 'parsed' return NOT_GIVEN @overload async def _completions_create( self, messages: list[ModelMessage], stream: Literal[True], model_settings: GroqModelSettings, model_request_parameters: ModelRequestParameters, ) -> AsyncStream[chat.ChatCompletionChunk]: pass @overload async def _completions_create( self, messages: list[ModelMessage], stream: Literal[False], model_settings: GroqModelSettings, model_request_parameters: ModelRequestParameters, ) -> chat.ChatCompletion: pass async def _completions_create( self, messages: list[ModelMessage], stream: bool, model_settings: GroqModelSettings, model_request_parameters: ModelRequestParameters, ) -> chat.ChatCompletion | AsyncStream[chat.ChatCompletionChunk]: tools, tool_choice = self._get_tool_choice(model_settings, model_request_parameters) native_tools, search_settings = self._get_native_tools(model_request_parameters) tools += native_tools groq_messages = await self._map_messages(messages, model_request_parameters) 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.get('supports_json_object_output', False) ): # pragma: no branch response_format = {'type': 'json_object'} extra_headers = dict(model_settings.get('extra_headers', {})) extra_headers.setdefault('User-Agent', get_user_agent()) # qwen3 truly disables reasoning by sending `reasoning_effort='none'` (in `extra_body`); `_translate_thinking` # then omits `reasoning_format`. The flag is computed once and shared with `_translate_thinking` so the two # stay aligned for the default path. An explicit `groq_reasoning_format` does still ride alongside # `reasoning_effort='none'` on the wire (it short-circuits `_translate_thinking`), but Groq accepts the pair # (HTTP 200) and lets `reasoning_effort='none'` win — reasoning is disabled and the format is ignored. disable_via_effort = model_request_parameters.thinking is False and self.profile.get( 'groq_supports_reasoning_disable', False ) extra_body = model_settings.get('extra_body') # `reasoning_effort` value sets are family-specific on Groq, so precedence is: # qwen3 disable (`'none'`) > explicit `groq_reasoning_effort` > unified `thinking` mapping > nothing. # The unified mapping only applies to graded families (gpt-oss: low/medium/high); qwen3's enable levels # have no gradation (only none/default) so unified thinking there just controls `reasoning_format` above. groq_reasoning_effort = model_settings.get('groq_reasoning_effort') if disable_via_effort and groq_reasoning_effort is not None: warnings.warn( "`thinking=False` disables reasoning on this Groq model via `reasoning_effort='none'`, " 'which overrides the `groq_reasoning_effort` setting; `groq_reasoning_effort` will be ignored.', UserWarning, ) effort = 'none' if disable_via_effort else groq_reasoning_effort if effort is None and self.profile.get('groq_supports_graded_reasoning_effort', False): thinking = model_request_parameters.thinking if thinking is True: effort = 'medium' elif thinking is not None and thinking is not False: effort = GROQ_GPT_OSS_REASONING_EFFORT_MAP[thinking] if effort is not None: # `reasoning_effort` isn't a named param in the Groq SDK, so it's passed via `extra_body`. # `ModelSettings.extra_body` is typed `object`, so narrowing it for the merge reads back as `Unknown`. merged_extra_body: dict[str, object] = {} if isinstance(extra_body, Mapping): merged_extra_body.update(extra_body) # pyright: ignore[reportUnknownArgumentType] merged_extra_body['reasoning_effort'] = effort extra_body = merged_extra_body with _map_api_errors(self.model_name, self._provider.model_id_namespace): return await self.client.chat.completions.create( model=self._model_name, messages=groq_messages, n=1, parallel_tool_calls=model_settings.get('parallel_tool_calls', NOT_GIVEN) if tools else NOT_GIVEN, tools=tools or NOT_GIVEN, 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), timeout=model_settings.get('timeout', NOT_GIVEN), seed=model_settings.get('seed', NOT_GIVEN), presence_penalty=model_settings.get('presence_penalty', NOT_GIVEN), reasoning_format=self._translate_thinking(model_settings, model_request_parameters, disable_via_effort), frequency_penalty=model_settings.get('frequency_penalty', NOT_GIVEN), logit_bias=model_settings.get('logit_bias', NOT_GIVEN), extra_headers=extra_headers, extra_body=extra_body, search_settings=search_settings, ) def _process_response(self, response: chat.ChatCompletion) -> ModelResponse: """Process a non-streamed response, and prepare a message to return.""" choice = response.choices[0] items: list[ModelResponsePart] = [] if choice.message.reasoning is not None: # NOTE: The `reasoning` field is only present if `groq_reasoning_format` is set to `parsed`. items.append(ThinkingPart(content=choice.message.reasoning)) if choice.message.executed_tools: for tool in choice.message.executed_tools: call_part, return_part = _map_executed_tool(tool, self.system) if call_part and return_part: # pragma: no branch items.append(call_part) items.append(return_part) if choice.message.content: # NOTE: The `` tag is only present if `groq_reasoning_format` is set to `raw`. items.extend( split_content_into_text_and_thinking( choice.message.content, self.profile.get('thinking_tags', DEFAULT_THINKING_TAGS) ) ) if choice.message.tool_calls is not None: for c in choice.message.tool_calls: items.append(ToolCallPart(tool_name=c.function.name, args=c.function.arguments, tool_call_id=c.id)) raw_finish_reason = choice.finish_reason provider_details: dict[str, Any] = {'finish_reason': raw_finish_reason} if response.created: # pragma: no branch provider_details['timestamp'] = number_to_datetime(response.created) finish_reason = _FINISH_REASON_MAP.get(raw_finish_reason) return ModelResponse( parts=items, usage=_map_usage(response, self._provider.name, self.base_url, response.model), model_name=response.model, provider_response_id=response.id, provider_name=self._provider.name, provider_url=self.base_url, finish_reason=finish_reason, provider_details=provider_details, ) async def _process_streamed_response( self, response: AsyncStream[chat.ChatCompletionChunk], model_request_parameters: ModelRequestParameters ) -> GroqStreamedResponse: """Process a streamed response, and prepare a streaming response to return.""" peekable_response: _utils.PeekableAsyncStream[ chat.ChatCompletionChunk, AsyncStream[chat.ChatCompletionChunk] ] = _utils.PeekableAsyncStream(response) with _map_api_errors(self.model_name, self._provider.model_id_namespace): first_chunk = await peekable_response.peek() if isinstance(first_chunk, _utils.Unset): raise UnexpectedModelBehavior( # pragma: no cover 'Streamed response ended without content or tool calls' ) return GroqStreamedResponse( model_request_parameters=model_request_parameters, _response=peekable_response, _model_name=first_chunk.model, _model_profile=self.profile, _provider_name=self._provider.name, _model_id_namespace=self._provider.model_id_namespace, _provider_url=self.base_url, _provider_timestamp=number_to_datetime(first_chunk.created), ) def _get_tool_choice( self, model_settings: GroqModelSettings, model_request_parameters: ModelRequestParameters, ) -> tuple[list[chat.ChatCompletionToolParam], ChatCompletionToolChoiceOptionParam | None]: """Determine which tools to send and the API tool_choice value. Returns: A tuple of (filtered_tools, tool_choice). """ resolved_tool_choice = resolve_tool_choice(model_settings, model_request_parameters) tool_defs = model_request_parameters.declared_tool_defs tool_choice: ChatCompletionToolChoiceOptionParam if resolved_tool_choice in ('auto', 'required', 'none'): # Use native 'none' mode to keep tool definitions cached while disabling tool calls tool_choice = resolved_tool_choice elif isinstance(resolved_tool_choice, tuple): tool_choice_mode, tool_names = resolved_tool_choice if tool_choice_mode == 'required' and len(tool_names) == 1: tool_choice = ChatCompletionNamedToolChoiceParam( type='function', function={'name': next(iter(tool_names))}, ) else: # Breaks caching, but Groq doesn't support limiting tools via API arg tool_defs = {k: v for k, v in tool_defs.items() if k in tool_names} tool_choice = tool_choice_mode else: assert_never(resolved_tool_choice) tools: list[chat.ChatCompletionToolParam] = [self._map_tool_definition(t) for t in tool_defs.values()] if not tools: return tools, None return tools, tool_choice def _get_native_tools( self, model_request_parameters: ModelRequestParameters ) -> tuple[list[chat.ChatCompletionToolParam], SearchSettings | NotGiven]: tools: list[chat.ChatCompletionToolParam] = [] search_settings: SearchSettings | NotGiven = NOT_GIVEN for tool in model_request_parameters.native_tools: if isinstance(tool, WebSearchTool): if not self.profile.get('groq_always_has_web_search_builtin_tool', False): raise UserError('`WebSearchTool` is not supported by Groq') # pragma: no cover # Compound models run web search implicitly, so we forward only the domain filters # (as `search_settings`) rather than emitting a tool definition. ss: SearchSettings = {} if tool.allowed_domains: ss['include_domains'] = tool.allowed_domains if tool.blocked_domains: ss['exclude_domains'] = tool.blocked_domains if ss: search_settings = ss else: # pragma: no cover raise UserError( f'`{tool.__class__.__name__}` is not supported by `GroqModel`. If it should be, please file an issue.' ) return tools, search_settings async def _map_messages( # noqa: C901 self, messages: list[ModelMessage], model_request_parameters: ModelRequestParameters ) -> list[chat.ChatCompletionMessageParam]: """Just maps a `pydantic_ai.Message` to a `groq.types.ChatCompletionMessageParam`.""" groq_messages: list[chat.ChatCompletionMessageParam] = [] for message in messages: if isinstance(message, ModelRequest): async for item in self._map_user_message(message): groq_messages.append(item) elif isinstance(message, ModelResponse): texts: list[str] = [] tool_calls: list[chat.ChatCompletionMessageToolCallParam] = [] for item in message.parts: if isinstance(item, TextPart): texts.append(item.content) elif isinstance(item, ToolCallPart): tool_calls.append(self._map_tool_call(item)) elif isinstance(item, ThinkingPart): start_tag, end_tag = self.profile.get('thinking_tags', DEFAULT_THINKING_TAGS) texts.append('\n'.join([start_tag, item.content, end_tag])) elif isinstance(item, NativeToolCallPart | NativeToolReturnPart): # pragma: no cover # These are not currently sent back pass elif isinstance(item, FilePart): # pragma: no cover # Files generated by models are not sent back to models that don't themselves generate files. pass elif isinstance(item, CompactionPart): # pragma: no cover # Compaction parts are not sent back to models that don't support compaction. pass elif isinstance(item, SpeechPart): # pragma: no cover # Unconverted realtime speech; `prepare_messages` turns these into `TextPart`s in `Model.prepare_messages`. raise _unconverted_speech_part_error() else: assert_never(item) message_param = chat.ChatCompletionAssistantMessageParam(role='assistant') if texts: # Note: model responses from this model should only have one text item, so the following # shouldn't merge multiple texts into one unless you switch models between runs: message_param['content'] = '\n\n'.join(texts) if tool_calls: message_param['tool_calls'] = tool_calls groq_messages.append(message_param) else: assert_never(message) if instruction_parts := self._get_instruction_parts(messages, model_request_parameters): system_prompt_count = next( (i for i, m in enumerate(groq_messages) if m.get('role') != 'system'), len(groq_messages) ) groq_messages[system_prompt_count:system_prompt_count] = [ chat.ChatCompletionSystemMessageParam(role='system', content=part.content) for part in instruction_parts ] return groq_messages @staticmethod def _map_tool_call(t: ToolCallPart) -> chat.ChatCompletionMessageToolCallParam: return chat.ChatCompletionMessageToolCallParam( id=_guard_tool_call_id(t=t), type='function', function={'name': t.tool_name, 'arguments': t.args_as_json_str()}, ) @staticmethod def _map_tool_definition(f: ToolDefinition) -> chat.ChatCompletionToolParam: return { 'type': 'function', 'function': { 'name': f.name, 'description': f.description or '', 'parameters': f.parameters_json_schema, }, } 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 async def _map_user_message(self, message: ModelRequest) -> AsyncIterable[chat.ChatCompletionMessageParam]: file_content: list[UserContent] = [] for part in message.parts: if isinstance(part, SystemPromptPart): yield chat.ChatCompletionSystemMessageParam(role='system', content=part.content) elif isinstance(part, UserPromptPart): yield await self._map_user_prompt(part) elif isinstance(part, ToolReturnPart): tool_text, tool_file_content = part.model_response_str_and_user_content() file_content.extend(tool_file_content) yield chat.ChatCompletionToolMessageParam( role='tool', tool_call_id=_guard_tool_call_id(t=part), content=tool_text, ) elif isinstance(part, RetryPromptPart): # pragma: no branch if part.tool_name is None: yield chat.ChatCompletionUserMessageParam(role='user', content=part.model_response()) else: yield chat.ChatCompletionToolMessageParam( role='tool', tool_call_id=_guard_tool_call_id(t=part), content=part.model_response(), ) if file_content: yield await self._map_user_prompt(UserPromptPart(content=file_content)) async def _map_user_prompt(self, part: UserPromptPart) -> chat.ChatCompletionUserMessageParam: content: str | list[chat.ChatCompletionContentPartParam] if isinstance(part.content, str): content = part.content else: content = [] for item in part.content: if isinstance(item, str | TextContent): text = item if isinstance(item, str) else item.content content.append(chat.ChatCompletionContentPartTextParam(text=text, type='text')) elif isinstance(item, ImageUrl): image_url_str = item.url if item.force_download: downloaded = await download_item(item, data_format='base64_uri') image_url_str = downloaded['data'] image_url: ImageURL = {'url': image_url_str} if metadata := item.vendor_metadata: image_url['detail'] = metadata.get('detail', 'auto') content.append(chat.ChatCompletionContentPartImageParam(image_url=image_url, type='image_url')) elif isinstance(item, BinaryContent): if item.is_image: image_url: ImageURL = {'url': item.data_uri} if metadata := item.vendor_metadata: image_url['detail'] = metadata.get('detail', 'auto') content.append(chat.ChatCompletionContentPartImageParam(image_url=image_url, type='image_url')) else: raise NotImplementedError('Only images are supported for BinaryContent in Groq user prompts') elif isinstance(item, DocumentUrl): raise NotImplementedError('DocumentUrl is not supported in Groq user prompts') elif isinstance(item, AudioUrl): raise NotImplementedError('AudioUrl is not supported in Groq user prompts') elif isinstance(item, VideoUrl): raise NotImplementedError('VideoUrl is not supported in Groq user prompts') elif isinstance(item, UploadedFile): raise NotImplementedError('UploadedFile is not supported in Groq user prompts') elif isinstance(item, CachePoint): pass else: assert_never(item) return chat.ChatCompletionUserMessageParam(role='user', content=content) @dataclass class GroqStreamedResponse(StreamedResponse): """Implementation of `StreamedResponse` for Groq models.""" _model_name: GroqModelName _model_profile: ModelProfile _response: _utils.PeekableAsyncStream[chat.ChatCompletionChunk, AsyncStream[chat.ChatCompletionChunk]] _provider_name: str _model_id_namespace: str _provider_url: str _provider_timestamp: datetime | None = None _timestamp: datetime = field(default_factory=_utils.now_utc) async def close_stream(self) -> None: await self._response.source.close() async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]: # noqa: C901 with _map_api_errors(self._model_name, self._model_id_namespace): try: executed_tool_call_id: str | None = None reasoning_index = 0 reasoning = False if self._provider_timestamp is not None: # pragma: no branch self.provider_details = {'timestamp': self._provider_timestamp} async for chunk in self._response: self._usage += _map_usage(chunk, self._provider_name, self._provider_url, self._model_name) if chunk.id: # pragma: no branch self.provider_response_id = chunk.id try: choice = chunk.choices[0] except IndexError: continue if raw_finish_reason := choice.finish_reason: self.provider_details = {**(self.provider_details or {}), 'finish_reason': raw_finish_reason} self.finish_reason = _FINISH_REASON_MAP.get(raw_finish_reason) if choice.delta.reasoning is not None: if not reasoning: reasoning_index += 1 reasoning = True # NOTE: The `reasoning` field is only present if `groq_reasoning_format` is set to `parsed`. for event in self._parts_manager.handle_thinking_delta( vendor_part_id=f'reasoning-{reasoning_index}', content=choice.delta.reasoning ): yield event else: reasoning = False if choice.delta.executed_tools: for tool in choice.delta.executed_tools: call_part, return_part = _map_executed_tool( tool, self.provider_name, streaming=True, tool_call_id=executed_tool_call_id ) if call_part: executed_tool_call_id = call_part.tool_call_id yield self._parts_manager.handle_part( vendor_part_id=f'executed_tools-{tool.index}-call', part=call_part ) if return_part: executed_tool_call_id = None yield self._parts_manager.handle_part( vendor_part_id=f'executed_tools-{tool.index}-return', part=return_part ) # Handle the text part of the response content = choice.delta.content if content: for event in self._parts_manager.handle_text_delta( vendor_part_id='content', content=content, thinking_tags=self._model_profile.get('thinking_tags', DEFAULT_THINKING_TAGS), ignore_leading_whitespace=self._model_profile.get( 'ignore_streamed_leading_whitespace', False ), ): yield 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.name if dtc.function is not None else None, args=dtc.function.arguments if dtc.function is not None else None, tool_call_id=dtc.id, ) if maybe_event is not None: yield maybe_event except APIError as e: # 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 if (failed_generation := _parse_tool_use_failed_error(e.body)) is not None: if isinstance(failed_generation, _GroqToolUseFailedGeneration): yield self._parts_manager.handle_tool_call_part( vendor_part_id='tool_use_failed', tool_name=failed_generation.name, args=failed_generation.arguments, ) elif failed_generation: # pragma: no cover # This branch is not covered because when streaming, the non-tool call text would already # have streamed before the `tool_use_failed` error which comes with `failed_generation=''`, # but we keep this here for (hypothetical?) cases where that field would not be empty. for event in self._parts_manager.handle_text_delta( vendor_part_id='tool_use_failed', content=failed_generation, thinking_tags=self._model_profile.get('thinking_tags', DEFAULT_THINKING_TAGS), ignore_leading_whitespace=self._model_profile.get( 'ignore_streamed_leading_whitespace', False ), ): yield event return raise @property def model_name(self) -> GroqModelName: """Get the model name of the response.""" return self._model_name @property def provider_name(self) -> str: """Get the provider name.""" return self._provider_name @property def provider_url(self) -> str: """Get the provider base URL.""" return self._provider_url @property def timestamp(self) -> datetime: """Get the timestamp of the response.""" return self._timestamp def _map_usage( completion: chat.ChatCompletionChunk | chat.ChatCompletion, provider: str, provider_url: str, model: str, ) -> usage.RequestUsage: response_usage = None if isinstance(completion, chat.ChatCompletion): response_usage = completion.usage elif completion.x_groq is not None: response_usage = completion.x_groq.usage if response_usage is None: return usage.RequestUsage() usage_data = response_usage.model_dump(exclude_none=True) details: dict[str, int] = { k: v for k, v in usage_data.items() if k not in {'prompt_tokens', 'completion_tokens', 'total_tokens'} if isinstance(v, int) } # Lift only `completion_tokens_details` (reasoning_tokens) into `details`: genai-prices # doesn't surface those, but it does map `prompt_tokens_details.cached_tokens` to # first-class `cache_read_tokens`, so lifting that too would double-report it. completion_tokens_details: dict[str, Any] = usage_data.get('completion_tokens_details') or {} details.update({k: v for k, v in completion_tokens_details.items() if isinstance(v, int)}) return usage.RequestUsage.extract( dict(model=model, usage=usage_data), provider=provider, provider_url=provider_url, provider_fallback='groq', details=details or None, ) 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: str 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 def _parse_tool_use_failed_error(body: Any) -> _GroqToolUseFailedGeneration | str | None: if not isinstance(body, dict): return None try: error = _GroqToolUseFailedError.model_validate(body) error = error.error except ValidationError: try: error = _GroqToolUseFailedInnerError.model_validate(body) except ValidationError: return None try: return _GroqToolUseFailedGeneration.model_validate_json(error.failed_generation) except ValidationError: return error.failed_generation def _map_executed_tool( tool: ExecutedTool, provider_name: str, streaming: bool = False, tool_call_id: str | None = None ) -> tuple[NativeToolCallPart | None, NativeToolReturnPart | None]: if tool.type == 'search': if tool.search_results and (tool.search_results.images or tool.search_results.results): results = tool.search_results.model_dump(mode='json') else: results = tool.output tool_call_id = tool_call_id or generate_tool_call_id() call_part = NativeToolCallPart( tool_name=WebSearchTool.kind, args=from_json(tool.arguments), provider_name=provider_name, tool_call_id=tool_call_id, ) return_part = NativeToolReturnPart( tool_name=WebSearchTool.kind, content=results, provider_name=provider_name, tool_call_id=tool_call_id, ) if streaming: if results: return None, return_part else: return call_part, None else: return call_part, return_part else: return None, None