scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

Add logic to include prompt information in the log line

Browse source
This commit is contained in:
Javier Garcia 2026-03-03 08:17:34 +01:00
parent 647f2f5d86
commit 4558df5b63
3 changed files with 355 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
litellm/integrations/braintrust_logging.py
View file

@ -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

111
tests/test_litellm/integrations/test_braintrust_logging.py
View file

@ -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')
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)

227
tests/test_litellm/integrations/test_braintrust_span_name.py
View file

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