mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-13 03:59:12 +00:00
8a103e26ab
* test: move routing, caching and callback tests out of local_testing Moves the 37 routing, caching and callback files in tests/local_testing into mirrored tests/unit paths and tests/integration sdk and security groups per the migration manifest, and deletes the 8 rows it marks for deletion * test: cover least-busy routing and tighten migrated router, cache and pass-through tests Add router-level least-busy tests for the three legacy rows that had no destination, move completion fallback and model alias tests to their mirror files, move router batch tests next to the router, and strengthen tests whose mutants survived (specific deployment, batch pair grouping, custom cache key, router cache ttl). Add a unit test for configured pass-through Authorization forwarding. * test: tighten migrated callback, proxy config and cache ttl tests Assert the full StandardLoggingPayload key set and the wire user on callback tests, add an async embedding failure callback test, drive the prometheus redis failure through a real RedisCache, use the real salt key in proxy config tests, pin the redis ttl value and rename fallback test helpers * test: fix semantic cache assertions and order-dependent router and callback tests Count semantic cache embeddings on the route, drain async cache writes before the second lookup, echo the requested model from the callback provider mock so response_cost does not depend on model_cost entries other tests register, and move the two-key authorization assertion to the auth failure test it belongs to * test: cover model_list fan-out and the usage-based routing tpm limit lost in the move * test: wait for the streamed cache entry in redis and drop deleted files from the hosted mock guard * test: restore sync callback, audio redaction and langfuse export coverage, keep redis-compat filter on the moved file * ci: run the remaining router-keyword legacy tests on one node * test: refuse the redis write in-process instead of resolving an invalid host --------- Co-authored-by: yuneng <yuneng@berri.ai>
411 lines
16 KiB
Python
411 lines
16 KiB
Python
import json
|
|
import os
|
|
import unittest
|
|
from datetime import datetime
|
|
from typing import Final
|
|
from unittest.mock import MagicMock, Mock, patch
|
|
|
|
import httpx
|
|
import litellm
|
|
import pytest
|
|
import respx
|
|
from litellm.litellm_core_utils import litellm_logging
|
|
from litellm.integrations.braintrust_logging import BraintrustLogger
|
|
|
|
|
|
class TestBraintrustLogger(unittest.TestCase):
|
|
@patch.dict(os.environ, {"BRAINTRUST_API_KEY": "test-env-api-key"})
|
|
@patch.dict(os.environ, {"BRAINTRUST_API_BASE": "https://test-env-api.com/v1"})
|
|
def test_init_with_env_var(self):
|
|
"""Test BraintrustLogger initialization with environment variable."""
|
|
logger = BraintrustLogger()
|
|
self.assertEqual(logger.api_key, "test-env-api-key")
|
|
self.assertEqual(logger.api_base, "https://test-env-api.com/v1")
|
|
self.assertEqual(logger.headers["Authorization"], "Bearer test-env-api-key")
|
|
self.assertEqual(logger.headers["Content-Type"], "application/json")
|
|
|
|
def test_init_with_explicit_params(self):
|
|
"""Test BraintrustLogger initialization with explicit parameters."""
|
|
logger = BraintrustLogger(
|
|
api_key="explicit-key", api_base="https://custom-api.com/v1"
|
|
)
|
|
self.assertEqual(logger.api_key, "explicit-key")
|
|
self.assertEqual(logger.api_base, "https://custom-api.com/v1")
|
|
self.assertEqual(logger.headers["Authorization"], "Bearer explicit-key")
|
|
|
|
@patch.dict(os.environ, {}, clear=True)
|
|
def test_init_missing_api_key(self):
|
|
"""Test BraintrustLogger initialization fails without API key."""
|
|
with self.assertRaises(Exception) as context:
|
|
BraintrustLogger()
|
|
self.assertIn("Missing keys=['BRAINTRUST_API_KEY']", str(context.exception))
|
|
|
|
def test_validate_environment_with_api_key(self):
|
|
"""Test validate_environment method with valid API key."""
|
|
logger = BraintrustLogger(api_key="test-key")
|
|
# Should not raise an exception
|
|
logger.validate_environment(api_key="test-key")
|
|
|
|
def test_validate_environment_missing_api_key(self):
|
|
"""Test validate_environment method with missing API key."""
|
|
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))
|
|
|
|
@patch("litellm.integrations.braintrust_logging.HTTPHandler")
|
|
def test_log_success_event_with_default_span_name(self, MockHTTPHandler):
|
|
"""Test log_success_event uses default span name when not 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"}})
|
|
# Mock the __getitem__ to support response_obj["choices"][0]["message"]
|
|
choice_mock.__getitem__ = Mock(return_value=message_mock)
|
|
|
|
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.HTTPHandler")
|
|
def test_log_success_event_with_custom_span_name(self, MockHTTPHandler):
|
|
"""Test log_success_event uses custom span name when 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": {"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.get_async_httpx_client")
|
|
async def test_async_log_success_event_with_default_span_name(
|
|
self, mock_get_http_handler
|
|
):
|
|
"""Test async_log_success_event uses default span name when not provided."""
|
|
# Mock async HTTP response
|
|
mock_response = Mock()
|
|
mock_response.json.return_value = {"id": "test-project-id"}
|
|
mock_http_handler = MagicMock()
|
|
mock_http_handler.post = MagicMock(return_value=mock_response)
|
|
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 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
|
|
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.get_async_httpx_client")
|
|
async def test_async_log_success_event_with_custom_span_name(
|
|
self, mock_get_http_handler
|
|
):
|
|
"""Test async_log_success_event uses custom span name when provided."""
|
|
# Mock async HTTP response
|
|
mock_response = Mock()
|
|
mock_response.json.return_value = {"id": "test-project-id"}
|
|
mock_http_handler = MagicMock()
|
|
mock_http_handler.post = MagicMock(return_value=mock_response)
|
|
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 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": {"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.HTTPHandler")
|
|
def test_span_name_with_multiple_metadata_fields(self, MockHTTPHandler):
|
|
"""Test that span_name works correctly alongside other metadata fields."""
|
|
# 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": {
|
|
"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")
|
|
|
|
|
|
async def _assert_completion_is_logged_by_braintrust(
|
|
project_id: str, metadata: dict[str, object]
|
|
) -> None:
|
|
logger: Final = BraintrustLogger(
|
|
api_key="test-key",
|
|
api_base="https://api.braintrustdata.com/v1",
|
|
)
|
|
logger.default_project_id = project_id
|
|
messages: Final = [{"role": "user", "content": "Summarize this request"}]
|
|
response: Final = litellm.ModelResponse(
|
|
id="chatcmpl-test",
|
|
model="gpt-4o-mini",
|
|
choices=[
|
|
{
|
|
"index": 0,
|
|
"message": {"role": "assistant", "content": "The request is summarized."},
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
usage=litellm.Usage(prompt_tokens=2, completion_tokens=5, total_tokens=7),
|
|
)
|
|
kwargs: Final = {
|
|
"model": "gpt-4o-mini",
|
|
"messages": messages,
|
|
"litellm_params": {"metadata": metadata},
|
|
"litellm_call_id": "call-test",
|
|
"standard_logging_object": {
|
|
"model": "gpt-4o-mini",
|
|
"messages": messages,
|
|
"response": response.model_dump(),
|
|
"metadata": metadata,
|
|
},
|
|
}
|
|
|
|
with respx.mock(assert_all_called=True) as router:
|
|
insert: Final = router.post(
|
|
f"https://api.braintrustdata.com/v1/project_logs/{project_id}/insert"
|
|
).mock(return_value=httpx.Response(200, json={"ok": True}))
|
|
fixed_time: Final = datetime(2024, 1, 1)
|
|
logger.log_success_event(kwargs, response, fixed_time, fixed_time)
|
|
|
|
assert insert.call_count == 1
|
|
event: Final = json.loads(insert.calls[0].request.content)["events"][0]
|
|
assert event["input"] == messages
|
|
assert event["output"][0]["message"]["content"] == "The request is summarized."
|
|
assert event["metadata"]["model"] == "gpt-4o-mini"
|
|
assert event["metadata"]["messages"] == messages
|
|
assert event["metadata"]["response"]["choices"][0]["message"]["content"] == "The request is summarized."
|
|
assert event["metrics"]["prompt_tokens"] == response.usage.prompt_tokens
|
|
assert event["metrics"]["completion_tokens"] == response.usage.completion_tokens
|
|
assert event["metrics"]["total_tokens"] == response.usage.total_tokens
|
|
assert event["span_attributes"]["name"] == "Chat Completion"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_braintrust_logging() -> None:
|
|
await _assert_completion_is_logged_by_braintrust(
|
|
"default-project",
|
|
{},
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_braintrust_logging_specific_project_id() -> None:
|
|
await _assert_completion_is_logged_by_braintrust(
|
|
"123",
|
|
{"project_id": "123", "requester_metadata": {"request_id": "request-123"}},
|
|
)
|
|
|
|
|
|
def test_braintrust_logger_initialization_reuses_instance(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("BRAINTRUST_API_KEY", "test-api-key")
|
|
monkeypatch.setenv("BRAINTRUST_API_BASE", "https://api.braintrustdata.com/v1")
|
|
first: Final = litellm_logging._init_custom_logger_compatible_class(
|
|
"braintrust",
|
|
internal_usage_cache=None,
|
|
llm_router=None,
|
|
)
|
|
second: Final = litellm_logging._init_custom_logger_compatible_class(
|
|
"braintrust",
|
|
internal_usage_cache=None,
|
|
llm_router=None,
|
|
)
|
|
|
|
assert isinstance(first, BraintrustLogger)
|
|
assert second is first
|