mirror of
https://github.com/tiennm99/litellm.git
synced 2026-06-18 07:33:58 +00:00
fix(langsmith.py): add langsmith_sampling_rate as a dynamic param
Closes LIT-879
This commit is contained in:
@@ -39,6 +39,7 @@ class LangsmithLogger(CustomBatchLogger):
|
||||
langsmith_api_key: Optional[str] = None,
|
||||
langsmith_project: Optional[str] = None,
|
||||
langsmith_base_url: Optional[str] = None,
|
||||
langsmith_sampling_rate: Optional[float] = None,
|
||||
**kwargs,
|
||||
):
|
||||
self.flush_lock = asyncio.Lock()
|
||||
@@ -49,7 +50,8 @@ class LangsmithLogger(CustomBatchLogger):
|
||||
langsmith_base_url=langsmith_base_url,
|
||||
)
|
||||
self.sampling_rate: float = (
|
||||
float(os.getenv("LANGSMITH_SAMPLING_RATE")) # type: ignore
|
||||
langsmith_sampling_rate
|
||||
or float(os.getenv("LANGSMITH_SAMPLING_RATE")) # type: ignore
|
||||
if os.getenv("LANGSMITH_SAMPLING_RATE") is not None
|
||||
and os.getenv("LANGSMITH_SAMPLING_RATE").strip().isdigit() # type: ignore
|
||||
else 1.0
|
||||
@@ -442,9 +444,9 @@ class LangsmithLogger(CustomBatchLogger):
|
||||
|
||||
Otherwise, use the default credentials.
|
||||
"""
|
||||
standard_callback_dynamic_params: Optional[
|
||||
StandardCallbackDynamicParams
|
||||
] = kwargs.get("standard_callback_dynamic_params", None)
|
||||
standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = (
|
||||
kwargs.get("standard_callback_dynamic_params", None)
|
||||
)
|
||||
if standard_callback_dynamic_params is not None:
|
||||
credentials = self.get_credentials_from_env(
|
||||
langsmith_api_key=standard_callback_dynamic_params.get(
|
||||
|
||||
@@ -0,0 +1,134 @@
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../.."))
|
||||
|
||||
from litellm.integrations.langsmith import LangsmithLogger
|
||||
|
||||
|
||||
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("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
|
||||
):
|
||||
"""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",
|
||||
langsmith_project="test-project",
|
||||
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
|
||||
):
|
||||
"""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!
|
||||
)
|
||||
|
||||
# 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):
|
||||
"""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):
|
||||
"""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):
|
||||
"""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("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):
|
||||
"""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("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):
|
||||
"""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}"
|
||||
|
||||
@patch("asyncio.create_task")
|
||||
def test_langsmith_sampling_rate_attribute_exists(self, mock_create_task):
|
||||
"""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"
|
||||
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}"
|
||||
Reference in New Issue
Block a user