mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
fix: support logging dynamic metadata values to braintrust
This commit is contained in:
Krrish Dholakia
parent
b6f6dc5c1c
commit
4a422ca897
2 changed files with 106 additions and 74 deletions
|
|
@ -159,7 +159,7 @@ class BraintrustLogger(CustomLogger):
|
|||
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 @@ class BraintrustLogger(CustomLogger):
|
|||
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 @@ class BraintrustLogger(CustomLogger):
|
|||
):
|
||||
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 @@ class BraintrustLogger(CustomLogger):
|
|||
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 @@ class BraintrustLogger(CustomLogger):
|
|||
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 @@ class BraintrustLogger(CustomLogger):
|
|||
):
|
||||
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
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ from litellm.integrations.braintrust_logging import BraintrustLogger
|
|||
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 @@ class TestBraintrustSpanName(unittest.TestCase):
|
|||
# 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 @@ class TestBraintrustSpanName(unittest.TestCase):
|
|||
# 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 @@ class TestBraintrustSpanName(unittest.TestCase):
|
|||
# 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 @@ class TestBraintrustSpanName(unittest.TestCase):
|
|||
"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']
|
||||
|
||||
# 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
|
||||
json_data = call_args.kwargs["json"]
|
||||
|
||||
@patch('litellm.integrations.braintrust_logging.get_async_httpx_client')
|
||||
# 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"]
|
||||
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
|
||||
|
||||
@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 @@ class TestBraintrustSpanName(unittest.TestCase):
|
|||
# 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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue