From faf6f1dcc3e5873bdea9fb7d272d9574596ff639 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Fri, 4 Sep 2026 14:29:10 -0400 Subject: [PATCH 1/2] fix(pydantic_ai): trace instructions and dynamic system prompts Capture resolved instructions and system prompts on agent spans and preserve their structured representation on model spans across sync, async, and streaming paths.\n\nAdd VCR regression coverage for both supported Pydantic AI versions. --- ...nstructions_and_dynamic_system_prompt.yaml | 104 +++++++++++++++ ...nstructions_and_dynamic_system_prompt.yaml | 120 ++++++++++++++++++ .../test_pydantic_ai_integration.py | 33 +++++ .../integrations/pydantic_ai/tracing.py | 99 ++++++++++++--- 4 files changed, 335 insertions(+), 21 deletions(-) create mode 100644 py/src/braintrust/integrations/pydantic_ai/cassettes/1.10.0/test_agent_with_instructions_and_dynamic_system_prompt.yaml create mode 100644 py/src/braintrust/integrations/pydantic_ai/cassettes/latest/test_agent_with_instructions_and_dynamic_system_prompt.yaml 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..685cc680 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 @@ -1066,6 +1066,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): 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}) From db86fde4b8e4a7fb015c9534a110d1b5b7324190 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Fri, 4 Sep 2026 14:49:16 -0400 Subject: [PATCH 2/2] no mocks no fakes --- .../test_pydantic_ai_integration.py | 206 ------------------ 1 file changed, 206 deletions(-) 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 685cc680..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 @@ -2650,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."""