pydantic-ai-slim 0.0.13__py3-none-any.whl → 0.0.15__py3-none-any.whl

pydantic_ai/models/mistral.py

@@ -8,6 +8,7 @@ from datetime import datetime, timezone
 from itertools import chain
 from typing import Any, Callable, Literal, Union
 
+import pydantic_core
 from httpx import AsyncClient as AsyncHTTPClient, Timeout
 from typing_extensions import assert_never
 
@@ -26,7 +27,7 @@ from ..messages import (
     ToolReturnPart,
     UserPromptPart,
 )
-from ..result import Cost
+from ..result import Usage
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
@@ -39,7 +40,6 @@ from . import (
 )
 
 try:
-    from json_repair import repair_json
     from mistralai import (
         UNSET,
         CompletionChunk as MistralCompletionChunk,
@@ -156,10 +156,10 @@ class MistralAgentModel(AgentModel):
 
     async def request(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> tuple[ModelResponse, Cost]:
+    ) -> tuple[ModelResponse, Usage]:
         """Make a non-streaming request to the model from Pydantic AI call."""
         response = await self._completions_create(messages, model_settings)
-        return self._process_response(response), _map_cost(response)
+        return self._process_response(response), _map_usage(response)
 
     @asynccontextmanager
     async def request_stream(
@@ -198,11 +198,10 @@ class MistralAgentModel(AgentModel):
         """Create a streaming completion request to the Mistral model."""
         response: MistralEventStreamAsync[MistralCompletionEvent] | None
         mistral_messages = list(chain(*(self._map_message(m) for m in messages)))
-
         model_settings = model_settings or {}
 
         if self.result_tools and self.function_tools or self.function_tools:
-            # Function Calling Mode
+            # Function Calling
             response = await self.client.chat.stream_async(
                 model=str(self.model_name),
                 messages=mistral_messages,
@@ -218,9 +217,9 @@ class MistralAgentModel(AgentModel):
         elif self.result_tools:
             # Json Mode
             parameters_json_schemas = [tool.parameters_json_schema for tool in self.result_tools]
-
             user_output_format_message = self._generate_user_output_format(parameters_json_schemas)
             mistral_messages.append(user_output_format_message)
+
             response = await self.client.chat.stream_async(
                 model=str(self.model_name),
                 messages=mistral_messages,
@@ -270,12 +269,13 @@ class MistralAgentModel(AgentModel):
     @staticmethod
     def _process_response(response: MistralChatCompletionResponse) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
+        assert response.choices, 'Unexpected empty response choice.'
+
         if response.created:
             timestamp = datetime.fromtimestamp(response.created, tz=timezone.utc)
         else:
             timestamp = _now_utc()
 
-        assert response.choices, 'Unexpected empty response choice.'
         choice = response.choices[0]
         content = choice.message.content
         tool_calls = choice.message.tool_calls
@@ -297,7 +297,7 @@ class MistralAgentModel(AgentModel):
         response: MistralEventStreamAsync[MistralCompletionEvent],
     ) -> EitherStreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
-        start_cost = Cost()
+        start_usage = Usage()
 
         # Iterate until we get either `tool_calls` or `content` from the first chunk.
         while True:
@@ -307,7 +307,7 @@ class MistralAgentModel(AgentModel):
             except StopAsyncIteration as e:
                 raise UnexpectedModelBehavior('Streamed response ended without content or tool calls') from e
 
-            start_cost += _map_cost(chunk)
+            start_usage += _map_usage(chunk)
 
             if chunk.created:
                 timestamp = datetime.fromtimestamp(chunk.created, tz=timezone.utc)
@@ -329,11 +329,11 @@ class MistralAgentModel(AgentModel):
                         response,
                         content,
                         timestamp,
-                        start_cost,
+                        start_usage,
                     )
 
                 elif content:
-                    return MistralStreamTextResponse(content, response, timestamp, start_cost)
+                    return MistralStreamTextResponse(content, response, timestamp, start_usage)
 
     @staticmethod
     def _map_to_mistral_tool_call(t: ToolCallPart) -> MistralToolCall:
@@ -474,7 +474,7 @@ class MistralStreamTextResponse(StreamTextResponse):
     _first: str | None
     _response: MistralEventStreamAsync[MistralCompletionEvent]
     _timestamp: datetime
-    _cost: Cost
+    _usage: Usage
     _buffer: list[str] = field(default_factory=list, init=False)
 
     async def __anext__(self) -> None:
@@ -484,7 +484,7 @@ class MistralStreamTextResponse(StreamTextResponse):
             return None
 
         chunk = await self._response.__anext__()
-        self._cost += _map_cost(chunk.data)
+        self._usage += _map_usage(chunk.data)
 
         try:
             choice = chunk.data.choices[0]
@@ -502,8 +502,8 @@ class MistralStreamTextResponse(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
@@ -518,11 +518,11 @@ class MistralStreamStructuredResponse(StreamStructuredResponse):
     _response: MistralEventStreamAsync[MistralCompletionEvent]
     _delta_content: str | None
     _timestamp: datetime
-    _cost: Cost
+    _usage: Usage
 
     async def __anext__(self) -> None:
         chunk = await self._response.__anext__()
-        self._cost += _map_cost(chunk.data)
+        self._usage += _map_usage(chunk.data)
 
         try:
             choice = chunk.data.choices[0]
@@ -546,39 +546,31 @@ class MistralStreamStructuredResponse(StreamStructuredResponse):
                 calls.append(tool)
 
         elif self._delta_content and self._result_tools:
-            # NOTE: Params set for the most efficient and fastest way.
-            output_json = repair_json(self._delta_content, return_objects=True, skip_json_loads=True)
-            assert isinstance(
-                output_json, dict
-            ), f'Expected repair_json as type dict, invalid type: {type(output_json)}'
+            output_json: dict[str, Any] | None = pydantic_core.from_json(
+                self._delta_content, allow_partial='trailing-strings'
+            )
 
             if output_json:
                 for result_tool in self._result_tools.values():
-                    # NOTE: Additional verification to prevent JSON validation to crash in `result.py`
+                    # NOTE: Additional verification to prevent JSON validation to crash in `_result.py`
                     # Ensures required parameters in the JSON schema are respected, especially for stream-based return types.
-                    # For example, `return_type=list[str]` expects a 'response' key with value type array of str.
-                    # when `{"response":` then `repair_json` sets `{"response": ""}` (type not found default str)
-                    # when `{"response": {` then `repair_json` sets `{"response": {}}` (type found)
-                    # This ensures it's corrected to `{"response": {}}` and other required parameters and type.
-                    if not self._validate_required_json_shema(output_json, result_tool.parameters_json_schema):
+                    # Example with BaseModel and required fields.
+                    if not self._validate_required_json_schema(output_json, result_tool.parameters_json_schema):
                         continue
 
-                    tool = ToolCallPart.from_dict(
-                        tool_name=result_tool.name,
-                        args_dict=output_json,
-                    )
+                    tool = ToolCallPart.from_raw_args(result_tool.name, output_json)
                     calls.append(tool)
 
         return ModelResponse(calls, timestamp=self._timestamp)
 
-    def cost(self) -> Cost:
-        return self._cost
+    def usage(self) -> Usage:
+        return self._usage
 
     def timestamp(self) -> datetime:
         return self._timestamp
 
     @staticmethod
-    def _validate_required_json_shema(json_dict: dict[str, Any], json_schema: dict[str, Any]) -> bool:
+    def _validate_required_json_schema(json_dict: dict[str, Any], json_schema: dict[str, Any]) -> bool:
         """Validate that all required parameters in the JSON schema are present in the JSON dictionary."""
         required_params = json_schema.get('required', [])
         properties = json_schema.get('properties', {})
@@ -602,7 +594,7 @@ class MistralStreamStructuredResponse(StreamStructuredResponse):
 
             if isinstance(json_dict[param], dict) and 'properties' in param_schema:
                 nested_schema = param_schema
-                if not MistralStreamStructuredResponse._validate_required_json_shema(json_dict[param], nested_schema):
+                if not MistralStreamStructuredResponse._validate_required_json_schema(json_dict[param], nested_schema):
                     return False
 
         return True
@@ -633,29 +625,20 @@ def _map_mistral_to_pydantic_tool_call(tool_call: MistralToolCall) -> ToolCallPa
     tool_call_id = tool_call.id or None
     func_call = tool_call.function
 
-    if isinstance(func_call.arguments, str):
-        return ToolCallPart.from_json(
-            tool_name=func_call.name,
-            args_json=func_call.arguments,
-            tool_call_id=tool_call_id,
-        )
-    else:
-        return ToolCallPart.from_dict(
-            tool_name=func_call.name, args_dict=func_call.arguments, tool_call_id=tool_call_id
-        )
+    return ToolCallPart.from_raw_args(func_call.name, func_call.arguments, tool_call_id)
 
 
-def _map_cost(response: MistralChatCompletionResponse | MistralCompletionChunk) -> Cost:
-    """Maps a Mistral Completion Chunk or Chat Completion Response to a Cost."""
+def _map_usage(response: MistralChatCompletionResponse | MistralCompletionChunk) -> Usage:
+    """Maps a Mistral Completion Chunk or Chat Completion Response to a Usage."""
     if response.usage:
-        return Cost(
+        return Usage(
             request_tokens=response.usage.prompt_tokens,
             response_tokens=response.usage.completion_tokens,
             total_tokens=response.usage.total_tokens,
             details=None,
         )
     else:
-        return Cost()
+        return Usage()
 
 
 def _map_content(content: MistralOptionalNullable[MistralContent]) -> str | None:

pydantic_ai/models/openai.py

@@ -13,7 +13,6 @@ 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,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -25,7 +24,7 @@ from ..messages import (
     ToolReturnPart,
     UserPromptPart,
 )
-from ..result import Cost
+from ..result import Usage
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
@@ -147,9 +146,9 @@ class OpenAIAgentModel(AgentModel):
 
     async def request(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> tuple[ModelResponse, result.Cost]:
+    ) -> tuple[ModelResponse, result.Usage]:
         response = await self._completions_create(messages, False, model_settings)
-        return self._process_response(response), _map_cost(response)
+        return self._process_response(response), _map_usage(response)
 
     @asynccontextmanager
     async def request_stream(
@@ -211,14 +210,14 @@ class OpenAIAgentModel(AgentModel):
             items.append(TextPart(choice.message.content))
         if choice.message.tool_calls is not None:
             for c in choice.message.tool_calls:
-                items.append(ToolCallPart.from_json(c.function.name, c.function.arguments, c.id))
+                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:
@@ -227,19 +226,19 @@ class OpenAIAgentModel(AgentModel):
                 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 OpenAIStreamTextResponse(delta.content, response, timestamp, start_cost)
+                    return OpenAIStreamTextResponse(delta.content, response, timestamp, start_usage)
                 elif delta.tool_calls is not None:
                     return OpenAIStreamStructuredResponse(
                         response,
                         {c.index: c for c in delta.tool_calls},
                         timestamp,
-                        start_cost,
+                        start_usage,
                     )
                 # else continue until we get either delta.content or delta.tool_calls
 
@@ -302,7 +301,7 @@ class OpenAIStreamTextResponse(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:
@@ -312,7 +311,7 @@ class OpenAIStreamTextResponse(StreamTextResponse):
             return None
 
         chunk = await self._response.__anext__()
-        self._cost += _map_cost(chunk)
+        self._usage += _map_usage(chunk)
         try:
             choice = chunk.choices[0]
         except IndexError:
@@ -328,8 +327,8 @@ class OpenAIStreamTextResponse(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
@@ -342,11 +341,11 @@ class OpenAIStreamStructuredResponse(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]
         except IndexError:
@@ -372,37 +371,36 @@ class OpenAIStreamStructuredResponse(StreamStructuredResponse):
         for c in self._delta_tool_calls.values():
             if f := c.function:
                 if f.name is not None and f.arguments is not None:
-                    items.append(ToolCallPart.from_json(f.name, f.arguments, c.id))
+                    items.append(ToolCallPart.from_raw_args(f.name, f.arguments, c.id))
 
         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 _map_tool_call(t: ToolCallPart) -> chat.ChatCompletionMessageToolCallParam:
-    assert isinstance(t.args, ArgsJson), f'Expected ArgsJson, got {t.args}'
     return chat.ChatCompletionMessageToolCallParam(
         id=_guard_tool_call_id(t=t, model_source='OpenAI'),
         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(response: chat.ChatCompletion | ChatCompletionChunk) -> result.Cost:
+def _map_usage(response: chat.ChatCompletion | ChatCompletionChunk) -> result.Usage:
     usage = response.usage
     if usage is None:
-        return result.Cost()
+        return result.Usage()
     else:
         details: dict[str, int] = {}
         if usage.completion_tokens_details is not None:
             details.update(usage.completion_tokens_details.model_dump(exclude_none=True))
         if usage.prompt_tokens_details is not None:
             details.update(usage.prompt_tokens_details.model_dump(exclude_none=True))
-        return result.Cost(
+        return result.Usage(
             request_tokens=usage.prompt_tokens,
             response_tokens=usage.completion_tokens,
             total_tokens=usage.total_tokens,

pydantic_ai/models/test.py

@@ -21,7 +21,7 @@ from ..messages import (
     ToolCallPart,
     ToolReturnPart,
 )
-from ..result import Cost
+from ..result import Usage
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
@@ -31,6 +31,7 @@ from . import (
     StreamStructuredResponse,
     StreamTextResponse,
 )
+from .function import _estimate_string_usage, _estimate_usage  # pyright: ignore[reportPrivateUsage]
 
 
 @dataclass
@@ -131,15 +132,17 @@ class TestAgentModel(AgentModel):
 
     async def request(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> tuple[ModelResponse, Cost]:
-        return self._request(messages, model_settings), Cost()
+    ) -> tuple[ModelResponse, Usage]:
+        model_response = self._request(messages, model_settings)
+        usage = _estimate_usage([*messages, model_response])
+        return model_response, usage
 
     @asynccontextmanager
     async def request_stream(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
     ) -> AsyncIterator[EitherStreamedResponse]:
         msg = self._request(messages, model_settings)
-        cost = Cost()
+        usage = _estimate_usage(messages)
 
         # TODO: Rework this once we make StreamTextResponse more general
         texts: list[str] = []
@@ -153,9 +156,9 @@ class TestAgentModel(AgentModel):
                 assert_never(item)
 
         if texts:
-            yield TestStreamTextResponse('\n\n'.join(texts), cost)
+            yield TestStreamTextResponse('\n\n'.join(texts), usage)
         else:
-            yield TestStreamStructuredResponse(msg, cost)
+            yield TestStreamStructuredResponse(msg, usage)
 
     def gen_tool_args(self, tool_def: ToolDefinition) -> Any:
         return _JsonSchemaTestData(tool_def.parameters_json_schema, self.seed).generate()
@@ -164,7 +167,7 @@ class TestAgentModel(AgentModel):
         # if there are tools, the first thing we want to do is call all of them
         if self.tool_calls and not any(isinstance(m, ModelResponse) for m in messages):
             return ModelResponse(
-                parts=[ToolCallPart.from_dict(name, self.gen_tool_args(args)) for name, args in self.tool_calls]
+                parts=[ToolCallPart.from_raw_args(name, self.gen_tool_args(args)) for name, args in self.tool_calls]
             )
 
         if messages:
@@ -176,7 +179,7 @@ class TestAgentModel(AgentModel):
             if new_retry_names:
                 return ModelResponse(
                     parts=[
-                        ToolCallPart.from_dict(name, self.gen_tool_args(args))
+                        ToolCallPart.from_raw_args(name, self.gen_tool_args(args))
                         for name, args in self.tool_calls
                         if name in new_retry_names
                     ]
@@ -202,10 +205,10 @@ class TestAgentModel(AgentModel):
             custom_result_args = self.result.right
             result_tool = self.result_tools[self.seed % len(self.result_tools)]
             if custom_result_args is not None:
-                return ModelResponse(parts=[ToolCallPart.from_dict(result_tool.name, custom_result_args)])
+                return ModelResponse(parts=[ToolCallPart.from_raw_args(result_tool.name, custom_result_args)])
             else:
                 response_args = self.gen_tool_args(result_tool)
-                return ModelResponse(parts=[ToolCallPart.from_dict(result_tool.name, response_args)])
+                return ModelResponse(parts=[ToolCallPart.from_raw_args(result_tool.name, response_args)])
 
 
 @dataclass
@@ -213,7 +216,7 @@ class TestStreamTextResponse(StreamTextResponse):
     """A text response that streams test data."""
 
     _text: str
-    _cost: Cost
+    _usage: Usage
     _iter: Iterator[str] = field(init=False)
     _timestamp: datetime = field(default_factory=_utils.now_utc)
     _buffer: list[str] = field(default_factory=list, init=False)
@@ -228,14 +231,17 @@ class TestStreamTextResponse(StreamTextResponse):
         self._iter = iter(words)
 
     async def __anext__(self) -> None:
-        self._buffer.append(_utils.sync_anext(self._iter))
+        next_str = _utils.sync_anext(self._iter)
+        response_tokens = _estimate_string_usage(next_str)
+        self._usage += Usage(response_tokens=response_tokens, total_tokens=response_tokens)
+        self._buffer.append(next_str)
 
     def get(self, *, final: bool = False) -> Iterable[str]:
         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
@@ -246,7 +252,7 @@ class TestStreamStructuredResponse(StreamStructuredResponse):
     """A structured response that streams test data."""
 
     _structured_response: ModelResponse
-    _cost: Cost
+    _usage: Usage
     _iter: Iterator[None] = field(default_factory=lambda: iter([None]))
     _timestamp: datetime = field(default_factory=_utils.now_utc, init=False)
 
@@ -256,8 +262,8 @@ class TestStreamStructuredResponse(StreamStructuredResponse):
     def get(self, *, final: bool = False) -> ModelResponse:
         return self._structured_response
 
-    def cost(self) -> Cost:
-        return self._cost
+    def usage(self) -> Usage:
+        return self._usage
 
     def timestamp(self) -> datetime:
         return self._timestamp