feat: adds new exception to shared usage logger

This commit is contained in:
hajdul88 2026-01-15 15:58:40 +01:00
parent 7ebe8563c5
commit b83af5f63f
3 changed files with 18 additions and 12 deletions

View file

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

View file

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

View file

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