diff --git a/docs/agents.md b/docs/agents.md index e32b30c30..a16aacf0a 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -148,6 +148,7 @@ async def main(): model_name='gpt-4o', timestamp=datetime.datetime(...), kind='response', + vendor_id=None, ) ), End(data=FinalResult(output='Paris', tool_name=None, tool_call_id=None)), @@ -212,6 +213,7 @@ async def main(): model_name='gpt-4o', timestamp=datetime.datetime(...), kind='response', + vendor_id=None, ) ), End(data=FinalResult(output='Paris', tool_name=None, tool_call_id=None)), @@ -808,6 +810,7 @@ with capture_run_messages() as messages: # (2)! model_name='gpt-4o', timestamp=datetime.datetime(...), kind='response', + vendor_id=None, ), ModelRequest( parts=[ @@ -834,6 +837,7 @@ with capture_run_messages() as messages: # (2)! model_name='gpt-4o', timestamp=datetime.datetime(...), kind='response', + vendor_id=None, ), ] """ diff --git a/docs/message-history.md b/docs/message-history.md index 608113471..2a8243c55 100644 --- a/docs/message-history.md +++ b/docs/message-history.md @@ -67,6 +67,7 @@ print(result.all_messages()) model_name='gpt-4o', timestamp=datetime.datetime(...), kind='response', + vendor_id=None, ), ] """ @@ -145,6 +146,7 @@ async def main(): model_name='gpt-4o', timestamp=datetime.datetime(...), kind='response', + vendor_id=None, ), ] """ @@ -204,6 +206,7 @@ print(result2.all_messages()) model_name='gpt-4o', timestamp=datetime.datetime(...), kind='response', + vendor_id=None, ), ModelRequest( parts=[ @@ -226,6 +229,7 @@ print(result2.all_messages()) model_name='gpt-4o', timestamp=datetime.datetime(...), kind='response', + vendor_id=None, ), ] """ @@ -332,6 +336,7 @@ print(result2.all_messages()) model_name='gpt-4o', timestamp=datetime.datetime(...), kind='response', + vendor_id=None, ), ModelRequest( parts=[ @@ -354,6 +359,7 @@ print(result2.all_messages()) model_name='gemini-1.5-pro', timestamp=datetime.datetime(...), kind='response', + vendor_id=None, ), ] """ diff --git a/docs/tools.md b/docs/tools.md index edd99664f..5696f7f04 100644 --- a/docs/tools.md +++ b/docs/tools.md @@ -99,6 +99,7 @@ print(dice_result.all_messages()) model_name='gemini-1.5-flash', timestamp=datetime.datetime(...), kind='response', + vendor_id=None, ), ModelRequest( parts=[ @@ -125,6 +126,7 @@ print(dice_result.all_messages()) model_name='gemini-1.5-flash', timestamp=datetime.datetime(...), kind='response', + vendor_id=None, ), ModelRequest( parts=[ @@ -149,6 +151,7 @@ print(dice_result.all_messages()) model_name='gemini-1.5-flash', timestamp=datetime.datetime(...), kind='response', + vendor_id=None, ), ] """ diff --git a/pydantic_ai_slim/pydantic_ai/agent.py b/pydantic_ai_slim/pydantic_ai/agent.py index 25307ed02..98178e592 100644 --- a/pydantic_ai_slim/pydantic_ai/agent.py +++ b/pydantic_ai_slim/pydantic_ai/agent.py @@ -554,6 +554,7 @@ async def main(): model_name='gpt-4o', timestamp=datetime.datetime(...), kind='response', + vendor_id=None, ) ), End(data=FinalResult(output='Paris', tool_name=None, tool_call_id=None)), @@ -1717,6 +1718,7 @@ async def main(): model_name='gpt-4o', timestamp=datetime.datetime(...), kind='response', + vendor_id=None, ) ), End(data=FinalResult(output='Paris', tool_name=None, tool_call_id=None)), @@ -1855,6 +1857,7 @@ async def main(): model_name='gpt-4o', timestamp=datetime.datetime(...), kind='response', + vendor_id=None, ) ), End(data=FinalResult(output='Paris', tool_name=None, tool_call_id=None)), diff --git a/pydantic_ai_slim/pydantic_ai/messages.py b/pydantic_ai_slim/pydantic_ai/messages.py index 896e952c9..f0a71b100 100644 --- a/pydantic_ai_slim/pydantic_ai/messages.py +++ b/pydantic_ai_slim/pydantic_ai/messages.py @@ -566,6 +566,9 @@ class ModelResponse: kind: Literal['response'] = 'response' """Message type identifier, this is available on all parts as a discriminator.""" + vendor_id: str | None = None + """Vendor ID as specified by the model provider. This can be used to track the specific request to the model.""" + def otel_events(self) -> list[Event]: """Return OpenTelemetry events for the response.""" result: list[Event] = [] diff --git a/pydantic_ai_slim/pydantic_ai/models/bedrock.py b/pydantic_ai_slim/pydantic_ai/models/bedrock.py index a8882ba62..5996d7b64 100644 --- a/pydantic_ai_slim/pydantic_ai/models/bedrock.py +++ b/pydantic_ai_slim/pydantic_ai/models/bedrock.py @@ -269,7 +269,8 @@ async def _process_response(self, response: ConverseResponseTypeDef) -> tuple[Mo response_tokens=response['usage']['outputTokens'], total_tokens=response['usage']['totalTokens'], ) - return ModelResponse(items, model_name=self.model_name), u + vendor_id = response.get('ResponseMetadata', {}).get('RequestId', None) + return ModelResponse(items, model_name=self.model_name, vendor_id=vendor_id), u @overload async def _messages_create( diff --git a/tests/test_agent.py b/tests/test_agent.py index 720047d89..1e5523f0e 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -1697,6 +1697,7 @@ def test_binary_content_all_messages_json(): { 'parts': [{'content': 'success (no tool calls)', 'part_kind': 'text'}], 'model_name': 'test', + 'vendor_id': None, 'timestamp': IsStr(), 'kind': 'response', },