diff --git a/cognee/shared/utils.py b/cognee/shared/utils.py index 7c76cfa59..3b6e7ba14 100644 --- a/cognee/shared/utils.py +++ b/cognee/shared/utils.py @@ -8,7 +8,8 @@ import http.server import socketserver from threading import Thread import pathlib -from uuid import uuid4, uuid5, NAMESPACE_OID +from typing import Union, Any, Dict, List +from uuid import uuid4, uuid5, NAMESPACE_OID, UUID from cognee.base_config import get_base_config from cognee.shared.logging_utils import get_logger @@ -58,7 +59,7 @@ def get_anonymous_id(): return anonymous_id -def _sanitize_nested_properties(obj, property_names: list[str]): +def _sanitize_nested_properties(obj: Any, property_names: list[str]) -> Any: """ Recursively replaces any property whose key matches one of `property_names` (e.g., ['url', 'path']) in a nested dict or list with a uuid5 hash @@ -78,7 +79,9 @@ def _sanitize_nested_properties(obj, property_names: list[str]): return obj -def send_telemetry(event_name: str, user_id, additional_properties: dict = {}): +def send_telemetry(event_name: str, user_id: Union[str, UUID], additional_properties: dict = {}): + if additional_properties is None: + additional_properties = {} if os.getenv("TELEMETRY_DISABLED"): return @@ -108,7 +111,7 @@ def send_telemetry(event_name: str, user_id, additional_properties: dict = {}): print(f"Error sending telemetry through proxy: {response.status_code}") -def embed_logo(p, layout_scale, logo_alpha, position): +def embed_logo(p: Any, layout_scale: float, logo_alpha: float, position: str): """ Embed a logo into the graph visualization as a watermark. """ @@ -138,7 +141,11 @@ def embed_logo(p, layout_scale, logo_alpha, position): def start_visualization_server( - host="0.0.0.0", port=8001, handler_class=http.server.SimpleHTTPRequestHandler + host: str = "0.0.0.0", + port: int = 8001, + handler_class: type[ + http.server.SimpleHTTPRequestHandler + ] = http.server.SimpleHTTPRequestHandler, ): """ Spin up a simple HTTP server in a background thread to serve files.