scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

fix(langsmith): avoid no running event loop during sync init (#23727)

Browse source 
* 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
This commit is contained in:
Miguel Miranda Dias 2026-03-17 06:34:15 +01:00 • committed by GitHub
parent 84b4af40fa
commit e9291a97c3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 129 additions and 39 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

43
litellm/integrations/langsmith.py
View file

@ -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,

125
tests/test_litellm/integrations/test_langsmith_init.py
View file

@ -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