cognee 0.5.1.dev0__py3-none-any.whl → 0.5.2__py3-none-any.whl

cognee/modules/search/models/SearchResultPayload.py

@@ -0,0 +1,67 @@
+from uuid import UUID
+from typing import Optional, Any, List, Union
+from pydantic import BaseModel, ConfigDict, field_serializer
+from pydantic.alias_generators import to_camel
+from cognee.modules.search.types.SearchType import SearchType
+
+
+class SearchResultPayload(BaseModel):
+    """Result payload from retriever classes."""
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        alias_generator=to_camel,
+        populate_by_name=True,
+    )
+
+    result_object: Any = None
+    context: Optional[Union[str, List[str]]] = None
+    completion: Optional[Union[str, List[str], List[dict]]] = None
+
+    # TODO: Add return_type info
+    search_type: SearchType
+    only_context: bool = False
+
+    dataset_name: Optional[str] = None
+    dataset_id: Optional[UUID] = None
+    dataset_tenant_id: Optional[UUID] = None
+
+    @field_serializer("result_object")
+    def serialize_complex_types(self, v: Any):
+        """
+        Custom serializer to handle complex types in result_object.
+        Transforms non-JSON-compatible types to their string representation.
+        """
+
+        # Helper to check if a value is a "simple" JSON-compatible type
+        def is_simple(item):
+            return isinstance(item, (int, float, dict, str, bool, type(None)))
+
+        if isinstance(v, list) and all(isinstance(item, dict) for item in v):
+            # Handle List of Dictionaries
+            return [
+                {key: (val if is_simple(val) else str(val)) for key, val in item.items()}
+                for item in v
+            ]
+        elif isinstance(v, list):
+            # Handle Lists
+            return [item if is_simple(item) else str(item) for item in v]
+        elif isinstance(v, dict):
+            # Handle Dictionaries
+            return {key: (val if is_simple(val) else str(val)) for key, val in v.items()}
+        else:
+            # Fallback for the object itself
+            return v if is_simple(v) else str(v)
+
+    @property
+    def result(self) -> Any:
+        """Function used to determine search_result for users request.
+        Return context if only_context is True, else return completion if it exists, else return result_object."""
+        if self.only_context:
+            return self.context
+        elif self.completion:
+            return self.completion
+        elif self.context:
+            return self.context
+        else:
+            return self.result_object

cognee/modules/search/types/SearchResult.py

@@ -1,6 +1,6 @@
 from uuid import UUID
 from pydantic import BaseModel
-from typing import Any, Dict, List, Optional
+from typing import Any, Optional
 
 
 class SearchResultDataset(BaseModel):
@@ -8,13 +8,6 @@ class SearchResultDataset(BaseModel):
     name: str
 
 
-class CombinedSearchResult(BaseModel):
-    result: Optional[Any]
-    context: Dict[str, Any]
-    graphs: Optional[Dict[str, Any]] = {}
-    datasets: Optional[List[SearchResultDataset]] = None
-
-
 class SearchResult(BaseModel):
     search_result: Any
     dataset_id: Optional[UUID]

cognee/modules/search/types/SearchType.py

@@ -1,7 +1,7 @@
 from enum import Enum
 
 
-class SearchType(Enum):
+class SearchType(str, Enum):
     SUMMARIES = "SUMMARIES"
     CHUNKS = "CHUNKS"
     RAG_COMPLETION = "RAG_COMPLETION"
@@ -13,7 +13,6 @@ class SearchType(Enum):
     GRAPH_COMPLETION_COT = "GRAPH_COMPLETION_COT"
     GRAPH_COMPLETION_CONTEXT_EXTENSION = "GRAPH_COMPLETION_CONTEXT_EXTENSION"
     FEELING_LUCKY = "FEELING_LUCKY"
-    FEEDBACK = "FEEDBACK"
     TEMPORAL = "TEMPORAL"
     CODING_RULES = "CODING_RULES"
     CHUNKS_LEXICAL = "CHUNKS_LEXICAL"

cognee/modules/search/types/__init__.py

@@ -1,2 +1,2 @@
 from .SearchType import SearchType
-from .SearchResult import SearchResult, SearchResultDataset, CombinedSearchResult
+from .SearchResult import SearchResult, SearchResultDataset

cognee/modules/search/utils/__init__.py

@@ -1,2 +1 @@
-from .prepare_search_result import prepare_search_result
-from .transform_context_to_graph import transform_context_to_graph
+from .transform_triplets_to_graph import transform_triplets_to_graph

cognee/modules/search/utils/transform_insights_to_graph.py

@@ -1,11 +1,11 @@
 from typing import Dict, List, Tuple
 
 
-def transform_insights_to_graph(context: List[Tuple[Dict, Dict, Dict]]):
+def transform_insights_to_graph(input_triplets: List[Tuple[Dict, Dict, Dict]]):
     nodes = {}
     edges = {}
 
-    for triplet in context:
+    for triplet in input_triplets:
         nodes[triplet[0]["id"]] = {
             "id": triplet[0]["id"],
             "label": triplet[0]["name"] if "name" in triplet[0] else triplet[0]["id"],

cognee/modules/search/utils/{transform_context_to_graph.py → transform_triplets_to_graph.py}

@@ -3,11 +3,11 @@ from typing import List
 from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
 
 
-def transform_context_to_graph(context: List[Edge]):
+def transform_triplets_to_graph(input_triplets: List[Edge]):
     nodes = {}
     edges = {}
 
-    for triplet in context:
+    for triplet in input_triplets:
         nodes[triplet.node1.id] = {
             "id": triplet.node1.id,
             "label": triplet.node1.attributes["name"]

cognee/modules/users/authentication/default/default_transport.py

@@ -1,12 +1,22 @@
 import os
 from fastapi_users.authentication import CookieTransport
 
+# Get cookie domain from environment variable
+# If not set or empty, use None to allow cookie to work on any domain
+cookie_domain = os.getenv("AUTH_TOKEN_COOKIE_DOMAIN")
+if cookie_domain == "":
+    cookie_domain = None
+
+# Note: Cookie expiration is automatically set by FastAPI Users based on JWT Strategy's lifetime_seconds
+# The JWT Strategy lifetime_seconds is configured in get_client_auth_backend.py
+# and reads from JWT_LIFETIME_SECONDS environment variable
+
 default_transport = CookieTransport(
     cookie_name=os.getenv("AUTH_TOKEN_COOKIE_NAME", "auth_token"),
     cookie_secure=False,
     cookie_httponly=True,
     cookie_samesite="Lax",
-    cookie_domain="localhost",
+    cookie_domain=cookie_domain,  # None allows cookie to work on any domain
 )
 
 default_transport.name = "cookie"

cognee/modules/users/authentication/get_api_auth_backend.py

@@ -16,8 +16,9 @@ def get_api_auth_backend():
 
     def get_jwt_strategy() -> JWTStrategy[models.UP, models.ID]:
         secret = os.getenv("FASTAPI_USERS_JWT_SECRET", "super_secret")
+        lifetime_seconds = int(os.getenv("JWT_LIFETIME_SECONDS", "3600"))
 
-        return APIJWTStrategy(secret, lifetime_seconds=36000)
+        return APIJWTStrategy(secret, lifetime_seconds=lifetime_seconds)
 
     auth_backend = AuthenticationBackend(
         name=transport.name,

cognee/modules/users/authentication/get_client_auth_backend.py

@@ -18,8 +18,9 @@ def get_client_auth_backend():
         from .default.default_jwt_strategy import DefaultJWTStrategy
 
         secret = os.getenv("FASTAPI_USERS_JWT_SECRET", "super_secret")
+        lifetime_seconds = int(os.getenv("JWT_LIFETIME_SECONDS", "3600"))
 
-        return DefaultJWTStrategy(secret, lifetime_seconds=3600)
+        return DefaultJWTStrategy(secret, lifetime_seconds=lifetime_seconds)
 
     auth_backend = AuthenticationBackend(
         name=transport.name,

cognee/modules/users/methods/create_user.py

@@ -1,18 +1,9 @@
-from uuid import UUID, uuid4
 from fastapi_users.exceptions import UserAlreadyExists
-from sqlalchemy.ext.asyncio import AsyncSession
 
 from cognee.infrastructure.databases.relational import get_relational_engine
-from cognee.modules.notebooks.models.Notebook import Notebook
-from cognee.modules.notebooks.methods.create_notebook import _create_tutorial_notebook
-from cognee.modules.users.exceptions import TenantNotFoundError
 from cognee.modules.users.get_user_manager import get_user_manager_context
 from cognee.modules.users.get_user_db import get_user_db_context
 from cognee.modules.users.models.User import UserCreate
-from cognee.modules.users.models.Tenant import Tenant
-
-from sqlalchemy import select
-from typing import Optional
 
 
 async def create_user(

cognee/modules/users/permissions/methods/has_user_management_permission.py

@@ -0,0 +1,29 @@
+from uuid import UUID
+
+from cognee.modules.users.permissions.methods.get_tenant import get_tenant
+from cognee.modules.users.exceptions import PermissionDeniedError
+
+
+async def has_user_management_permission(requester_id: UUID, tenant_id: UUID) -> bool:
+    """
+    Check if requester is allowed to manage users for a tenant.
+    Args:
+        requester_id: Id of the user making the request
+        tenant_id: Id of the tenant
+
+    Returns:
+        True if requester has permission
+
+    Raises:
+        PermissionDeniedError: If requester is not authorized.
+        TenantNotFoundError: If the tenant does not exist.
+    """
+    tenant = await get_tenant(tenant_id)
+
+    # TODO: extend to support admin roles
+    if tenant.owner_id != requester_id:
+        raise PermissionDeniedError(
+            message="User is not authorized to manage users for this tenant"
+        )
+
+    return True

cognee/modules/visualization/cognee_network_visualization.py

@@ -92,7 +92,7 @@ async def cognee_network_visualization(graph_data, destination_file_path: str =
         }
         links_list.append(link_data)
 
-    html_template = """
+    html_template = r"""
     <!DOCTYPE html>
     <html>
     <head>

cognee/run_migrations.py

@@ -0,0 +1,48 @@
+import os
+import sys
+import subprocess
+from pathlib import Path
+import importlib.resources as pkg_resources
+
+# Assuming your package is named 'cognee' and the migrations are under 'cognee/alembic'
+# This is a placeholder for the path logic.
+MIGRATIONS_PACKAGE = "cognee"
+MIGRATIONS_DIR_NAME = "alembic"
+
+
+async def run_migrations():
+    """
+    Finds the Alembic configuration within the installed package and
+    programmatically executes 'alembic upgrade head'.
+    """
+    # 1. Locate the base path of the installed package.
+    # This reliably finds the root directory of the installed 'cognee' package.
+    # We look for the parent of the 'migrations' directory.
+    package_root = str(pkg_resources.files(MIGRATIONS_PACKAGE))
+
+    # 2. Define the paths for config and scripts
+    alembic_ini_path = os.path.join(package_root, "alembic.ini")
+    script_location_path = os.path.join(package_root, MIGRATIONS_DIR_NAME)
+
+    if not os.path.exists(alembic_ini_path):
+        raise FileNotFoundError(
+            f"Error: alembic.ini not found at expected locations for package '{MIGRATIONS_PACKAGE}'."
+        )
+    if not os.path.exists(script_location_path):
+        raise FileNotFoundError(
+            f"Error: Migrations directory not found at expected locations for package '{MIGRATIONS_PACKAGE}'."
+        )
+
+    migration_result = subprocess.run(
+        ["python", "-m", "alembic", "upgrade", "head"],
+        capture_output=True,
+        text=True,
+        cwd=Path(package_root),
+    )
+
+    if migration_result.returncode != 0:
+        migration_output = migration_result.stderr + migration_result.stdout
+        print(f"Migration failed with unexpected error: {migration_output}")
+        sys.exit(1)
+
+    print("Migration completed successfully.")

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

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

@@ -0,0 +1,332 @@
+import asyncio
+import inspect
+import os
+from datetime import datetime, timezone
+from functools import singledispatch, wraps
+from typing import Any, Callable, Optional
+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
+
+logger = get_logger("usage_logger")
+
+
+@singledispatch
+def _sanitize_value(value: Any) -> Any:
+    """Default handler for JSON serialization - converts to string."""
+    try:
+        str_repr = str(value)
+        if str_repr.startswith("<") and str_repr.endswith(">"):
+            return f"<cannot be serialized: {type(value).__name__}>"
+        return str_repr
+    except Exception:
+        return f"<cannot be serialized: {type(value).__name__}>"
+
+
+@_sanitize_value.register(type(None))
+def _(value: None) -> None:
+    """Handle None values - returns None as-is."""
+    return None
+
+
+@_sanitize_value.register(str)
+@_sanitize_value.register(int)
+@_sanitize_value.register(float)
+@_sanitize_value.register(bool)
+def _(value: str | int | float | bool) -> str | int | float | bool:
+    """Handle primitive types - returns value as-is since they're JSON-serializable."""
+    return value
+
+
+@_sanitize_value.register(UUID)
+def _(value: UUID) -> str:
+    """Convert UUID to string representation."""
+    return str(value)
+
+
+@_sanitize_value.register(datetime)
+def _(value: datetime) -> str:
+    """Convert datetime to ISO format string."""
+    return value.isoformat()
+
+
+@_sanitize_value.register(list)
+@_sanitize_value.register(tuple)
+def _(value: list | tuple) -> list:
+    """Recursively sanitize list or tuple elements."""
+    return [_sanitize_value(v) for v in value]
+
+
+@_sanitize_value.register(dict)
+def _(value: dict) -> dict:
+    """Recursively sanitize dictionary keys and values."""
+    sanitized = {}
+    for k, v in value.items():
+        key_str = k if isinstance(k, str) else _sanitize_dict_key(k)
+        sanitized[key_str] = _sanitize_value(v)
+    return sanitized
+
+
+def _sanitize_dict_key(key: Any) -> str:
+    """Convert a non-string dict key to a string."""
+    sanitized_key = _sanitize_value(key)
+    if isinstance(sanitized_key, str):
+        if sanitized_key.startswith("<cannot be serialized"):
+            return f"<key:{type(key).__name__}>"
+        return sanitized_key
+    return str(sanitized_key)
+
+
+def _get_param_names(func: Callable) -> list[str]:
+    """Get parameter names from function signature."""
+    try:
+        return list(inspect.signature(func).parameters.keys())
+    except Exception:
+        return []
+
+
+def _get_param_defaults(func: Callable) -> dict[str, Any]:
+    """Get parameter defaults from function signature."""
+    try:
+        sig = inspect.signature(func)
+        defaults = {}
+        for param_name, param in sig.parameters.items():
+            if param.default != inspect.Parameter.empty:
+                defaults[param_name] = param.default
+        return defaults
+    except Exception:
+        return {}
+
+
+def _extract_user_id(args: tuple, kwargs: dict, param_names: list[str]) -> Optional[str]:
+    """Extract user_id from function arguments if available."""
+    try:
+        if "user" in kwargs and kwargs["user"] is not None:
+            user = kwargs["user"]
+            if hasattr(user, "id"):
+                return str(user.id)
+
+        for i, param_name in enumerate(param_names):
+            if i < len(args) and param_name == "user":
+                user = args[i]
+                if user is not None and hasattr(user, "id"):
+                    return str(user.id)
+        return None
+    except Exception:
+        return None
+
+
+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."""
+    params = {}
+
+    for key, value in kwargs.items():
+        if key != "user":
+            params[key] = _sanitize_value(value)
+
+    if param_names:
+        for i, param_name in enumerate(param_names):
+            if i < len(args) and param_name != "user" and param_name not in kwargs:
+                params[param_name] = _sanitize_value(args[i])
+    else:
+        for i, arg_value in enumerate(args):
+            params[f"arg_{i}"] = _sanitize_value(arg_value)
+
+    if param_names:
+        defaults = _get_param_defaults(func)
+        for param_name in param_names:
+            if param_name != "user" and param_name not in params and param_name in defaults:
+                params[param_name] = _sanitize_value(defaults[param_name])
+
+    return params
+
+
+async def _log_usage_async(
+    function_name: str,
+    log_type: str,
+    user_id: Optional[str],
+    parameters: dict,
+    result: Any,
+    success: bool,
+    error: Optional[str],
+    duration_ms: float,
+    start_time: datetime,
+    end_time: datetime,
+):
+    """Asynchronously log function usage to Redis.
+
+    Args:
+        function_name: Name of the function being logged.
+        log_type: Type of log entry (e.g., "api_endpoint", "mcp_tool", "function").
+        user_id: User identifier, or None to use "unknown".
+        parameters: Dictionary of function parameters (sanitized).
+        result: Function return value (will be sanitized).
+        success: Whether the function executed successfully.
+        error: Error message if function failed, None otherwise.
+        duration_ms: Execution duration in milliseconds.
+        start_time: Function start timestamp.
+        end_time: Function end timestamp.
+
+    Note:
+        This function silently handles errors to avoid disrupting the original
+        function execution. Logs are written to Redis with TTL from config.
+    """
+    try:
+        logger.debug(f"Starting to log usage for {function_name} at {start_time.isoformat()}")
+        config = get_cache_config()
+        if not config.usage_logging:
+            logger.debug("Usage logging disabled, skipping log")
+            return
+
+        logger.debug(f"Getting cache engine for {function_name}")
+        cache_engine = get_cache_engine()
+        if cache_engine is None:
+            logger.warning(
+                f"Cache engine not available for usage logging (function: {function_name})"
+            )
+            return
+
+        logger.debug(f"Cache engine obtained for {function_name}")
+
+        if user_id is None:
+            user_id = "unknown"
+            logger.debug(f"No user_id provided, using 'unknown' for {function_name}")
+
+        log_entry = {
+            "timestamp": start_time.isoformat(),
+            "type": log_type,
+            "function_name": function_name,
+            "user_id": user_id,
+            "parameters": parameters,
+            "result": _sanitize_value(result),
+            "success": success,
+            "error": error,
+            "duration_ms": round(duration_ms, 2),
+            "start_time": start_time.isoformat(),
+            "end_time": end_time.isoformat(),
+            "metadata": {
+                "cognee_version": cognee_version,
+                "environment": os.getenv("ENV", "prod"),
+            },
+        }
+
+        logger.debug(f"Calling log_usage for {function_name}, user_id={user_id}")
+        await cache_engine.log_usage(
+            user_id=user_id,
+            log_entry=log_entry,
+            ttl=config.usage_logging_ttl,
+        )
+        logger.info(f"Successfully logged usage for {function_name} (user_id={user_id})")
+    except Exception as e:
+        logger.error(f"Failed to log usage for {function_name}: {str(e)}", exc_info=True)
+
+
+def log_usage(function_name: Optional[str] = None, log_type: str = "function"):
+    """
+    Decorator to log function usage to Redis.
+
+    This decorator is completely transparent - it doesn't change function behavior.
+    It logs function name, parameters, result, timing, and user (if available).
+
+    Args:
+        function_name: Optional name for the function (defaults to func.__name__)
+        log_type: Type of log entry (e.g., "api_endpoint", "mcp_tool")
+
+    Usage:
+        @log_usage(function_name="MCP my_mcp_tool", log_type="mcp_tool")
+        async def my_mcp_tool(...):
+            # mcp code
+
+        @log_usage(function_name="POST API /v1/add", log_type="api_endpoint")
+        async def add(...):
+            # endpoint code
+    """
+
+    def decorator(func: Callable) -> Callable:
+        """Inner decorator that wraps the function with usage logging.
+
+        Args:
+            func: The async function to wrap with usage logging.
+
+        Returns:
+            Callable: The wrapped function with usage logging enabled.
+
+        Raises:
+            UsageLoggerError: If the function is not async.
+        """
+        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):
+            """Wrapper function that executes the original function and logs usage.
+
+            This wrapper:
+            - Extracts user ID and parameters from function arguments
+            - Executes the original function
+            - Captures result, success status, and any errors
+            - Logs usage information asynchronously without blocking
+
+            Args:
+                *args: Positional arguments passed to the original function.
+                **kwargs: Keyword arguments passed to the original function.
+
+            Returns:
+                Any: The return value of the original function.
+
+            Raises:
+                Any exception raised by the original function (re-raised after logging).
+            """
+            config = get_cache_config()
+            if not config.usage_logging:
+                return await func(*args, **kwargs)
+
+            start_time = datetime.now(timezone.utc)
+
+            param_names = _get_param_names(func)
+            user_id = _extract_user_id(args, kwargs, param_names)
+            parameters = _extract_parameters(args, kwargs, param_names, func)
+
+            result = None
+            success = True
+            error = None
+
+            try:
+                result = await func(*args, **kwargs)
+                return result
+            except Exception as e:
+                success = False
+                error = str(e)
+                raise
+            finally:
+                end_time = datetime.now(timezone.utc)
+                duration_ms = (end_time - start_time).total_seconds() * 1000
+
+                try:
+                    await _log_usage_async(
+                        function_name=function_name or func.__name__,
+                        log_type=log_type,
+                        user_id=user_id,
+                        parameters=parameters,
+                        result=result,
+                        success=success,
+                        error=error,
+                        duration_ms=duration_ms,
+                        start_time=start_time,
+                        end_time=end_time,
+                    )
+                except Exception as e:
+                    logger.error(
+                        f"Failed to log usage for {function_name or func.__name__}: {str(e)}",
+                        exc_info=True,
+                    )
+
+        return async_wrapper
+
+    return decorator