litellm/tests/test_litellm/integrations/test_langsmith_init.py
mrinal 4c179f2f59 test(langsmith): type the flush race test and cancel its periodic task
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-09-15 21:17:36 +00:00

576 lines
22 KiB
Python

import asyncio
import os
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import litellm
from litellm.integrations.langsmith import LangsmithLogger
from litellm.types.integrations.langsmith import LangsmithQueueObject
@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()
@pytest.mark.asyncio
async def test_langsmith_init_starts_periodic_flush_with_running_loop(self):
"""Test that init schedules periodic flush when a running loop exists."""
logger = LangsmithLogger(
langsmith_api_key="test-key", langsmith_project="test-project", flush_interval=0.01
)
batch_sent = asyncio.Event()
logger.async_send_batch = AsyncMock(side_effect=batch_sent.set)
logger.log_queue.append({"id": "run-id"})
flush_task = logger._flush_task
assert isinstance(flush_task, asyncio.Task)
await asyncio.wait_for(batch_sent.wait(), timeout=5)
flush_task.cancel()
with pytest.raises(asyncio.CancelledError):
await flush_task
@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
@pytest.mark.asyncio
async def test_events_appended_during_flush_are_not_dropped():
logger = LangsmithLogger(langsmith_api_key="test-key", langsmith_project="test-project")
try:
sent_batches: Final[list[list[dict[str, str]]]] = []
late_event: Final = LangsmithQueueObject(
credentials=logger.default_credentials, data={"id": "late"}
)
async def fake_post(
url: str, json: dict[str, list[dict[str, str]]], headers: dict[str, str]
) -> MagicMock:
if not sent_batches:
logger.log_queue.append(late_event)
sent_batches.append(json["post"])
response = MagicMock()
response.status_code = 200
response.raise_for_status = MagicMock()
return response
logger.async_httpx_client = MagicMock(post=AsyncMock(side_effect=fake_post))
logger.log_queue = [
LangsmithQueueObject(credentials=logger.default_credentials, data={"id": "a"}),
LangsmithQueueObject(credentials=logger.default_credentials, data={"id": "b"}),
]
await logger.flush_queue()
assert [e["id"] for e in sent_batches[0]] == ["a", "b"]
assert logger.log_queue == [late_event]
await logger.flush_queue()
assert [e["id"] for e in sent_batches[1]] == ["late"]
assert logger.log_queue == []
finally:
if logger._flush_task is not None:
logger._flush_task.cancel()
await asyncio.gather(logger._flush_task, return_exceptions=True)