cognee 0.2.4__py3-none-any.whl → 0.3.0.dev0__py3-none-any.whl

cognee/modules/search/types/SearchResult.py

@@ -0,0 +1,21 @@
+from uuid import UUID
+from pydantic import BaseModel
+from typing import Any, Dict, List, Optional
+
+
+class SearchResultDataset(BaseModel):
+    id: UUID
+    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]
+    dataset_name: Optional[str]

cognee/modules/search/types/SearchType.py

@@ -15,3 +15,5 @@ class SearchType(Enum):
     GRAPH_COMPLETION_CONTEXT_EXTENSION = "GRAPH_COMPLETION_CONTEXT_EXTENSION"
     FEELING_LUCKY = "FEELING_LUCKY"
     FEEDBACK = "FEEDBACK"
+    TEMPORAL = "TEMPORAL"
+    CODING_RULES = "CODING_RULES"

cognee/modules/search/types/__init__.py

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

cognee/modules/search/utils/__init__.py

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

cognee/modules/search/utils/prepare_search_result.py

@@ -0,0 +1,41 @@
+from typing import List, cast
+
+from cognee.modules.graph.utils import resolve_edges_to_text
+from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
+from cognee.modules.search.utils.transform_context_to_graph import transform_context_to_graph
+
+
+async def prepare_search_result(search_result):
+    result, context, datasets = search_result
+
+    graphs = None
+    result_graph = None
+    context_texts = {}
+
+    if isinstance(context, List) and len(context) > 0 and isinstance(context[0], Edge):
+        context_graph = transform_context_to_graph(context)
+
+        graphs = {
+            "*": context_graph,
+        }
+        context_texts = {
+            "*": await resolve_edges_to_text(context),
+        }
+    elif isinstance(context, str):
+        context_texts = {
+            "*": context,
+        }
+    elif isinstance(context, List) and len(context) > 0 and isinstance(context[0], str):
+        context_texts = {
+            "*": "\n".join(cast(List[str], context)),
+        }
+
+    if isinstance(result, List) and len(result) > 0 and isinstance(result[0], Edge):
+        result_graph = transform_context_to_graph(result)
+
+    return {
+        "result": result_graph or result,
+        "graphs": graphs,
+        "context": context_texts,
+        "datasets": datasets,
+    }

cognee/modules/search/utils/transform_context_to_graph.py

@@ -0,0 +1,38 @@
+from typing import List
+
+from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
+
+
+def transform_context_to_graph(context: List[Edge]):
+    nodes = {}
+    edges = {}
+
+    for triplet in context:
+        nodes[triplet.node1.id] = {
+            "id": triplet.node1.id,
+            "label": triplet.node1.attributes["name"]
+            if "name" in triplet.node1.attributes
+            else triplet.node1.id,
+            "type": triplet.node1.attributes["type"],
+            "attributes": triplet.node2.attributes,
+        }
+        nodes[triplet.node2.id] = {
+            "id": triplet.node2.id,
+            "label": triplet.node2.attributes["name"]
+            if "name" in triplet.node2.attributes
+            else triplet.node2.id,
+            "type": triplet.node2.attributes["type"],
+            "attributes": triplet.node2.attributes,
+        }
+        edges[
+            f"{triplet.node1.id}_{triplet.attributes['relationship_name']}_{triplet.node2.id}"
+        ] = {
+            "source": triplet.node1.id,
+            "target": triplet.node2.id,
+            "label": triplet.attributes["relationship_name"],
+        }
+
+    return {
+        "nodes": list(nodes.values()),
+        "edges": list(edges.values()),
+    }

cognee/modules/sync/__init__.py

@@ -0,0 +1 @@
+# Sync module for tracking sync operations

cognee/modules/sync/methods/__init__.py

@@ -0,0 +1,23 @@
+from .create_sync_operation import create_sync_operation
+from .get_sync_operation import (
+    get_sync_operation,
+    get_user_sync_operations,
+    get_running_sync_operations_for_user,
+)
+from .update_sync_operation import (
+    update_sync_operation,
+    mark_sync_started,
+    mark_sync_completed,
+    mark_sync_failed,
+)
+
+__all__ = [
+    "create_sync_operation",
+    "get_sync_operation",
+    "get_user_sync_operations",
+    "get_running_sync_operations_for_user",
+    "update_sync_operation",
+    "mark_sync_started",
+    "mark_sync_completed",
+    "mark_sync_failed",
+]

cognee/modules/sync/methods/create_sync_operation.py

@@ -0,0 +1,53 @@
+from uuid import UUID
+from typing import Optional, List
+from datetime import datetime, timezone
+from cognee.modules.sync.models import SyncOperation, SyncStatus
+from cognee.infrastructure.databases.relational import get_relational_engine
+
+
+async def create_sync_operation(
+    run_id: str,
+    dataset_ids: List[UUID],
+    dataset_names: List[str],
+    user_id: UUID,
+    total_records_to_sync: Optional[int] = None,
+    total_records_to_download: Optional[int] = None,
+    total_records_to_upload: Optional[int] = None,
+) -> SyncOperation:
+    """
+    Create a new sync operation record in the database.
+
+    Args:
+        run_id: Unique public identifier for this sync operation
+        dataset_ids: List of dataset UUIDs being synced
+        dataset_names: List of dataset names being synced
+        user_id: UUID of the user who initiated the sync
+        total_records_to_sync: Total number of records to sync (if known)
+        total_records_to_download: Total number of records to download (if known)
+        total_records_to_upload: Total number of records to upload (if known)
+
+    Returns:
+        SyncOperation: The created sync operation record
+    """
+    db_engine = get_relational_engine()
+
+    sync_operation = SyncOperation(
+        run_id=run_id,
+        dataset_ids=[
+            str(uuid) for uuid in dataset_ids
+        ],  # Convert UUIDs to strings for JSON storage
+        dataset_names=dataset_names,
+        user_id=user_id,
+        status=SyncStatus.STARTED,
+        total_records_to_sync=total_records_to_sync,
+        total_records_to_download=total_records_to_download,
+        total_records_to_upload=total_records_to_upload,
+        created_at=datetime.now(timezone.utc),
+    )
+
+    async with db_engine.get_async_session() as session:
+        session.add(sync_operation)
+        await session.commit()
+        await session.refresh(sync_operation)
+
+    return sync_operation

cognee/modules/sync/methods/get_sync_operation.py

@@ -0,0 +1,107 @@
+from uuid import UUID
+from typing import List, Optional
+from sqlalchemy import select, desc, and_
+from cognee.modules.sync.models import SyncOperation, SyncStatus
+from cognee.infrastructure.databases.relational import get_relational_engine
+
+
+async def get_sync_operation(run_id: str) -> Optional[SyncOperation]:
+    """
+    Get a sync operation by its run_id.
+
+    Args:
+        run_id: The public run_id of the sync operation
+
+    Returns:
+        SyncOperation: The sync operation record, or None if not found
+    """
+    db_engine = get_relational_engine()
+
+    async with db_engine.get_async_session() as session:
+        query = select(SyncOperation).where(SyncOperation.run_id == run_id)
+        result = await session.execute(query)
+        return result.scalars().first()
+
+
+async def get_user_sync_operations(
+    user_id: UUID, limit: int = 50, offset: int = 0
+) -> List[SyncOperation]:
+    """
+    Get sync operations for a specific user, ordered by most recent first.
+
+    Args:
+        user_id: UUID of the user
+        limit: Maximum number of records to return
+        offset: Number of records to skip
+
+    Returns:
+        List[SyncOperation]: List of sync operations for the user
+    """
+    db_engine = get_relational_engine()
+
+    async with db_engine.get_async_session() as session:
+        query = (
+            select(SyncOperation)
+            .where(SyncOperation.user_id == user_id)
+            .order_by(desc(SyncOperation.created_at))
+            .limit(limit)
+            .offset(offset)
+        )
+        result = await session.execute(query)
+        return list(result.scalars().all())
+
+
+async def get_sync_operations_by_dataset(
+    dataset_id: UUID, limit: int = 50, offset: int = 0
+) -> List[SyncOperation]:
+    """
+    Get sync operations for a specific dataset.
+
+    Args:
+        dataset_id: UUID of the dataset
+        limit: Maximum number of records to return
+        offset: Number of records to skip
+
+    Returns:
+        List[SyncOperation]: List of sync operations for the dataset
+    """
+    db_engine = get_relational_engine()
+
+    async with db_engine.get_async_session() as session:
+        query = (
+            select(SyncOperation)
+            .where(SyncOperation.dataset_id == dataset_id)
+            .order_by(desc(SyncOperation.created_at))
+            .limit(limit)
+            .offset(offset)
+        )
+        result = await session.execute(query)
+        return list(result.scalars().all())
+
+
+async def get_running_sync_operations_for_user(user_id: UUID) -> List[SyncOperation]:
+    """
+    Get all currently running sync operations for a specific user.
+    Checks for operations with STARTED or IN_PROGRESS status.
+
+    Args:
+        user_id: UUID of the user
+
+    Returns:
+        List[SyncOperation]: List of running sync operations for the user
+    """
+    db_engine = get_relational_engine()
+
+    async with db_engine.get_async_session() as session:
+        query = (
+            select(SyncOperation)
+            .where(
+                and_(
+                    SyncOperation.user_id == user_id,
+                    SyncOperation.status.in_([SyncStatus.STARTED, SyncStatus.IN_PROGRESS]),
+                )
+            )
+            .order_by(desc(SyncOperation.created_at))
+        )
+        result = await session.execute(query)
+        return list(result.scalars().all())

cognee/modules/sync/methods/update_sync_operation.py

@@ -0,0 +1,248 @@
+import asyncio
+from typing import Optional, List
+from datetime import datetime, timezone
+from sqlalchemy import select
+from sqlalchemy.exc import SQLAlchemyError, DisconnectionError, OperationalError, TimeoutError
+from cognee.modules.sync.models import SyncOperation, SyncStatus
+from cognee.infrastructure.databases.relational import get_relational_engine
+from cognee.shared.logging_utils import get_logger
+from cognee.infrastructure.utils.calculate_backoff import calculate_backoff
+
+logger = get_logger("sync.db_operations")
+
+
+async def _retry_db_operation(operation_func, run_id: str, max_retries: int = 3):
+    """
+    Retry database operations with exponential backoff for transient failures.
+
+    Args:
+        operation_func: Async function to retry
+        run_id: Run ID for logging context
+        max_retries: Maximum number of retry attempts
+
+    Returns:
+        Result of the operation function
+
+    Raises:
+        Exception: Re-raises the last exception if all retries fail
+    """
+    attempt = 0
+    last_exception = None
+
+    while attempt < max_retries:
+        try:
+            return await operation_func()
+        except (DisconnectionError, OperationalError, TimeoutError) as e:
+            attempt += 1
+            last_exception = e
+
+            if attempt >= max_retries:
+                logger.error(
+                    f"Database operation failed after {max_retries} attempts for run_id {run_id}: {str(e)}"
+                )
+                break
+
+            backoff_time = calculate_backoff(attempt - 1)  # calculate_backoff is 0-indexed
+            logger.warning(
+                f"Database operation failed for run_id {run_id}, retrying in {backoff_time:.2f}s (attempt {attempt}/{max_retries}): {str(e)}"
+            )
+            await asyncio.sleep(backoff_time)
+
+        except Exception as e:
+            # Non-transient errors should not be retried
+            logger.error(f"Non-retryable database error for run_id {run_id}: {str(e)}")
+            raise
+
+    # If we get here, all retries failed
+    raise last_exception
+
+
+async def update_sync_operation(
+    run_id: str,
+    status: Optional[SyncStatus] = None,
+    progress_percentage: Optional[int] = None,
+    records_downloaded: Optional[int] = None,
+    total_records_to_sync: Optional[int] = None,
+    total_records_to_download: Optional[int] = None,
+    total_records_to_upload: Optional[int] = None,
+    records_uploaded: Optional[int] = None,
+    bytes_downloaded: Optional[int] = None,
+    bytes_uploaded: Optional[int] = None,
+    dataset_sync_hashes: Optional[dict] = None,
+    error_message: Optional[str] = None,
+    retry_count: Optional[int] = None,
+    started_at: Optional[datetime] = None,
+    completed_at: Optional[datetime] = None,
+) -> Optional[SyncOperation]:
+    """
+    Update a sync operation record with new status/progress information.
+
+    Args:
+        run_id: The public run_id of the sync operation to update
+        status: New status for the operation
+        progress_percentage: Progress percentage (0-100)
+        records_downloaded: Number of records downloaded so far
+        total_records_to_sync: Total number of records that need to be synced
+        total_records_to_download: Total number of records to download from cloud
+        total_records_to_upload: Total number of records to upload to cloud
+        records_uploaded: Number of records uploaded so far
+        bytes_downloaded: Total bytes downloaded from cloud
+        bytes_uploaded: Total bytes uploaded to cloud
+        dataset_sync_hashes: Dict mapping dataset_id -> {uploaded: [hashes], downloaded: [hashes]}
+        error_message: Error message if operation failed
+        retry_count: Number of retry attempts
+        started_at: When the actual processing started
+        completed_at: When the operation completed (success or failure)
+
+    Returns:
+        SyncOperation: The updated sync operation record, or None if not found
+    """
+
+    async def _perform_update():
+        db_engine = get_relational_engine()
+
+        async with db_engine.get_async_session() as session:
+            try:
+                # Find the sync operation
+                query = select(SyncOperation).where(SyncOperation.run_id == run_id)
+                result = await session.execute(query)
+                sync_operation = result.scalars().first()
+
+                if not sync_operation:
+                    logger.warning(f"Sync operation not found for run_id: {run_id}")
+                    return None
+
+                # Log what we're updating for debugging
+                updates = []
+                if status is not None:
+                    updates.append(f"status={status.value}")
+                if progress_percentage is not None:
+                    updates.append(f"progress={progress_percentage}%")
+                if records_downloaded is not None:
+                    updates.append(f"downloaded={records_downloaded}")
+                if records_uploaded is not None:
+                    updates.append(f"uploaded={records_uploaded}")
+                if total_records_to_sync is not None:
+                    updates.append(f"total_sync={total_records_to_sync}")
+                if total_records_to_download is not None:
+                    updates.append(f"total_download={total_records_to_download}")
+                if total_records_to_upload is not None:
+                    updates.append(f"total_upload={total_records_to_upload}")
+
+                if updates:
+                    logger.debug(f"Updating sync operation {run_id}: {', '.join(updates)}")
+
+                # Update fields that were provided
+                if status is not None:
+                    sync_operation.status = status
+
+                if progress_percentage is not None:
+                    sync_operation.progress_percentage = max(0, min(100, progress_percentage))
+
+                if records_downloaded is not None:
+                    sync_operation.records_downloaded = records_downloaded
+
+                if records_uploaded is not None:
+                    sync_operation.records_uploaded = records_uploaded
+
+                if total_records_to_sync is not None:
+                    sync_operation.total_records_to_sync = total_records_to_sync
+
+                if total_records_to_download is not None:
+                    sync_operation.total_records_to_download = total_records_to_download
+
+                if total_records_to_upload is not None:
+                    sync_operation.total_records_to_upload = total_records_to_upload
+
+                if bytes_downloaded is not None:
+                    sync_operation.bytes_downloaded = bytes_downloaded
+
+                if bytes_uploaded is not None:
+                    sync_operation.bytes_uploaded = bytes_uploaded
+
+                if dataset_sync_hashes is not None:
+                    sync_operation.dataset_sync_hashes = dataset_sync_hashes
+
+                if error_message is not None:
+                    sync_operation.error_message = error_message
+
+                if retry_count is not None:
+                    sync_operation.retry_count = retry_count
+
+                if started_at is not None:
+                    sync_operation.started_at = started_at
+
+                if completed_at is not None:
+                    sync_operation.completed_at = completed_at
+
+                # Auto-set completion timestamp for terminal statuses
+                if (
+                    status in [SyncStatus.COMPLETED, SyncStatus.FAILED, SyncStatus.CANCELLED]
+                    and completed_at is None
+                ):
+                    sync_operation.completed_at = datetime.now(timezone.utc)
+
+                # Auto-set started timestamp when moving to IN_PROGRESS
+                if status == SyncStatus.IN_PROGRESS and sync_operation.started_at is None:
+                    sync_operation.started_at = datetime.now(timezone.utc)
+
+                await session.commit()
+                await session.refresh(sync_operation)
+
+                logger.debug(f"Successfully updated sync operation {run_id}")
+                return sync_operation
+
+            except SQLAlchemyError as e:
+                logger.error(
+                    f"Database error updating sync operation {run_id}: {str(e)}", exc_info=True
+                )
+                await session.rollback()
+                raise
+            except Exception as e:
+                logger.error(
+                    f"Unexpected error updating sync operation {run_id}: {str(e)}", exc_info=True
+                )
+                await session.rollback()
+                raise
+
+    # Use retry logic for the database operation
+    return await _retry_db_operation(_perform_update, run_id)
+
+
+async def mark_sync_started(run_id: str) -> Optional[SyncOperation]:
+    """Convenience method to mark a sync operation as started."""
+    return await update_sync_operation(
+        run_id=run_id, status=SyncStatus.IN_PROGRESS, started_at=datetime.now(timezone.utc)
+    )
+
+
+async def mark_sync_completed(
+    run_id: str,
+    records_downloaded: int = 0,
+    records_uploaded: int = 0,
+    bytes_downloaded: int = 0,
+    bytes_uploaded: int = 0,
+    dataset_sync_hashes: Optional[dict] = None,
+) -> Optional[SyncOperation]:
+    """Convenience method to mark a sync operation as completed successfully."""
+    return await update_sync_operation(
+        run_id=run_id,
+        status=SyncStatus.COMPLETED,
+        progress_percentage=100,
+        records_downloaded=records_downloaded,
+        records_uploaded=records_uploaded,
+        bytes_downloaded=bytes_downloaded,
+        bytes_uploaded=bytes_uploaded,
+        dataset_sync_hashes=dataset_sync_hashes,
+        completed_at=datetime.now(timezone.utc),
+    )
+
+
+async def mark_sync_failed(run_id: str, error_message: str) -> Optional[SyncOperation]:
+    """Convenience method to mark a sync operation as failed."""
+    return await update_sync_operation(
+        run_id=run_id,
+        status=SyncStatus.FAILED,
+        error_message=error_message,
+        completed_at=datetime.now(timezone.utc),
+    )