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

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.

cognee/tasks/chunks/__init__.py

@@ -1,3 +1,12 @@
+"""
+Text chunking and chunk management tasks.
+
+This module provides functionality for splitting text into chunks using
+different strategies (word, sentence, paragraph, or row-based) and for
+cleaning up disconnected or obsolete chunks to support downstream
+processing and knowledge graph workflows.
+"""
+
 from .chunk_by_word import chunk_by_word
 from .chunk_by_sentence import chunk_by_sentence
 from .chunk_by_paragraph import chunk_by_paragraph

cognee/tasks/cleanup/cleanup_unused_data.py

@@ -0,0 +1,172 @@
+"""
+Task for automatically deleting unused data from the memify pipeline.
+
+This task identifies and removes entire documents that haven't
+been accessed by retrievers for a specified period, helping maintain system
+efficiency and storage optimization through whole-document removal.
+"""
+
+import json
+from datetime import datetime, timezone, timedelta
+from typing import Optional, Dict, Any
+from uuid import UUID
+import os
+from cognee.infrastructure.databases.graph import get_graph_engine
+from cognee.infrastructure.databases.vector import get_vector_engine
+from cognee.infrastructure.databases.relational import get_relational_engine
+from cognee.modules.data.models import Data, DatasetData
+from cognee.shared.logging_utils import get_logger
+from sqlalchemy import select, or_
+import cognee
+import sqlalchemy as sa
+from cognee.modules.graph.cognee_graph.CogneeGraph import CogneeGraph
+
+logger = get_logger(__name__)
+
+
+async def cleanup_unused_data(
+    minutes_threshold: Optional[int], dry_run: bool = True, user_id: Optional[UUID] = None
+) -> Dict[str, Any]:
+    """
+    Identify and remove unused data from the memify pipeline.
+
+    Parameters
+    ----------
+    minutes_threshold : int
+        Minutes since last access to consider data unused
+    dry_run : bool
+        If True, only report what would be deleted without actually deleting (default: True)
+    user_id : UUID, optional
+        Limit cleanup to specific user's data (default: None)
+
+    Returns
+    -------
+    Dict[str, Any]
+        Cleanup results with status, counts, and timestamp
+    """
+    # Check 1: Environment variable must be enabled
+    if os.getenv("ENABLE_LAST_ACCESSED", "false").lower() != "true":
+        logger.warning("Cleanup skipped: ENABLE_LAST_ACCESSED is not enabled.")
+        return {
+            "status": "skipped",
+            "reason": "ENABLE_LAST_ACCESSED not enabled",
+            "unused_count": 0,
+            "deleted_count": {},
+            "cleanup_date": datetime.now(timezone.utc).isoformat(),
+        }
+
+    # Check 2: Verify tracking has actually been running
+    db_engine = get_relational_engine()
+    async with db_engine.get_async_session() as session:
+        # Count records with non-NULL last_accessed
+        tracked_count = await session.execute(
+            select(sa.func.count(Data.id)).where(Data.last_accessed.isnot(None))
+        )
+        tracked_records = tracked_count.scalar()
+
+        if tracked_records == 0:
+            logger.warning(
+                "Cleanup skipped: No records have been tracked yet. "
+                "ENABLE_LAST_ACCESSED may have been recently enabled. "
+                "Wait for retrievers to update timestamps before running cleanup."
+            )
+            return {
+                "status": "skipped",
+                "reason": "No tracked records found - tracking may be newly enabled",
+                "unused_count": 0,
+                "deleted_count": {},
+                "cleanup_date": datetime.now(timezone.utc).isoformat(),
+            }
+
+    logger.info(
+        "Starting cleanup task",
+        minutes_threshold=minutes_threshold,
+        dry_run=dry_run,
+        user_id=str(user_id) if user_id else None,
+    )
+
+    # Calculate cutoff timestamp
+    cutoff_date = datetime.now(timezone.utc) - timedelta(minutes=minutes_threshold)
+
+    # Document-level approach (recommended)
+    return await _cleanup_via_sql(cutoff_date, dry_run, user_id)
+
+
+async def _cleanup_via_sql(
+    cutoff_date: datetime, dry_run: bool, user_id: Optional[UUID] = None
+) -> Dict[str, Any]:
+    """
+    SQL-based cleanup: Query Data table for unused documents and use cognee.delete().
+
+    Parameters
+    ----------
+    cutoff_date : datetime
+        Cutoff date for last_accessed filtering
+    dry_run : bool
+        If True, only report what would be deleted
+    user_id : UUID, optional
+        Filter by user ID if provided
+
+    Returns
+    -------
+    Dict[str, Any]
+        Cleanup results
+    """
+    db_engine = get_relational_engine()
+
+    async with db_engine.get_async_session() as session:
+        # Query for Data records with old last_accessed timestamps
+        query = (
+            select(Data, DatasetData)
+            .join(DatasetData, Data.id == DatasetData.data_id)
+            .where(or_(Data.last_accessed < cutoff_date, Data.last_accessed.is_(None)))
+        )
+
+        if user_id:
+            from cognee.modules.data.models import Dataset
+
+            query = query.join(Dataset, DatasetData.dataset_id == Dataset.id).where(
+                Dataset.owner_id == user_id
+            )
+
+        result = await session.execute(query)
+        unused_data = result.all()
+
+    logger.info(f"Found {len(unused_data)} unused documents in SQL")
+
+    if dry_run:
+        return {
+            "status": "dry_run",
+            "unused_count": len(unused_data),
+            "deleted_count": {"data_items": 0, "documents": 0},
+            "cleanup_date": datetime.now(timezone.utc).isoformat(),
+            "preview": {"documents": len(unused_data)},
+        }
+
+    # Delete each document using cognee.delete()
+    deleted_count = 0
+    from cognee.modules.users.methods import get_default_user
+
+    user = await get_default_user() if user_id is None else None
+
+    for data, dataset_data in unused_data:
+        try:
+            await cognee.delete(
+                data_id=data.id,
+                dataset_id=dataset_data.dataset_id,
+                mode="hard",  # Use hard mode to also remove orphaned entities
+                user=user,
+            )
+            deleted_count += 1
+            logger.info(f"Deleted document {data.id} from dataset {dataset_data.dataset_id}")
+        except Exception as e:
+            logger.error(f"Failed to delete document {data.id}: {e}")
+
+    logger.info("Cleanup completed", deleted_count=deleted_count)
+
+    return {
+        "status": "completed",
+        "unused_count": len(unused_data),
+        "deleted_count": {"data_items": deleted_count, "documents": deleted_count},
+        "cleanup_date": datetime.now(timezone.utc).isoformat(),
+    }

cognee/tasks/graph/__init__.py

@@ -1,2 +1,9 @@
+"""
+Graph extraction and manipulation tasks.
+
+This module provides tasks for extracting knowledge graphs from data,
+building relationships between entities, and managing graph structures.
+"""
+
 from .extract_graph_from_data import extract_graph_from_data
 from .extract_graph_from_code import extract_graph_from_code

cognee/tasks/ingestion/data_item.py

@@ -0,0 +1,8 @@
+from dataclasses import dataclass
+from typing import Any, Optional
+
+
+@dataclass
+class DataItem:
+    data: Any
+    label: Optional[str] = None

cognee/tasks/ingestion/ingest_data.py

@@ -20,6 +20,7 @@ from cognee.modules.data.methods import (
 
 from .save_data_item_to_storage import save_data_item_to_storage
 from .data_item_to_text_file import data_item_to_text_file
+from .data_item import DataItem
 
 
 async def ingest_data(
@@ -78,8 +79,16 @@ async def ingest_data(
         dataset_data_map = {str(data.id): True for data in dataset_data}
 
         for data_item in data:
+            # Support for DataItem (custom label + data wrapper)
+            current_label = None
+            underlying_data = data_item
+
+            if isinstance(data_item, DataItem):
+                underlying_data = data_item.data
+                current_label = data_item.label
+
             # Get file path of data item or create a file if it doesn't exist
-            original_file_path = await save_data_item_to_storage(data_item)
+            original_file_path = await save_data_item_to_storage(underlying_data)
             # Transform file path to be OS usable
             actual_file_path = get_data_file_path(original_file_path)
 
@@ -139,6 +148,7 @@ async def ingest_data(
                 data_point.external_metadata = ext_metadata
                 data_point.node_set = json.dumps(node_set) if node_set else None
                 data_point.tenant_id = user.tenant_id if user.tenant_id else None
+                data_point.label = current_label
 
                 # Check if data is already in dataset
                 if str(data_point.id) in dataset_data_map:
@@ -169,6 +179,7 @@ async def ingest_data(
                     tenant_id=user.tenant_id if user.tenant_id else None,
                     pipeline_status={},
                     token_count=-1,
+                    label=current_label,
                 )
 
                 new_datapoints.append(data_point)

cognee/tasks/ingestion/save_data_item_to_storage.py

@@ -9,6 +9,7 @@ from cognee.shared.logging_utils import get_logger
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 from cognee.tasks.web_scraper.utils import fetch_page_content
+from cognee.tasks.ingestion.data_item import DataItem
 
 
 logger = get_logger()
@@ -95,5 +96,9 @@ async def save_data_item_to_storage(data_item: Union[BinaryIO, str, Any]) -> str
         # data is text, save it to data storage and return the file path
         return await save_data_to_file(data_item)
 
+    if isinstance(data_item, DataItem):
+        # If instance is DataItem use the underlying data
+        return await save_data_item_to_storage(data_item.data)
+
     # data is not a supported type
     raise IngestionError(message=f"Data type not supported: {type(data_item)}")

cognee/tasks/memify/__init__.py

@@ -1,3 +1,11 @@
+"""
+Memory and subgraph extraction tasks.
+
+This module provides tasks for extracting subgraphs, document chunks, and
+user session data, as well as initiating session cognification workflows,
+to support memory enrichment and downstream knowledge graph processing.
+"""
+
 from .extract_subgraph import extract_subgraph
 from .extract_subgraph_chunks import extract_subgraph_chunks
 from .cognify_session import cognify_session