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