This commit is contained in:
Vinnie-Singleton-NN 2025-04-24 14:06:41 -04:00 committed by GitHub
commit 257b81bebf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 35 additions and 1 deletions

View file

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

View file

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