mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
ee68813530
A request carrying a session/trace header fans the header value into litellm metadata as both trace_id and session_id. LangSmith then rejected the whole ingest batch: a root run's trace_id must equal the run id embedded in dotted_order (400), and a run-body session_id must reference an existing tracer session (404/422). Override caller trace_id on runs that post as roots and drop session_id only when it mirrors trace_id, so deliberate child-run and valid tracer-session fields still pass through.
534 lines
20 KiB
Python
534 lines
20 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"
|
|
|
|
|
|
class TestLangsmithRootRunIdConsistency:
|
|
"""Regression tests for LIT-5878 / #37269.
|
|
|
|
A request that carries a session/trace header (e.g. x-claude-code-session-id)
|
|
fans the header value out into litellm metadata as both trace_id and
|
|
session_id. LangSmith then rejected the whole ingest batch twice over:
|
|
a root run whose trace_id does not match the run id embedded in dotted_order
|
|
(400), and a run-body session_id that does not reference an existing tracer
|
|
session (404, or 422 for non-UUID values).
|
|
"""
|
|
|
|
def _prepare(self, request_metadata):
|
|
payload = {
|
|
"id": "slp-1",
|
|
"response": {"choices": []},
|
|
"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,
|
|
}
|
|
logger = LangsmithLogger(
|
|
langsmith_api_key="test-key",
|
|
langsmith_project="test-project",
|
|
)
|
|
return logger._prepare_log_data(
|
|
kwargs={
|
|
"litellm_params": {"metadata": request_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",
|
|
},
|
|
)
|
|
|
|
def test_header_derived_ids_yield_self_consistent_root_run(self):
|
|
header_value = "ed29c3bb-44fa-4eec-9b7b-fecaa3e82d64"
|
|
data = self._prepare({"trace_id": header_value, "session_id": header_value})
|
|
|
|
assert data["trace_id"] == data["id"]
|
|
assert data["trace_id"] != header_value
|
|
assert data["dotted_order"].endswith(data["id"])
|
|
assert len(data["dotted_order"]) == 22 + len(data["id"])
|
|
assert "session_id" not in data
|
|
|
|
def test_distinct_session_id_is_still_forwarded(self):
|
|
data = self._prepare({"session_id": "11111111-2222-3333-4444-555555555555"})
|
|
|
|
assert data["session_id"] == "11111111-2222-3333-4444-555555555555"
|
|
|
|
def test_trace_id_only_root_run_is_overridden(self):
|
|
data = self._prepare({"trace_id": "ed29c3bb-44fa-4eec-9b7b-fecaa3e82d64"})
|
|
|
|
assert data["trace_id"] == data["id"]
|
|
assert data["trace_id"] != "ed29c3bb-44fa-4eec-9b7b-fecaa3e82d64"
|
|
assert data["dotted_order"].endswith(data["id"])
|
|
|
|
def test_root_run_without_caller_ids_is_self_consistent(self):
|
|
data = self._prepare({})
|
|
|
|
assert data["trace_id"] == data["id"]
|
|
assert data["dotted_order"].endswith(data["id"])
|
|
|
|
def test_child_run_keeps_caller_trace_id(self):
|
|
data = self._prepare(
|
|
{
|
|
"trace_id": "trace-1",
|
|
"parent_run_id": "parent-1",
|
|
"run_id": "child-1",
|
|
}
|
|
)
|
|
|
|
assert data["trace_id"] == "trace-1"
|
|
assert data["id"] == "child-1"
|
|
assert data["parent_run_id"] == "parent-1"
|
|
|
|
def test_caller_supplied_dotted_order_and_trace_id_are_untouched(self):
|
|
dotted = "20260820T000000000000Ztrace-1.20260820T000001000000Zrun-1"
|
|
data = self._prepare(
|
|
{
|
|
"trace_id": "trace-1",
|
|
"run_id": "run-1",
|
|
"dotted_order": dotted,
|
|
}
|
|
)
|
|
|
|
assert data["trace_id"] == "trace-1"
|
|
assert data["dotted_order"] == dotted
|