pydantic-ai-slim 0.0.12__py3-none-any.whl → 0.0.14__py3-none-any.whl

pydantic_ai/models/groq.py

@@ -4,23 +4,28 @@ from collections.abc import AsyncIterator, Iterable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
+from itertools import chain
 from typing import Literal, overload
 
 from httpx import AsyncClient as AsyncHTTPClient
 from typing_extensions import assert_never
 
 from .. import UnexpectedModelBehavior, _utils, result
+from .._utils import guard_tool_call_id as _guard_tool_call_id
 from ..messages import (
-    ArgsJson,
-    Message,
-    ModelAnyResponse,
-    ModelStructuredResponse,
-    ModelTextResponse,
-    RetryPrompt,
-    ToolCall,
-    ToolReturn,
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    ModelResponsePart,
+    RetryPromptPart,
+    SystemPromptPart,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
 )
-from ..result import Cost
+from ..result import Usage
+from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
     AgentModel,
@@ -40,10 +45,11 @@ try:
 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[groq]'`"
+        "you can use the `groq` optional group — `pip install 'pydantic-ai-slim[groq]'`"
     ) from _import_error
 
 GroqModelName = Literal[
+    'llama-3.3-70b-versatile',
     'llama-3.1-70b-versatile',
     'llama3-groq-70b-8192-tool-use-preview',
     'llama3-groq-8b-8192-tool-use-preview',
@@ -149,28 +155,34 @@ class GroqAgentModel(AgentModel):
     allow_text_result: bool
     tools: list[chat.ChatCompletionToolParam]
 
-    async def request(self, messages: list[Message]) -> tuple[ModelAnyResponse, result.Cost]:
-        response = await self._completions_create(messages, False)
-        return self._process_response(response), _map_cost(response)
+    async def request(
+        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+    ) -> tuple[ModelResponse, result.Usage]:
+        response = await self._completions_create(messages, False, model_settings)
+        return self._process_response(response), _map_usage(response)
 
     @asynccontextmanager
-    async def request_stream(self, messages: list[Message]) -> AsyncIterator[EitherStreamedResponse]:
-        response = await self._completions_create(messages, True)
+    async def request_stream(
+        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+    ) -> AsyncIterator[EitherStreamedResponse]:
+        response = await self._completions_create(messages, True, model_settings)
         async with response:
             yield await self._process_streamed_response(response)
 
     @overload
     async def _completions_create(
-        self, messages: list[Message], stream: Literal[True]
+        self, messages: list[ModelMessage], stream: Literal[True], model_settings: ModelSettings | None
     ) -> AsyncStream[ChatCompletionChunk]:
         pass
 
     @overload
-    async def _completions_create(self, messages: list[Message], stream: Literal[False]) -> chat.ChatCompletion:
+    async def _completions_create(
+        self, messages: list[ModelMessage], stream: Literal[False], model_settings: ModelSettings | None
+    ) -> chat.ChatCompletion:
         pass
 
     async def _completions_create(
-        self, messages: list[Message], stream: bool
+        self, messages: list[ModelMessage], stream: bool, model_settings: ModelSettings | None
     ) -> chat.ChatCompletion | AsyncStream[ChatCompletionChunk]:
         # standalone function to make it easier to override
         if not self.tools:
@@ -180,37 +192,42 @@ class GroqAgentModel(AgentModel):
         else:
             tool_choice = 'auto'
 
-        groq_messages = [self._map_message(m) for m in messages]
+        groq_messages = list(chain(*(self._map_message(m) for m in messages)))
+
+        model_settings = model_settings or {}
+
         return await self.client.chat.completions.create(
             model=str(self.model_name),
             messages=groq_messages,
-            temperature=0.0,
             n=1,
             parallel_tool_calls=True if self.tools else NOT_GIVEN,
             tools=self.tools or NOT_GIVEN,
             tool_choice=tool_choice or NOT_GIVEN,
             stream=stream,
+            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),
         )
 
     @staticmethod
-    def _process_response(response: chat.ChatCompletion) -> ModelAnyResponse:
+    def _process_response(response: chat.ChatCompletion) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
         timestamp = datetime.fromtimestamp(response.created, tz=timezone.utc)
         choice = response.choices[0]
+        items: list[ModelResponsePart] = []
+        if choice.message.content is not None:
+            items.append(TextPart(choice.message.content))
         if choice.message.tool_calls is not None:
-            return ModelStructuredResponse(
-                [ToolCall.from_json(c.function.name, c.function.arguments, c.id) for c in choice.message.tool_calls],
-                timestamp=timestamp,
-            )
-        else:
-            assert choice.message.content is not None, choice
-            return ModelTextResponse(choice.message.content, timestamp=timestamp)
+            for c in choice.message.tool_calls:
+                items.append(ToolCallPart.from_raw_args(c.function.name, c.function.arguments, c.id))
+        return ModelResponse(items, timestamp=timestamp)
 
     @staticmethod
     async def _process_streamed_response(response: AsyncStream[ChatCompletionChunk]) -> EitherStreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
         timestamp: datetime | None = None
-        start_cost = Cost()
+        start_usage = Usage()
         # the first chunk may contain enough information so we iterate until we get either `tool_calls` or `content`
         while True:
             try:
@@ -218,62 +235,70 @@ class GroqAgentModel(AgentModel):
             except StopAsyncIteration as e:
                 raise UnexpectedModelBehavior('Streamed response ended without content or tool calls') from e
             timestamp = timestamp or datetime.fromtimestamp(chunk.created, tz=timezone.utc)
-            start_cost += _map_cost(chunk)
+            start_usage += _map_usage(chunk)
 
             if chunk.choices:
                 delta = chunk.choices[0].delta
 
                 if delta.content is not None:
-                    return GroqStreamTextResponse(delta.content, response, timestamp, start_cost)
+                    return GroqStreamTextResponse(delta.content, response, timestamp, start_usage)
                 elif delta.tool_calls is not None:
                     return GroqStreamStructuredResponse(
                         response,
                         {c.index: c for c in delta.tool_calls},
                         timestamp,
-                        start_cost,
+                        start_usage,
                     )
 
-    @staticmethod
-    def _map_message(message: Message) -> chat.ChatCompletionMessageParam:
+    @classmethod
+    def _map_message(cls, message: ModelMessage) -> Iterable[chat.ChatCompletionMessageParam]:
         """Just maps a `pydantic_ai.Message` to a `groq.types.ChatCompletionMessageParam`."""
-        if message.role == 'system':
-            # SystemPrompt ->
-            return chat.ChatCompletionSystemMessageParam(role='system', content=message.content)
-        elif message.role == 'user':
-            # UserPrompt ->
-            return chat.ChatCompletionUserMessageParam(role='user', content=message.content)
-        elif message.role == 'tool-return':
-            # ToolReturn ->
-            return chat.ChatCompletionToolMessageParam(
-                role='tool',
-                tool_call_id=_guard_tool_id(message),
-                content=message.model_response_str(),
-            )
-        elif message.role == 'retry-prompt':
-            # RetryPrompt ->
-            if message.tool_name is None:
-                return chat.ChatCompletionUserMessageParam(role='user', content=message.model_response())
-            else:
-                return chat.ChatCompletionToolMessageParam(
-                    role='tool',
-                    tool_call_id=_guard_tool_id(message),
-                    content=message.model_response(),
-                )
-        elif message.role == 'model-text-response':
-            # ModelTextResponse ->
-            return chat.ChatCompletionAssistantMessageParam(role='assistant', content=message.content)
-        elif message.role == 'model-structured-response':
-            assert (
-                message.role == 'model-structured-response'
-            ), f'Expected role to be "llm-tool-calls", got {message.role}'
-            # ModelStructuredResponse ->
-            return chat.ChatCompletionAssistantMessageParam(
-                role='assistant',
-                tool_calls=[_map_tool_call(t) for t in message.calls],
-            )
+        if isinstance(message, ModelRequest):
+            yield from cls._map_user_message(message)
+        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(_map_tool_call(item))
+                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
+            yield message_param
         else:
             assert_never(message)
 
+    @classmethod
+    def _map_user_message(cls, message: ModelRequest) -> Iterable[chat.ChatCompletionMessageParam]:
+        for part in message.parts:
+            if isinstance(part, SystemPromptPart):
+                yield chat.ChatCompletionSystemMessageParam(role='system', content=part.content)
+            elif isinstance(part, UserPromptPart):
+                yield chat.ChatCompletionUserMessageParam(role='user', content=part.content)
+            elif isinstance(part, ToolReturnPart):
+                yield chat.ChatCompletionToolMessageParam(
+                    role='tool',
+                    tool_call_id=_guard_tool_call_id(t=part, model_source='Groq'),
+                    content=part.model_response_str(),
+                )
+            elif isinstance(part, RetryPromptPart):
+                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, model_source='Groq'),
+                        content=part.model_response(),
+                    )
+
 
 @dataclass
 class GroqStreamTextResponse(StreamTextResponse):
@@ -282,7 +307,7 @@ class GroqStreamTextResponse(StreamTextResponse):
     _first: str | None
     _response: AsyncStream[ChatCompletionChunk]
     _timestamp: datetime
-    _cost: result.Cost
+    _usage: result.Usage
     _buffer: list[str] = field(default_factory=list, init=False)
 
     async def __anext__(self) -> None:
@@ -292,7 +317,7 @@ class GroqStreamTextResponse(StreamTextResponse):
             return None
 
         chunk = await self._response.__anext__()
-        self._cost = _map_cost(chunk)
+        self._usage = _map_usage(chunk)
 
         try:
             choice = chunk.choices[0]
@@ -309,8 +334,8 @@ class GroqStreamTextResponse(StreamTextResponse):
         yield from self._buffer
         self._buffer.clear()
 
-    def cost(self) -> Cost:
-        return self._cost
+    def usage(self) -> Usage:
+        return self._usage
 
     def timestamp(self) -> datetime:
         return self._timestamp
@@ -323,11 +348,11 @@ class GroqStreamStructuredResponse(StreamStructuredResponse):
     _response: AsyncStream[ChatCompletionChunk]
     _delta_tool_calls: dict[int, ChoiceDeltaToolCall]
     _timestamp: datetime
-    _cost: result.Cost
+    _usage: result.Usage
 
     async def __anext__(self) -> None:
         chunk = await self._response.__anext__()
-        self._cost = _map_cost(chunk)
+        self._usage = _map_usage(chunk)
 
         try:
             choice = chunk.choices[0]
@@ -349,38 +374,31 @@ class GroqStreamStructuredResponse(StreamStructuredResponse):
             else:
                 self._delta_tool_calls[new.index] = new
 
-    def get(self, *, final: bool = False) -> ModelStructuredResponse:
-        calls: list[ToolCall] = []
+    def get(self, *, final: bool = False) -> ModelResponse:
+        items: list[ModelResponsePart] = []
         for c in self._delta_tool_calls.values():
             if f := c.function:
                 if f.name is not None and f.arguments is not None:
-                    calls.append(ToolCall.from_json(f.name, f.arguments, c.id))
+                    items.append(ToolCallPart.from_raw_args(f.name, f.arguments, c.id))
 
-        return ModelStructuredResponse(calls, timestamp=self._timestamp)
+        return ModelResponse(items, timestamp=self._timestamp)
 
-    def cost(self) -> Cost:
-        return self._cost
+    def usage(self) -> Usage:
+        return self._usage
 
     def timestamp(self) -> datetime:
         return self._timestamp
 
 
-def _guard_tool_id(t: ToolCall | ToolReturn | RetryPrompt) -> str:
-    """Type guard that checks a `tool_id` is not None both for static typing and runtime."""
-    assert t.tool_id is not None, f'Groq requires `tool_id` to be set: {t}'
-    return t.tool_id
-
-
-def _map_tool_call(t: ToolCall) -> chat.ChatCompletionMessageToolCallParam:
-    assert isinstance(t.args, ArgsJson), f'Expected ArgsJson, got {t.args}'
+def _map_tool_call(t: ToolCallPart) -> chat.ChatCompletionMessageToolCallParam:
     return chat.ChatCompletionMessageToolCallParam(
-        id=_guard_tool_id(t),
+        id=_guard_tool_call_id(t=t, model_source='Groq'),
         type='function',
-        function={'name': t.tool_name, 'arguments': t.args.args_json},
+        function={'name': t.tool_name, 'arguments': t.args_as_json_str()},
     )
 
 
-def _map_cost(completion: ChatCompletionChunk | ChatCompletion) -> result.Cost:
+def _map_usage(completion: ChatCompletionChunk | ChatCompletion) -> result.Usage:
     usage = None
     if isinstance(completion, ChatCompletion):
         usage = completion.usage
@@ -388,9 +406,9 @@ def _map_cost(completion: ChatCompletionChunk | ChatCompletion) -> result.Cost:
         usage = completion.x_groq.usage
 
     if usage is None:
-        return result.Cost()
+        return result.Usage()
 
-    return result.Cost(
+    return result.Usage(
         request_tokens=usage.prompt_tokens,
         response_tokens=usage.completion_tokens,
         total_tokens=usage.total_tokens,