test(braintrust-logging): add span_name tests for events

Add tests to verify custom and default span_name in BraintrustLogger,
including async, metadata merging, and span name behavior.

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()