From 3e4281c18adfc2f33252637452f3d94a9630e906 Mon Sep 17 00:00:00 2001 From: Javier Garcia <112962975+javiergarciapleo@users.noreply.github.com> Date: Thu, 23 Oct 2025 10:09:25 +0200 Subject: [PATCH 1/2] Add logic to include span attributes in async process --- .../integrations/test_braintrust_span_name.py | 70 +++++++++++++++++++ 1 file changed, 70 insertions(+) diff --git a/tests/test_litellm/integrations/test_braintrust_span_name.py b/tests/test_litellm/integrations/test_braintrust_span_name.py index 7050a6d355f1..ac1f5237f6c7 100644 --- a/tests/test_litellm/integrations/test_braintrust_span_name.py +++ b/tests/test_litellm/integrations/test_braintrust_span_name.py @@ -294,6 +294,76 @@ def test_span_attributes_with_multiple_metadata_fields(self, MockHTTPHandler): self.assertEqual(event_metadata['user_id'], 'user123') self.assertEqual(event_metadata['session_id'], 'session456') + @patch("litellm.integrations.braintrust_logging.get_async_httpx_client") + async def test_async_span_attributes_with_multiple_metadata_fields(self, mock_get_http_handler): + """Test async logging with custom span name.""" + # Mock async HTTP response + mock_http_handler = MagicMock() + mock_http_handler.post = MagicMock(return_value=Mock()) + mock_get_http_handler.return_value = mock_http_handler + + # Setup + logger = BraintrustLogger(api_key="test-key") + logger.default_project_id = "test-project-id" + + # Create a properly structured mock response + response_obj = litellm.ModelResponse( + id="test-id", + object="chat.completion", + created=1234567890, + model="gpt-3.5-turbo", + choices=[ + { + "index": 0, + "message": {"role": "assistant", "content": "test response"}, + "finish_reason": "stop", + } + ], + usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + ) + + kwargs = { + "litellm_call_id": "test-call-id", + "messages": [{"role": "user", "content": "test"}], + "litellm_params": { + "metadata": { + "span_name": "Async Custom Operation", + "span_id": "span_id", + "root_span_id": "root_span_id", + "span_parents": "span_parent1,span_parent2", + "project_id": "custom-project", + "user_id": "user123", + "session_id": "session456" + } + }, + "model": "gpt-3.5-turbo", + "response_cost": 0.001, + } + + # Execute + await logger.async_log_success_event( + kwargs, response_obj, datetime.now(), datetime.now() + ) + + # Verify + call_args = mock_http_handler.post.call_args + self.assertIsNotNone(call_args) + json_data = call_args.kwargs["json"] + self.assertEqual( + json_data["events"][0]["span_attributes"]["name"], "Async Custom Operation" + ) + # Check span name + self.assertEqual(json_data['events'][0]['span_attributes']['name'], 'Multi Metadata Test') + self.assertEqual(json_data['events'][0]['span_id'], 'span_id') + self.assertEqual(json_data['events'][0]['root_span_id'], 'root_span_id') + self.assertEqual(json_data['events'][0]['span_parents'][0], 'span_parent1') + self.assertEqual(json_data['events'][0]['span_parents'][1], 'span_parent2') + + # Check that other metadata is preserved + event_metadata = json_data['events'][0]['metadata'] + self.assertEqual(event_metadata['user_id'], 'user123') + self.assertEqual(event_metadata['session_id'], 'session456') + if __name__ == "__main__": unittest.main() From ed286001467a0ace1d2e3bf36d05736dd9151358 Mon Sep 17 00:00:00 2001 From: Javier Garcia <112962975+javiergarciapleo@users.noreply.github.com> Date: Thu, 23 Oct 2025 10:12:52 +0200 Subject: [PATCH 2/2] Commit change --- litellm/integrations/braintrust_logging.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/litellm/integrations/braintrust_logging.py b/litellm/integrations/braintrust_logging.py index 364fa3f5defd..4ab373139f35 100644 --- a/litellm/integrations/braintrust_logging.py +++ b/litellm/integrations/braintrust_logging.py @@ -350,6 +350,20 @@ async def async_log_success_event( # noqa: PLR0915 # Allow metadata override for span name span_name = dynamic_metadata.get("span_name", "Chat Completion") + + # Span parents is a special case + span_parents = dynamic_metadata.get("span_parents") + + # Convert comma-separated string to list if present + if span_parents: + span_parents = [s.strip() for s in span_parents.split(",") if s.strip()] + + # Add optional span attributes only if present + span_attributes = { + "span_id": dynamic_metadata.get("span_id"), + "root_span_id": dynamic_metadata.get("root_span_id"), + "span_parents": span_parents, + } request_data = { "id": litellm_call_id, @@ -359,6 +373,12 @@ async def async_log_success_event( # noqa: PLR0915 "tags": tags, "span_attributes": {"name": span_name, "type": "llm"}, } + + # Only add those that are not None (or falsy) + for key, value in span_attributes.items(): + if value: + request_data[key] = value + if choices is not None: request_data["output"] = [choice.dict() for choice in choices] else: