feat: adds new exception to shared usage logger
This commit is contained in:
parent
7ebe8563c5
commit
b83af5f63f
3 changed files with 18 additions and 12 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue