gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

ruff

Browse source
This commit is contained in:
hajdul88 2026-01-15 15:21:24 +01:00
parent 8dc358da39
commit 7ebe8563c5
1 changed files with 10 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
cognee/shared/usage_logger.py
View file

@ -116,11 +116,11 @@ def _extract_user_id(args: tuple, kwargs: dict, param_names: list[str]) -> Optio
def _extract_parameters(args: tuple, kwargs: dict, param_names: list[str], func: Callable) -> dict: def _extract_parameters(args: tuple, kwargs: dict, param_names: list[str], func: Callable) -> dict:
"""Extract function parameters - captures all parameters including defaults, sanitizes for JSON.""" """Extract function parameters - captures all parameters including defaults, sanitizes for JSON."""
params = {} params = {}
for key, value in kwargs.items(): for key, value in kwargs.items():
if key != "user": if key != "user":
params[key] = _sanitize_value(value) params[key] = _sanitize_value(value)
if param_names: if param_names:
for i, param_name in enumerate(param_names): for i, param_name in enumerate(param_names):
if i < len(args) and param_name != "user" and param_name not in kwargs: if i < len(args) and param_name != "user" and param_name not in kwargs:
@ -128,13 +128,13 @@ def _extract_parameters(args: tuple, kwargs: dict, param_names: list[str], func:
else: else:
for i, arg_value in enumerate(args): for i, arg_value in enumerate(args):
params[f"arg_{i}"] = _sanitize_value(arg_value) params[f"arg_{i}"] = _sanitize_value(arg_value)
if param_names: if param_names:
defaults = _get_param_defaults(func) defaults = _get_param_defaults(func)
for param_name in param_names: for param_name in param_names:
if param_name != "user" and param_name not in params and param_name in defaults: if param_name != "user" and param_name not in params and param_name in defaults:
params[param_name] = _sanitize_value(defaults[param_name]) params[param_name] = _sanitize_value(defaults[param_name])
return params return params
@ -288,12 +288,13 @@ if __name__ == "__main__":
self.id = user_id self.id = user_id
@log_usage(function_name="example_with_user", log_type="example") @log_usage(function_name="example_with_user", log_type="example")
async def example_with_user(data: str, user: MockUser, wrong_param=datetime.utcnow().isoformat()): async def example_with_user(
data: str, user: MockUser, wrong_param=datetime.utcnow().isoformat()
):
"""Example function with user parameter.""" """Example function with user parameter."""
await asyncio.sleep(0.05) await asyncio.sleep(0.05)
return float("nan") return float("nan")
@log_usage(function_name="returns_cycle", log_type="function") @log_usage(function_name="returns_cycle", log_type="function")
async def returns_cycle(): async def returns_cycle():
a = [] a = []
@ -314,7 +315,9 @@ if __name__ == "__main__":
# Example 2: Function with user # Example 2: Function with user
print("\n2. Running example function with user:") print("\n2. Running example function with user:")
mock_user = MockUser("example-user-123") mock_user = MockUser("example-user-123")
result2 = await example_with_user("sample_data", user=mock_user, wrong_param=datetime.utcnow().isoformat()) result2 = await example_with_user(
"sample_data", user=mock_user, wrong_param=datetime.utcnow().isoformat()
)
result3 = await example_with_user("sample_data", user=mock_user) result3 = await example_with_user("sample_data", user=mock_user)
print(f" Result: {result2}") print(f" Result: {result2}")
@ -353,7 +356,5 @@ if __name__ == "__main__":
print(" - CACHE_HOST=localhost") print(" - CACHE_HOST=localhost")
print(" - CACHE_PORT=6379") print(" - CACHE_PORT=6379")
# Run the example # Run the example
asyncio.run(run_example()) asyncio.run(run_example())