fix: support logging dynamic metadata values to braintrust

Author: krrishdholakia

litellm/integrations/braintrust_logging.py

@@ -159,7 +159,7 @@ def log_success_event(  # noqa: PLR0915
                 output = response_obj["data"]
 
             litellm_params = kwargs.get("litellm_params", {}) or {}
-            dynamic_metadata = litellm_params.get("dynamic_metadata", {}) or {}
+            dynamic_metadata = litellm_params.get("metadata", {}) or {}
 
             # Get project_id from metadata or create default if needed
             project_id = dynamic_metadata.get("project_id")
@@ -175,6 +175,7 @@ def log_success_event(  # noqa: PLR0915
                 project_id = self.default_project_id
 
             tags = []
+
             if isinstance(dynamic_metadata, dict):
                 for key, value in dynamic_metadata.items():
                     # generate langfuse tags - Default Tags sent to Langfuse from LiteLLM Proxy
@@ -185,6 +186,11 @@ def log_success_event(  # noqa: PLR0915
                     ):
                         tags.append(f"{key}:{value}")
 
+                    if (
+                        isinstance(value, str) and key not in standard_logging_object
+                    ):  # support logging dynamic metadata to braintrust
+                        standard_logging_object[key] = value
+
             cost = kwargs.get("response_cost", None)
 
             metrics: Optional[dict] = None
@@ -265,9 +271,7 @@ async def async_log_success_event(  # noqa: PLR0915
                 output = response_obj["data"]
 
             litellm_params = kwargs.get("litellm_params", {})
-            dynamic_metadata = litellm_params.get("dynamic_metadata", {}) or {}
-
-            clean_metadata = {}
+            dynamic_metadata = litellm_params.get("metadata", {}) or {}
 
             # Get project_id from metadata or create default if needed
             project_id = dynamic_metadata.get("project_id")
@@ -285,6 +289,7 @@ async def async_log_success_event(  # noqa: PLR0915
                 project_id = self.default_project_id
 
             tags = []
+
             if isinstance(dynamic_metadata, dict):
                 for key, value in dynamic_metadata.items():
                     # generate langfuse tags - Default Tags sent to Langfuse from LiteLLM Proxy
@@ -295,6 +300,11 @@ async def async_log_success_event(  # noqa: PLR0915
                     ):
                         tags.append(f"{key}:{value}")
 
+                    if (
+                        isinstance(value, str) and key not in standard_logging_object
+                    ):  # support logging dynamic metadata to braintrust
+                        standard_logging_object[key] = value
+
             cost = kwargs.get("response_cost", None)
 
             metrics: Optional[dict] = None

tests/test_litellm/integrations/test_braintrust_span_name.py

@@ -11,7 +11,7 @@
 class TestBraintrustSpanName(unittest.TestCase):
     """Test custom span_name functionality in Braintrust logging."""
 
-    @patch('litellm.integrations.braintrust_logging.HTTPHandler')
+    @patch("litellm.integrations.braintrust_logging.HTTPHandler")
     def test_default_span_name(self, MockHTTPHandler):
         """Test that default span name is 'Chat Completion' when not provided."""
         # Mock HTTP response
@@ -22,39 +22,43 @@ def test_default_span_name(self, MockHTTPHandler):
         # 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}
+            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
+            "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')
+        json_data = call_args.kwargs["json"]
+        self.assertEqual(
+            json_data["events"][0]["span_attributes"]["name"], "Chat Completion"
+        )
 
-    @patch('litellm.integrations.braintrust_logging.HTTPHandler')
+    @patch("litellm.integrations.braintrust_logging.HTTPHandler")
     def test_custom_span_name(self, MockHTTPHandler):
         """Test that custom span name is used when provided in metadata."""
         # Mock HTTP response
@@ -65,39 +69,43 @@ def test_custom_span_name(self, MockHTTPHandler):
         # 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}
+            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
+            "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')
+        json_data = call_args.kwargs["json"]
+        self.assertEqual(
+            json_data["events"][0]["span_attributes"]["name"], "Custom Operation"
+        )
 
-    @patch('litellm.integrations.braintrust_logging.HTTPHandler')
+    @patch("litellm.integrations.braintrust_logging.HTTPHandler")
     def test_span_name_with_other_metadata(self, MockHTTPHandler):
         """Test that span_name works alongside other metadata fields."""
         # Mock HTTP response
@@ -108,21 +116,23 @@ def test_span_name_with_other_metadata(self, MockHTTPHandler):
         # 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}
+            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"}],
@@ -132,34 +142,40 @@ def test_span_name_with_other_metadata(self, MockHTTPHandler):
                     "project_id": "custom-project",
                     "user_id": "user123",
                     "session_id": "session456",
-                    "environment": "production"
+                    "environment": "production",
                 }
             },
             "model": "gpt-3.5-turbo",
-            "response_cost": 0.001
+            "response_cost": 0.001,
+            "standard_logging_object": {
+                "user_id": "user123",
+            },
         }
-        
+
         # 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']
-        
+        json_data = call_args.kwargs["json"]
+
         # Check span name
-        self.assertEqual(json_data['events'][0]['span_attributes']['name'], 'Multi Metadata Test')
-        
+        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')
-        
+        event_metadata = json_data["events"][0]["metadata"]
+        print(event_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
+        self.assertIn("span_name", event_metadata)  # span_name is also kept in metadata
 
-    @patch('litellm.integrations.braintrust_logging.get_async_httpx_client')
+    @patch("litellm.integrations.braintrust_logging.get_async_httpx_client")
     async def test_async_custom_span_name(self, mock_get_http_handler):
         """Test async logging with custom span name."""
         # Mock async HTTP response
@@ -170,38 +186,44 @@ async def test_async_custom_span_name(self, mock_get_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}
+            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
+            "response_cost": 0.001,
         }
-        
+
         # Execute
-        await logger.async_log_success_event(kwargs, response_obj, datetime.now(), datetime.now())
-        
+        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')
+        json_data = call_args.kwargs["json"]
+        self.assertEqual(
+            json_data["events"][0]["span_attributes"]["name"], "Async Custom Operation"
+        )
 
 
 if __name__ == "__main__":
-    unittest.main()
+    unittest.main()