Merge pull request #13573 from Ivy-Interactive/feature/braintrust-span-name-metadata

Feature/braintrust span name metadata

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

@@ -71,6 +71,10 @@ curl -X POST 'http://0.0.0.0:4000/chat/completions' \
 
 It is recommended that you include the `project_id` or `project_name` to ensure your traces are being written out to the correct Braintrust project.
 
+### Custom Span Names
+
+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".
+
 <Tabs>
 <TabItem value="sdk" label="SDK">
 
@@ -84,7 +88,9 @@ response = litellm.completion(
     "project_id": "1234",
     # passing project_name will try to find a project with that name, or create one if it doesn't exist
     # if both project_id and project_name are passed, project_id will be used
-    # "project_name": "my-special-project"
+    # "project_name": "my-special-project",
+    # custom span name for this operation (default: "Chat Completion")
+    "span_name": "User Greeting Handler"
   }
 )
 ```
@@ -99,6 +105,7 @@ response = litellm.completion(
   ],
   metadata={
     "project_id": "1234",
+    "span_name": "Custom Operation",
     "item1": "an item",
     "item2": "another item"
   }
@@ -121,7 +128,8 @@ curl -X POST 'http://0.0.0.0:4000/chat/completions' \
         { "role": "user", "content": "What time is it now? Use your tool"}
     ],
     "metadata": {
-        "project_id": "my-special-project"
+        "project_id": "my-special-project",
+        "span_name": "Tool Usage Request"
     }
 }'
 ```
@@ -146,7 +154,8 @@ response = client.chat.completions.create(
     ],
     extra_body={ # pass in any provider-specific param, if not supported by openai, https://docs.litellm.ai/docs/completion/input#provider-specific-params
         "metadata": { # 👈 use for logging additional params (e.g. to braintrust)
-            "project_id": "my-special-project"
+            "project_id": "my-special-project",
+            "span_name": "Poetry Generation"
         }
     }
 )
@@ -168,3 +177,7 @@ Here's everything you can pass in metadata for a braintrust request
 `braintrust_*` - If you are adding metadata from _proxy request headers_, any metadata field starting with `braintrust_` will be passed as metadata to the logging request. If you are using the SDK, just pass your metadata like normal (e.g., `metadata={"project_name": "my-test-project", "item1": "an item", "item2": "another item"}`)
 
 `project_id` - Set the project id for a braintrust call. Default is `litellm`.
+
+`project_name` - Set the project name for a braintrust call. Will try to find a project with that name, or create one if it doesn't exist. If both `project_id` and `project_name` are passed, `project_id` will be used.
+
+`span_name` - Set a custom span name for the operation. Default is `"Chat Completion"`. Use this to provide more descriptive names for different types of operations in your application (e.g., "User Query", "Document Summary", "Code Generation").

litellm/integrations/braintrust_logging.py

@@ -274,12 +274,15 @@ class BraintrustLogger(CustomLogger):
                     "end": end_time.timestamp(),
                 }
 
+            # Allow metadata override for span name
+            span_name = metadata.get("span_name", "Chat Completion")
+            
             request_data = {
                 "id": litellm_call_id,
                 "input": prompt["messages"],
                 "metadata": clean_metadata,
                 "tags": tags,
-                "span_attributes": {"name": "Chat Completion", "type": "llm"},
+                "span_attributes": {"name": span_name, "type": "llm"},
             }
             if choices is not None:
                 request_data["output"] = [choice.dict() for choice in choices]
@@ -426,13 +429,16 @@ class BraintrustLogger(CustomLogger):
                         - api_call_start_time.timestamp()
                     )
 
+            # Allow metadata override for span name
+            span_name = metadata.get("span_name", "Chat Completion")
+            
             request_data = {
                 "id": litellm_call_id,
                 "input": prompt["messages"],
                 "output": output,
                 "metadata": clean_metadata,
                 "tags": tags,
-                "span_attributes": {"name": "Chat Completion", "type": "llm"},
+                "span_attributes": {"name": span_name, "type": "llm"},
             }
             if choices is not None:
                 request_data["output"] = [choice.dict() for choice in choices]

tests/test_litellm/integrations/test_braintrust_logging.py

@@ -1,7 +1,9 @@
 import os
 import unittest
-from unittest.mock import patch
+from datetime import datetime
+from unittest.mock import MagicMock, Mock, patch
 
+import litellm
 from litellm.integrations.braintrust_logging import BraintrustLogger
 
 class TestBraintrustLogger(unittest.TestCase):
@@ -40,4 +42,244 @@ class TestBraintrustLogger(unittest.TestCase):
         with patch.dict(os.environ, {}, clear=True):
             with self.assertRaises(Exception) as context:
                 BraintrustLogger(api_key=None)
-            self.assertIn("Missing keys=['BRAINTRUST_API_KEY']", str(context.exception))
+            self.assertIn("Missing keys=['BRAINTRUST_API_KEY']", str(context.exception))
+
+    @patch('litellm.integrations.braintrust_logging.global_braintrust_sync_http_handler')
+    def test_log_success_event_with_default_span_name(self, mock_http_handler):
+        """Test log_success_event uses default span name when not provided."""
+        # Setup
+        logger = BraintrustLogger(api_key="test-key")
+        logger.default_project_id = "test-project-id"
+        
+        mock_response = Mock()
+        mock_response.json.return_value = {"id": "test-project-id"}
+        mock_http_handler.post.return_value = mock_response
+        
+        # Create a mock response object
+        message_mock = Mock()
+        message_mock.json = Mock(return_value={"content": "test"})
+        
+        choice_mock = Mock()
+        choice_mock.message = message_mock
+        choice_mock.dict = Mock(return_value={"message": {"content": "test"}})
+        
+        response_obj = Mock(spec=litellm.ModelResponse)
+        response_obj.choices = [choice_mock]
+        # Mock the __getitem__ to support response_obj["choices"]
+        response_obj.__getitem__ = Mock(return_value=[choice_mock])
+        response_obj.usage = litellm.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": {}},
+            "model": "gpt-3.5-turbo",
+            "response_cost": 0.001
+        }
+        
+        # Execute
+        logger.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'], 'Chat Completion')
+
+    @patch('litellm.integrations.braintrust_logging.global_braintrust_sync_http_handler')
+    def test_log_success_event_with_custom_span_name(self, mock_http_handler):
+        """Test log_success_event uses custom span name when provided."""
+        # Setup
+        logger = BraintrustLogger(api_key="test-key")
+        logger.default_project_id = "test-project-id"
+        
+        mock_response = Mock()
+        mock_response.json.return_value = {"id": "test-project-id"}
+        mock_http_handler.post.return_value = mock_response
+        
+        # Create a mock response object
+        message_mock = Mock()
+        message_mock.json = Mock(return_value={"content": "test"})
+        
+        choice_mock = Mock()
+        choice_mock.message = message_mock
+        choice_mock.dict = Mock(return_value={"message": {"content": "test"}})
+        
+        response_obj = Mock(spec=litellm.ModelResponse)
+        response_obj.choices = [choice_mock]
+        response_obj.__getitem__ = Mock(return_value=[choice_mock])
+        response_obj.usage = litellm.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": "Custom Operation"}},
+            "model": "gpt-3.5-turbo",
+            "response_cost": 0.001
+        }
+        
+        # Execute
+        logger.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'], 'Custom Operation')
+
+    @patch('litellm.integrations.braintrust_logging.global_braintrust_http_handler')
+    async def test_async_log_success_event_with_default_span_name(self, mock_http_handler):
+        """Test async_log_success_event uses default span name when not provided."""
+        # Setup
+        logger = BraintrustLogger(api_key="test-key")
+        logger.default_project_id = "test-project-id"
+        
+        mock_response = Mock()
+        mock_response.json.return_value = {"id": "test-project-id"}
+        mock_http_handler.post = MagicMock(return_value=mock_response)
+        
+        # Create a mock response object
+        message_mock = Mock()
+        message_mock.json = Mock(return_value={"content": "test"})
+        
+        choice_mock = Mock()
+        choice_mock.message = message_mock
+        choice_mock.dict = Mock(return_value={"message": {"content": "test"}})
+        
+        response_obj = Mock(spec=litellm.ModelResponse)
+        response_obj.choices = [choice_mock]
+        response_obj.__getitem__ = Mock(return_value=[choice_mock])
+        response_obj.usage = litellm.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": {}},
+            "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'], 'Chat Completion')
+
+    @patch('litellm.integrations.braintrust_logging.global_braintrust_http_handler')
+    async def test_async_log_success_event_with_custom_span_name(self, mock_http_handler):
+        """Test async_log_success_event uses custom span name when provided."""
+        # Setup
+        logger = BraintrustLogger(api_key="test-key")
+        logger.default_project_id = "test-project-id"
+        
+        mock_response = Mock()
+        mock_response.json.return_value = {"id": "test-project-id"}
+        mock_http_handler.post = MagicMock(return_value=mock_response)
+        
+        # Create a mock response object
+        message_mock = Mock()
+        message_mock.json = Mock(return_value={"content": "test"})
+        
+        choice_mock = Mock()
+        choice_mock.message = message_mock
+        choice_mock.dict = Mock(return_value={"message": {"content": "test"}})
+        
+        response_obj = Mock(spec=litellm.ModelResponse)
+        response_obj.choices = [choice_mock]
+        response_obj.__getitem__ = Mock(return_value=[choice_mock])
+        response_obj.usage = litellm.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"}},
+            "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')
+
+    @patch('litellm.integrations.braintrust_logging.global_braintrust_sync_http_handler')
+    def test_span_name_with_multiple_metadata_fields(self, mock_http_handler):
+        """Test that span_name works correctly alongside other metadata fields."""
+        # Setup
+        logger = BraintrustLogger(api_key="test-key")
+        logger.default_project_id = "test-project-id"
+        
+        mock_response = Mock()
+        mock_response.json.return_value = {"id": "test-project-id"}
+        mock_http_handler.post.return_value = mock_response
+        
+        # Create a mock response object
+        message_mock = Mock()
+        message_mock.json = Mock(return_value={"content": "test"})
+        
+        choice_mock = Mock()
+        choice_mock.message = message_mock
+        choice_mock.dict = Mock(return_value={"message": {"content": "test"}})
+        
+        response_obj = Mock(spec=litellm.ModelResponse)
+        response_obj.choices = [choice_mock]
+        response_obj.__getitem__ = Mock(return_value=[choice_mock])
+        response_obj.usage = litellm.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": "Multi Metadata Test",
+                    "project_id": "custom-project",
+                    "user_id": "user123",
+                    "session_id": "session456"
+                }
+            },
+            "model": "gpt-3.5-turbo",
+            "response_cost": 0.001
+        }
+        
+        # Execute
+        logger.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']
+        
+        # Check span name
+        self.assertEqual(json_data['events'][0]['span_attributes']['name'], 'Multi Metadata Test')
+        
+        # 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')

tests/test_litellm/integrations/test_braintrust_span_name.py

@@ -0,0 +1,199 @@
+import json
+import os
+import unittest
+from datetime import datetime
+from unittest.mock import MagicMock, Mock, patch
+
+import litellm
+from litellm.integrations.braintrust_logging import BraintrustLogger
+
+
+class TestBraintrustSpanName(unittest.TestCase):
+    """Test custom span_name functionality in Braintrust logging."""
+
+    @patch('litellm.integrations.braintrust_logging.global_braintrust_sync_http_handler')
+    def test_default_span_name(self, mock_http_handler):
+        """Test that default span name is 'Chat Completion' when not provided."""
+        # Setup
+        logger = BraintrustLogger(api_key="test-key")
+        logger.default_project_id = "test-project-id"
+        
+        # Mock HTTP response
+        mock_http_handler.post.return_value = Mock()
+        
+        # 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": {}},
+            "model": "gpt-3.5-turbo",
+            "response_cost": 0.001
+        }
+        
+        # Execute
+        logger.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'], 'Chat Completion')
+
+    @patch('litellm.integrations.braintrust_logging.global_braintrust_sync_http_handler')
+    def test_custom_span_name(self, mock_http_handler):
+        """Test that custom span name is used when provided in metadata."""
+        # Setup
+        logger = BraintrustLogger(api_key="test-key")
+        logger.default_project_id = "test-project-id"
+        
+        # Mock HTTP response
+        mock_http_handler.post.return_value = Mock()
+        
+        # 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": "Custom Operation"}},
+            "model": "gpt-3.5-turbo",
+            "response_cost": 0.001
+        }
+        
+        # Execute
+        logger.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'], 'Custom Operation')
+
+    @patch('litellm.integrations.braintrust_logging.global_braintrust_sync_http_handler')
+    def test_span_name_with_other_metadata(self, mock_http_handler):
+        """Test that span_name works alongside other metadata fields."""
+        # Setup
+        logger = BraintrustLogger(api_key="test-key")
+        logger.default_project_id = "test-project-id"
+        
+        # Mock HTTP response
+        mock_http_handler.post.return_value = Mock()
+        
+        # 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": "Multi Metadata Test",
+                    "project_id": "custom-project",
+                    "user_id": "user123",
+                    "session_id": "session456",
+                    "environment": "production"
+                }
+            },
+            "model": "gpt-3.5-turbo",
+            "response_cost": 0.001
+        }
+        
+        # Execute
+        logger.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']
+        
+        # Check span name
+        self.assertEqual(json_data['events'][0]['span_attributes']['name'], 'Multi Metadata Test')
+        
+        # Check that other metadata is preserved (except for filtered keys)
+        event_metadata = json_data['events'][0]['metadata']
+        self.assertEqual(event_metadata['user_id'], 'user123')
+        self.assertEqual(event_metadata['session_id'], 'session456')
+        self.assertEqual(event_metadata['environment'], 'production')
+        
+        # Span name should be in span_attributes, not in metadata
+        self.assertIn('span_name', event_metadata)  # span_name is also kept in metadata
+
+    @patch('litellm.integrations.braintrust_logging.global_braintrust_http_handler')
+    async def test_async_custom_span_name(self, mock_http_handler):
+        """Test async logging with custom span name."""
+        # Setup
+        logger = BraintrustLogger(api_key="test-key")
+        logger.default_project_id = "test-project-id"
+        
+        # Mock async HTTP response
+        mock_http_handler.post = MagicMock(return_value=Mock())
+        
+        # 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"}},
+            "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')
+
+
+if __name__ == "__main__":
+    unittest.main()