Skip to content

Commit cf07dab

Browse files
authored
perf(pydantic_ai): simplify trace shaping (#357)
Avoid pre-serializing Pydantic AI values in tracing helpers. Keep integration-side shaping focused on readable span payloads and materializing binary content into Braintrust attachments, while leaving general serialization to Braintrust logging. Update private helper names and tests from serialize_* to shape_* to reflect the narrower responsibility. ``` I did a quick local microbenchmark against an approximation of the previous helper behavior using the pydantic-ai latest nox env. Results: ┌────────────────────────────────────┬──────────┬─────────┬────────────────────────┐ │ case │ old │ new │ speedup │ ├────────────────────────────────────┼──────────┼─────────┼────────────────────────┤ │ binary message shaping │ 24.83 µs │ 5.45 µs │ 4.6x │ ├────────────────────────────────────┼──────────┼─────────┼────────────────────────┤ │ text messages x5 │ 51.79 µs │ 7.31 µs │ 7.1x │ ├────────────────────────────────────┼──────────┼─────────┼────────────────────────┤ │ nested binary content part │ 15.01 µs │ 4.86 µs │ 3.1x │ ├────────────────────────────────────┼──────────┼─────────┼────────────────────────┤ │ model response text │ 13.84 µs │ 1.94 µs │ 7.1x │ ├────────────────────────────────────┼──────────┼─────────┼────────────────────────┤ │ model settings copy vs passthrough │ 0.79 µs │ ~0 µs │ effectively eliminated │ └────────────────────────────────────┴──────────┴─────────┴────────────────────────┘ ```
1 parent 91124f8 commit cf07dab

2 files changed

Lines changed: 151 additions & 147 deletions

File tree

py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py

Lines changed: 20 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -2333,19 +2333,19 @@ def test_setup_pydantic_ai_is_idempotent_across_new_patch_points():
23332333
assert agent_graph_module.ToolManager.__dict__[tool_method_name] is tool_method
23342334

23352335

2336-
def test_serialize_content_part_with_binary_content():
2337-
"""Unit test to verify _serialize_content_part handles BinaryContent correctly.
2336+
def test_shape_content_part_with_binary_content():
2337+
"""Unit test to verify _shape_content_part handles BinaryContent correctly.
23382338
2339-
This tests the direct serialization of BinaryContent objects and verifies
2339+
This tests the direct shaping of BinaryContent objects and verifies
23402340
they are converted to Braintrust Attachment objects.
23412341
"""
2342-
from braintrust.integrations.pydantic_ai.tracing import _serialize_content_part
2342+
from braintrust.integrations.pydantic_ai.tracing import _shape_content_part
23432343
from braintrust.logger import Attachment
23442344
from pydantic_ai.models.function import BinaryContent
23452345

23462346
# Test 1: Direct BinaryContent serialization
23472347
binary = BinaryContent(data=b"test pdf data", media_type="application/pdf")
2348-
result = _serialize_content_part(binary)
2348+
result = _shape_content_part(binary)
23492349

23502350
assert result is not None, "Should serialize BinaryContent"
23512351
assert result["type"] == "binary", "Should have type 'binary'"
@@ -2356,14 +2356,14 @@ def test_serialize_content_part_with_binary_content():
23562356
assert result["attachment"]._reference["content_type"] == "application/pdf"
23572357

23582358

2359-
def test_serialize_content_part_with_user_prompt_part():
2360-
"""Unit test to verify _serialize_content_part handles UserPromptPart with nested BinaryContent.
2359+
def test_shape_content_part_with_user_prompt_part():
2360+
"""Unit test to verify _shape_content_part handles UserPromptPart with nested BinaryContent.
23612361
23622362
This is the critical test for the bug: when a UserPromptPart has a content list
2363-
containing BinaryContent, we need to recursively serialize the content items
2363+
containing BinaryContent, we need to recursively shape the content items
23642364
so that BinaryContent is converted to Braintrust Attachment.
23652365
"""
2366-
from braintrust.integrations.pydantic_ai.tracing import _serialize_content_part
2366+
from braintrust.integrations.pydantic_ai.tracing import _shape_content_part
23672367
from braintrust.logger import Attachment
23682368
from pydantic_ai.messages import UserPromptPart
23692369
from pydantic_ai.models.function import BinaryContent
@@ -2373,18 +2373,18 @@ def test_serialize_content_part_with_user_prompt_part():
23732373
binary = BinaryContent(data=pdf_data, media_type="application/pdf")
23742374
user_prompt_part = UserPromptPart(content=[binary, "What is in this document?"])
23752375

2376-
# Serialize the UserPromptPart
2377-
result = _serialize_content_part(user_prompt_part)
2376+
# Shape the UserPromptPart
2377+
result = _shape_content_part(user_prompt_part)
23782378

2379-
# Verify the result is a dict with serialized content
2379+
# Verify the result is a dict with shaped content
23802380
assert isinstance(result, dict), f"Should return dict, got {type(result)}"
23812381
assert "content" in result, f"Should have 'content' key. Keys: {result.keys()}"
23822382

23832383
content = result["content"]
23842384
assert isinstance(content, list), f"Content should be a list, got {type(content)}"
23852385
assert len(content) == 2, f"Should have 2 content items, got {len(content)}"
23862386

2387-
# CRITICAL: First item should be serialized BinaryContent with Attachment
2387+
# CRITICAL: First item should be shaped BinaryContent with Attachment
23882388
binary_item = content[0]
23892389
assert isinstance(binary_item, dict), f"Binary item should be dict, got {type(binary_item)}"
23902390
assert binary_item.get("type") == "binary", f"Binary item should have type='binary'. Got: {binary_item}"
@@ -2398,13 +2398,13 @@ def test_serialize_content_part_with_user_prompt_part():
23982398
assert content[1] == "What is in this document?"
23992399

24002400

2401-
def test_serialize_messages_with_binary_content():
2402-
"""Unit test to verify _serialize_messages handles ModelRequest with BinaryContent in parts.
2401+
def test_shape_messages_with_binary_content():
2402+
"""Unit test to verify _shape_messages handles ModelRequest with BinaryContent in parts.
24032403
2404-
This tests the full message serialization path that's used for the chat span,
2404+
This tests the full message shaping path that's used for the chat span,
24052405
ensuring that nested BinaryContent in UserPromptPart is properly converted.
24062406
"""
2407-
from braintrust.integrations.pydantic_ai.tracing import _serialize_messages
2407+
from braintrust.integrations.pydantic_ai.tracing import _shape_messages
24082408
from braintrust.logger import Attachment
24092409
from pydantic_ai.messages import ModelRequest, UserPromptPart
24102410
from pydantic_ai.models.function import BinaryContent
@@ -2415,9 +2415,9 @@ def test_serialize_messages_with_binary_content():
24152415
user_prompt_part = UserPromptPart(content=[binary, "What is in this document?"])
24162416
model_request = ModelRequest(parts=[user_prompt_part])
24172417

2418-
# Serialize the messages
2418+
# Shape the messages
24192419
messages = [model_request]
2420-
result = _serialize_messages(messages)
2420+
result = _shape_messages(messages)
24212421

24222422
# Verify structure
24232423
assert len(result) == 1, f"Should have 1 message, got {len(result)}"
@@ -2435,7 +2435,7 @@ def test_serialize_messages_with_binary_content():
24352435
assert isinstance(content, list), f"Content should be list, got {type(content)}"
24362436
assert len(content) == 2, f"Should have 2 content items, got {len(content)}"
24372437

2438-
# CRITICAL: First content item should be serialized BinaryContent with Attachment
2438+
# CRITICAL: First content item should be shaped BinaryContent with Attachment
24392439
binary_item = content[0]
24402440
assert isinstance(binary_item, dict), f"Binary item should be dict, got {type(binary_item)}"
24412441
assert binary_item.get("type") == "binary", f"Binary item should have type='binary'. Got: {binary_item}"

0 commit comments

Comments
 (0)