diff --git a/agentops/sdk/__init__.py b/agentops/sdk/__init__.py index 1b0779dd5..4e74162e2 100644 --- a/agentops/sdk/__init__.py +++ b/agentops/sdk/__init__.py @@ -11,7 +11,11 @@ from agentops.sdk.decorators import agent, operation, session, task, workflow # from agentops.sdk.traced import TracedObject # Merged into TracedObject from agentops.sdk.types import TracingConfig +# Initialize SDK first +from agentops.sdk.decorators.sentry_manager import set_opt_out_sentry +# Initialize sentry settings +set_opt_out_sentry(False) # Enable Sentry error tracking __all__ = [ # Core components "TracingCore", diff --git a/agentops/sdk/decorators/sentry_manager.py b/agentops/sdk/decorators/sentry_manager.py new file mode 100644 index 000000000..3a3d7f439 --- /dev/null +++ b/agentops/sdk/decorators/sentry_manager.py @@ -0,0 +1,141 @@ +import sentry_sdk +import functools +from contextlib import contextmanager +from typing import Optional, Dict +import inspect + +# Dictionary to store per-module Sentry settings +_module_sentry_settings: Dict[str, bool] = {} +# Global default setting +_default_opt_out = False + + +def set_opt_out_sentry(value: bool, module_name: Optional[str] = None): + """ + Function to enable or disable Sentry error tracking globally or for specific modules. + + Args: + value (bool): Set to True to disable Sentry, False to enable it. + module_name (Optional[str]): If provided, sets Sentry setting for specific module. + If None, sets the global default. + + Example: + # Disable Sentry for current module: + set_opt_out_sentry(True, __name__) + + # Disable Sentry globally: + set_opt_out_sentry(True) + + # Enable Sentry for specific module: + set_opt_out_sentry(False, "my_module") + """ + global _default_opt_out, _module_sentry_settings + + if module_name is None: + _default_opt_out = value + # Initialize Sentry globally based on default setting + if not _default_opt_out: + sentry_sdk.init( + dsn="", + traces_sample_rate=1.0, + send_default_pii=True, + ) + print("Global Sentry error tracking is enabled.") + else: + sentry_sdk.init() # Initialize with empty DSN to effectively disable + print("Global Sentry error tracking is disabled.") + else: + _module_sentry_settings[module_name] = value + print(f"Sentry error tracking is {'disabled' if value else 'enabled'} for module: {module_name}") + + +def is_sentry_enabled(module_name: Optional[str] = None) -> bool: + """ + Check if Sentry error tracking is enabled for a specific module or globally. + + Args: + module_name (Optional[str]): Module name to check. If None, checks global setting. + + Returns: + bool: True if Sentry is enabled, False if disabled + """ + if module_name is None: + return not _default_opt_out + return not _module_sentry_settings.get(module_name, _default_opt_out) + + +def track_errors(func=None, *, module_override: Optional[str] = None): + """ + Decorator to automatically track errors in Sentry. + Respects per-module and global Sentry settings. + + Args: + func: The function to decorate + module_override: Optional module name to override the automatic module detection + + Usage: + @track_errors + def your_function(): + # Your code here + + @track_errors(module_override="custom_module") + def your_function(): + # Your code here + """ + + def decorator(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + # Get the module name from the function or use override + mod_name = module_override or func.__module__ + try: + return func(*args, **kwargs) + except Exception as e: + if is_sentry_enabled(mod_name): + sentry_sdk.capture_exception(e) + raise + + return wrapper + + if func is None: + return decorator + return decorator(func) + + +@contextmanager +def track_errors_context(module_name: Optional[str] = None): + """ + Context manager to track errors in Sentry. + Respects per-module and global Sentry settings. + + Args: + module_name: Optional module name to override the automatic module detection + + Usage: + with track_errors_context(): + # Your code here + + with track_errors_context("custom_module"): + # Your code here + """ + try: + yield + except Exception as e: + # If no module_name provided, get it from the caller's frame + if module_name is None: + frame = inspect.currentframe() + if frame and frame.f_back: + module_name = frame.f_back.f_globals.get("__name__") + + if is_sentry_enabled(module_name): + sentry_sdk.capture_exception(e) + raise + + +def capture_error(exception): + """Captures the error in Sentry if it's enabled for the calling module.""" + frame = inspect.currentframe() + if frame and frame.f_back: + module_name = frame.f_back.f_globals.get("__name__") + if is_sentry_enabled(module_name): + sentry_sdk.capture_exception(exception) diff --git a/tests/test_sentry_manager.py b/tests/test_sentry_manager.py new file mode 100644 index 000000000..81add76a3 --- /dev/null +++ b/tests/test_sentry_manager.py @@ -0,0 +1,222 @@ +import pytest +from unittest.mock import patch +from agentops.sdk.decorators.sentry_manager import ( + set_opt_out_sentry, + is_sentry_enabled, + track_errors, + track_errors_context, + capture_error, +) + + +# Mock sentry_sdk for testing +@pytest.fixture +def mock_sentry(): + with patch("agentops.sdk.decorators.sentry_manager.sentry_sdk") as mock_sdk: + yield mock_sdk + + +# Test Global Sentry Settings +class TestGlobalSentrySettings: + def test_default_sentry_enabled(self): + """Test that Sentry is enabled by default""" + assert is_sentry_enabled() is True + + def test_global_opt_out(self): + """Test global opt-out functionality""" + set_opt_out_sentry(True) + assert is_sentry_enabled() is False + # Reset to default + set_opt_out_sentry(False) + + def test_global_opt_in(self): + """Test global opt-in functionality""" + set_opt_out_sentry(False) + assert is_sentry_enabled() is True + + +# Test Module-Level Settings +class TestModuleLevelSettings: + def test_module_specific_setting(self): + """Test module-specific Sentry settings""" + test_module = "test_module" + # Enable globally but disable for specific module + set_opt_out_sentry(False) # Global enable + set_opt_out_sentry(True, test_module) # Module disable + assert is_sentry_enabled() is True # Global still enabled + assert is_sentry_enabled(test_module) is False # Module disabled + + def test_module_override_global(self): + """Test that module settings override global settings""" + test_module = "test_module" + # Disable globally but enable for specific module + set_opt_out_sentry(True) # Global disable + set_opt_out_sentry(False, test_module) # Module enable + assert is_sentry_enabled() is False # Global still disabled + assert is_sentry_enabled(test_module) is True # Module enabled + + +# Test Error Tracking Decorator +class TestErrorTrackingDecorator: + def test_decorator_with_enabled_sentry(self, mock_sentry): + """Test that errors are captured when Sentry is enabled""" + set_opt_out_sentry(False) # Enable Sentry + + @track_errors + def raise_error(): + raise ValueError("Test error") + + with pytest.raises(ValueError): + raise_error() + + mock_sentry.capture_exception.assert_called_once() + + def test_decorator_with_disabled_sentry(self, mock_sentry): + """Test that errors are not captured when Sentry is disabled""" + set_opt_out_sentry(True) # Disable Sentry + + @track_errors + def raise_error(): + raise ValueError("Test error") + + with pytest.raises(ValueError): + raise_error() + + mock_sentry.capture_exception.assert_not_called() + + def test_decorator_with_module_override(self, mock_sentry): + """Test decorator with explicit module override""" + test_module = "test_module" + set_opt_out_sentry(True) # Global disable + set_opt_out_sentry(False, test_module) # Module enable + + @track_errors(module_override=test_module) + def raise_error(): + raise ValueError("Test error") + + with pytest.raises(ValueError): + raise_error() + + mock_sentry.capture_exception.assert_called_once() + + +# Test Context Manager +class TestContextManager: + def test_context_manager_with_enabled_sentry(self, mock_sentry): + """Test that errors are captured in context when Sentry is enabled""" + set_opt_out_sentry(False) # Enable Sentry + + with pytest.raises(ValueError): + with track_errors_context(): + raise ValueError("Test error") + + mock_sentry.capture_exception.assert_called_once() + + def test_context_manager_with_disabled_sentry(self, mock_sentry): + """Test that errors are not captured in context when Sentry is disabled""" + set_opt_out_sentry(True) # Disable Sentry + + with pytest.raises(ValueError): + with track_errors_context(): + raise ValueError("Test error") + + mock_sentry.capture_exception.assert_not_called() + + def test_context_manager_with_module_name(self, mock_sentry): + """Test context manager with specific module name""" + test_module = "test_module" + set_opt_out_sentry(True) # Global disable + set_opt_out_sentry(False, test_module) # Module enable + + with pytest.raises(ValueError): + with track_errors_context(module_name=test_module): + raise ValueError("Test error") + + mock_sentry.capture_exception.assert_called_once() + + def test_nested_context_managers(self, mock_sentry): + """Test nested context managers with different module settings""" + outer_module = "outer_module" + inner_module = "inner_module" + + set_opt_out_sentry(False) # Global enable + set_opt_out_sentry(True, outer_module) # Outer module disable + set_opt_out_sentry(False, inner_module) # Inner module enable + + with pytest.raises(ValueError): + with track_errors_context(module_name=outer_module): + # This error should not be captured (outer module disabled) + with track_errors_context(module_name=inner_module): + # This error should be captured (inner module enabled) + raise ValueError("Test error") + + # Only inner context should capture the error + assert mock_sentry.capture_exception.call_count == 1 + + +# Test Direct Error Capture +class TestDirectErrorCapture: + def test_capture_error_enabled(self, mock_sentry): + """Test direct error capture when Sentry is enabled""" + set_opt_out_sentry(False) + error = ValueError("Test error") + capture_error(error) + mock_sentry.capture_exception.assert_called_once_with(error) + + def test_capture_error_disabled(self, mock_sentry): + """Test direct error capture when Sentry is disabled""" + set_opt_out_sentry(True) + error = ValueError("Test error") + capture_error(error) + mock_sentry.capture_exception.assert_not_called() + + +# Test Real-world Scenarios +class TestRealWorldScenarios: + def test_multiple_functions_different_modules(self, mock_sentry): + """Test multiple functions with different module settings""" + module1 = "module1" + module2 = "module2" + + set_opt_out_sentry(True) # Global disable + set_opt_out_sentry(False, module1) # Enable module1 + set_opt_out_sentry(True, module2) # Disable module2 + + @track_errors(module_override=module1) + def function1(): + raise ValueError("Error 1") + + @track_errors(module_override=module2) + def function2(): + raise ValueError("Error 2") + + # Should capture error (module1 enabled) + with pytest.raises(ValueError): + function1() + assert mock_sentry.capture_exception.call_count == 1 + + # Should not capture error (module2 disabled) + with pytest.raises(ValueError): + function2() + assert mock_sentry.capture_exception.call_count == 1 # Count shouldn't increase + + def test_error_propagation(self, mock_sentry): + """Test that errors properly propagate through multiple layers""" + set_opt_out_sentry(False) + + @track_errors + def inner_function(): + raise ValueError("Inner error") + + @track_errors + def outer_function(): + try: + inner_function() + except ValueError: + raise RuntimeError("Outer error") + + with pytest.raises(RuntimeError): + outer_function() + + # Should capture both errors + assert mock_sentry.capture_exception.call_count == 2