From e9291a97c32303a2ca18313bede88a75a52ce258 Mon Sep 17 00:00:00 2001 From: Miguel Miranda Dias <7780875+pandego@users.noreply.github.com> Date: Tue, 17 Mar 2026 06:34:15 +0100 Subject: [PATCH] fix(langsmith): avoid no running event loop during sync init (#23727) * fix(langsmith): skip periodic flush task without event loop * fix(langsmith): lazily start periodic flush task * test(langsmith): tighten flush task coverage * test(langsmith): cover lazy failure flush startup * refactor(langsmith): keep flush startup private --- litellm/integrations/langsmith.py | 43 ++++-- .../integrations/test_langsmith_init.py | 125 ++++++++++++++---- 2 files changed, 129 insertions(+), 39 deletions(-) diff --git a/litellm/integrations/langsmith.py b/litellm/integrations/langsmith.py index 03845af521d..df2e3c1e2b7 100644 --- a/litellm/integrations/langsmith.py +++ b/litellm/integrations/langsmith.py @@ -83,7 +83,26 @@ class LangsmithLogger(CustomBatchLogger): if _batch_size: self.batch_size = int(_batch_size) self.log_queue: List[LangsmithQueueObject] = [] - asyncio.create_task(self.periodic_flush()) + self._flush_task: Optional[asyncio.Task[Any]] = self._start_periodic_flush_task() + + def _start_periodic_flush_task(self) -> Optional[asyncio.Task[Any]]: + """Start the periodic flush task only when an event loop is already running.""" + try: + loop = asyncio.get_running_loop() + except RuntimeError: + verbose_logger.debug( + "Langsmith logger init: no running event loop, skipping periodic flush task startup" + ) + return None + + return loop.create_task(self.periodic_flush()) + + def _ensure_periodic_flush_task(self) -> None: + # This helper is intentionally synchronous. In asyncio's cooperative + # execution model, there is no await between the check and assignment, + # so one caller cannot interleave here and create a duplicate task. + if self._flush_task is None or self._flush_task.done(): + self._flush_task = self._start_periodic_flush_task() def get_credentials_from_env( self, @@ -255,6 +274,7 @@ class LangsmithLogger(CustomBatchLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): try: + self._ensure_periodic_flush_task() sampling_rate = self._get_sampling_rate_to_use_for_request(kwargs=kwargs) random_sample = random.random() if random_sample > sampling_rate: @@ -296,17 +316,18 @@ class LangsmithLogger(CustomBatchLogger): ) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - sampling_rate = self._get_sampling_rate_to_use_for_request(kwargs=kwargs) - random_sample = random.random() - if random_sample > sampling_rate: - verbose_logger.info( - "Skipping Langsmith logging. Sampling rate={}, random_sample={}".format( - sampling_rate, random_sample - ) - ) - return # Skip logging - verbose_logger.info("Langsmith Failure Event Logging!") try: + self._ensure_periodic_flush_task() + sampling_rate = self._get_sampling_rate_to_use_for_request(kwargs=kwargs) + random_sample = random.random() + if random_sample > sampling_rate: + verbose_logger.info( + "Skipping Langsmith logging. Sampling rate={}, random_sample={}".format( + sampling_rate, random_sample + ) + ) + return # Skip logging + verbose_logger.info("Langsmith Failure Event Logging!") credentials = self._get_credentials_to_use_for_request(kwargs=kwargs) data = self._prepare_log_data( kwargs=kwargs, diff --git a/tests/test_litellm/integrations/test_langsmith_init.py b/tests/test_litellm/integrations/test_langsmith_init.py index 9f7db4095bc..edc827033f0 100644 --- a/tests/test_litellm/integrations/test_langsmith_init.py +++ b/tests/test_litellm/integrations/test_langsmith_init.py @@ -16,13 +16,9 @@ class TestLangsmithLoggerInit: Note: The current implementation has some edge cases in the sampling rate logic. """ - @patch("asyncio.create_task") @patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "1"}, clear=False) - def test_langsmith_sampling_rate_parameter_respected_with_valid_env( - self, mock_create_task - ): + 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.""" - # When there's a valid integer in env var, the parameter should be used due to 'or' logic sampling_rate = 0.5 logger = LangsmithLogger( langsmith_api_key="test-key", @@ -30,58 +26,47 @@ class TestLangsmithLoggerInit: langsmith_sampling_rate=sampling_rate, ) - # With the current 'or' logic and valid env var, the parameter should be used assert ( logger.sampling_rate == sampling_rate ), f"Expected sampling_rate to be {sampling_rate}, got {logger.sampling_rate}" - @patch("asyncio.create_task") @patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "1"}, clear=False) - def test_langsmith_sampling_rate_zero_parameter_falls_back_to_env( - self, mock_create_task - ): + 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.""" - # This demonstrates the current behavior where 0.0 is falsy and falls back to env logger = LangsmithLogger( langsmith_api_key="test-key", langsmith_project="test-project", - langsmith_sampling_rate=0.0, # This is falsy! + langsmith_sampling_rate=0.0, ) - # Due to current 'or' logic, 0.0 falls back to env var assert ( logger.sampling_rate == 1.0 ), f"Expected sampling_rate to fall back to 1.0 from env, got {logger.sampling_rate}" - @patch("asyncio.create_task") @patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "1"}, clear=False) - def test_langsmith_sampling_rate_from_integer_env_var(self, mock_create_task): + 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" ) - # Should use env var since it's a valid integer assert ( logger.sampling_rate == 1.0 ), f"Expected sampling_rate to be 1.0 from env var, got {logger.sampling_rate}" - @patch("asyncio.create_task") @patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "0.8"}, clear=False) - def test_langsmith_sampling_rate_decimal_env_var_ignored(self, mock_create_task): + 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" ) - # Decimal env vars are ignored due to isdigit() check, falls back to 1.0 assert ( logger.sampling_rate == 1.0 ), f"Expected sampling_rate to default to 1.0 (decimal env ignored), got {logger.sampling_rate}" - @patch("asyncio.create_task") @patch.dict(os.environ, {}, clear=True) - def test_langsmith_sampling_rate_default_value(self, mock_create_task): + 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" @@ -91,9 +76,8 @@ class TestLangsmithLoggerInit: logger.sampling_rate == 1.0 ), f"Expected default sampling_rate to be 1.0, got {logger.sampling_rate}" - @patch("asyncio.create_task") @patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "invalid"}, clear=False) - def test_langsmith_sampling_rate_invalid_env_var_defaults(self, mock_create_task): + 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" @@ -103,9 +87,8 @@ class TestLangsmithLoggerInit: logger.sampling_rate == 1.0 ), f"Expected sampling_rate to default to 1.0 with invalid env var, got {logger.sampling_rate}" - @patch("asyncio.create_task") @patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": ""}, clear=False) - def test_langsmith_sampling_rate_empty_env_var_defaults(self, mock_create_task): + 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" @@ -115,14 +98,12 @@ class TestLangsmithLoggerInit: logger.sampling_rate == 1.0 ), f"Expected sampling_rate to default to 1.0 with empty env var, got {logger.sampling_rate}" - @patch("asyncio.create_task") - def test_langsmith_sampling_rate_attribute_exists(self, mock_create_task): + 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" ) - # Verify the attribute exists and is a float assert hasattr( logger, "sampling_rate" ), "LangsmithLogger should have sampling_rate attribute" @@ -132,3 +113,91 @@ class TestLangsmithLoggerInit: 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