pydantic-ai-slim 0.0.24__py3-none-any.whl → 0.0.25__py3-none-any.whl

pydantic_ai/messages.py

@@ -1,5 +1,6 @@
 from __future__ import annotations as _annotations
 
+import uuid
 from dataclasses import dataclass, field, replace
 from datetime import datetime
 from typing import Annotated, Any, Literal, Union, cast, overload
@@ -445,3 +446,33 @@ class PartDeltaEvent:
 
 ModelResponseStreamEvent = Annotated[Union[PartStartEvent, PartDeltaEvent], pydantic.Discriminator('event_kind')]
 """An event in the model response stream, either starting a new part or applying a delta to an existing one."""
+
+
+@dataclass
+class FunctionToolCallEvent:
+    """An event indicating the start to a call to a function tool."""
+
+    part: ToolCallPart
+    """The (function) tool call to make."""
+    call_id: str = field(init=False)
+    """An ID used for matching details about the call to its result. If present, defaults to the part's tool_call_id."""
+    event_kind: Literal['function_tool_call'] = 'function_tool_call'
+    """Event type identifier, used as a discriminator."""
+
+    def __post_init__(self):
+        self.call_id = self.part.tool_call_id or str(uuid.uuid4())
+
+
+@dataclass
+class FunctionToolResultEvent:
+    """An event indicating the result of a function tool call."""
+
+    result: ToolReturnPart | RetryPromptPart
+    """The result of the call to the function tool."""
+    call_id: str
+    """An ID used to match the result to its original call."""
+    event_kind: Literal['function_tool_result'] = 'function_tool_result'
+    """Event type identifier, used as a discriminator."""
+
+
+HandleResponseEvent = Annotated[Union[FunctionToolCallEvent, FunctionToolResultEvent], pydantic.Discriminator('kind')]

pydantic_ai/models/__init__.py

@@ -234,6 +234,8 @@ class StreamedResponse(ABC):
 
         This method should be implemented by subclasses to translate the vendor-specific stream of events into
         pydantic_ai-format events.
+
+        It should use the `_parts_manager` to handle deltas, and should update the `_usage` attributes as it goes.
         """
         raise NotImplementedError()
         # noinspection PyUnreachableCode
@@ -362,7 +364,6 @@ def infer_model(model: Model | KnownModelName) -> Model:
         raise UserError(f'Unknown model: {model}')
 
 
-@cache
 def cached_async_http_client(timeout: int = 600, connect: int = 5) -> httpx.AsyncClient:
     """Cached HTTPX async client so multiple agents and calls can share the same client.
 
@@ -373,6 +374,16 @@ def cached_async_http_client(timeout: int = 600, connect: int = 5) -> httpx.Asyn
     The default timeouts match those of OpenAI,
     see <https://github.com/openai/openai-python/blob/v1.54.4/src/openai/_constants.py#L9>.
     """
+    client = _cached_async_http_client(timeout=timeout, connect=connect)
+    if client.is_closed:
+        # This happens if the context manager is used, so we need to create a new client.
+        _cached_async_http_client.cache_clear()
+        client = _cached_async_http_client(timeout=timeout, connect=connect)
+    return client
+
+
+@cache
+def _cached_async_http_client(timeout: int = 600, connect: int = 5) -> httpx.AsyncClient:
     return httpx.AsyncClient(
         timeout=httpx.Timeout(timeout=timeout, connect=connect),
         headers={'User-Agent': get_user_agent()},

pydantic_ai/models/anthropic.py

@@ -272,64 +272,56 @@ class AnthropicModel(Model):
         anthropic_messages: list[MessageParam] = []
         for m in messages:
             if isinstance(m, ModelRequest):
-                for part in m.parts:
-                    if isinstance(part, SystemPromptPart):
-                        system_prompt += part.content
-                    elif isinstance(part, UserPromptPart):
-                        anthropic_messages.append(MessageParam(role='user', content=part.content))
-                    elif isinstance(part, ToolReturnPart):
-                        anthropic_messages.append(
-                            MessageParam(
-                                role='user',
-                                content=[
-                                    ToolResultBlockParam(
-                                        tool_use_id=_guard_tool_call_id(t=part, model_source='Anthropic'),
-                                        type='tool_result',
-                                        content=part.model_response_str(),
-                                        is_error=False,
-                                    )
-                                ],
-                            )
+                user_content_params: list[ToolResultBlockParam | TextBlockParam] = []
+                for request_part in m.parts:
+                    if isinstance(request_part, SystemPromptPart):
+                        system_prompt += request_part.content
+                    elif isinstance(request_part, UserPromptPart):
+                        text_block_param = TextBlockParam(type='text', text=request_part.content)
+                        user_content_params.append(text_block_param)
+                    elif isinstance(request_part, ToolReturnPart):
+                        tool_result_block_param = ToolResultBlockParam(
+                            tool_use_id=_guard_tool_call_id(t=request_part, model_source='Anthropic'),
+                            type='tool_result',
+                            content=request_part.model_response_str(),
+                            is_error=False,
                         )
-                    elif isinstance(part, RetryPromptPart):
-                        if part.tool_name is None:
-                            anthropic_messages.append(MessageParam(role='user', content=part.model_response()))
+                        user_content_params.append(tool_result_block_param)
+                    elif isinstance(request_part, RetryPromptPart):
+                        if request_part.tool_name is None:
+                            retry_param = TextBlockParam(type='text', text=request_part.model_response())
                         else:
-                            anthropic_messages.append(
-                                MessageParam(
-                                    role='user',
-                                    content=[
-                                        ToolResultBlockParam(
-                                            tool_use_id=_guard_tool_call_id(t=part, model_source='Anthropic'),
-                                            type='tool_result',
-                                            content=part.model_response(),
-                                            is_error=True,
-                                        ),
-                                    ],
-                                )
+                            retry_param = ToolResultBlockParam(
+                                tool_use_id=_guard_tool_call_id(t=request_part, model_source='Anthropic'),
+                                type='tool_result',
+                                content=request_part.model_response(),
+                                is_error=True,
                             )
+                        user_content_params.append(retry_param)
+                anthropic_messages.append(
+                    MessageParam(
+                        role='user',
+                        content=user_content_params,
+                    )
+                )
             elif isinstance(m, ModelResponse):
-                content: list[TextBlockParam | ToolUseBlockParam] = []
-                for item in m.parts:
-                    if isinstance(item, TextPart):
-                        content.append(TextBlockParam(text=item.content, type='text'))
+                assistant_content_params: list[TextBlockParam | ToolUseBlockParam] = []
+                for response_part in m.parts:
+                    if isinstance(response_part, TextPart):
+                        assistant_content_params.append(TextBlockParam(text=response_part.content, type='text'))
                     else:
-                        assert isinstance(item, ToolCallPart)
-                        content.append(self._map_tool_call(item))
-                anthropic_messages.append(MessageParam(role='assistant', content=content))
+                        tool_use_block_param = ToolUseBlockParam(
+                            id=_guard_tool_call_id(t=response_part, model_source='Anthropic'),
+                            type='tool_use',
+                            name=response_part.tool_name,
+                            input=response_part.args_as_dict(),
+                        )
+                        assistant_content_params.append(tool_use_block_param)
+                anthropic_messages.append(MessageParam(role='assistant', content=assistant_content_params))
             else:
                 assert_never(m)
         return system_prompt, anthropic_messages
 
-    @staticmethod
-    def _map_tool_call(t: ToolCallPart) -> ToolUseBlockParam:
-        return ToolUseBlockParam(
-            id=_guard_tool_call_id(t=t, model_source='Anthropic'),
-            type='tool_use',
-            name=t.tool_name,
-            input=t.args_as_dict(),
-        )
-
     @staticmethod
     def _map_tool_definition(f: ToolDefinition) -> ToolParam:
         return {

pydantic_ai/models/cohere.py

@@ -124,7 +124,7 @@ class CohereModel(Model):
             assert api_key is None, 'Cannot provide both `cohere_client` and `api_key`'
             self.client = cohere_client
         else:
-            self.client = AsyncClientV2(api_key=api_key, httpx_client=http_client)  # type: ignore
+            self.client = AsyncClientV2(api_key=api_key, httpx_client=http_client)
 
     async def request(
         self,

pydantic_ai/models/function.py

@@ -109,9 +109,9 @@ class FunctionModel(Model):
             model_settings,
         )
 
-        assert (
-            self.stream_function is not None
-        ), 'FunctionModel must receive a `stream_function` to support streamed requests'
+        assert self.stream_function is not None, (
+            'FunctionModel must receive a `stream_function` to support streamed requests'
+        )
 
         response_stream = PeekableAsyncStream(self.stream_function(messages, agent_info))
 

pydantic_ai/models/gemini.py

@@ -254,7 +254,7 @@ class GeminiModel(Model):
         async for chunk in aiter_bytes:
             content.extend(chunk)
             responses = _gemini_streamed_response_ta.validate_json(
-                content,
+                _ensure_decodeable(content),
                 experimental_allow_partial='trailing-strings',
             )
             if responses:
@@ -370,7 +370,7 @@ class GeminiStreamedResponse(StreamedResponse):
             self._content.extend(chunk)
 
             gemini_responses = _gemini_streamed_response_ta.validate_json(
-                self._content,
+                _ensure_decodeable(self._content),
                 experimental_allow_partial='trailing-strings',
             )
 
@@ -774,3 +774,19 @@ class _GeminiJsonSchema:
 
         if items_schema := schema.get('items'):  # pragma: no branch
             self._simplify(items_schema, refs_stack)
+
+
+def _ensure_decodeable(content: bytearray) -> bytearray:
+    """Trim any invalid unicode point bytes off the end of a bytearray.
+
+    This is necessary before attempting to parse streaming JSON bytes.
+
+    This is a temporary workaround until https://github.com/pydantic/pydantic-core/issues/1633 is resolved
+    """
+    while True:
+        try:
+            content.decode()
+        except UnicodeDecodeError:
+            content = content[:-1]  # this will definitely succeed before we run out of bytes
+        else:
+            return content

pydantic_ai/models/instrumented.py

@@ -0,0 +1,225 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Iterator
+from contextlib import asynccontextmanager, contextmanager
+from dataclasses import dataclass, field
+from functools import partial
+from typing import Any, Callable, Literal
+
+import logfire_api
+from opentelemetry._events import Event, EventLogger, EventLoggerProvider, get_event_logger_provider
+from opentelemetry.trace import Tracer, TracerProvider, get_tracer_provider
+
+from ..messages import (
+    ModelMessage,
+    ModelRequest,
+    ModelRequestPart,
+    ModelResponse,
+    RetryPromptPart,
+    SystemPromptPart,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
+)
+from ..settings import ModelSettings
+from ..usage import Usage
+from . import KnownModelName, Model, ModelRequestParameters, StreamedResponse
+from .wrapper import WrapperModel
+
+MODEL_SETTING_ATTRIBUTES: tuple[
+    Literal[
+        'max_tokens',
+        'top_p',
+        'seed',
+        'temperature',
+        'presence_penalty',
+        'frequency_penalty',
+    ],
+    ...,
+] = (
+    'max_tokens',
+    'top_p',
+    'seed',
+    'temperature',
+    'presence_penalty',
+    'frequency_penalty',
+)
+
+NOT_GIVEN = object()
+
+
+@dataclass
+class InstrumentedModel(WrapperModel):
+    """Model which is instrumented with logfire."""
+
+    tracer: Tracer = field(repr=False)
+    event_logger: EventLogger = field(repr=False)
+
+    def __init__(
+        self,
+        wrapped: Model | KnownModelName,
+        tracer_provider: TracerProvider | None = None,
+        event_logger_provider: EventLoggerProvider | None = None,
+    ):
+        super().__init__(wrapped)
+        tracer_provider = tracer_provider or get_tracer_provider()
+        event_logger_provider = event_logger_provider or get_event_logger_provider()
+        self.tracer = tracer_provider.get_tracer('pydantic-ai')
+        self.event_logger = event_logger_provider.get_event_logger('pydantic-ai')
+
+    @classmethod
+    def from_logfire(
+        cls,
+        wrapped: Model | KnownModelName,
+        logfire_instance: logfire_api.Logfire = logfire_api.DEFAULT_LOGFIRE_INSTANCE,
+    ) -> InstrumentedModel:
+        if hasattr(logfire_instance.config, 'get_event_logger_provider'):
+            event_provider = logfire_instance.config.get_event_logger_provider()
+        else:
+            event_provider = None
+        tracer_provider = logfire_instance.config.get_tracer_provider()
+        return cls(wrapped, tracer_provider, event_provider)
+
+    async def request(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> tuple[ModelResponse, Usage]:
+        with self._instrument(messages, model_settings) as finish:
+            response, usage = await super().request(messages, model_settings, model_request_parameters)
+            finish(response, usage)
+            return response, usage
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> AsyncIterator[StreamedResponse]:
+        with self._instrument(messages, model_settings) as finish:
+            response_stream: StreamedResponse | None = None
+            try:
+                async with super().request_stream(
+                    messages, model_settings, model_request_parameters
+                ) as response_stream:
+                    yield response_stream
+            finally:
+                if response_stream:
+                    finish(response_stream.get(), response_stream.usage())
+
+    @contextmanager
+    def _instrument(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+    ) -> Iterator[Callable[[ModelResponse, Usage], None]]:
+        operation = 'chat'
+        model_name = self.model_name
+        span_name = f'{operation} {model_name}'
+        system = getattr(self.wrapped, 'system', '') or self.wrapped.__class__.__name__.removesuffix('Model').lower()
+        system = {'google-gla': 'gemini', 'google-vertex': 'vertex_ai', 'mistral': 'mistral_ai'}.get(system, system)
+        # TODO Missing attributes:
+        #  - server.address: requires a Model.base_url abstract method or similar
+        #  - server.port: to parse from the base_url
+        #  - error.type: unclear if we should do something here or just always rely on span exceptions
+        #  - gen_ai.request.stop_sequences/top_k: model_settings doesn't include these
+        attributes: dict[str, Any] = {
+            'gen_ai.operation.name': operation,
+            'gen_ai.system': system,
+            'gen_ai.request.model': model_name,
+        }
+
+        if model_settings:
+            for key in MODEL_SETTING_ATTRIBUTES:
+                if (value := model_settings.get(key, NOT_GIVEN)) is not NOT_GIVEN:
+                    attributes[f'gen_ai.request.{key}'] = value
+
+        emit_event = partial(self._emit_event, system)
+
+        with self.tracer.start_as_current_span(span_name, attributes=attributes) as span:
+            if span.is_recording():
+                for message in messages:
+                    if isinstance(message, ModelRequest):
+                        for part in message.parts:
+                            event_name, body = _request_part_body(part)
+                            if event_name:
+                                emit_event(event_name, body)
+                    elif isinstance(message, ModelResponse):
+                        for body in _response_bodies(message):
+                            emit_event('gen_ai.assistant.message', body)
+
+            def finish(response: ModelResponse, usage: Usage):
+                if not span.is_recording():
+                    return
+
+                for response_body in _response_bodies(response):
+                    if response_body:
+                        emit_event(
+                            'gen_ai.choice',
+                            {
+                                # TODO finish_reason
+                                'index': 0,
+                                'message': response_body,
+                            },
+                        )
+                span.set_attributes(
+                    {
+                        k: v
+                        for k, v in {
+                            # TODO finish_reason (https://github.com/open-telemetry/semantic-conventions/issues/1277), id
+                            #  https://github.com/pydantic/pydantic-ai/issues/886
+                            'gen_ai.response.model': response.model_name or model_name,
+                            'gen_ai.usage.input_tokens': usage.request_tokens,
+                            'gen_ai.usage.output_tokens': usage.response_tokens,
+                        }.items()
+                        if v is not None
+                    }
+                )
+
+            yield finish
+
+    def _emit_event(self, system: str, event_name: str, body: dict[str, Any]) -> None:
+        self.event_logger.emit(Event(event_name, body=body, attributes={'gen_ai.system': system}))
+
+
+def _request_part_body(part: ModelRequestPart) -> tuple[str, dict[str, Any]]:
+    if isinstance(part, SystemPromptPart):
+        return 'gen_ai.system.message', {'content': part.content, 'role': 'system'}
+    elif isinstance(part, UserPromptPart):
+        return 'gen_ai.user.message', {'content': part.content, 'role': 'user'}
+    elif isinstance(part, ToolReturnPart):
+        return 'gen_ai.tool.message', {'content': part.content, 'role': 'tool', 'id': part.tool_call_id}
+    elif isinstance(part, RetryPromptPart):
+        if part.tool_name is None:
+            return 'gen_ai.user.message', {'content': part.model_response(), 'role': 'user'}
+        else:
+            return 'gen_ai.tool.message', {'content': part.model_response(), 'role': 'tool', 'id': part.tool_call_id}
+    else:
+        return '', {}
+
+
+def _response_bodies(message: ModelResponse) -> list[dict[str, Any]]:
+    body: dict[str, Any] = {'role': 'assistant'}
+    result = [body]
+    for part in message.parts:
+        if isinstance(part, ToolCallPart):
+            body.setdefault('tool_calls', []).append(
+                {
+                    'id': part.tool_call_id,
+                    'type': 'function',  # TODO https://github.com/pydantic/pydantic-ai/issues/888
+                    'function': {
+                        'name': part.tool_name,
+                        'arguments': part.args,
+                    },
+                }
+            )
+        elif isinstance(part, TextPart):
+            if body.get('content'):
+                body = {'role': 'assistant'}
+                result.append(body)
+            body['content'] = part.content
+
+    return result

pydantic_ai/models/mistral.py

@@ -134,9 +134,6 @@ class MistralModel(Model):
             api_key = os.getenv('MISTRAL_API_KEY') if api_key is None else api_key
             self.client = Mistral(api_key=api_key, async_client=http_client or cached_async_http_client())
 
-    def name(self) -> str:
-        return f'mistral:{self._model_name}'
-
     async def request(
         self,
         messages: list[ModelMessage],

pydantic_ai/models/openai.py

@@ -119,9 +119,9 @@ class OpenAIModel(Model):
         """
         self._model_name = model_name
         # This is a workaround for the OpenAI client requiring an API key, whilst locally served,
-        # openai compatible models do not always need an API key.
+        # openai compatible models do not always need an API key, but a placeholder (non-empty) key is required.
         if api_key is None and 'OPENAI_API_KEY' not in os.environ and base_url is not None and openai_client is None:
-            api_key = ''
+            api_key = 'api-key-not-set'
 
         if openai_client is not None:
             assert http_client is None, 'Cannot provide both `openai_client` and `http_client`'
@@ -135,9 +135,6 @@ class OpenAIModel(Model):
         self.system_prompt_role = system_prompt_role
         self._system = system
 
-    def name(self) -> str:
-        return f'openai:{self._model_name}'
-
     async def request(
         self,
         messages: list[ModelMessage],

pydantic_ai/models/test.py

@@ -130,15 +130,15 @@ class TestModel(Model):
 
     def _get_result(self, model_request_parameters: ModelRequestParameters) -> _TextResult | _FunctionToolResult:
         if self.custom_result_text is not None:
-            assert (
-                model_request_parameters.allow_text_result
-            ), 'Plain response not allowed, but `custom_result_text` is set.'
+            assert model_request_parameters.allow_text_result, (
+                'Plain response not allowed, but `custom_result_text` is set.'
+            )
             assert self.custom_result_args is None, 'Cannot set both `custom_result_text` and `custom_result_args`.'
             return _TextResult(self.custom_result_text)
         elif self.custom_result_args is not None:
-            assert (
-                model_request_parameters.result_tools is not None
-            ), 'No result tools provided, but `custom_result_args` is set.'
+            assert model_request_parameters.result_tools is not None, (
+                'No result tools provided, but `custom_result_args` is set.'
+            )
             result_tool = model_request_parameters.result_tools[0]
 
             if k := result_tool.outer_typed_dict_key:

pydantic_ai/models/wrapper.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import Any
+
+from ..messages import ModelMessage, ModelResponse
+from ..settings import ModelSettings
+from ..usage import Usage
+from . import KnownModelName, Model, ModelRequestParameters, StreamedResponse, infer_model
+
+
+@dataclass(init=False)
+class WrapperModel(Model):
+    """Model which wraps another model."""
+
+    wrapped: Model
+
+    def __init__(self, wrapped: Model | KnownModelName):
+        self.wrapped = infer_model(wrapped)
+
+    async def request(self, *args: Any, **kwargs: Any) -> tuple[ModelResponse, Usage]:
+        return await self.wrapped.request(*args, **kwargs)
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> AsyncIterator[StreamedResponse]:
+        async with self.wrapped.request_stream(messages, model_settings, model_request_parameters) as response_stream:
+            yield response_stream
+
+    @property
+    def model_name(self) -> str:
+        return self.wrapped.model_name
+
+    @property
+    def system(self) -> str | None:
+        return self.wrapped.system
+
+    def __getattr__(self, item: str):
+        return getattr(self.wrapped, item)