mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
6a0d03914c
* test: drop the cwd-relative sys.path.insert calls from the test suite
TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.
Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.
Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.
* test: drop the duplicate imports the sys.path sweep exposed to F811
* test(pre-call-utils): restore the os import the new bedrock tests need
434 lines
17 KiB
Python
434 lines
17 KiB
Python
import os
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
|
|
import litellm
|
|
from litellm.integrations.langsmith import LangsmithLogger
|
|
|
|
|
|
@pytest.fixture
|
|
def reset_redact_flag():
|
|
"""Reset redact_user_api_key_info between tests so global state doesn't leak."""
|
|
original = litellm.redact_user_api_key_info
|
|
yield
|
|
litellm.redact_user_api_key_info = original
|
|
|
|
|
|
class TestLangsmithLoggerInit:
|
|
"""Test cases for LangSmith logger initialization, particularly sampling rate handling.
|
|
|
|
These tests verify that the sampling_rate attribute is set during initialization.
|
|
Note: The current implementation has some edge cases in the sampling rate logic.
|
|
"""
|
|
|
|
@patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "1"}, clear=False)
|
|
def test_langsmith_sampling_rate_parameter_respected_with_valid_env(self):
|
|
"""Test that langsmith_sampling_rate parameter is properly set when env var condition is met."""
|
|
sampling_rate = 0.5
|
|
logger = LangsmithLogger(
|
|
langsmith_api_key="test-key",
|
|
langsmith_project="test-project",
|
|
langsmith_sampling_rate=sampling_rate,
|
|
)
|
|
|
|
assert (
|
|
logger.sampling_rate == sampling_rate
|
|
), f"Expected sampling_rate to be {sampling_rate}, got {logger.sampling_rate}"
|
|
|
|
@patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "1"}, clear=False)
|
|
def test_langsmith_sampling_rate_zero_parameter_falls_back_to_env(self):
|
|
"""Test that 0.0 parameter falls back to env var due to falsy value."""
|
|
logger = LangsmithLogger(
|
|
langsmith_api_key="test-key",
|
|
langsmith_project="test-project",
|
|
langsmith_sampling_rate=0.0,
|
|
)
|
|
|
|
assert (
|
|
logger.sampling_rate == 1.0
|
|
), f"Expected sampling_rate to fall back to 1.0 from env, got {logger.sampling_rate}"
|
|
|
|
@patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "1"}, clear=False)
|
|
def test_langsmith_sampling_rate_from_integer_env_var(self):
|
|
"""Test that sampling rate uses environment variable when parameter not provided and env var is integer."""
|
|
logger = LangsmithLogger(
|
|
langsmith_api_key="test-key", langsmith_project="test-project"
|
|
)
|
|
|
|
assert (
|
|
logger.sampling_rate == 1.0
|
|
), f"Expected sampling_rate to be 1.0 from env var, got {logger.sampling_rate}"
|
|
|
|
@patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "0.8"}, clear=False)
|
|
def test_langsmith_sampling_rate_decimal_env_var_ignored(self):
|
|
"""Test that decimal environment variables are ignored due to isdigit() check."""
|
|
logger = LangsmithLogger(
|
|
langsmith_api_key="test-key", langsmith_project="test-project"
|
|
)
|
|
|
|
assert (
|
|
logger.sampling_rate == 1.0
|
|
), f"Expected sampling_rate to default to 1.0 (decimal env ignored), got {logger.sampling_rate}"
|
|
|
|
@patch.dict(os.environ, {}, clear=True)
|
|
def test_langsmith_sampling_rate_default_value(self):
|
|
"""Test that sampling rate defaults to 1.0 when no parameter or env var provided."""
|
|
logger = LangsmithLogger(
|
|
langsmith_api_key="test-key", langsmith_project="test-project"
|
|
)
|
|
|
|
assert (
|
|
logger.sampling_rate == 1.0
|
|
), f"Expected default sampling_rate to be 1.0, got {logger.sampling_rate}"
|
|
|
|
@patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "invalid"}, clear=False)
|
|
def test_langsmith_sampling_rate_invalid_env_var_defaults(self):
|
|
"""Test that invalid environment variable falls back to default value."""
|
|
logger = LangsmithLogger(
|
|
langsmith_api_key="test-key", langsmith_project="test-project"
|
|
)
|
|
|
|
assert (
|
|
logger.sampling_rate == 1.0
|
|
), f"Expected sampling_rate to default to 1.0 with invalid env var, got {logger.sampling_rate}"
|
|
|
|
@patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": ""}, clear=False)
|
|
def test_langsmith_sampling_rate_empty_env_var_defaults(self):
|
|
"""Test that empty environment variable falls back to default value."""
|
|
logger = LangsmithLogger(
|
|
langsmith_api_key="test-key", langsmith_project="test-project"
|
|
)
|
|
|
|
assert (
|
|
logger.sampling_rate == 1.0
|
|
), f"Expected sampling_rate to default to 1.0 with empty env var, got {logger.sampling_rate}"
|
|
|
|
def test_langsmith_sampling_rate_attribute_exists(self):
|
|
"""Test that the sampling_rate attribute is always set on the logger instance."""
|
|
logger = LangsmithLogger(
|
|
langsmith_api_key="test-key", langsmith_project="test-project"
|
|
)
|
|
|
|
assert hasattr(
|
|
logger, "sampling_rate"
|
|
), "LangsmithLogger should have sampling_rate attribute"
|
|
assert isinstance(
|
|
logger.sampling_rate, float
|
|
), f"sampling_rate should be a float, got {type(logger.sampling_rate)}"
|
|
assert (
|
|
logger.sampling_rate >= 0.0
|
|
), f"sampling_rate should be non-negative, got {logger.sampling_rate}"
|
|
|
|
@patch.object(LangsmithLogger, "_start_periodic_flush_task", return_value=None)
|
|
def test_langsmith_init_skips_periodic_flush_without_running_loop(
|
|
self, mock_start_periodic_flush_task
|
|
):
|
|
"""Test that sync initialization leaves the periodic flush task unset."""
|
|
logger = LangsmithLogger(
|
|
langsmith_api_key="test-key", langsmith_project="test-project"
|
|
)
|
|
|
|
assert logger is not None
|
|
mock_start_periodic_flush_task.assert_called_once()
|
|
assert logger._flush_task is None
|
|
|
|
@patch(
|
|
"asyncio.get_running_loop", side_effect=RuntimeError("no running event loop")
|
|
)
|
|
def test_start_periodic_flush_task_returns_none_without_running_loop(
|
|
self, mock_get_running_loop
|
|
):
|
|
"""Test that helper returns None when no running event loop exists."""
|
|
with patch.object(
|
|
LangsmithLogger, "_start_periodic_flush_task", return_value=None
|
|
):
|
|
logger = LangsmithLogger(
|
|
langsmith_api_key="test-key",
|
|
langsmith_project="test-project",
|
|
)
|
|
|
|
mock_get_running_loop.reset_mock()
|
|
|
|
assert logger._start_periodic_flush_task() is None
|
|
mock_get_running_loop.assert_called_once()
|
|
|
|
@patch("asyncio.get_running_loop")
|
|
def test_langsmith_init_starts_periodic_flush_with_running_loop(
|
|
self, mock_get_running_loop
|
|
):
|
|
"""Test that init schedules periodic flush when a running loop exists."""
|
|
mock_loop = MagicMock()
|
|
mock_task = MagicMock()
|
|
mock_loop.create_task.return_value = mock_task
|
|
mock_get_running_loop.return_value = mock_loop
|
|
|
|
logger = LangsmithLogger(
|
|
langsmith_api_key="test-key", langsmith_project="test-project"
|
|
)
|
|
|
|
assert logger._flush_task == mock_task
|
|
mock_loop.create_task.assert_called_once()
|
|
scheduled_coro = mock_loop.create_task.call_args.args[0]
|
|
scheduled_coro.close()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_log_success_event_lazily_starts_periodic_flush(self):
|
|
"""Test that async logging lazily starts periodic flush after sync init."""
|
|
with patch.object(
|
|
LangsmithLogger, "_start_periodic_flush_task", return_value=None
|
|
):
|
|
logger = LangsmithLogger(
|
|
langsmith_api_key="test-key",
|
|
langsmith_project="test-project",
|
|
)
|
|
logger._get_sampling_rate_to_use_for_request = MagicMock(return_value=1.0)
|
|
logger._get_credentials_to_use_for_request = MagicMock(
|
|
return_value=logger.default_credentials
|
|
)
|
|
logger._prepare_log_data = MagicMock(return_value={"id": "run-id"})
|
|
logger._start_periodic_flush_task = MagicMock(return_value=MagicMock())
|
|
|
|
await logger.async_log_success_event({}, {}, None, None)
|
|
|
|
logger._start_periodic_flush_task.assert_called_once()
|
|
assert len(logger.log_queue) == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_log_failure_event_lazily_starts_periodic_flush(self):
|
|
"""Test that async failure logging lazily starts periodic flush after sync init."""
|
|
with patch.object(
|
|
LangsmithLogger, "_start_periodic_flush_task", return_value=None
|
|
):
|
|
logger = LangsmithLogger(
|
|
langsmith_api_key="test-key",
|
|
langsmith_project="test-project",
|
|
)
|
|
logger._get_sampling_rate_to_use_for_request = MagicMock(return_value=1.0)
|
|
logger._get_credentials_to_use_for_request = MagicMock(
|
|
return_value=logger.default_credentials
|
|
)
|
|
logger._prepare_log_data = MagicMock(return_value={"id": "run-id"})
|
|
logger._start_periodic_flush_task = MagicMock(return_value=MagicMock())
|
|
|
|
await logger.async_log_failure_event({}, {}, None, None)
|
|
|
|
logger._start_periodic_flush_task.assert_called_once()
|
|
assert len(logger.log_queue) == 1
|
|
|
|
|
|
class TestLangsmithPrepareLogData:
|
|
"""Regression test for #24001: _prepare_log_data must inject
|
|
usage_metadata into outputs so LangSmith's Cost column is populated."""
|
|
|
|
@patch("asyncio.create_task")
|
|
@patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "1"}, clear=False)
|
|
def test_outputs_contain_usage_metadata(self, mock_create_task):
|
|
logger = LangsmithLogger(
|
|
langsmith_api_key="test-key",
|
|
langsmith_project="test-project",
|
|
)
|
|
|
|
payload = {
|
|
"id": "test-id",
|
|
"response": {"choices": [{"message": {"content": "hi"}}]},
|
|
"metadata": {},
|
|
"startTime": 1.0,
|
|
"endTime": 2.0,
|
|
"request_tags": [],
|
|
"error_str": None,
|
|
"status": "success",
|
|
"response_cost": 0.0042,
|
|
"prompt_tokens": 100,
|
|
"completion_tokens": 50,
|
|
"total_tokens": 150,
|
|
}
|
|
|
|
kwargs = {
|
|
"litellm_params": {"metadata": {}},
|
|
"standard_logging_object": payload,
|
|
}
|
|
|
|
credentials = {
|
|
"LANGSMITH_API_KEY": "test-key",
|
|
"LANGSMITH_PROJECT": "test-project",
|
|
"LANGSMITH_BASE_URL": "https://api.smith.langchain.com",
|
|
}
|
|
|
|
data = logger._prepare_log_data(
|
|
kwargs=kwargs,
|
|
response_obj=None,
|
|
start_time=1.0,
|
|
end_time=2.0,
|
|
credentials=credentials,
|
|
)
|
|
|
|
assert "usage_metadata" in data["outputs"]
|
|
um = data["outputs"]["usage_metadata"]
|
|
assert um["total_cost"] == 0.0042
|
|
assert um["input_tokens"] == 100
|
|
assert um["output_tokens"] == 50
|
|
assert um["total_tokens"] == 150
|
|
|
|
|
|
class TestLangsmithRedactUserApiKeyInfo:
|
|
"""Verify litellm.redact_user_api_key_info is honored for LangSmith."""
|
|
|
|
def _logger(self):
|
|
return LangsmithLogger(
|
|
langsmith_api_key="test-key",
|
|
langsmith_project="test-project",
|
|
)
|
|
|
|
def _metadata_with_user_api_key_fields(self):
|
|
return {
|
|
"user_api_key_hash": "abc123",
|
|
"user_api_key_alias": "engineer-key",
|
|
"user_api_key_user_id": "default_user_id",
|
|
"user_api_key_team_id": "team-uuid",
|
|
"user_api_key_team_alias": "GNT",
|
|
"user_api_key_request_route": "/chat/completions",
|
|
"user_api_key_spend": 1.64,
|
|
"model": "gpt-4",
|
|
"requester_metadata": {
|
|
"user_api_key_team_id": "team-uuid",
|
|
"user_api_key_user_id": "default_user_id",
|
|
"session_id": "sess-1",
|
|
},
|
|
}
|
|
|
|
def test_redact_disabled_keeps_user_api_key_fields(self, reset_redact_flag):
|
|
"""Flag off: user_api_key_* fields are preserved (no behavior change)."""
|
|
litellm.redact_user_api_key_info = False
|
|
logger = self._logger()
|
|
metadata = self._metadata_with_user_api_key_fields()
|
|
|
|
extra = logger._build_extra_metadata(metadata)
|
|
|
|
assert extra["user_api_key_hash"] == "abc123"
|
|
assert extra["user_api_key_team_id"] == "team-uuid"
|
|
assert extra["requester_metadata"]["user_api_key_user_id"] == "default_user_id"
|
|
|
|
def test_redact_enabled_strips_top_level_user_api_key_fields(
|
|
self, reset_redact_flag
|
|
):
|
|
"""Flag on: top-level user_api_key_* keys removed; other keys preserved."""
|
|
litellm.redact_user_api_key_info = True
|
|
logger = self._logger()
|
|
metadata = self._metadata_with_user_api_key_fields()
|
|
|
|
extra = logger._build_extra_metadata(metadata)
|
|
|
|
for key in (
|
|
"user_api_key_hash",
|
|
"user_api_key_alias",
|
|
"user_api_key_user_id",
|
|
"user_api_key_team_id",
|
|
"user_api_key_team_alias",
|
|
"user_api_key_request_route",
|
|
"user_api_key_spend",
|
|
):
|
|
assert key not in extra, f"{key} should be redacted at top level"
|
|
assert extra["model"] == "gpt-4"
|
|
|
|
def test_redact_enabled_strips_nested_requester_metadata(self, reset_redact_flag):
|
|
"""Flag on: nested requester_metadata.user_api_key_* removed; session_id still lifted."""
|
|
litellm.redact_user_api_key_info = True
|
|
logger = self._logger()
|
|
metadata = self._metadata_with_user_api_key_fields()
|
|
|
|
extra = logger._build_extra_metadata(metadata)
|
|
|
|
nested = extra["requester_metadata"]
|
|
assert "user_api_key_team_id" not in nested
|
|
assert "user_api_key_user_id" not in nested
|
|
assert nested["session_id"] == "sess-1"
|
|
assert extra["session_id"] == "sess-1"
|
|
|
|
def test_redact_enabled_strips_user_api_key_info_from_inputs(self, reset_redact_flag):
|
|
"""
|
|
Regression (LIT-4306): `inputs` is the whole StandardLoggingPayload, so
|
|
`redact_user_api_key_info` has to cover `inputs.metadata` the same way it
|
|
covers `extra` - including the nested `requester_metadata` copy. Before
|
|
the fix `extra` was redacted and `inputs` shipped every user_api_key_*
|
|
field verbatim.
|
|
"""
|
|
litellm.redact_user_api_key_info = True
|
|
logger = self._logger()
|
|
metadata = self._metadata_with_user_api_key_fields()
|
|
metadata["user_api_key_auth_metadata"] = {"priority": "high"}
|
|
payload = {
|
|
"id": "run-1",
|
|
"response": {"choices": []},
|
|
"metadata": metadata,
|
|
"startTime": 1.0,
|
|
"endTime": 2.0,
|
|
"request_tags": [],
|
|
"error_str": None,
|
|
"status": "success",
|
|
"response_cost": 0.0,
|
|
"prompt_tokens": 1,
|
|
"completion_tokens": 1,
|
|
"total_tokens": 2,
|
|
}
|
|
credentials = {
|
|
"LANGSMITH_API_KEY": "test-key",
|
|
"LANGSMITH_PROJECT": "test-project",
|
|
"LANGSMITH_BASE_URL": "https://api.smith.langchain.com",
|
|
}
|
|
|
|
data = logger._prepare_log_data(
|
|
kwargs={"litellm_params": {"metadata": metadata}, "standard_logging_object": payload},
|
|
response_obj=None,
|
|
start_time=1.0,
|
|
end_time=2.0,
|
|
credentials=credentials,
|
|
)
|
|
|
|
inputs_metadata = data["inputs"]["metadata"]
|
|
assert [k for k in inputs_metadata if k.startswith("user_api_key")] == []
|
|
assert [k for k in inputs_metadata["requester_metadata"] if k.startswith("user_api_key")] == []
|
|
# inputs and extra must agree - they go through the same redaction now
|
|
assert [k for k in data["extra"] if k.startswith("user_api_key")] == []
|
|
# non-identity payload is untouched
|
|
assert inputs_metadata["model"] == "gpt-4"
|
|
assert inputs_metadata["requester_metadata"]["session_id"] == "sess-1"
|
|
assert data["inputs"]["total_tokens"] == 2
|
|
# the shared standard_logging_object other loggers read is not mutated
|
|
assert "user_api_key_hash" in payload["metadata"]
|
|
assert "user_api_key_user_id" in payload["metadata"]["requester_metadata"]
|
|
|
|
def test_redact_disabled_keeps_user_api_key_info_in_inputs(self, reset_redact_flag):
|
|
"""Flag off: the identity fields stay. The flag governs them, not this fix."""
|
|
litellm.redact_user_api_key_info = False
|
|
logger = self._logger()
|
|
metadata = self._metadata_with_user_api_key_fields()
|
|
payload = {
|
|
"id": "run-1",
|
|
"response": {"choices": []},
|
|
"metadata": metadata,
|
|
"startTime": 1.0,
|
|
"endTime": 2.0,
|
|
"request_tags": [],
|
|
"error_str": None,
|
|
"status": "success",
|
|
"response_cost": 0.0,
|
|
"prompt_tokens": 1,
|
|
"completion_tokens": 1,
|
|
"total_tokens": 2,
|
|
}
|
|
|
|
data = logger._prepare_log_data(
|
|
kwargs={"litellm_params": {"metadata": metadata}, "standard_logging_object": payload},
|
|
response_obj=None,
|
|
start_time=1.0,
|
|
end_time=2.0,
|
|
credentials={
|
|
"LANGSMITH_API_KEY": "test-key",
|
|
"LANGSMITH_PROJECT": "test-project",
|
|
"LANGSMITH_BASE_URL": "https://api.smith.langchain.com",
|
|
},
|
|
)
|
|
|
|
assert data["inputs"]["metadata"]["user_api_key_hash"] == "abc123"
|