cognee 0.3.6__py3-none-any.whl → 0.3.7.dev1__py3-none-any.whl

cognee/api/v1/sync/routers/get_sync_router.py

@@ -12,6 +12,7 @@ from cognee.modules.sync.methods import get_running_sync_operations_for_user, ge
 from cognee.shared.utils import send_telemetry
 from cognee.shared.logging_utils import get_logger
 from cognee.api.v1.sync import SyncResponse
+from cognee import __version__ as cognee_version
 from cognee.context_global_variables import set_database_global_context_variables
 
 logger = get_logger()
@@ -99,6 +100,7 @@ def get_sync_router() -> APIRouter:
             user.id,
             additional_properties={
                 "endpoint": "POST /v1/sync",
+                "cognee_version": cognee_version,
                 "dataset_ids": [str(id) for id in request.dataset_ids]
                 if request.dataset_ids
                 else "*",
@@ -205,6 +207,7 @@ def get_sync_router() -> APIRouter:
             user.id,
             additional_properties={
                 "endpoint": "GET /v1/sync/status",
+                "cognee_version": cognee_version,
             },
         )
 

cognee/api/v1/ui/ui.py

@@ -502,22 +502,46 @@ def start_ui(
 
     if start_mcp:
         logger.info("Starting Cognee MCP server with Docker...")
-        cwd = os.getcwd()
-        env_file = os.path.join(cwd, ".env")
         try:
+            image = "cognee/cognee-mcp:main"
+            subprocess.run(["docker", "pull", image], check=True)
+
+            import uuid
+
+            container_name = f"cognee-mcp-{uuid.uuid4().hex[:8]}"
+
+            docker_cmd = [
+                "docker",
+                "run",
+                "--name",
+                container_name,
+                "-p",
+                f"{mcp_port}:8000",
+                "--rm",
+                "-e",
+                "TRANSPORT_MODE=sse",
+            ]
+
+            if start_backend:
+                docker_cmd.extend(
+                    [
+                        "-e",
+                        f"API_URL=http://localhost:{backend_port}",
+                    ]
+                )
+                logger.info(
+                    f"Configuring MCP to connect to backend API at http://localhost:{backend_port}"
+                )
+                logger.info("(localhost will be auto-converted to host.docker.internal)")
+            else:
+                cwd = os.getcwd()
+                env_file = os.path.join(cwd, ".env")
+                docker_cmd.extend(["--env-file", env_file])
+
+            docker_cmd.append(image)
+
             mcp_process = subprocess.Popen(
-                [
-                    "docker",
-                    "run",
-                    "-p",
-                    f"{mcp_port}:8000",
-                    "--rm",
-                    "--env-file",
-                    env_file,
-                    "-e",
-                    "TRANSPORT_MODE=sse",
-                    "cognee/cognee-mcp:main",
-                ],
+                docker_cmd,
                 stdout=subprocess.PIPE,
                 stderr=subprocess.PIPE,
                 preexec_fn=os.setsid if hasattr(os, "setsid") else None,
@@ -526,8 +550,13 @@ def start_ui(
             _stream_process_output(mcp_process, "stdout", "[MCP]", "\033[34m")  # Blue
             _stream_process_output(mcp_process, "stderr", "[MCP]", "\033[34m")  # Blue
 
-            pid_callback(mcp_process.pid)
-            logger.info(f"✓ Cognee MCP server starting on http://127.0.0.1:{mcp_port}/sse")
+            # Pass both PID and container name using a tuple
+            pid_callback((mcp_process.pid, container_name))
+
+            mode_info = "API mode" if start_backend else "direct mode"
+            logger.info(
+                f"✓ Cognee MCP server starting on http://127.0.0.1:{mcp_port}/sse ({mode_info})"
+            )
         except Exception as e:
             logger.error(f"Failed to start MCP server with Docker: {str(e)}")
     # Start backend server if requested

cognee/api/v1/update/routers/get_update_router.py

@@ -9,6 +9,7 @@ from cognee.shared.logging_utils import get_logger
 from cognee.modules.users.models import User
 from cognee.modules.users.methods import get_authenticated_user
 from cognee.shared.utils import send_telemetry
+from cognee import __version__ as cognee_version
 from cognee.modules.pipelines.models.PipelineRunInfo import (
     PipelineRunErrored,
 )
@@ -64,6 +65,7 @@ def get_update_router() -> APIRouter:
                 "dataset_id": str(dataset_id),
                 "data_id": str(data_id),
                 "node_set": str(node_set),
+                "cognee_version": cognee_version,
             },
         )
 
@@ -75,7 +77,7 @@ def get_update_router() -> APIRouter:
                 data=data,
                 dataset_id=dataset_id,
                 user=user,
-                node_set=node_set,
+                node_set=node_set if node_set else None,
             )
 
             # If any cognify run errored return JSONResponse with proper error status code

cognee/api/v1/update/update.py

@@ -1,5 +1,5 @@
 from uuid import UUID
-from typing import Union, BinaryIO, List, Optional
+from typing import Union, BinaryIO, List, Optional, Any
 
 from cognee.modules.users.models import User
 from cognee.api.v1.delete import delete
@@ -10,12 +10,12 @@ from cognee.api.v1.cognify import cognify
 async def update(
     data_id: UUID,
     data: Union[BinaryIO, list[BinaryIO], str, list[str]],
+    dataset_id: UUID,
     user: User = None,
     node_set: Optional[List[str]] = None,
-    dataset_id: Optional[UUID] = None,
     vector_db_config: dict = None,
     graph_db_config: dict = None,
-    preferred_loaders: List[str] = None,
+    preferred_loaders: dict[str, dict[str, Any]] = None,
     incremental_loading: bool = True,
 ):
     """

cognee/api/v1/users/routers/get_visualize_router.py

@@ -8,6 +8,7 @@ from cognee.modules.users.models import User
 
 from cognee.context_global_variables import set_database_global_context_variables
 from cognee.shared.utils import send_telemetry
+from cognee import __version__ as cognee_version
 
 logger = get_logger()
 
@@ -46,6 +47,7 @@ def get_visualize_router() -> APIRouter:
             additional_properties={
                 "endpoint": "GET /v1/visualize",
                 "dataset_id": str(dataset_id),
+                "cognee_version": cognee_version,
             },
         )
 

cognee/cli/_cognee.py

@@ -175,19 +175,59 @@ def main() -> int:
     # Handle UI flag
     if hasattr(args, "start_ui") and args.start_ui:
         spawned_pids = []
+        docker_container = None
 
         def signal_handler(signum, frame):
             """Handle Ctrl+C and other termination signals"""
-            nonlocal spawned_pids
-            fmt.echo("\nShutting down UI server...")
+            nonlocal spawned_pids, docker_container
 
+            try:
+                fmt.echo("\nShutting down UI server...")
+            except (BrokenPipeError, OSError):
+                pass
+
+            # First, stop Docker container if running
+            if docker_container:
+                try:
+                    result = subprocess.run(
+                        ["docker", "stop", docker_container],
+                        capture_output=True,
+                        timeout=10,
+                        check=False,
+                    )
+                    try:
+                        if result.returncode == 0:
+                            fmt.success(f"✓ Docker container {docker_container} stopped.")
+                        else:
+                            fmt.warning(
+                                f"Could not stop container {docker_container}: {result.stderr.decode()}"
+                            )
+                    except (BrokenPipeError, OSError):
+                        pass
+                except subprocess.TimeoutExpired:
+                    try:
+                        fmt.warning(
+                            f"Timeout stopping container {docker_container}, forcing removal..."
+                        )
+                    except (BrokenPipeError, OSError):
+                        pass
+                    subprocess.run(
+                        ["docker", "rm", "-f", docker_container], capture_output=True, check=False
+                    )
+                except Exception:
+                    pass
+
+            # Then, stop regular processes
             for pid in spawned_pids:
                 try:
                     if hasattr(os, "killpg"):
                         # Unix-like systems: Use process groups
                         pgid = os.getpgid(pid)
                         os.killpg(pgid, signal.SIGTERM)
-                        fmt.success(f"✓ Process group {pgid} (PID {pid}) terminated.")
+                        try:
+                            fmt.success(f"✓ Process group {pgid} (PID {pid}) terminated.")
+                        except (BrokenPipeError, OSError):
+                            pass
                     else:
                         # Windows: Use taskkill to terminate process and its children
                         subprocess.run(
@@ -195,24 +235,35 @@ def main() -> int:
                             capture_output=True,
                             check=False,
                         )
-                        fmt.success(f"✓ Process {pid} and its children terminated.")
-                except (OSError, ProcessLookupError, subprocess.SubprocessError) as e:
-                    fmt.warning(f"Could not terminate process {pid}: {e}")
+                        try:
+                            fmt.success(f"✓ Process {pid} and its children terminated.")
+                        except (BrokenPipeError, OSError):
+                            pass
+                except (OSError, ProcessLookupError, subprocess.SubprocessError):
+                    pass
 
             sys.exit(0)
 
         signal.signal(signal.SIGINT, signal_handler)  # Ctrl+C
         signal.signal(signal.SIGTERM, signal_handler)  # Termination request
+        if hasattr(signal, "SIGHUP"):
+            signal.signal(signal.SIGHUP, signal_handler)
 
         try:
             from cognee import start_ui
 
             fmt.echo("Starting cognee UI...")
 
-            # Callback to capture PIDs of all spawned processes
-            def pid_callback(pid):
-                nonlocal spawned_pids
-                spawned_pids.append(pid)
+            # Callback to capture PIDs and Docker container of all spawned processes
+            def pid_callback(pid_or_tuple):
+                nonlocal spawned_pids, docker_container
+                # Handle both regular PIDs and (PID, container_name) tuples
+                if isinstance(pid_or_tuple, tuple):
+                    pid, container_name = pid_or_tuple
+                    spawned_pids.append(pid)
+                    docker_container = container_name
+                else:
+                    spawned_pids.append(pid_or_tuple)
 
             frontend_port = 3000
             start_backend, backend_port = True, 8000

cognee/cli/commands/add_command.py

@@ -70,11 +70,11 @@ After adding data, use `cognee cognify` to process it into knowledge graphs.
                     await cognee.add(data=data_to_add, dataset_name=args.dataset_name)
                     fmt.success(f"Successfully added data to dataset '{args.dataset_name}'")
                 except Exception as e:
-                    raise CliCommandInnerException(f"Failed to add data: {str(e)}")
+                    raise CliCommandInnerException(f"Failed to add data: {str(e)}") from e
 
             asyncio.run(run_add())
 
         except Exception as e:
             if isinstance(e, CliCommandInnerException):
-                raise CliCommandException(str(e), error_code=1)
-            raise CliCommandException(f"Error adding data: {str(e)}", error_code=1)
+                raise CliCommandException(str(e), error_code=1) from e
+            raise CliCommandException(f"Failed to add data: {str(e)}", error_code=1) from e

cognee/cli/commands/cognify_command.py

@@ -107,7 +107,7 @@ After successful cognify processing, use `cognee search` to query the knowledge
                     )
                     return result
                 except Exception as e:
-                    raise CliCommandInnerException(f"Failed to cognify: {str(e)}")
+                    raise CliCommandInnerException(f"Failed to cognify: {str(e)}") from e
 
             result = asyncio.run(run_cognify())
 
@@ -124,5 +124,5 @@ After successful cognify processing, use `cognee search` to query the knowledge
 
         except Exception as e:
             if isinstance(e, CliCommandInnerException):
-                raise CliCommandException(str(e), error_code=1)
-            raise CliCommandException(f"Error during cognification: {str(e)}", error_code=1)
+                raise CliCommandException(str(e), error_code=1) from e
+            raise CliCommandException(f"Error during cognification: {str(e)}", error_code=1) from e

cognee/cli/commands/config_command.py

@@ -79,8 +79,10 @@ Configuration changes will affect how cognee processes and stores data.
 
         except Exception as e:
             if isinstance(e, CliCommandInnerException):
-                raise CliCommandException(str(e), error_code=1)
-            raise CliCommandException(f"Error managing configuration: {str(e)}", error_code=1)
+                raise CliCommandException(str(e), error_code=1) from e
+            raise CliCommandException(
+                f"Error managing configuration: {str(e)}", error_code=1
+            ) from e
 
     def _handle_get(self, args: argparse.Namespace) -> None:
         try:
@@ -122,7 +124,7 @@ Configuration changes will affect how cognee processes and stores data.
                     fmt.note("Configuration viewing not fully implemented yet")
 
         except Exception as e:
-            raise CliCommandInnerException(f"Failed to get configuration: {str(e)}")
+            raise CliCommandInnerException(f"Failed to get configuration: {str(e)}") from e
 
     def _handle_set(self, args: argparse.Namespace) -> None:
         try:
@@ -141,7 +143,7 @@ Configuration changes will affect how cognee processes and stores data.
                 fmt.error(f"Failed to set configuration key '{args.key}'")
 
         except Exception as e:
-            raise CliCommandInnerException(f"Failed to set configuration: {str(e)}")
+            raise CliCommandInnerException(f"Failed to set configuration: {str(e)}") from e
 
     def _handle_unset(self, args: argparse.Namespace) -> None:
         try:
@@ -189,7 +191,7 @@ Configuration changes will affect how cognee processes and stores data.
                 fmt.note("Use 'cognee config list' to see all available configuration options")
 
         except Exception as e:
-            raise CliCommandInnerException(f"Failed to unset configuration: {str(e)}")
+            raise CliCommandInnerException(f"Failed to unset configuration: {str(e)}") from e
 
     def _handle_list(self, args: argparse.Namespace) -> None:
         try:
@@ -209,7 +211,7 @@ Configuration changes will affect how cognee processes and stores data.
             fmt.echo("  cognee config reset         - Reset all to defaults")
 
         except Exception as e:
-            raise CliCommandInnerException(f"Failed to list configuration: {str(e)}")
+            raise CliCommandInnerException(f"Failed to list configuration: {str(e)}") from e
 
     def _handle_reset(self, args: argparse.Namespace) -> None:
         try:
@@ -222,4 +224,4 @@ Configuration changes will affect how cognee processes and stores data.
             fmt.echo("This would reset all settings to their default values")
 
         except Exception as e:
-            raise CliCommandInnerException(f"Failed to reset configuration: {str(e)}")
+            raise CliCommandInnerException(f"Failed to reset configuration: {str(e)}") from e

cognee/cli/commands/delete_command.py

@@ -100,7 +100,7 @@ Be careful with deletion operations as they are irreversible.
                     else:
                         await cognee.delete(dataset_name=args.dataset_name, user_id=args.user_id)
                 except Exception as e:
-                    raise CliCommandInnerException(f"Failed to delete: {str(e)}")
+                    raise CliCommandInnerException(f"Failed to delete: {str(e)}") from e
 
             asyncio.run(run_delete())
             # This success message may be inaccurate due to the underlying bug, but we leave it for now.
@@ -108,5 +108,5 @@ Be careful with deletion operations as they are irreversible.
 
         except Exception as e:
             if isinstance(e, CliCommandInnerException):
-                raise CliCommandException(str(e), error_code=1)
-            raise CliCommandException(f"Error deleting data: {str(e)}", error_code=1)
+                raise CliCommandException(str(e), error_code=1) from e
+            raise CliCommandException(f"Error deleting data: {str(e)}", error_code=1) from e

cognee/cli/commands/search_command.py

@@ -31,10 +31,6 @@ Search Types & Use Cases:
     Traditional RAG using document chunks without graph structure.
     Best for: Direct document retrieval, specific fact-finding.
 
-**INSIGHTS**:
-    Structured entity relationships and semantic connections.
-    Best for: Understanding concept relationships, knowledge mapping.
-
 **CHUNKS**:
     Raw text segments that match the query semantically.
     Best for: Finding specific passages, citations, exact content.
@@ -108,7 +104,7 @@ Search Types & Use Cases:
                     )
                     return results
                 except Exception as e:
-                    raise CliCommandInnerException(f"Failed to search: {str(e)}")
+                    raise CliCommandInnerException(f"Failed to search: {str(e)}") from e
 
             results = asyncio.run(run_search())
 
@@ -145,5 +141,5 @@ Search Types & Use Cases:
 
         except Exception as e:
             if isinstance(e, CliCommandInnerException):
-                raise CliCommandException(str(e), error_code=1)
-            raise CliCommandException(f"Error searching: {str(e)}", error_code=1)
+                raise CliCommandException(str(e), error_code=1) from e
+            raise CliCommandException(f"Error searching: {str(e)}", error_code=1) from e

cognee/cli/config.py

@@ -19,7 +19,6 @@ COMMAND_DESCRIPTIONS = {
 SEARCH_TYPE_CHOICES = [
     "GRAPH_COMPLETION",
     "RAG_COMPLETION",
-    "INSIGHTS",
     "CHUNKS",
     "SUMMARIES",
     "CODE",

cognee/context_global_variables.py

@@ -12,6 +12,11 @@ from cognee.modules.users.methods import get_user
 #       for different async tasks, threads and processes
 vector_db_config = ContextVar("vector_db_config", default=None)
 graph_db_config = ContextVar("graph_db_config", default=None)
+session_user = ContextVar("session_user", default=None)
+
+
+async def set_session_user_context_variable(user):
+    session_user.set(user)
 
 
 async def set_database_global_context_variables(dataset: Union[str, UUID], user_id: UUID):

cognee/exceptions/exceptions.py

@@ -56,7 +56,7 @@ class CogneeValidationError(CogneeApiError):
         self,
         message: str = "A validation error occurred.",
         name: str = "CogneeValidationError",
-        status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+        status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
         log=True,
         log_level="ERROR",
     ):

cognee/infrastructure/databases/cache/__init__.py

@@ -0,0 +1,2 @@
+from .get_cache_engine import get_cache_engine
+from .config import get_cache_config

cognee/infrastructure/databases/cache/cache_db_interface.py

@@ -0,0 +1,79 @@
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+
+
+class CacheDBInterface(ABC):
+    """
+    Abstract base class for distributed cache coordination systems (e.g., Redis, Memcached).
+    Provides a common interface for lock acquisition, release, and context-managed locking.
+    """
+
+    def __init__(self, host: str, port: int, lock_key: str):
+        self.host = host
+        self.port = port
+        self.lock_key = lock_key
+        self.lock = None
+
+    @abstractmethod
+    def acquire_lock(self):
+        """
+        Acquire a lock on the given key.
+        Must be implemented by subclasses.
+        """
+        pass
+
+    @abstractmethod
+    def release_lock(self):
+        """
+        Release the lock if it is held.
+        Must be implemented by subclasses.
+        """
+        pass
+
+    @contextmanager
+    def hold_lock(self):
+        """
+        Context manager for safely acquiring and releasing the lock.
+        """
+        self.acquire()
+        try:
+            yield
+        finally:
+            self.release()
+
+    @abstractmethod
+    async def add_qa(
+        self,
+        user_id: str,
+        session_id: str,
+        question: str,
+        context: str,
+        answer: str,
+        ttl: int | None = 86400,
+    ):
+        """
+        Add a Q/A/context triplet to a cache session.
+        """
+
+        pass
+
+    @abstractmethod
+    async def get_latest_qa(self, user_id: str, session_id: str, last_n: int = 5):
+        """
+        Retrieve the most recent Q/A/context triplets for a session.
+        """
+        pass
+
+    @abstractmethod
+    async def get_all_qas(self, user_id: str, session_id: str):
+        """
+        Retrieve all Q/A/context triplets for the given session.
+        """
+        pass
+
+    @abstractmethod
+    async def close(self):
+        """
+        Gracefully close any async connections.
+        """
+        pass

cognee/infrastructure/databases/cache/config.py

@@ -0,0 +1,44 @@
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from functools import lru_cache
+from typing import Optional
+
+
+class CacheConfig(BaseSettings):
+    """
+    Configuration for distributed cache systems (e.g., Redis), used for locking or coordination.
+
+    Attributes:
+    - shared_kuzu_lock: Shared kuzu lock logic on/off.
+    - cache_host: Hostname of the cache service.
+    - cache_port: Port number for the cache service.
+    - agentic_lock_expire: Automatic lock expiration time (in seconds).
+    - agentic_lock_timeout: Maximum time (in seconds) to wait for the lock release.
+    """
+
+    caching: bool = False
+    shared_kuzu_lock: bool = False
+    cache_host: str = "localhost"
+    cache_port: int = 6379
+    cache_username: Optional[str] = None
+    cache_password: Optional[str] = None
+    agentic_lock_expire: int = 240
+    agentic_lock_timeout: int = 300
+
+    model_config = SettingsConfigDict(env_file=".env", extra="allow")
+
+    def to_dict(self) -> dict:
+        return {
+            "caching": self.caching,
+            "shared_kuzu_lock": self.shared_kuzu_lock,
+            "cache_host": self.cache_host,
+            "cache_port": self.cache_port,
+            "cache_username": self.cache_username,
+            "cache_password": self.cache_password,
+            "agentic_lock_expire": self.agentic_lock_expire,
+            "agentic_lock_timeout": self.agentic_lock_timeout,
+        }
+
+
+@lru_cache
+def get_cache_config():
+    return CacheConfig()

cognee/infrastructure/databases/cache/get_cache_engine.py

@@ -0,0 +1,67 @@
+"""Factory to get the appropriate cache coordination engine (e.g., Redis)."""
+
+from functools import lru_cache
+from typing import Optional
+from cognee.infrastructure.databases.cache.config import get_cache_config
+from cognee.infrastructure.databases.cache.cache_db_interface import CacheDBInterface
+
+config = get_cache_config()
+
+
+@lru_cache
+def create_cache_engine(
+    cache_host: str,
+    cache_port: int,
+    cache_username: str,
+    cache_password: str,
+    lock_key: str,
+    agentic_lock_expire: int = 240,
+    agentic_lock_timeout: int = 300,
+):
+    """
+    Factory function to instantiate a cache coordination backend (currently Redis).
+
+    Parameters:
+    -----------
+    - cache_host: Hostname or IP of the cache server.
+    - cache_port: Port number to connect to.
+    - cache_username: Username to authenticate with.
+    - cache_password: Password to authenticate with.
+    - lock_key: Identifier used for the locking resource.
+    - agentic_lock_expire: Duration to hold the lock after acquisition.
+    - agentic_lock_timeout: Max time to wait for the lock before failing.
+
+    Returns:
+    --------
+    - CacheDBInterface: An instance of the appropriate cache adapter. :TODO: Now we support only Redis. later if we add more here we can split the logic
+    """
+    if config.caching:
+        from cognee.infrastructure.databases.cache.redis.RedisAdapter import RedisAdapter
+
+        return RedisAdapter(
+            host=cache_host,
+            port=cache_port,
+            username=cache_username,
+            password=cache_password,
+            lock_name=lock_key,
+            timeout=agentic_lock_expire,
+            blocking_timeout=agentic_lock_timeout,
+        )
+    else:
+        return None
+
+
+def get_cache_engine(lock_key: Optional[str] = None) -> CacheDBInterface:
+    """
+    Returns a cache adapter instance using current context configuration.
+    """
+
+    return create_cache_engine(
+        cache_host=config.cache_host,
+        cache_port=config.cache_port,
+        cache_username=config.cache_username,
+        cache_password=config.cache_password,
+        lock_key=lock_key,
+        agentic_lock_expire=config.agentic_lock_expire,
+        agentic_lock_timeout=config.agentic_lock_timeout,
+    )