diff --git a/py/src/braintrust/integrations/pydantic_ai/cassettes/1.10.0/test_agent_with_instructions_and_dynamic_system_prompt.yaml b/py/src/braintrust/integrations/pydantic_ai/cassettes/1.10.0/test_agent_with_instructions_and_dynamic_system_prompt.yaml new file mode 100644 index 00000000..39d7ce5a --- /dev/null +++ b/py/src/braintrust/integrations/pydantic_ai/cassettes/1.10.0/test_agent_with_instructions_and_dynamic_system_prompt.yaml @@ -0,0 +1,104 @@ +interactions: +- request: + body: '{"messages":[{"content":"Answer with only the number requested by the user.","role":"system"},{"role":"system","content":"The + user is currently taking a math quiz."},{"role":"user","content":"What is 2+2? + Answer with just the number."}],"model":"gpt-4o-mini","max_completion_tokens":100,"stream":false}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '302' + content-type: + - application/json + host: + - api.openai.com + user-agent: + - pydantic-ai/1.10.0 + x-stainless-arch: + - arm64 + x-stainless-async: + - async:asyncio + x-stainless-lang: + - python + x-stainless-os: + - MacOS + x-stainless-package-version: + - 3.8.0 + x-stainless-read-timeout: + - '600' + x-stainless-retry-count: + - '0' + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.14.6 + method: POST + uri: https://api.openai.com/v1/chat/completions + response: + body: + string: "{\n \"id\": \"chatcmpl-EKSb4SbtYB0fiLGdkgdeIqTZXUUOz\",\n \"object\": + \"chat.completion\",\n \"created\": 1788545194,\n \"model\": \"gpt-4o-mini-2024-07-18\",\n + \ \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\": + \"assistant\",\n \"content\": \"4\",\n \"refusal\": null,\n + \ \"annotations\": []\n },\n \"logprobs\": null,\n \"finish_reason\": + \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\": 47,\n \"completion_tokens\": + 1,\n \"total_tokens\": 48,\n \"prompt_tokens_details\": {\n \"cached_tokens\": + 0,\n \"audio_tokens\": 0\n },\n \"completion_tokens_details\": + {\n \"reasoning_tokens\": 0,\n \"audio_tokens\": 0,\n \"accepted_prediction_tokens\": + 0,\n \"rejected_prediction_tokens\": 0\n }\n },\n \"service_tier\": + \"default\",\n \"system_fingerprint\": \"fp_b85f3645a5\"\n}\n" + headers: + access-control-expose-headers: + - X-Request-ID + - CF-Ray + - CF-Ray + alt-svc: + - h3=":443"; ma=86400 + cf-cache-status: + - DYNAMIC + cf-ray: + - a35eedc85be9cceb-YYZ + connection: + - keep-alive + content-length: + - '807' + content-type: + - application/json + date: + - Fri, 04 Sep 2026 18:06:35 GMT + openai-processing-ms: + - '515' + openai-version: + - '2020-10-01' + server: + - cloudflare + strict-transport-security: + - max-age=31536000; includeSubDomains; preload + transfer-encoding: + - chunked + x-content-type-options: + - nosniff + x-openai-proxy-wasm: + - v0.1 + x-ratelimit-limit-requests: + - '30000' + x-ratelimit-limit-tokens: + - '150000000' + x-ratelimit-remaining-requests: + - '29999' + x-ratelimit-remaining-tokens: + - '149999962' + x-ratelimit-reset-requests: + - 2ms + x-ratelimit-reset-tokens: + - 0s + x-request-id: + - req_e67f592357a64de29dda6a4c82b4e1a4 + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/pydantic_ai/cassettes/latest/test_agent_with_instructions_and_dynamic_system_prompt.yaml b/py/src/braintrust/integrations/pydantic_ai/cassettes/latest/test_agent_with_instructions_and_dynamic_system_prompt.yaml new file mode 100644 index 00000000..5118e95d --- /dev/null +++ b/py/src/braintrust/integrations/pydantic_ai/cassettes/latest/test_agent_with_instructions_and_dynamic_system_prompt.yaml @@ -0,0 +1,120 @@ +interactions: +- request: + body: '{"input":[{"role":"system","content":"The user is currently taking a math + quiz."},{"role":"user","content":"What is 2+2? Answer with just the number."}],"instructions":"Answer + with only the number requested by the user.","max_output_tokens":100,"model":"gpt-4o-mini","stream":false}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate, zstd + connection: + - keep-alive + content-length: + - '282' + content-type: + - application/json + host: + - api.openai.com + user-agent: + - pydantic-ai/2.36.0 + x-stainless-arch: + - arm64 + x-stainless-async: + - async:asyncio + x-stainless-lang: + - python + x-stainless-os: + - MacOS + x-stainless-package-version: + - 3.8.0 + x-stainless-read-timeout: + - '600' + x-stainless-retry-count: + - '0' + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.14.6 + method: POST + uri: https://api.openai.com/v1/responses + response: + body: + string: "{\n \"id\": \"resp_006792c93537c683006a9b07f51b4087d29047361acacaf77b\",\n + \ \"object\": \"response\",\n \"created_at\": 1788545013,\n \"status\": + \"completed\",\n \"background\": false,\n \"billing\": {\n \"payer\": + \"developer\"\n },\n \"completed_at\": 1788545015,\n \"error\": null,\n + \ \"frequency_penalty\": 0.0,\n \"incomplete_details\": null,\n \"instructions\": + \"Answer with only the number requested by the user.\",\n \"max_output_tokens\": + 100,\n \"max_tool_calls\": null,\n \"model\": \"gpt-4o-mini-2024-07-18\",\n + \ \"moderation\": null,\n \"output\": [\n {\n \"id\": \"msg_006792c93537c683006a9b07f743fc87d2b7cf465113ee4972\",\n + \ \"type\": \"message\",\n \"status\": \"completed\",\n \"content\": + [\n {\n \"type\": \"output_text\",\n \"annotations\": + [],\n \"logprobs\": [],\n \"text\": \"4\"\n }\n ],\n + \ \"role\": \"assistant\"\n }\n ],\n \"parallel_tool_calls\": true,\n + \ \"presence_penalty\": 0.0,\n \"previous_response_id\": null,\n \"prompt_cache_key\": + null,\n \"prompt_cache_retention\": \"in_memory\",\n \"reasoning\": {\n + \ \"context\": null,\n \"effort\": null,\n \"summary\": null\n },\n + \ \"safety_identifier\": null,\n \"service_tier\": \"default\",\n \"store\": + true,\n \"temperature\": 1.0,\n \"text\": {\n \"format\": {\n \"type\": + \"text\"\n },\n \"verbosity\": \"medium\"\n },\n \"tool_choice\": + \"auto\",\n \"tool_usage\": {\n \"image_gen\": {\n \"input_tokens\": + 0,\n \"input_tokens_details\": {\n \"image_tokens\": 0,\n \"text_tokens\": + 0\n },\n \"output_tokens\": 0,\n \"output_tokens_details\": + {\n \"image_tokens\": 0,\n \"text_tokens\": 0\n },\n \"total_tokens\": + 0\n },\n \"web_search\": {\n \"num_requests\": 0\n }\n },\n + \ \"tools\": [],\n \"top_logprobs\": 0,\n \"top_p\": 1.0,\n \"truncation\": + \"disabled\",\n \"usage\": {\n \"input_tokens\": 47,\n \"input_tokens_details\": + {\n \"cache_write_tokens\": 0,\n \"cached_tokens\": 0\n },\n + \ \"output_tokens\": 2,\n \"output_tokens_details\": {\n \"reasoning_tokens\": + 0\n },\n \"total_tokens\": 49\n },\n \"user\": null,\n \"metadata\": + {}\n}" + headers: + access-control-expose-headers: + - X-Request-ID + - CF-Ray + - CF-Ray + alt-svc: + - h3=":443"; ma=86400 + cf-cache-status: + - DYNAMIC + cf-ray: + - a35ee9567d4caaa4-YYZ + connection: + - keep-alive + content-length: + - '2017' + content-type: + - application/json + date: + - Fri, 04 Sep 2026 18:03:35 GMT + openai-processing-ms: + - '2402' + openai-version: + - '2020-10-01' + server: + - cloudflare + strict-transport-security: + - max-age=31536000; includeSubDomains; preload + transfer-encoding: + - chunked + x-content-type-options: + - nosniff + x-ratelimit-limit-requests: + - '30000' + x-ratelimit-limit-tokens: + - '150000000' + x-ratelimit-remaining-requests: + - '29999' + x-ratelimit-remaining-tokens: + - '149999932' + x-ratelimit-reset-requests: + - 2ms + x-ratelimit-reset-tokens: + - 0s + x-request-id: + - req_3c6a1c7bbb5e4b7ba59fb6c27f5ef4b4 + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py b/py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py index 42f4052a..ef8bf770 100644 --- a/py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py +++ b/py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py @@ -2,7 +2,6 @@ # pyright: reportUnknownMemberType=false # pyright: reportUnknownParameterType=false # pyright: reportPrivateUsage=false -import asyncio import inspect import time @@ -1066,6 +1065,39 @@ async def test_agent_with_system_prompt_in_metadata(memory_logger): assert agent_span["metadata"]["provider"] == "openai" +@pytest.mark.vcr +@pytest.mark.asyncio +async def test_agent_with_instructions_and_dynamic_system_prompt(memory_logger): + """Resolved instructions and system prompts should appear on agent and model spans.""" + assert not memory_logger.pop() + + instructions = "Answer with only the number requested by the user." + dynamic_system_prompt = "The user is currently taking a math quiz." + agent = Agent(MODEL, instructions=instructions, model_settings=ModelSettings(max_tokens=100)) + + @agent.system_prompt + def add_dynamic_system_prompt(): + return dynamic_system_prompt + + result = await agent.run(TEST_PROMPT) + assert "4" in str(result.output) + + spans = memory_logger.pop() + assert len(spans) == 2, f"Expected 2 spans (agent_run + chat), got {len(spans)}" + + agent_span = next(span for span in spans if span["span_attributes"]["type"] == SpanTypeAttribute.TASK) + chat_span = next(span for span in spans if span["span_attributes"]["type"] == SpanTypeAttribute.LLM) + + assert agent_span["input"]["instructions"] == instructions + assert agent_span["input"]["system_prompt"] == dynamic_system_prompt + + request = chat_span["input"]["messages"][0] + assert request["instructions"] == instructions + assert any( + part["part_kind"] == "system-prompt" and part["content"] == dynamic_system_prompt for part in request["parts"] + ) + + @pytest.mark.vcr @pytest.mark.asyncio async def test_agent_with_message_history(memory_logger): @@ -2617,211 +2649,6 @@ def test_shape_messages_with_binary_content(): assert content[1] == "What is in this document?" -@pytest.mark.asyncio -async def test_streaming_wrappers_capture_time_to_first_token(): - """Unit test verifying all streaming wrappers capture time_to_first_token. - - This test uses mocks to verify the internal wrapper logic without requiring - API calls. It ensures that _first_token_time is tracked correctly in: - - _AgentStreamWrapper (async agent streaming) - - _DirectStreamWrapper (async direct API streaming) - - _AgentStreamResultSyncProxy (sync agent streaming) - - _DirectStreamWrapperSync (sync direct API streaming) - """ - from unittest.mock import AsyncMock, MagicMock, Mock - - from braintrust.integrations.pydantic_ai.tracing import ( - _AgentStreamResultSyncProxy, - _AgentStreamWrapper, - _DirectStreamIteratorProxy, - _DirectStreamIteratorSyncProxy, - _DirectStreamWrapper, - _DirectStreamWrapperSync, - _StreamResultProxy, - ) - - # Test 1: _AgentStreamWrapper captures first token time - print("\n--- Testing _AgentStreamWrapper ---") - - class MockStreamResult: - async def stream_text(self, delta=True): - for i in range(3): - await asyncio.sleep(0.001) - yield f"token{i} " - - def usage(self): - usage_mock = Mock(input_tokens=50, output_tokens=20, total_tokens=70) - usage_mock.cache_read_tokens = None - usage_mock.cache_write_tokens = None - return usage_mock - - mock_stream_result = MockStreamResult() - wrapper = _AgentStreamWrapper( - stream_cm=AsyncMock(), - span_name="test_stream", - input_data={"prompt": "test"}, - metadata={"model": "gpt-4o"}, - ) - - wrapper.span_cm = MagicMock() - wrapper.span_cm.__enter__ = MagicMock() - wrapper.start_time = time.time() - wrapper.stream_result = mock_stream_result - - proxy = _StreamResultProxy(mock_stream_result, wrapper) - - assert wrapper._first_token_time is None - - chunk_count = 0 - async for text in proxy.stream_text(delta=True): - chunk_count += 1 - if chunk_count == 1: - assert wrapper._first_token_time is not None - assert wrapper._first_token_time > wrapper.start_time - - assert chunk_count == 3 - assert wrapper._first_token_time is not None - print("✓ _AgentStreamWrapper captures first token time") - - # Test 2: _DirectStreamWrapper captures first token time - print("\n--- Testing _DirectStreamWrapper ---") - - class MockStream: - def __init__(self): - self.chunks = [] - - async def __anext__(self): - if len(self.chunks) < 3: - await asyncio.sleep(0.001) - chunk = Mock(delta=Mock(content_delta=f"chunk{len(self.chunks)}")) - self.chunks.append(chunk) - return chunk - raise StopAsyncIteration - - def __aiter__(self): - return self - - def get(self): - usage_mock = Mock(input_tokens=50, output_tokens=20, total_tokens=70) - usage_mock.cache_read_tokens = None - usage_mock.cache_write_tokens = None - return Mock(usage=usage_mock) - - mock_stream = MockStream() - direct_wrapper = _DirectStreamWrapper( - stream_cm=AsyncMock(), - span_name="test_direct_stream", - input_data={"messages": []}, - metadata={"model": "gpt-4o"}, - ) - - direct_wrapper.span_cm = MagicMock() - direct_wrapper.start_time = time.time() - direct_wrapper.stream = mock_stream - - proxy = _DirectStreamIteratorProxy(mock_stream, direct_wrapper) - - assert direct_wrapper._first_token_time is None - - chunk_count = 0 - async for chunk in proxy: - chunk_count += 1 - if chunk_count == 1: - assert direct_wrapper._first_token_time is not None - assert direct_wrapper._first_token_time > direct_wrapper.start_time - - assert chunk_count == 3 - assert direct_wrapper._first_token_time is not None - print("✓ _DirectStreamWrapper captures first token time") - - # Test 3: _AgentStreamResultSyncProxy captures first token time - print("\n--- Testing _AgentStreamResultSyncProxy ---") - - class MockSyncStreamResult: - def stream_text(self, delta=True): - for i in range(3): - time.sleep(0.001) - yield f"token{i} " - - def usage(self): - usage_mock = Mock(input_tokens=50, output_tokens=20, total_tokens=70) - usage_mock.cache_read_tokens = None - usage_mock.cache_write_tokens = None - return usage_mock - - mock_sync_result = MockSyncStreamResult() - sync_proxy = _AgentStreamResultSyncProxy( - stream_result=mock_sync_result, - span=MagicMock(), - span_cm=MagicMock(), - start_time=time.time(), - ) - - assert sync_proxy._first_token_time is None - - chunk_count = 0 - for text in sync_proxy.stream_text(delta=True): - chunk_count += 1 - if chunk_count == 1: - assert sync_proxy._first_token_time is not None - - assert chunk_count == 3 - assert sync_proxy._first_token_time is not None - print("✓ _AgentStreamResultSyncProxy captures first token time") - - # Test 4: _DirectStreamWrapperSync captures first token time - print("\n--- Testing _DirectStreamWrapperSync ---") - - class MockSyncStream: - def __init__(self): - self.chunks = [] - - def __iter__(self): - return self - - def __next__(self): - if len(self.chunks) < 3: - time.sleep(0.001) - chunk = Mock(delta=Mock(content_delta=f"chunk{len(self.chunks)}")) - self.chunks.append(chunk) - return chunk - raise StopIteration - - def get(self): - usage_mock = Mock(input_tokens=50, output_tokens=20, total_tokens=70) - usage_mock.cache_read_tokens = None - usage_mock.cache_write_tokens = None - return Mock(usage=usage_mock) - - mock_sync_stream = MockSyncStream() - sync_wrapper = _DirectStreamWrapperSync( - stream_cm=MagicMock(), - span_name="test_sync_stream", - input_data={"messages": []}, - metadata={"model": "gpt-4o"}, - ) - - sync_wrapper.start_time = time.time() - sync_wrapper.stream = mock_sync_stream - - sync_proxy = _DirectStreamIteratorSyncProxy(mock_sync_stream, sync_wrapper) - - assert sync_wrapper._first_token_time is None - - chunk_count = 0 - for chunk in sync_proxy: - chunk_count += 1 - if chunk_count == 1: - assert sync_wrapper._first_token_time is not None - assert sync_wrapper._first_token_time > sync_wrapper.start_time - - assert chunk_count == 3 - assert sync_wrapper._first_token_time is not None - print("✓ _DirectStreamWrapperSync captures first token time") - - print("\n✅ All streaming wrapper unit tests passed!") - - @pytest.mark.asyncio async def test_attachment_preserved_in_model_settings(memory_logger): """Test that attachments in model_settings are preserved through serialization.""" diff --git a/py/src/braintrust/integrations/pydantic_ai/tracing.py b/py/src/braintrust/integrations/pydantic_ai/tracing.py index 3058d104..65b4f8f2 100644 --- a/py/src/braintrust/integrations/pydantic_ai/tracing.py +++ b/py/src/braintrust/integrations/pydantic_ai/tracing.py @@ -82,6 +82,28 @@ def _maybe_create_tool_spans_from_messages(result: Any) -> None: _create_tool_spans_from_messages(result) +def _log_agent_result(span: Any, input_data: dict[str, Any], result: Any, output: Any, metrics: Any) -> None: + resolved_input = input_data + if result is not None: + resolved_input = dict(input_data) + for message in reversed(result.new_messages()): + instructions = getattr(message, "instructions", None) + if instructions: + resolved_input["instructions"] = instructions + break + + system_prompts = [ + part.content + for message in result.all_messages() + for part in getattr(message, "parts", ()) + if getattr(part, "part_kind", None) == "system-prompt" and getattr(part, "content", None) + ] + if system_prompts: + resolved_input["system_prompt"] = "\n\n".join(system_prompts) + + span.log(input=resolved_input, output=output, metrics=metrics) + + async def _agent_run_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any): input_data, metadata = _build_agent_input_and_metadata(args, kwargs, instance) @@ -99,8 +121,13 @@ async def _agent_run_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any _maybe_create_tool_spans_from_messages(result) - output = _shape_result_output(result) - agent_span.log(output=output, metrics=_wrapper_span_metrics(start_time, end_time)) + _log_agent_result( + agent_span, + input_data, + result, + _shape_result_output(result), + _wrapper_span_metrics(start_time, end_time), + ) return result finally: _reset_tool_trace_capture(tool_trace_token) @@ -123,8 +150,13 @@ def _agent_run_sync_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any) _maybe_create_tool_spans_from_messages(result) - output = _shape_result_output(result) - agent_span.log(output=output, metrics=_wrapper_span_metrics(start_time, end_time)) + _log_agent_result( + agent_span, + input_data, + result, + _shape_result_output(result), + _wrapper_span_metrics(start_time, end_time), + ) return result finally: _reset_tool_trace_capture(tool_trace_token) @@ -185,6 +217,7 @@ def _agent_run_stream_sync_wrapper(wrapped: Any, instance: Any, args: Any, kwarg span, span_cm, start_time, + input_data, tool_trace_token, ) except Exception: @@ -480,7 +513,7 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): "event_count": self._event_count, } output = _shape_result_output(self._final_result) if self._final_result is not None else None - self.agent_span.log(output=output, metrics=metrics) + _log_agent_result(self.agent_span, self.input_data, self._final_result, output, metrics) if self.span_cm: if asyncio.current_task() is self._enter_task: @@ -584,10 +617,12 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): _maybe_create_tool_spans_from_messages(self.stream_result) - output = _shape_stream_output(self.stream_result) - self.span_cm.log( - output=output, - metrics=_wrapper_span_metrics(self.start_time, end_time, self._first_token_time), + _log_agent_result( + self.span_cm, + self.input_data, + self.stream_result, + _shape_stream_output(self.stream_result), + _wrapper_span_metrics(self.start_time, end_time, self._first_token_time), ) if self.span_cm: @@ -721,12 +756,14 @@ def __init__( span: Any, span_cm: Any, start_time: float, + input_data: dict[str, Any] | None = None, tool_trace_token: Any = None, ): self._stream_result = stream_result self._span = span self._span_cm = span_cm self._start_time = start_time + self._input_data = input_data or {} self._logged = False self._finalize_on_del = True self._first_token_time = None @@ -767,10 +804,12 @@ def _finalize(self): _maybe_create_tool_spans_from_messages(self._stream_result) - output = _shape_stream_output(self._stream_result) - self._span.log( - output=output, - metrics=_wrapper_span_metrics(self._start_time, end_time, self._first_token_time), + _log_agent_result( + self._span, + self._input_data, + self._stream_result, + _shape_stream_output(self._stream_result), + _wrapper_span_metrics(self._start_time, end_time, self._first_token_time), ) self._logged = True finally: @@ -976,8 +1015,17 @@ def _msg_timestamp(msg: Any) -> float | None: _MISSING = object() -_MESSAGE_FIELDS = ("kind", "role", "timestamp", "state") -_PART_FIELDS = ("kind", "part_kind", "tool_name", "tool_call_id") +_MESSAGE_FIELDS = ( + "kind", + "role", + "timestamp", + "state", + "instructions", + "run_id", + "conversation_id", + "metadata", +) +_PART_FIELDS = ("kind", "part_kind", "tool_name", "tool_call_id", "timestamp", "dynamic_ref", "id", "args") _RESPONSE_FIELDS = ( "kind", "model_name", @@ -1057,8 +1105,13 @@ def _shape_message(message: Any) -> Any: parts = _field_value(message, "parts") if not parts: return message - if not any(_has_binary_leaf(part) for part in parts): - # Let Braintrust's dataclass/Pydantic serializer preserve every field. + has_prompt_data = _field_value(message, "instructions") not in (_MISSING, None, "") or any( + _field_value(part, "part_kind") == "system-prompt" for part in parts + ) + if not has_prompt_data and not any(_has_binary_leaf(part) for part in parts): + # Let Braintrust's serializer preserve every top-level field. It does not recurse + # into dataclasses nested in a list, so prompt-bearing parts have to be shaped + # here or they reach the span as repr strings; binary leaves need materializing. return message return _shape_object( message, fields=_MESSAGE_FIELDS, overrides={"parts": [_shape_content_part(part) for part in parts]} @@ -1074,12 +1127,16 @@ def _shape_content_part(part: Any) -> Any: return attachment_payload content = _field_value(part, "content") - if content is _MISSING or not _has_binary_leaf(content): + if content is _MISSING: return part - shaped_content = ( - [_shape_content_part(item) for item in content] if isinstance(content, list) else _shape_content_part(content) - ) + shaped_content = content + if _has_binary_leaf(content): + shaped_content = ( + [_shape_content_part(item) for item in content] + if isinstance(content, list) + else _shape_content_part(content) + ) return _shape_object(part, fields=_PART_FIELDS, overrides={"content": shaped_content})