From 4558df5b6319e7cfe57c3a805bdbfe6bc989b98b Mon Sep 17 00:00:00 2001 From: Javier Garcia <112962975+javiergarciapleo@users.noreply.github.com> Date: Tue, 3 Mar 2026 08:17:34 +0100 Subject: [PATCH] Add logic to include prompt information in the log line --- litellm/integrations/braintrust_logging.py | 18 ++ .../integrations/test_braintrust_logging.py | 111 ++++++++- .../integrations/test_braintrust_span_name.py | 227 ++++++++++++++++++ 3 files changed, 355 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/braintrust_logging.py b/litellm/integrations/braintrust_logging.py index 364fa3f5def..e3133467eb5 100644 --- a/litellm/integrations/braintrust_logging.py +++ b/litellm/integrations/braintrust_logging.py @@ -187,6 +187,15 @@ class BraintrustLogger(CustomLogger): ): # support logging dynamic metadata to braintrust standard_logging_object[key] = value + prompt_id = dynamic_metadata.get("prompt_id") + prompt_slug = dynamic_metadata.get("prompt_slug") + if prompt_id is not None or prompt_slug is not None: + standard_logging_object["prompt"] = { + "id": prompt_id, + "project_id": project_id, + "slug": prompt_slug, + } + cost = kwargs.get("response_cost", None) metrics: Optional[dict] = None @@ -321,6 +330,15 @@ class BraintrustLogger(CustomLogger): ): # support logging dynamic metadata to braintrust standard_logging_object[key] = value + prompt_id = dynamic_metadata.get("prompt_id") + prompt_slug = dynamic_metadata.get("prompt_slug") + if prompt_id is not None or prompt_slug is not None: + standard_logging_object["prompt"] = { + "id": prompt_id, + "project_id": project_id, + "slug": prompt_slug, + } + cost = kwargs.get("response_cost", None) metrics: Optional[dict] = None diff --git a/tests/test_litellm/integrations/test_braintrust_logging.py b/tests/test_litellm/integrations/test_braintrust_logging.py index cb227148ed9..374fd238ddd 100644 --- a/tests/test_litellm/integrations/test_braintrust_logging.py +++ b/tests/test_litellm/integrations/test_braintrust_logging.py @@ -303,4 +303,113 @@ class TestBraintrustLogger(unittest.TestCase): # 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') \ No newline at end of file + self.assertEqual(event_metadata['session_id'], 'session456') + + @patch('litellm.integrations.braintrust_logging.HTTPHandler') + def test_log_success_event_with_prompt_id(self, MockHTTPHandler): + """Test log_success_event sets prompt fields in metadata when prompt_id is provided.""" + # Mock HTTP response + mock_response = Mock() + mock_response.json.return_value = {"id": "test-project-id"} + mock_http_handler = Mock() + mock_http_handler.post.return_value = mock_response + MockHTTPHandler.return_value = mock_http_handler + + # Setup + logger = BraintrustLogger(api_key="test-key") + logger.default_project_id = "test-project-id" + + # 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"}}) + choice_mock.__getitem__ = Mock(return_value=message_mock) + + 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": { + "prompt_id": "my-prompt-id", + "prompt_slug": "my-prompt-slug", + "project_id": "my-project-id", + } + }, + "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'] + event_metadata = json_data['events'][0]['metadata'] + self.assertIn('prompt', event_metadata) + self.assertEqual(event_metadata['prompt']['id'], 'my-prompt-id') + self.assertEqual(event_metadata['prompt']['project_id'], 'my-project-id') + self.assertEqual(event_metadata['prompt']['slug'], 'my-prompt-slug') + + @patch('litellm.integrations.braintrust_logging.HTTPHandler') + def test_log_success_event_without_prompt_id(self, MockHTTPHandler): + """Test log_success_event does not set prompt in metadata when prompt_id is absent.""" + # Mock HTTP response + mock_response = Mock() + mock_response.json.return_value = {"id": "test-project-id"} + mock_http_handler = Mock() + mock_http_handler.post.return_value = mock_response + MockHTTPHandler.return_value = mock_http_handler + + # Setup + logger = BraintrustLogger(api_key="test-key") + logger.default_project_id = "test-project-id" + + # 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"}}) + choice_mock.__getitem__ = Mock(return_value=message_mock) + + 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 + 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'] + event_metadata = json_data['events'][0]['metadata'] + self.assertNotIn('prompt', event_metadata) \ No newline at end of file diff --git a/tests/test_litellm/integrations/test_braintrust_span_name.py b/tests/test_litellm/integrations/test_braintrust_span_name.py index 7050a6d355f..a34baa7b520 100644 --- a/tests/test_litellm/integrations/test_braintrust_span_name.py +++ b/tests/test_litellm/integrations/test_braintrust_span_name.py @@ -294,6 +294,233 @@ class TestBraintrustSpanName(unittest.TestCase): self.assertEqual(event_metadata['user_id'], 'user123') self.assertEqual(event_metadata['session_id'], 'session456') + @patch("litellm.integrations.braintrust_logging.HTTPHandler") + def test_prompt_id_is_set_in_metadata(self, MockHTTPHandler): + """Test that prompt fields from dynamic_metadata are stored in event metadata.""" + # Mock HTTP response + mock_http_handler = Mock() + mock_http_handler.post.return_value = Mock() + MockHTTPHandler.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": { + "prompt_id": "my-prompt-id", + "prompt_slug": "my-prompt-slug", + "project_id": "my-project-id", + } + }, + "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"] + event_metadata = json_data["events"][0]["metadata"] + self.assertIn("prompt", event_metadata) + self.assertEqual(event_metadata["prompt"]["id"], "my-prompt-id") + self.assertEqual(event_metadata["prompt"]["project_id"], "my-project-id") + self.assertEqual(event_metadata["prompt"]["slug"], "my-prompt-slug") + + @patch("litellm.integrations.braintrust_logging.HTTPHandler") + def test_prompt_id_absent_when_not_provided(self, MockHTTPHandler): + """Test that prompt key is absent from event metadata when prompt_id is not in dynamic_metadata.""" + # Mock HTTP response + mock_http_handler = Mock() + mock_http_handler.post.return_value = Mock() + MockHTTPHandler.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": {}}, + "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"] + event_metadata = json_data["events"][0]["metadata"] + self.assertNotIn("prompt", event_metadata) + + @patch("litellm.integrations.braintrust_logging.HTTPHandler") + def test_prompt_id_alongside_other_metadata(self, MockHTTPHandler): + """Test that prompt fields work correctly alongside other metadata fields.""" + # Mock HTTP response + mock_http_handler = Mock() + mock_http_handler.post.return_value = Mock() + MockHTTPHandler.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": { + "prompt_id": "my-prompt-id", + "prompt_slug": "my-prompt-slug", + "project_id": "my-project-id", + "span_name": "Custom Operation", + "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"] + event_metadata = json_data["events"][0]["metadata"] + + # Check prompt is set correctly + self.assertIn("prompt", event_metadata) + self.assertEqual(event_metadata["prompt"]["id"], "my-prompt-id") + self.assertEqual(event_metadata["prompt"]["project_id"], "my-project-id") + self.assertEqual(event_metadata["prompt"]["slug"], "my-prompt-slug") + + # Check span name is also set correctly + self.assertEqual( + json_data["events"][0]["span_attributes"]["name"], "Custom Operation" + ) + + # Check other metadata is preserved + 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_prompt_id_is_set_in_metadata(self, mock_get_http_handler): + """Test async logging sets prompt fields in metadata when prompt info is provided.""" + # 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": { + "prompt_id": "async-prompt-id", + "prompt_slug": "async-prompt-slug", + "project_id": "async-project-id", + } + }, + "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"] + event_metadata = json_data["events"][0]["metadata"] + self.assertIn("prompt", event_metadata) + self.assertEqual(event_metadata["prompt"]["id"], "async-prompt-id") + self.assertEqual(event_metadata["prompt"]["project_id"], "async-project-id") + self.assertEqual(event_metadata["prompt"]["slug"], "async-prompt-slug") + if __name__ == "__main__": unittest.main()