diff --git a/cognee/shared/exceptions/__init__.py b/cognee/shared/exceptions/__init__.py index 5e2ae6875..776a803b4 100644 --- a/cognee/shared/exceptions/__init__.py +++ b/cognee/shared/exceptions/__init__.py @@ -4,6 +4,4 @@ Custom exceptions for the Cognee API. This module defines a set of exceptions for handling various shared utility errors """ -from .exceptions import ( - IngestionError, -) +from .exceptions import IngestionError, UsageLoggerError diff --git a/cognee/shared/exceptions/exceptions.py b/cognee/shared/exceptions/exceptions.py index 3740f677d..c00a39b9e 100644 --- a/cognee/shared/exceptions/exceptions.py +++ b/cognee/shared/exceptions/exceptions.py @@ -1,4 +1,4 @@ -from cognee.exceptions import CogneeValidationError +from cognee.exceptions import CogneeConfigurationError, CogneeValidationError from fastapi import status @@ -10,3 +10,13 @@ class IngestionError(CogneeValidationError): status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, ): super().__init__(message, name, status_code) + + +class UsageLoggerError(CogneeConfigurationError): + def __init__( + self, + message: str = "Usage logging configuration is invalid.", + name: str = "UsageLoggerError", + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + ): + super().__init__(message, name, status_code) diff --git a/cognee/shared/usage_logger.py b/cognee/shared/usage_logger.py index 3456991e0..bf7b50b64 100644 --- a/cognee/shared/usage_logger.py +++ b/cognee/shared/usage_logger.py @@ -8,6 +8,7 @@ from uuid import UUID from cognee.infrastructure.databases.cache.config import get_cache_config from cognee.infrastructure.databases.cache.get_cache_engine import get_cache_engine +from cognee.shared.exceptions import UsageLoggerError from cognee.shared.logging_utils import get_logger from cognee import __version__ as cognee_version @@ -223,6 +224,11 @@ def log_usage(function_name: Optional[str] = None, log_type: str = "function"): """ def decorator(func: Callable) -> Callable: + if not inspect.iscoroutinefunction(func): + raise UsageLoggerError( + f"@log_usage requires an async function. Got {func.__name__} which is not async." + ) + @wraps(func) async def async_wrapper(*args, **kwargs): config = get_cache_config() @@ -295,12 +301,6 @@ if __name__ == "__main__": await asyncio.sleep(0.05) return float("nan") - @log_usage(function_name="returns_cycle", log_type="function") - async def returns_cycle(): - a = [] - a.append(a) - return a - async def run_example(): """Run example demonstrations.""" print("Usage Logger Example") @@ -323,8 +323,6 @@ if __name__ == "__main__": print(f" Result: {result2}") await asyncio.sleep(0.2) # Wait for async logging to complete - await returns_cycle() - # Example 3: Retrieve logs (if cache engine is available) print("\n3. Retrieving usage logs:") try: