Merge branch 'main' into update-versins

Author: medaminezghal

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -815,14 +815,21 @@ def capture_run_messages() -> Iterator[list[_messages.ModelMessage]]:
         If you call `run`, `run_sync`, or `run_stream` more than once within a single `capture_run_messages` context,
         `messages` will represent the messages exchanged during the first call only.
     """
+    token = None
+    messages: list[_messages.ModelMessage] = []
+
+    # Try to reuse existing message context if available
     try:
-        yield _messages_ctx_var.get().messages
+        messages = _messages_ctx_var.get().messages
     except LookupError:
-        messages: list[_messages.ModelMessage] = []
+        # No existing context, create a new one
         token = _messages_ctx_var.set(_RunMessages(messages))
-        try:
-            yield messages
-        finally:
+
+    try:
+        yield messages
+    finally:
+        # Clean up context if we created it
+        if token is not None:
             _messages_ctx_var.reset(token)
 
 

pydantic_ai_slim/pydantic_ai/agent.py

@@ -2218,12 +2218,18 @@ def _set_output_tool_return(self, return_content: str) -> list[_messages.ModelMe
         """
         if not self._output_tool_name:
             raise ValueError('Cannot set output tool return content when the return type is `str`.')
-        messages = deepcopy(self._state.message_history)
+
+        messages = self._state.message_history
         last_message = messages[-1]
-        for part in last_message.parts:
+        for idx, part in enumerate(last_message.parts):
             if isinstance(part, _messages.ToolReturnPart) and part.tool_name == self._output_tool_name:
-                part.content = return_content
-                return messages
+                # Only do deepcopy when we have to modify
+                copied_messages = list(messages)
+                copied_last = deepcopy(last_message)
+                copied_last.parts[idx].content = return_content  # type: ignore[misc]
+                copied_messages[-1] = copied_last
+                return copied_messages
+
         raise LookupError(f'No tool call found with tool name {self._output_tool_name!r}.')
 
     @overload

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -230,10 +230,16 @@
         'o1-mini-2024-09-12',
         'o1-preview',
         'o1-preview-2024-09-12',
+        'o1-pro',
+        'o1-pro-2025-03-19',
         'o3',
         'o3-2025-04-16',
+        'o3-deep-research',
+        'o3-deep-research-2025-06-26',
         'o3-mini',
         'o3-mini-2025-01-31',
+        'o3-pro',
+        'o3-pro-2025-06-10',
         'openai:chatgpt-4o-latest',
         'openai:codex-mini-latest',
         'openai:gpt-3.5-turbo',
@@ -283,12 +289,22 @@
         'openai:o1-mini-2024-09-12',
         'openai:o1-preview',
         'openai:o1-preview-2024-09-12',
+        'openai:o1-pro',
+        'openai:o1-pro-2025-03-19',
         'openai:o3',
         'openai:o3-2025-04-16',
+        'openai:o3-deep-research',
+        'openai:o3-deep-research-2025-06-26',
         'openai:o3-mini',
         'openai:o3-mini-2025-01-31',
         'openai:o4-mini',
         'openai:o4-mini-2025-04-16',
+        'openai:o4-mini-deep-research',
+        'openai:o4-mini-deep-research-2025-06-26',
+        'openai:o3-pro',
+        'openai:o3-pro-2025-06-10',
+        'openai:computer-use-preview',
+        'openai:computer-use-preview-2025-03-11',
         'test',
     ],
 )

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -50,7 +50,7 @@
 
 try:
     from openai import NOT_GIVEN, APIStatusError, AsyncOpenAI, AsyncStream, NotGiven
-    from openai.types import ChatModel, chat, responses
+    from openai.types import AllModels, chat, responses
     from openai.types.chat import (
         ChatCompletionChunk,
         ChatCompletionContentPartImageParam,
@@ -80,7 +80,7 @@
     'OpenAIModelName',
 )
 
-OpenAIModelName = Union[str, ChatModel]
+OpenAIModelName = Union[str, AllModels]
 """
 Possible OpenAI model names.
 

pydantic_ai_slim/pydantic_ai/usage.py

@@ -57,13 +57,19 @@ def __add__(self, other: Usage) -> Usage:
 
     def opentelemetry_attributes(self) -> dict[str, int]:
         """Get the token limits as OpenTelemetry attributes."""
-        result = {
-            'gen_ai.usage.input_tokens': self.request_tokens,
-            'gen_ai.usage.output_tokens': self.response_tokens,
-        }
-        for key, value in (self.details or {}).items():
-            result[f'gen_ai.usage.details.{key}'] = value  # pragma: no cover
-        return {k: v for k, v in result.items() if v}
+        result: dict[str, int] = {}
+        if self.request_tokens:
+            result['gen_ai.usage.input_tokens'] = self.request_tokens
+        if self.response_tokens:
+            result['gen_ai.usage.output_tokens'] = self.response_tokens
+        details = self.details
+        if details:
+            prefix = 'gen_ai.usage.details.'
+            for key, value in details.items():
+                # Skipping check for value since spec implies all detail values are relevant
+                if value:
+                    result[prefix + key] = value
+        return result
 
     def has_values(self) -> bool:
         """Whether any values are set and non-zero."""

tests/test_agent.py

@@ -2673,6 +2673,18 @@ def test_double_capture_run_messages() -> None:
     )
 
 
+def test_capture_run_messages_with_user_exception_does_not_contain_internal_errors() -> None:
+    """Test that user exceptions within capture_run_messages context have clean stack traces."""
+    agent = Agent('test')
+
+    try:
+        with capture_run_messages():
+            agent.run_sync('Hello')
+            raise ZeroDivisionError('division by zero')
+    except Exception as e:
+        assert e.__context__ is None
+
+
 def test_dynamic_false_no_reevaluate():
     """When dynamic is false (default), the system prompt is not reevaluated
     i.e: SystemPromptPart(

tests/test_usage_limits.py

@@ -176,6 +176,15 @@ async def delegate_to_other_agent2(ctx: RunContext[None], sentence: str) -> int:
     # confirm the usage from result2 is the sum of the usage from result1
     assert result2.usage() == functools.reduce(operator.add, run_1_usages)
 
+    result1_usage = result1.usage()
+    result1_usage.details = {'custom1': 10, 'custom2': 20, 'custom3': 0}
+    assert result1_usage.opentelemetry_attributes() == {
+        'gen_ai.usage.input_tokens': 103,
+        'gen_ai.usage.output_tokens': 13,
+        'gen_ai.usage.details.custom1': 10,
+        'gen_ai.usage.details.custom2': 20,
+    }
+
 
 async def test_multi_agent_usage_sync():
     """As in `test_multi_agent_usage_async`, with a sync tool."""