Fixes for excluding content (#2180)

Author: adtyavrdhn

pydantic_ai_slim/pydantic_ai/messages.py

@@ -411,9 +411,9 @@ class UserPromptPart:
     """Part type identifier, this is available on all parts as a discriminator."""
 
     def otel_event(self, settings: InstrumentationSettings) -> Event:
-        content: str | list[dict[str, Any] | str]
+        content: str | list[dict[str, Any] | str] | dict[str, Any]
         if isinstance(self.content, str):
-            content = self.content
+            content = self.content if settings.include_content else {'kind': 'text'}
         else:
             content = []
             for part in self.content:
@@ -743,7 +743,7 @@ def new_event_body():
                         'type': 'function',  # TODO https://github.com/pydantic/pydantic-ai/issues/888
                         'function': {
                             'name': part.tool_name,
-                            'arguments': part.args,
+                            **({'arguments': part.args} if settings.include_content else {}),
                         },
                     }
                 )

tests/models/test_instrumented.py

@@ -863,6 +863,7 @@ def test_messages_without_content(document_content: BinaryContent):
         ModelRequest(parts=[ToolReturnPart('tool', 'tool_return_content', 'tool_call_1')]),
         ModelRequest(parts=[RetryPromptPart('retry_prompt', tool_name='tool', tool_call_id='tool_call_2')]),
         ModelRequest(parts=[UserPromptPart(content=['user_prompt2', document_content])]),
+        ModelRequest(parts=[UserPromptPart('simple text prompt')]),
     ]
     settings = InstrumentationSettings(include_content=False)
     assert [InstrumentedModel.event_to_dict(e) for e in settings.messages_to_otel_events(messages)] == snapshot(
@@ -896,7 +897,7 @@ def test_messages_without_content(document_content: BinaryContent):
                     {
                         'id': IsStr(),
                         'type': 'function',
-                        'function': {'name': 'my_tool', 'arguments': {'a': 13, 'b': 4}},
+                        'function': {'name': 'my_tool'},
                     }
                 ],
                 'gen_ai.message.index': 3,
@@ -922,5 +923,11 @@ def test_messages_without_content(document_content: BinaryContent):
                 'gen_ai.message.index': 6,
                 'event.name': 'gen_ai.user.message',
             },
+            {
+                'content': {'kind': 'text'},
+                'role': 'user',
+                'gen_ai.message.index': 7,
+                'event.name': 'gen_ai.user.message',
+            },
         ]
     )

tests/test_logfire.py

@@ -415,6 +415,114 @@ class MyOutput:
     )
 
 
+@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
+def test_instructions_with_structured_output_exclude_content(get_logfire_summary: Callable[[], LogfireSummary]) -> None:
+    @dataclass
+    class MyOutput:
+        content: str
+
+    settings: InstrumentationSettings = InstrumentationSettings(include_content=False)
+
+    my_agent = Agent(model=TestModel(), instructions='Here are some instructions', instrument=settings)
+
+    result = my_agent.run_sync('Hello', output_type=MyOutput)
+    assert result.output == snapshot(MyOutput(content='a'))
+
+    summary = get_logfire_summary()
+    assert summary.attributes[0] == snapshot(
+        {
+            'model_name': 'test',
+            'agent_name': 'my_agent',
+            'logfire.msg': 'my_agent run',
+            'logfire.span_type': 'span',
+            'gen_ai.usage.input_tokens': 51,
+            'gen_ai.usage.output_tokens': 5,
+            'all_messages_events': IsJson(
+                snapshot(
+                    [
+                        {
+                            'content': 'Here are some instructions',
+                            'role': 'system',
+                            'event.name': 'gen_ai.system.message',
+                        },
+                        {
+                            'content': {'kind': 'text'},
+                            'role': 'user',
+                            'gen_ai.message.index': 0,
+                            'event.name': 'gen_ai.user.message',
+                        },
+                        {
+                            'role': 'assistant',
+                            'tool_calls': [
+                                {
+                                    'id': IsStr(),
+                                    'type': 'function',
+                                    'function': {'name': 'final_result'},
+                                }
+                            ],
+                            'gen_ai.message.index': 1,
+                            'event.name': 'gen_ai.assistant.message',
+                        },
+                        {
+                            'role': 'tool',
+                            'id': IsStr(),
+                            'name': 'final_result',
+                            'gen_ai.message.index': 2,
+                            'event.name': 'gen_ai.tool.message',
+                        },
+                    ]
+                )
+            ),
+            'final_result': '{"content": "a"}',
+            'logfire.json_schema': IsJson(
+                snapshot(
+                    {
+                        'type': 'object',
+                        'properties': {'all_messages_events': {'type': 'array'}, 'final_result': {'type': 'object'}},
+                    }
+                )
+            ),
+        }
+    )
+    chat_span_attributes = summary.attributes[1]
+    assert chat_span_attributes['events'] == snapshot(
+        IsJson(
+            snapshot(
+                [
+                    {
+                        'content': 'Here are some instructions',
+                        'role': 'system',
+                        'gen_ai.system': 'test',
+                        'event.name': 'gen_ai.system.message',
+                    },
+                    {
+                        'event.name': 'gen_ai.user.message',
+                        'content': {'kind': 'text'},
+                        'role': 'user',
+                        'gen_ai.message.index': 0,
+                        'gen_ai.system': 'test',
+                    },
+                    {
+                        'event.name': 'gen_ai.choice',
+                        'index': 0,
+                        'message': {
+                            'role': 'assistant',
+                            'tool_calls': [
+                                {
+                                    'id': IsStr(),
+                                    'type': 'function',
+                                    'function': {'name': 'final_result'},
+                                }
+                            ],
+                        },
+                        'gen_ai.system': 'test',
+                    },
+                ]
+            )
+        )
+    )
+
+
 def test_instrument_all():
     model = TestModel()
     agent = Agent()