From f69e9b36368a1f6838388346f4c13de1b9a6c96e Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Mon, 2 Nov 2026 11:02:31 -0500 Subject: [PATCH 1/2] fix(pydantic-ai): trace fallback model attempts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit > 🤖 **This text was AI-generated.** When a Pydantic AI agent retries a failed request with a model instance, the fallback can run without a model span. The agent span can also retain the original model and provider after a different model answers. ```text Before: failed primary span → untraced fallback; agent shows primary After: failed primary span → fallback span; agent shows answering model ``` Wrap models resolved for retry and update agent metadata from the final response. Keep the provider reported by the response when a model ID uses an alias such as `openai-responses`. Closes #1042 --- ...st_model_fallback_traces_each_attempt.yaml | 315 ++++++++++++++++++ .../integrations/pydantic_ai/integration.py | 2 + .../integrations/pydantic_ai/patchers.py | 8 + .../test_pydantic_ai_integration.py | 51 +++ .../integrations/pydantic_ai/tracing.py | 17 +- 5 files changed, 391 insertions(+), 2 deletions(-) create mode 100644 py/src/braintrust/integrations/pydantic_ai/cassettes/latest/test_model_fallback_traces_each_attempt.yaml diff --git a/py/src/braintrust/integrations/pydantic_ai/cassettes/latest/test_model_fallback_traces_each_attempt.yaml b/py/src/braintrust/integrations/pydantic_ai/cassettes/latest/test_model_fallback_traces_each_attempt.yaml new file mode 100644 index 000000000..098daf30d --- /dev/null +++ b/py/src/braintrust/integrations/pydantic_ai/cassettes/latest/test_model_fallback_traces_each_attempt.yaml @@ -0,0 +1,315 @@ +interactions: +- request: + body: '{"model":"braintrust-nonexistent-model-for-fallback-test","stream":false,"messages":[{"role":"user","content":"What + is 2+2? Answer with just the number."}],"max_completion_tokens":50}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate, zstd + connection: + - keep-alive + content-length: + - '183' + content-type: + - application/json + host: + - api.openai.com + user-agent: + - pydantic-ai/2.64.0 + x-stainless-arch: + - arm64 + x-stainless-async: + - async:asyncio + x-stainless-lang: + - python + x-stainless-os: + - MacOS + x-stainless-package-version: + - 3.42.1 + x-stainless-read-timeout: + - '600' + x-stainless-retry-count: + - '0' + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.14.7 + method: POST + uri: https://api.openai.com/v1/chat/completions + response: + body: + string: "{\n \"error\": {\n \"message\": \"The model `braintrust-nonexistent-model-for-fallback-test` + does not exist or you do not have access to it.\",\n \"type\": \"invalid_request_error\",\n + \ \"param\": null,\n \"code\": \"model_not_found\"\n }\n}\n" + headers: + access-control-expose-headers: + - CF-Ray + - CF-Ray + alt-svc: + - h3=":443"; ma=86400 + cf-cache-status: + - DYNAMIC + cf-ray: + - a54451a5ddf20c26-YYZ + connection: + - keep-alive + content-length: + - '252' + content-type: + - application/json; charset=utf-8 + date: + - Mon, 02 Nov 2026 15:54:38 GMT + server: + - cloudflare + strict-transport-security: + - max-age=31536000; includeSubDomains; preload + transfer-encoding: + - chunked + vary: + - Origin + x-content-type-options: + - nosniff + x-openai-proxy-wasm: + - v0.1 + x-request-id: + - req_9411a089d66c46589834fc5f6de172f2 + status: + code: 404 + message: Not Found +- request: + body: '{"model":"gpt-4o-mini","stream":false,"input":[{"role":"user","content":"What + is 2+2? Answer with just the number."}],"max_output_tokens":50}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate, zstd + connection: + - keep-alive + content-length: + - '141' + content-type: + - application/json + host: + - api.openai.com + user-agent: + - pydantic-ai/2.64.0 + x-stainless-arch: + - arm64 + x-stainless-async: + - async:asyncio + x-stainless-lang: + - python + x-stainless-os: + - MacOS + x-stainless-package-version: + - 3.42.1 + x-stainless-read-timeout: + - '600' + x-stainless-retry-count: + - '0' + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.14.7 + method: POST + uri: https://api.openai.com/v1/responses + response: + body: + string: "{\n \"id\": \"resp_0a7f9ac6a30962e5016ae8b23ead9887d19955021bb7918a7d\",\n + \ \"object\": \"response\",\n \"created_at\": 1793634878,\n \"status\": + \"completed\",\n \"access_programs\": null,\n \"background\": false,\n \"billing\": + {\n \"payer\": \"developer\"\n },\n \"completed_at\": 1793634879,\n \"error\": + null,\n \"frequency_penalty\": 0.0,\n \"incomplete_details\": null,\n \"instructions\": + null,\n \"max_output_tokens\": 50,\n \"max_tool_calls\": null,\n \"model\": + \"gpt-4o-mini-2024-07-18\",\n \"moderation\": null,\n \"output\": [\n {\n + \ \"id\": \"msg_0a7f9ac6a30962e5016ae8b23f205487d1b43f0e83abb01c48\",\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\": + false,\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\": 20,\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\": 22\n },\n \"user\": null,\n \"metadata\": + {}\n}" + headers: + access-control-expose-headers: + - X-Request-ID, Server-Timing-Total + - CF-Ray + - CF-Ray + alt-svc: + - h3=":443"; ma=86400 + cf-cache-status: + - DYNAMIC + cf-ray: + - a54451a71dfb39fb-YYZ + connection: + - keep-alive + content-length: + - '1996' + content-type: + - application/json + date: + - Mon, 02 Nov 2026 15:54:39 GMT + openai-processing-ms: + - '537' + 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: + - '149999960' + x-ratelimit-reset-requests: + - 2ms + x-ratelimit-reset-tokens: + - 0s + x-request-id: + - req_5faffd3b9bb84d9e91f63ebe4a65ccb5 + status: + code: 200 + message: OK +- request: + body: '{"model":"gpt-4o-mini","stream":false,"input":[{"role":"user","content":"What + is 2+2? Answer with just the number."},{"role":"assistant","content":"4"},{"role":"user","content":"What + is 3+3? Answer with just the number."}],"max_output_tokens":50}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate, zstd + connection: + - keep-alive + content-length: + - '246' + content-type: + - application/json + host: + - api.openai.com + user-agent: + - pydantic-ai/2.64.0 + x-stainless-arch: + - arm64 + x-stainless-async: + - async:asyncio + x-stainless-lang: + - python + x-stainless-os: + - MacOS + x-stainless-package-version: + - 3.42.1 + x-stainless-read-timeout: + - '600' + x-stainless-retry-count: + - '0' + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.14.7 + method: POST + uri: https://api.openai.com/v1/responses + response: + body: + string: "{\n \"id\": \"resp_0fa2baed56a0624d016ae8b2f2423887d19b9ac3c8f933e0c5\",\n + \ \"object\": \"response\",\n \"created_at\": 1793635058,\n \"status\": + \"completed\",\n \"access_programs\": null,\n \"background\": false,\n \"billing\": + {\n \"payer\": \"developer\"\n },\n \"completed_at\": 1793635058,\n \"error\": + null,\n \"frequency_penalty\": 0.0,\n \"incomplete_details\": null,\n \"instructions\": + null,\n \"max_output_tokens\": 50,\n \"max_tool_calls\": null,\n \"model\": + \"gpt-4o-mini-2024-07-18\",\n \"moderation\": null,\n \"output\": [\n {\n + \ \"id\": \"msg_0fa2baed56a0624d016ae8b2f2d77087d1bdc4946521192873\",\n + \ \"type\": \"message\",\n \"status\": \"completed\",\n \"content\": + [\n {\n \"type\": \"output_text\",\n \"annotations\": + [],\n \"logprobs\": [],\n \"text\": \"6\"\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\": + false,\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\": 42,\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\": 44\n },\n \"user\": null,\n \"metadata\": + {}\n}" + headers: + access-control-expose-headers: + - X-Request-ID, Server-Timing-Total + - CF-Ray + - CF-Ray + alt-svc: + - h3=":443"; ma=86400 + cf-cache-status: + - DYNAMIC + cf-ray: + - a5445608bd32ab36-YYZ + connection: + - keep-alive + content-length: + - '1996' + content-type: + - application/json + date: + - Mon, 02 Nov 2026 15:57:38 GMT + openai-processing-ms: + - '655' + 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: + - '149999937' + x-ratelimit-reset-requests: + - 2ms + x-ratelimit-reset-tokens: + - 0s + x-request-id: + - req_1685ccb80c8e4439bc8f8f76ccebae83 + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/pydantic_ai/integration.py b/py/src/braintrust/integrations/pydantic_ai/integration.py index 93987f701..791e4b374 100644 --- a/py/src/braintrust/integrations/pydantic_ai/integration.py +++ b/py/src/braintrust/integrations/pydantic_ai/integration.py @@ -9,6 +9,7 @@ DirectModelRequestStreamSyncPatcher, DirectModelRequestSyncPatcher, DirectPrepareModelPatcher, + RetryModelPatcher, StreamedResponseSyncStartProducerPatcher, SyncStreamBridgeInitPatcher, ToolManagerFunctionToolPatcher, @@ -25,6 +26,7 @@ class PydanticAIIntegration(BaseIntegration): SyncStreamBridgeInitPatcher, StreamedResponseSyncStartProducerPatcher, AgentPatcher, + RetryModelPatcher, DirectPrepareModelPatcher, DirectModelRequestPatcher, DirectModelRequestSyncPatcher, diff --git a/py/src/braintrust/integrations/pydantic_ai/patchers.py b/py/src/braintrust/integrations/pydantic_ai/patchers.py index 56440103e..ab023de2e 100644 --- a/py/src/braintrust/integrations/pydantic_ai/patchers.py +++ b/py/src/braintrust/integrations/pydantic_ai/patchers.py @@ -20,6 +20,7 @@ _create_direct_model_request_wrapper, _create_start_producer_wrapper, _direct_prepare_model_wrapper, + _resolve_retry_model_wrapper, _sync_stream_bridge_init_wrapper, _tool_manager_call_function_tool_wrapper, _tool_manager_execute_function_tool_wrapper, @@ -97,6 +98,13 @@ class AgentPatcher(CompositeFunctionWrapperPatcher): ) +class RetryModelPatcher(FunctionWrapperPatcher): + name = "pydantic_ai.resolve_retry_model" + target_module = "pydantic_ai._agent_graph" + target_path = "ModelRequestNode._resolve_retry_model" + wrapper = _resolve_retry_model_wrapper + + class DirectPrepareModelPatcher(FunctionWrapperPatcher): name = "pydantic_ai.direct.prepare_model" target_module = "pydantic_ai.direct" 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 5051d60a1..d7b8af5fd 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 @@ -2440,6 +2440,57 @@ async def test_agent_with_prefill(memory_logger): _assert_metrics_are_valid(agent_span["metrics"], start, end) +@pytest.mark.vcr(cassette_name="test_model_fallback_traces_each_attempt") +@pytest.mark.asyncio +@pytest.mark.parametrize("candidate_kind", ["instance", "string"]) +async def test_model_fallback_traces_each_attempt(memory_logger, candidate_kind): + """A failed provider request and the model that answers each get their own span.""" + try: + from pydantic_ai.capabilities import ModelFallback + except ImportError: + pytest.skip("ModelFallback is unavailable in this Pydantic AI version") + from pydantic_ai.models.openai import OpenAIChatModel, OpenAIResponsesModel + + fallback = "openai-responses:gpt-4o-mini" if candidate_kind == "string" else OpenAIResponsesModel("gpt-4o-mini") + agent = Agent( + OpenAIChatModel("braintrust-nonexistent-model-for-fallback-test"), + capabilities=[ModelFallback(fallback, sticky=True)], + model_settings=ModelSettings(max_tokens=50), + ) + result = await agent.run(TEST_PROMPT) + assert "4" in result.output + + spans = memory_logger.pop() + agent_span = _find_agent_span(spans) + model_spans = [span for span in spans if span["span_attributes"].get("type") == SpanTypeAttribute.LLM] + assert agent_span is not None + assert len(model_spans) == 2 + failed = next( + span + for span in model_spans + if span["metadata"].get("model") == "braintrust-nonexistent-model-for-fallback-test" + ) + answered = next(span for span in model_spans if span["metadata"].get("model") == "gpt-4o-mini") + assert failed.get("error") + assert not answered.get("error") + assert agent_span["metadata"]["model"] == "gpt-4o-mini" + assert agent_span["metadata"]["provider"] == "openai" + for span in model_spans: + assert span["metadata"]["provider"] == "openai" + assert span["span_parents"] == [agent_span["span_id"]] + assert span["context"]["span_origin"]["instrumentation"]["name"] == "pydantic-ai-auto" + + continued = await agent.run("What is 3+3? Answer with just the number.", message_history=result.all_messages()) + assert "6" in continued.output + spans = memory_logger.pop() + agent_span = _find_agent_span(spans) + model_spans = [span for span in spans if span["span_attributes"].get("type") == SpanTypeAttribute.LLM] + assert len(model_spans) == 1 + assert model_spans[0]["metadata"]["model"] == "gpt-4o-mini" + assert agent_span["metadata"]["model"] == "gpt-4o-mini" + assert agent_span["metadata"]["provider"] == "openai" + + @pytest.mark.vcr @pytest.mark.asyncio async def test_agent_with_error_handling(memory_logger): diff --git a/py/src/braintrust/integrations/pydantic_ai/tracing.py b/py/src/braintrust/integrations/pydantic_ai/tracing.py index d03c33789..cc96cecdc 100644 --- a/py/src/braintrust/integrations/pydantic_ai/tracing.py +++ b/py/src/braintrust/integrations/pydantic_ai/tracing.py @@ -57,6 +57,11 @@ async def _agent_resolve_model_selection_wrapper(wrapped: Any, instance: Any, ar return _wrap_model_instance(await wrapped(*args, **kwargs)) +async def _resolve_retry_model_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any): + model, model_id = await wrapped(*args, **kwargs) + return _wrap_model_instance(model), model_id + + def _direct_prepare_model_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any): return _wrap_model_instance(wrapped(*args, **kwargs)) @@ -84,9 +89,17 @@ def _maybe_create_tool_spans_from_messages(result: Any) -> None: def _log_agent_result(span: Any, input_data: dict[str, Any], result: Any, output: Any, metrics: Any) -> None: resolved_input = input_data + metadata = {} if result is not None: resolved_input = dict(input_data) - for message in reversed(result.new_messages()): + new_messages = result.new_messages() + for message in reversed(new_messages): + model_id = getattr(message, "model_id", None) + if model_id: + model_name, provider = _parse_model_string(model_id) + metadata = _build_model_metadata(model_name, getattr(message, "provider_name", None) or provider) + break + for message in reversed(new_messages): instructions = getattr(message, "instructions", None) if instructions: resolved_input["instructions"] = instructions @@ -101,7 +114,7 @@ def _log_agent_result(span: Any, input_data: dict[str, Any], result: Any, output if system_prompts: resolved_input["system_prompt"] = "\n\n".join(system_prompts) - span.log(input=resolved_input, output=output, metrics=metrics) + span.log(input=resolved_input, output=output, metrics=metrics, metadata=metadata) async def _agent_run_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any): From 03f19e4c10b8d38d35716cb2fc62caebddb602fc Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Mon, 2 Nov 2026 11:14:55 -0500 Subject: [PATCH 2/2] fix(pydantic-ai): close failed streaming model spans --- ...fallback_stream_closes_failed_attempt.yaml | 212 ++++++++++++++++++ .../test_pydantic_ai_integration.py | 36 +++ .../integrations/pydantic_ai/tracing.py | 6 +- 3 files changed, 253 insertions(+), 1 deletion(-) create mode 100644 py/src/braintrust/integrations/pydantic_ai/cassettes/latest/test_model_fallback_stream_closes_failed_attempt.yaml diff --git a/py/src/braintrust/integrations/pydantic_ai/cassettes/latest/test_model_fallback_stream_closes_failed_attempt.yaml b/py/src/braintrust/integrations/pydantic_ai/cassettes/latest/test_model_fallback_stream_closes_failed_attempt.yaml new file mode 100644 index 000000000..ffed32a53 --- /dev/null +++ b/py/src/braintrust/integrations/pydantic_ai/cassettes/latest/test_model_fallback_stream_closes_failed_attempt.yaml @@ -0,0 +1,212 @@ +interactions: +- request: + body: '{"model":"braintrust-nonexistent-model-for-fallback-test","stream":true,"messages":[{"role":"user","content":"What + is 2+2? Answer with just the number."}],"max_completion_tokens":50,"stream_options":{"include_usage":true}}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate, zstd + connection: + - keep-alive + content-length: + - '222' + content-type: + - application/json + host: + - api.openai.com + user-agent: + - pydantic-ai/2.64.0 + x-stainless-arch: + - arm64 + x-stainless-async: + - async:asyncio + x-stainless-lang: + - python + x-stainless-os: + - MacOS + x-stainless-package-version: + - 3.42.1 + x-stainless-read-timeout: + - '600' + x-stainless-retry-count: + - '0' + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.14.7 + method: POST + uri: https://api.openai.com/v1/chat/completions + response: + body: + string: "{\n \"error\": {\n \"message\": \"The model `braintrust-nonexistent-model-for-fallback-test` + does not exist or you do not have access to it.\",\n \"type\": \"invalid_request_error\",\n + \ \"param\": null,\n \"code\": \"model_not_found\"\n }\n}\n" + headers: + access-control-expose-headers: + - CF-Ray + - CF-Ray + alt-svc: + - h3=":443"; ma=86400 + cf-cache-status: + - DYNAMIC + cf-ray: + - a5446e1f2826c109-YYZ + connection: + - keep-alive + content-length: + - '252' + content-type: + - application/json; charset=utf-8 + date: + - Mon, 02 Nov 2026 16:14:05 GMT + server: + - cloudflare + strict-transport-security: + - max-age=31536000; includeSubDomains; preload + transfer-encoding: + - chunked + vary: + - Origin + x-content-type-options: + - nosniff + x-openai-proxy-wasm: + - v0.1 + x-request-id: + - req_a68fe429713b4458b64e8acc0964f41c + status: + code: 404 + message: Not Found +- request: + body: '{"model":"gpt-4o-mini","stream":true,"input":[{"role":"user","content":"What + is 2+2? Answer with just the number."}],"max_output_tokens":50}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate, zstd + connection: + - keep-alive + content-length: + - '140' + content-type: + - application/json + host: + - api.openai.com + user-agent: + - pydantic-ai/2.64.0 + x-stainless-arch: + - arm64 + x-stainless-async: + - async:asyncio + x-stainless-lang: + - python + x-stainless-os: + - MacOS + x-stainless-package-version: + - 3.42.1 + x-stainless-read-timeout: + - '600' + x-stainless-retry-count: + - '0' + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.14.7 + method: POST + uri: https://api.openai.com/v1/responses + response: + body: + string: 'event: response.created + + data: {"type":"response.created","response":{"id":"resp_05d53df3d23eb9eb016ae8b6cd80fc87d1a59c2ec7f5ae47e9","object":"response","created_at":1793636045,"status":"in_progress","access_programs":null,"background":false,"completed_at":null,"error":null,"frequency_penalty":0.0,"incomplete_details":null,"instructions":null,"max_output_tokens":50,"max_tool_calls":null,"model":"gpt-4o-mini-2024-07-18","moderation":null,"output":[],"parallel_tool_calls":true,"presence_penalty":0.0,"previous_response_id":null,"prompt_cache_key":null,"prompt_cache_retention":"in_memory","reasoning":{"context":null,"effort":null,"summary":null},"safety_identifier":null,"service_tier":"auto","store":false,"temperature":1.0,"text":{"format":{"type":"text"},"verbosity":"medium"},"tool_choice":"auto","tool_usage":{"image_gen":{"input_tokens":0,"input_tokens_details":{"image_tokens":0,"text_tokens":0},"output_tokens":0,"output_tokens_details":{"image_tokens":0,"text_tokens":0},"total_tokens":0},"web_search":{"num_requests":0}},"tools":[],"top_logprobs":0,"top_p":1.0,"truncation":"disabled","usage":null,"user":null,"metadata":{}},"sequence_number":0} + + + event: response.in_progress + + data: {"type":"response.in_progress","response":{"id":"resp_05d53df3d23eb9eb016ae8b6cd80fc87d1a59c2ec7f5ae47e9","object":"response","created_at":1793636045,"status":"in_progress","access_programs":null,"background":false,"completed_at":null,"error":null,"frequency_penalty":0.0,"incomplete_details":null,"instructions":null,"max_output_tokens":50,"max_tool_calls":null,"model":"gpt-4o-mini-2024-07-18","moderation":null,"output":[],"parallel_tool_calls":true,"presence_penalty":0.0,"previous_response_id":null,"prompt_cache_key":null,"prompt_cache_retention":"in_memory","reasoning":{"context":null,"effort":null,"summary":null},"safety_identifier":null,"service_tier":"auto","store":false,"temperature":1.0,"text":{"format":{"type":"text"},"verbosity":"medium"},"tool_choice":"auto","tool_usage":{"image_gen":{"input_tokens":0,"input_tokens_details":{"image_tokens":0,"text_tokens":0},"output_tokens":0,"output_tokens_details":{"image_tokens":0,"text_tokens":0},"total_tokens":0},"web_search":{"num_requests":0}},"tools":[],"top_logprobs":0,"top_p":1.0,"truncation":"disabled","usage":null,"user":null,"metadata":{}},"sequence_number":1} + + + event: response.output_item.added + + data: {"type":"response.output_item.added","item":{"id":"msg_05d53df3d23eb9eb016ae8b6cdf79887d1a57a49654be7108f","type":"message","status":"in_progress","content":[],"role":"assistant"},"output_index":0,"sequence_number":2} + + + event: response.content_part.added + + data: {"type":"response.content_part.added","content_index":0,"item_id":"msg_05d53df3d23eb9eb016ae8b6cdf79887d1a57a49654be7108f","output_index":0,"part":{"type":"output_text","annotations":[],"logprobs":[],"text":""},"sequence_number":3} + + + event: response.output_text.delta + + data: {"type":"response.output_text.delta","content_index":0,"delta":"4","item_id":"msg_05d53df3d23eb9eb016ae8b6cdf79887d1a57a49654be7108f","logprobs":[],"obfuscation":"fGGZTujoCcguntJ","output_index":0,"sequence_number":4} + + + event: response.output_text.done + + data: {"type":"response.output_text.done","content_index":0,"item_id":"msg_05d53df3d23eb9eb016ae8b6cdf79887d1a57a49654be7108f","logprobs":[],"output_index":0,"sequence_number":5,"text":"4"} + + + event: response.content_part.done + + data: {"type":"response.content_part.done","content_index":0,"item_id":"msg_05d53df3d23eb9eb016ae8b6cdf79887d1a57a49654be7108f","output_index":0,"part":{"type":"output_text","annotations":[],"logprobs":[],"text":"4"},"sequence_number":6} + + + event: response.output_item.done + + data: {"type":"response.output_item.done","item":{"id":"msg_05d53df3d23eb9eb016ae8b6cdf79887d1a57a49654be7108f","type":"message","status":"completed","content":[{"type":"output_text","annotations":[],"logprobs":[],"text":"4"}],"role":"assistant"},"output_index":0,"sequence_number":7} + + + event: response.completed + + data: {"type":"response.completed","response":{"id":"resp_05d53df3d23eb9eb016ae8b6cd80fc87d1a59c2ec7f5ae47e9","object":"response","created_at":1793636045,"status":"completed","access_programs":null,"background":false,"completed_at":1793636045,"error":null,"frequency_penalty":0.0,"incomplete_details":null,"instructions":null,"max_output_tokens":50,"max_tool_calls":null,"model":"gpt-4o-mini-2024-07-18","moderation":null,"output":[{"id":"msg_05d53df3d23eb9eb016ae8b6cdf79887d1a57a49654be7108f","type":"message","status":"completed","content":[{"type":"output_text","annotations":[],"logprobs":[],"text":"4"}],"role":"assistant"}],"parallel_tool_calls":true,"presence_penalty":0.0,"previous_response_id":null,"prompt_cache_key":null,"prompt_cache_retention":"in_memory","reasoning":{"context":null,"effort":null,"summary":null},"safety_identifier":null,"service_tier":"default","store":false,"temperature":1.0,"text":{"format":{"type":"text"},"verbosity":"medium"},"tool_choice":"auto","tool_usage":{"image_gen":{"input_tokens":0,"input_tokens_details":{"image_tokens":0,"text_tokens":0},"output_tokens":0,"output_tokens_details":{"image_tokens":0,"text_tokens":0},"total_tokens":0},"web_search":{"num_requests":0}},"tools":[],"top_logprobs":0,"top_p":1.0,"truncation":"disabled","usage":{"input_tokens":20,"input_tokens_details":{"cache_write_tokens":0,"cached_tokens":0},"output_tokens":2,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":22},"user":null,"metadata":{}},"sequence_number":8} + + + ' + headers: + access-control-expose-headers: + - X-Request-ID, Server-Timing-Total + - CF-Ray + - CF-Ray + alt-svc: + - h3=":443"; ma=86400 + cf-cache-status: + - DYNAMIC + cf-ray: + - a5446e2378cfaabf-YYZ + connection: + - keep-alive + content-type: + - text/event-stream; charset=utf-8 + date: + - Mon, 02 Nov 2026 16:14:06 GMT + openai-processing-ms: + - '479' + 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: + - '149999962' + x-ratelimit-reset-requests: + - 2ms + x-ratelimit-reset-tokens: + - 0s + x-request-id: + - req_5fe15e08046d4c29a3e5fe363b1cd3f1 + 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 d7b8af5fd..1a07933ce 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 @@ -2491,6 +2491,42 @@ async def test_model_fallback_traces_each_attempt(memory_logger, candidate_kind) assert agent_span["metadata"]["provider"] == "openai" +@pytest.mark.vcr +@pytest.mark.asyncio +async def test_model_fallback_stream_closes_failed_attempt(memory_logger): + try: + from pydantic_ai.capabilities import ModelFallback + except ImportError: + pytest.skip("ModelFallback is unavailable in this Pydantic AI version") + from pydantic_ai.models.openai import OpenAIChatModel, OpenAIResponsesModel + + agent = Agent( + OpenAIChatModel("braintrust-nonexistent-model-for-fallback-test"), + capabilities=[ModelFallback(OpenAIResponsesModel("gpt-4o-mini"))], + model_settings=ModelSettings(max_tokens=50), + ) + async with agent.run_stream(TEST_PROMPT) as result: + output = await result.get_output() + assert "4" in output + + spans = memory_logger.pop() + agent_span = _find_agent_span(spans) + model_spans = [span for span in spans if span["span_attributes"].get("type") == SpanTypeAttribute.LLM] + assert agent_span is not None + assert len(model_spans) == 2 + failed = next( + span + for span in model_spans + if span["metadata"].get("model") == "braintrust-nonexistent-model-for-fallback-test" + ) + answered = next(span for span in model_spans if span["metadata"].get("model") == "gpt-4o-mini") + assert failed.get("error") + assert failed["metrics"].get("end") is not None + assert not answered.get("error") + assert answered["metrics"].get("end") is not None + assert failed["span_parents"] == answered["span_parents"] == [agent_span["span_id"]] + + @pytest.mark.vcr @pytest.mark.asyncio async def test_agent_with_error_handling(memory_logger): diff --git a/py/src/braintrust/integrations/pydantic_ai/tracing.py b/py/src/braintrust/integrations/pydantic_ai/tracing.py index cc96cecdc..c4ee2c145 100644 --- a/py/src/braintrust/integrations/pydantic_ai/tracing.py +++ b/py/src/braintrust/integrations/pydantic_ai/tracing.py @@ -770,7 +770,11 @@ async def __aenter__(self): self.span_cm.__enter__() self.start_time = time.time() - self.stream = await self.stream_cm.__aenter__() + try: + self.stream = await self.stream_cm.__aenter__() + except BaseException: + self.span_cm.__exit__(*sys.exc_info()) + raise return _DirectStreamIteratorProxy(self.stream, self)