diff --git a/cognee/shared/usage_logger.py b/cognee/shared/usage_logger.py index 33d2dac10..1c07aff47 100644 --- a/cognee/shared/usage_logger.py +++ b/cognee/shared/usage_logger.py @@ -29,6 +29,7 @@ def _sanitize_value(value: Any) -> Any: @_sanitize_value.register(type(None)) def _(value: None) -> None: + """Handle None values - returns None as-is.""" return None @@ -37,27 +38,32 @@ def _(value: None) -> None: @_sanitize_value.register(float) @_sanitize_value.register(bool) def _(value: str | int | float | bool) -> str | int | float | bool: + """Handle primitive types - returns value as-is since they're JSON-serializable.""" return value @_sanitize_value.register(UUID) def _(value: UUID) -> str: + """Convert UUID to string representation.""" return str(value) @_sanitize_value.register(datetime) def _(value: datetime) -> str: + """Convert datetime to ISO format string.""" return value.isoformat() @_sanitize_value.register(list) @_sanitize_value.register(tuple) def _(value: list | tuple) -> list: + """Recursively sanitize list or tuple elements.""" return [_sanitize_value(v) for v in value] @_sanitize_value.register(dict) def _(value: dict) -> dict: + """Recursively sanitize dictionary keys and values.""" sanitized = {} for k, v in value.items(): key_str = k if isinstance(k, str) else _sanitize_dict_key(k) @@ -151,7 +157,23 @@ async def _log_usage_async( start_time: datetime, end_time: datetime, ): - """Asynchronously log function usage to Redis.""" + """Asynchronously log function usage to Redis. + + Args: + function_name: Name of the function being logged. + log_type: Type of log entry (e.g., "api_endpoint", "mcp_tool", "function"). + user_id: User identifier, or None to use "unknown". + parameters: Dictionary of function parameters (sanitized). + result: Function return value (will be sanitized). + success: Whether the function executed successfully. + error: Error message if function failed, None otherwise. + start_time: Function start timestamp. + end_time: Function end timestamp. + + Note: + This function silently handles errors to avoid disrupting the original + function execution. Logs are written to Redis with TTL from config. + """ try: logger.debug(f"Starting to log usage for {function_name} at {start_time.isoformat()}") config = get_cache_config() @@ -224,6 +246,17 @@ def log_usage(function_name: Optional[str] = None, log_type: str = "function"): """ def decorator(func: Callable) -> Callable: + """Inner decorator that wraps the function with usage logging. + + Args: + func: The async function to wrap with usage logging. + + Returns: + Callable: The wrapped function with usage logging enabled. + + Raises: + UsageLoggerError: If the function is not async. + """ if not inspect.iscoroutinefunction(func): raise UsageLoggerError( f"@log_usage requires an async function. Got {func.__name__} which is not async." @@ -231,6 +264,24 @@ def log_usage(function_name: Optional[str] = None, log_type: str = "function"): @wraps(func) async def async_wrapper(*args, **kwargs): + """Wrapper function that executes the original function and logs usage. + + This wrapper: + - Extracts user ID and parameters from function arguments + - Executes the original function + - Captures result, success status, and any errors + - Logs usage information asynchronously without blocking + + Args: + *args: Positional arguments passed to the original function. + **kwargs: Keyword arguments passed to the original function. + + Returns: + Any: The return value of the original function. + + Raises: + Any exception raised by the original function (re-raised after logging). + """ config = get_cache_config() if not config.usage_logging: return await func(*args, **kwargs)