Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions litellm/integrations/braintrust_logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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:
Expand Down
70 changes: 70 additions & 0 deletions tests/test_litellm/integrations/test_braintrust_span_name.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Loading