diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 77d4fd7d5d..f8b3e1ea45 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2601,9 +2601,15 @@ def set_callbacks(callback_list, function_id=None): # noqa: PLR0915 if "SENTRY_API_TRACE_RATE" in os.environ else "1.0" ) + sentry_sample_rate = ( + os.environ.get("SENTRY_API_SAMPLE_RATE") + if "SENTRY_API_SAMPLE_RATE" in os.environ + else "1.0" + ) sentry_sdk_instance.init( dsn=os.environ.get("SENTRY_DSN"), traces_sample_rate=float(sentry_trace_rate), # type: ignore + sample_rate=float(sentry_sample_rate), ) capture_exception = sentry_sdk_instance.capture_exception add_breadcrumb = sentry_sdk_instance.add_breadcrumb diff --git a/tests/litellm/litellm_core_utils/test_litellm_logging.py b/tests/litellm/litellm_core_utils/test_litellm_logging.py index eedc15dd36..767b1ff617 100644 --- a/tests/litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/litellm/litellm_core_utils/test_litellm_logging.py @@ -11,7 +11,7 @@ sys.path.insert( import time -from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging +from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging, set_callbacks @pytest.fixture @@ -32,3 +32,31 @@ def test_get_masked_api_base(logging_obj): masked_api_base = logging_obj._get_masked_api_base(api_base) assert masked_api_base == "https://api.openai.com/v1" assert type(masked_api_base) == str + + +def test_sentry_sample_rate(): + existing_sample_rate = os.getenv("SENTRY_API_SAMPLE_RATE") + try: + # test with default value by removing the environment variable + if existing_sample_rate: + del os.environ["SENTRY_API_SAMPLE_RATE"] + + set_callbacks(["sentry"]) + # Check if the default sample rate is set to 1.0 + assert os.environ.get("SENTRY_API_SAMPLE_RATE") == "1.0" + + # test with custom value + os.environ["SENTRY_API_SAMPLE_RATE"] = "0.5" + + set_callbacks(["sentry"]) + # Check if the custom sample rate is set correctly + assert os.environ.get("SENTRY_API_SAMPLE_RATE") == "0.5" + except Exception as e: + print(f"Error: {e}") + finally: + # Restore the original environment variable + if existing_sample_rate: + os.environ["SENTRY_API_SAMPLE_RATE"] = existing_sample_rate + else: + if "SENTRY_API_SAMPLE_RATE" in os.environ: + del os.environ["SENTRY_API_SAMPLE_RATE"] \ No newline at end of file