feat: adds new exception to shared usage logger

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 (
+from .exceptions import IngestionError, UsageLoggerError
-    IngestionError,
-)

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)

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: