Skip to content

Commit b0a3a7c

Browse files
Add details in docs (#15721)
* Add details in docs * add logic to set span attributes and unit tests * Restore html files * Remove html files * Remove html files
1 parent 1cfc462 commit b0a3a7c

File tree

3 files changed

+96
-0
lines changed

3 files changed

+96
-0
lines changed

docs/my-website/docs/observability/braintrust.md

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,12 @@ It is recommended that you include the `project_id` or `project_name` to ensure
7575

7676
You can customize the span name in Braintrust logging by passing `span_name` in the metadata. By default, the span name is set to "Chat Completion".
7777

78+
### Custom Span Attributes
79+
80+
You can customize the span id, root span name and span parents in Braintrust logging by passing `span_id`, `root_span_id` and `span_parents` in the metadata.
81+
`span_parents` should be a string containing a list of span ids, joined by ,
82+
83+
7884
<Tabs>
7985
<TabItem value="sdk" label="SDK">
8086

litellm/integrations/braintrust_logging.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -206,6 +206,20 @@ def log_success_event( # noqa: PLR0915
206206

207207
# Allow metadata override for span name
208208
span_name = dynamic_metadata.get("span_name", "Chat Completion")
209+
210+
# Span parents is a special case
211+
span_parents = dynamic_metadata.get("span_parents")
212+
213+
# Convert comma-separated string to list if present
214+
if span_parents:
215+
span_parents = [s.strip() for s in span_parents.split(",") if s.strip()]
216+
217+
# Add optional span attributes only if present
218+
span_attributes = {
219+
"span_id": dynamic_metadata.get("span_id"),
220+
"root_span_id": dynamic_metadata.get("root_span_id"),
221+
"span_parents": span_parents,
222+
}
209223

210224
request_data = {
211225
"id": litellm_call_id,
@@ -214,6 +228,12 @@ def log_success_event( # noqa: PLR0915
214228
"tags": tags,
215229
"span_attributes": {"name": span_name, "type": "llm"},
216230
}
231+
232+
# Only add those that are not None (or falsy)
233+
for key, value in span_attributes.items():
234+
if value:
235+
request_data[key] = value
236+
217237
if choices is not None:
218238
request_data["output"] = [choice.dict() for choice in choices]
219239
else:

tests/test_litellm/integrations/test_braintrust_span_name.py

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -224,6 +224,76 @@ async def test_async_custom_span_name(self, mock_get_http_handler):
224224
json_data["events"][0]["span_attributes"]["name"], "Async Custom Operation"
225225
)
226226

227+
@patch('litellm.integrations.braintrust_logging.HTTPHandler')
228+
def test_span_attributes_with_multiple_metadata_fields(self, MockHTTPHandler):
229+
"""Test that span_name works correctly alongside other metadata fields."""
230+
# Mock HTTP response
231+
mock_response = Mock()
232+
mock_response.json.return_value = {"id": "test-project-id"}
233+
mock_http_handler = Mock()
234+
mock_http_handler.post.return_value = mock_response
235+
MockHTTPHandler.return_value = mock_http_handler
236+
237+
# Setup
238+
logger = BraintrustLogger(api_key="test-key")
239+
logger.default_project_id = "test-project-id"
240+
241+
# Create a mock response object
242+
message_mock = Mock()
243+
message_mock.json = Mock(return_value={"content": "test"})
244+
245+
choice_mock = Mock()
246+
choice_mock.message = message_mock
247+
choice_mock.dict = Mock(return_value={"message": {"content": "test"}})
248+
choice_mock.__getitem__ = Mock(return_value=message_mock)
249+
250+
response_obj = Mock(spec=litellm.ModelResponse)
251+
response_obj.choices = [choice_mock]
252+
response_obj.__getitem__ = Mock(return_value=[choice_mock])
253+
response_obj.usage = litellm.Usage(
254+
prompt_tokens=10,
255+
completion_tokens=20,
256+
total_tokens=30
257+
)
258+
259+
kwargs = {
260+
"litellm_call_id": "test-call-id",
261+
"messages": [{"role": "user", "content": "test"}],
262+
"litellm_params": {
263+
"metadata": {
264+
"span_name": "Multi Metadata Test",
265+
"span_id": "span_id",
266+
"root_span_id": "root_span_id",
267+
"span_parents": "span_parent1,span_parent2",
268+
"project_id": "custom-project",
269+
"user_id": "user123",
270+
"session_id": "session456"
271+
}
272+
},
273+
"model": "gpt-3.5-turbo",
274+
"response_cost": 0.001
275+
}
276+
277+
# Execute
278+
logger.log_success_event(kwargs, response_obj, datetime.now(), datetime.now())
279+
280+
# Verify
281+
call_args = mock_http_handler.post.call_args
282+
self.assertIsNotNone(call_args)
283+
json_data = call_args.kwargs['json']
284+
285+
# Check span name
286+
self.assertEqual(json_data['events'][0]['span_attributes']['name'], 'Multi Metadata Test')
287+
self.assertEqual(json_data['events'][0]['span_id'], 'span_id')
288+
self.assertEqual(json_data['events'][0]['root_span_id'], 'root_span_id')
289+
self.assertEqual(json_data['events'][0]['span_parents'][0], 'span_parent1')
290+
self.assertEqual(json_data['events'][0]['span_parents'][1], 'span_parent2')
291+
292+
# Check that other metadata is preserved
293+
event_metadata = json_data['events'][0]['metadata']
294+
self.assertEqual(event_metadata['user_id'], 'user123')
295+
self.assertEqual(event_metadata['session_id'], 'session456')
296+
227297

228298
if __name__ == "__main__":
229299
unittest.main()

0 commit comments

Comments
 (0)