fix(pydantic-ai): record audio duration usage

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

@@ -180,6 +180,35 @@ def test_extract_response_metrics_leaf_fields():
     # pylint: enable=unsupported-membership-test,unsubscriptable-object
 
 
+def test_audio_seconds_is_recorded_as_metadata():
+    """Keep Pydantic AI's nonstandard audio duration out of canonical metrics."""
+    from braintrust.integrations.pydantic_ai.tracing import _extract_response_metadata, _extract_response_metrics
+    from pydantic_ai.messages import ModelResponse, TextPart
+    from pydantic_ai.usage import RequestUsage
+
+    usage = RequestUsage(input_tokens=10, input_audio_tokens=7)
+    supports_audio_seconds = hasattr(usage, "audio_seconds")
+    if supports_audio_seconds:
+        setattr(usage, "audio_seconds", 2.5)
+
+    response = ModelResponse(
+        parts=[TextPart(content="spoken response")],
+        usage=usage,
+    )
+    metrics = _extract_response_metrics(response, start_time=1.0, end_time=2.0)
+    metadata = _extract_response_metadata(response)
+
+    assert metrics is not None
+    # pylint: disable=unsupported-membership-test,unsubscriptable-object
+    assert "audio_seconds" not in metrics
+    assert metrics["prompt_audio_tokens"] == 7
+    if supports_audio_seconds:
+        assert metadata == {"pydantic_ai_usage": {"audio_seconds": 2.5}}
+    else:
+        assert metadata is None
+    # pylint: enable=unsupported-membership-test,unsubscriptable-object
+
+
 @pytest.mark.vcr
 @pytest.mark.asyncio
 async def test_direct_model_request_creates_nested_chat_span_without_class_scan(memory_logger, direct):

py/src/braintrust/integrations/pydantic_ai/tracing.py

@@ -480,8 +480,9 @@ async def model_request_wrapper(wrapped: Any, instance: Any, args: Any, kwargs:
 
             output = _shape_model_response(result)
             metrics = _extract_response_metrics(result, start_time, end_time)
+            response_metadata = _extract_response_metadata(result)
 
-            span.log(output=output, metrics=metrics)
+            span.log(output=output, metadata=response_metadata, metrics=metrics)
             return result
 
     def model_request_stream_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any):
@@ -774,9 +775,11 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
                         metrics = _extract_response_metrics(
                             final_response, self.start_time, end_time, self._first_token_time
                         )
+                        response_metadata = _extract_response_metadata(final_response)
                     else:
                         metrics = _wrapper_span_metrics(self.start_time, end_time, self._first_token_time)
-                    self.span_cm.log(output=output, metrics=metrics)
+                        response_metadata = None
+                    self.span_cm.log(output=output, metadata=response_metadata, metrics=metrics)
                 except Exception as e:
                     logger.debug(f"Failed to extract stream output/metrics: {e}")
 
@@ -1427,6 +1430,15 @@ def _extract_response_metrics(
     return metrics if metrics else None
 
 
+def _extract_response_metadata(response: Any) -> dict[str, Any] | None:
+    usage = getattr(response, "usage", None)
+    audio_seconds = getattr(usage, "audio_seconds", None)
+    if audio_seconds is None:
+        return None
+
+    return {"pydantic_ai_usage": {"audio_seconds": float(audio_seconds)}}
+
+
 class _ContextPropagatingAsyncContextManager(AbstractAsyncContextManager):
     def __init__(self, cm: Any, ctx: contextvars.Context):
         self._cm = cm