cognee 1.4.0.dev1__tar.gz → 1.4.0.dev3__tar.gz

{cognee-1.4.0.dev1 → cognee-1.4.0.dev3}/Dockerfile

@@ -34,7 +34,7 @@ COPY README.md pyproject.toml uv.lock entrypoint.sh ./
 
 # Install the project's dependencies using the lockfile and settings
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv sync --extra debug --extra api --extra postgres --extra neo4j --extra llama-index --extra ollama --extra mistral --extra groq --extra anthropic --frozen --no-install-project --no-dev --no-editable
+    uv sync --extra debug --extra api --extra postgres --extra neo4j --extra llama-index --extra ollama --extra mistral --extra groq --extra anthropic --extra aws --frozen --no-install-project --no-dev --no-editable
 
 # Then, add the rest of the project source code and install it
 # Installing separately from its dependencies allows optimal layer caching
@@ -46,7 +46,7 @@ COPY ./cognee_db_workers /app/cognee_db_workers
 # imported at module load by alembic/versions/b9274c27a25a_kuzu_11_migration.py.
 COPY ./kuzu /app/kuzu
 RUN --mount=type=cache,target=/root/.cache/uv \
-uv sync --extra debug --extra api --extra postgres --extra neo4j --extra llama-index --extra ollama --extra mistral --extra groq --extra anthropic --frozen --no-dev --no-editable
+uv sync --extra debug --extra api --extra postgres --extra neo4j --extra llama-index --extra ollama --extra mistral --extra groq --extra anthropic --extra aws --frozen --no-dev --no-editable
 
 FROM python:3.12-slim-bookworm
 

{cognee-1.4.0.dev1 → cognee-1.4.0.dev3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognee
-Version: 1.4.0.dev1
+Version: 1.4.0.dev3
 Summary: Cognee - is a library for enriching LLM context with a semantic layer for better understanding and reasoning.
 Project-URL: Homepage, https://www.cognee.ai
 Project-URL: Repository, https://github.com/topoteretes/cognee

{cognee-1.4.0.dev1 → cognee-1.4.0.dev3}/cognee/api/v1/datasets/datasets.py

@@ -3,7 +3,6 @@ from uuid import UUID
 from typing import Optional
 
 from cognee.context_global_variables import set_database_global_context_variables
-from cognee.infrastructure.locks import dataset_lock
 from cognee.modules.users.models import User
 from cognee.modules.users.methods import get_default_user
 from cognee.modules.users.exceptions import PermissionDeniedError
@@ -113,34 +112,31 @@ class datasets:
         if not dataset:
             raise UnauthorizedDataAccessError(f"Dataset {dataset_id} not accessible.")
 
-        # Same per-dataset lock as pipeline runs: wait for any in-flight pipeline
-        # on this dataset and exclude concurrent deletes.
-        async with dataset_lock(dataset.id):
-            async with set_database_global_context_variables(dataset.id, dataset.owner_id):
-                await delete_dataset_nodes_and_edges(dataset_id, user.id)
-
-                dataset_data = await get_dataset_data(dataset.id)
-
-                # Delete dataset record first while DatasetData junction rows still exist,
-                # so pipeline_status cleanup can find related Data records.
-                result = await delete_dataset(dataset)
-
-                # Delete individual data records; use return_exceptions so all are attempted
-                # even if some fail.
-                if dataset_data:
-                    results = await asyncio.gather(
-                        *[delete_data(data, dataset_id) for data in dataset_data],
-                        return_exceptions=True,
+        async with set_database_global_context_variables(dataset.id, dataset.owner_id):
+            await delete_dataset_nodes_and_edges(dataset_id, user.id)
+
+            dataset_data = await get_dataset_data(dataset.id)
+
+            # Delete dataset record first while DatasetData junction rows still exist,
+            # so pipeline_status cleanup can find related Data records.
+            result = await delete_dataset(dataset)
+
+            # Delete individual data records; use return_exceptions so all are attempted
+            # even if some fail.
+            if dataset_data:
+                results = await asyncio.gather(
+                    *[delete_data(data, dataset_id) for data in dataset_data],
+                    return_exceptions=True,
+                )
+                deletion_errors = [r for r in results if isinstance(r, Exception)]
+                if deletion_errors:
+                    logger.error(
+                        "Failed to delete %d/%d data items from dataset %s: %s",
+                        len(deletion_errors),
+                        len(dataset_data),
+                        dataset_id,
+                        deletion_errors,
                     )
-                    deletion_errors = [r for r in results if isinstance(r, Exception)]
-                    if deletion_errors:
-                        logger.error(
-                            "Failed to delete %d/%d data items from dataset %s: %s",
-                            len(deletion_errors),
-                            len(dataset_data),
-                            dataset_id,
-                            deletion_errors,
-                        )
 
         return result
 
@@ -165,43 +161,38 @@ class datasets:
         if not dataset:
             raise UnauthorizedDataAccessError(f"Dataset {dataset_id} not accessible.")
 
-        # Same per-dataset lock as pipeline runs: wait for any in-flight pipeline
-        # on this dataset and exclude concurrent deletes.
-        async with dataset_lock(dataset.id):
-            dataset_data = [
-                data for data in await get_dataset_data(dataset.id) if data.id == data_id
-            ]
+        dataset_data = [data for data in await get_dataset_data(dataset.id) if data.id == data_id]
 
-            data = dataset_data[0] if len(dataset_data) > 0 else None
+        data = dataset_data[0] if len(dataset_data) > 0 else None
 
-            if not data:
-                # If data is not found in the system, user is using a custom graph model.
-                async with set_database_global_context_variables(dataset_id, dataset.owner_id):
-                    await delete_data_nodes_and_edges(dataset_id, data_id, user.id)
+        if not data:
+            # If data is not found in the system, user is using a custom graph model.
+            async with set_database_global_context_variables(dataset_id, dataset.owner_id):
+                await delete_data_nodes_and_edges(dataset_id, data_id, user.id)
 
-                    dataset_data = await get_dataset_data(dataset.id)
-                    if not dataset_data and delete_dataset_if_empty:
-                        await delete_dataset(dataset)
+                dataset_data = await get_dataset_data(dataset.id)
+                if not dataset_data and delete_dataset_if_empty:
+                    await delete_dataset(dataset)
 
-                return {"status": "success"}
+            return {"status": "success"}
 
-            if not any(ds.id == dataset_id for ds in data.datasets):
-                raise UnauthorizedDataAccessError(f"Data {data_id} not accessible.")
+        if not any(ds.id == dataset_id for ds in data.datasets):
+            raise UnauthorizedDataAccessError(f"Data {data_id} not accessible.")
 
-            async with set_database_global_context_variables(dataset_id, dataset.owner_id):
-                # Delete mode is exclusive: ledger rows imply the relational-ledger
-                # path; only ledger-free data probes the graph marker to distinguish
-                # graph-provenance data from legacy data without graph ownership.
-                if await has_data_related_nodes(dataset_id, data_id):
-                    await delete_data_nodes_and_edges(dataset_id, data_id, user.id)
-                elif not await try_delete_data_by_graph_provenance(dataset_id, data_id):
-                    await legacy_delete(data, "soft")
+        async with set_database_global_context_variables(dataset_id, dataset.owner_id):
+            # Delete mode is exclusive: ledger rows imply the relational-ledger
+            # path; only ledger-free data probes the graph marker to distinguish
+            # graph-provenance data from legacy data without graph ownership.
+            if await has_data_related_nodes(dataset_id, data_id):
+                await delete_data_nodes_and_edges(dataset_id, data_id, user.id)
+            elif not await try_delete_data_by_graph_provenance(dataset_id, data_id):
+                await legacy_delete(data, "soft")
 
-                await delete_data(data, dataset_id)
+            await delete_data(data, dataset_id)
 
-                dataset_data = await get_dataset_data(dataset.id)
-                if not dataset_data and delete_dataset_if_empty:
-                    await delete_dataset(dataset)
+            dataset_data = await get_dataset_data(dataset.id)
+            if not dataset_data and delete_dataset_if_empty:
+                await delete_dataset(dataset)
 
         return {"status": "success"}
 

{cognee-1.4.0.dev1 → cognee-1.4.0.dev3}/cognee/api/v1/datasets/routers/get_datasets_router.py

@@ -54,16 +54,6 @@ class DataDTO(OutDTO):
     dataset_id: UUID
 
 
-class DatasetGraphSummaryDTO(OutDTO):
-    dataset_id: UUID
-    pipeline_run_id: Optional[UUID] = None
-    num_nodes: int
-    num_edges: int
-    # None while pipeline_run_id is set means the last count attempt degraded
-    # (graph store unavailable) and wasn't cached — retried on the next poll.
-    computed_at: Optional[datetime] = None
-
-
 class GraphNodeDTO(OutDTO):
     id: UUID
     label: str
@@ -487,162 +477,6 @@ def get_datasets_router() -> APIRouter:
         except Exception as error:
             return JSONResponse(status_code=409, content={"error": str(error)})
 
-    @router.get("/graph-summary", response_model=List[DatasetGraphSummaryDTO])
-    async def get_datasets_graph_summary(
-        dataset_ids: Annotated[
-            List[UUID],
-            Query(
-                alias="dataset_ids",
-                description=(
-                    "Dataset UUIDs to summarize (from GET /api/v1/datasets)."
-                    " Omit to summarize every dataset you can read."
-                ),
-                examples=[["b8a7c3de-4f5a-4b6c-8d9e-0f1a2b3c4d5e"]],
-            ),
-        ] = [],
-        user: User = Depends(get_authenticated_user),
-    ):
-        """
-        Get node/edge counts per dataset, cached per cognify run.
-
-        Counts are computed once per dataset's latest cognify run and cached in
-        GraphMetrics, keyed by pipeline_run_id — orders of magnitude cheaper on
-        repeat polls than GET /{dataset_id}/graph, which does a full traversal.
-
-        ## Query Parameters
-        - **dataset_ids** (List[UUID], optional): Dataset UUIDs to summarize.
-          If omitted, summarizes every dataset the user has read permission on.
-
-        ## Response
-        Returns a list of summaries containing:
-        - **datasetId**: The dataset's UUID
-        - **pipelineRunId**: The dataset's latest cognify run, or null if it
-          has never been cognified
-        - **numNodes** / **numEdges**: Graph size for that run
-        - **computedAt**: When the count was cached, or null if the last
-          attempt degraded (graph store unavailable) and wasn't cached
-        """
-        send_telemetry(
-            "Datasets API Endpoint Invoked",
-            user.id,
-            additional_properties={
-                "endpoint": "GET /v1/datasets/graph-summary",
-                "dataset_ids": [str(dataset_id) for dataset_id in dataset_ids],
-                "cognee_version": cognee_version,
-            },
-        )
-
-        from datetime import timezone
-        from sqlalchemy import select, func
-        from sqlalchemy.orm import aliased
-        from cognee.context_global_variables import set_database_global_context_variables
-        from cognee.infrastructure.databases.graph import get_graph_engine
-        from cognee.modules.data.models import GraphMetrics
-        from cognee.modules.pipelines.models import PipelineRun
-
-        authorized_datasets = await get_authorized_existing_datasets(dataset_ids, "read", user)
-
-        if not authorized_datasets:
-            return []
-
-        db_engine = get_relational_engine()
-
-        async with db_engine.get_async_session() as session:
-            ranked_runs = (
-                select(
-                    PipelineRun,
-                    func.row_number()
-                    .over(
-                        partition_by=PipelineRun.dataset_id,
-                        order_by=PipelineRun.created_at.desc(),
-                    )
-                    .label("rn"),
-                )
-                .filter(PipelineRun.dataset_id.in_([dataset.id for dataset in authorized_datasets]))
-                .filter(PipelineRun.pipeline_name == "cognify_pipeline")
-                .subquery()
-            )
-            aliased_run = aliased(PipelineRun, ranked_runs)
-            latest_runs = (
-                (await session.execute(select(aliased_run).filter(ranked_runs.c.rn == 1)))
-                .scalars()
-                .all()
-            )
-            latest_run_by_dataset = {run.dataset_id: run for run in latest_runs}
-
-        summaries = []
-        for dataset in authorized_datasets:
-            latest_run = latest_run_by_dataset.get(dataset.id)
-            if latest_run is None:
-                summaries.append(
-                    DatasetGraphSummaryDTO(dataset_id=dataset.id, num_nodes=0, num_edges=0)
-                )
-                continue
-
-            async with db_engine.get_async_session() as session:
-                cached = (
-                    (
-                        await session.execute(
-                            select(GraphMetrics).where(
-                                GraphMetrics.id == latest_run.pipeline_run_id
-                            )
-                        )
-                    )
-                    .scalars()
-                    .first()
-                )
-
-            if cached is not None:
-                summaries.append(
-                    DatasetGraphSummaryDTO(
-                        dataset_id=dataset.id,
-                        pipeline_run_id=latest_run.pipeline_run_id,
-                        num_nodes=cached.num_nodes or 0,
-                        num_edges=cached.num_edges or 0,
-                        computed_at=cached.created_at,
-                    )
-                )
-                continue
-
-            try:
-                async with set_database_global_context_variables(dataset.id, dataset.owner_id):
-                    graph_engine = await get_graph_engine()
-                    graph_metrics = await graph_engine.get_graph_metrics(include_optional=False)
-
-                async with db_engine.get_async_session() as session:
-                    session.add(
-                        GraphMetrics(
-                            id=latest_run.pipeline_run_id,
-                            num_nodes=graph_metrics.get("num_nodes"),
-                            num_edges=graph_metrics.get("num_edges"),
-                        )
-                    )
-                    await session.commit()
-
-                summaries.append(
-                    DatasetGraphSummaryDTO(
-                        dataset_id=dataset.id,
-                        pipeline_run_id=latest_run.pipeline_run_id,
-                        num_nodes=graph_metrics.get("num_nodes") or 0,
-                        num_edges=graph_metrics.get("num_edges") or 0,
-                        computed_at=datetime.now(timezone.utc),
-                    )
-                )
-            except Exception as error:
-                logger.warning(
-                    "Failed to compute graph metrics for dataset %s: %s", dataset.id, error
-                )
-                summaries.append(
-                    DatasetGraphSummaryDTO(
-                        dataset_id=dataset.id,
-                        pipeline_run_id=latest_run.pipeline_run_id,
-                        num_nodes=0,
-                        num_edges=0,
-                    )
-                )
-
-        return summaries
-
     @router.get("/{dataset_id}/data/{data_id}/raw", response_class=FileResponse)
     async def get_raw_data(
         dataset_id: UUID = PathParam(

{cognee-1.4.0.dev1 → cognee-1.4.0.dev3}/cognee/api/v1/forget/forget.py

@@ -2,7 +2,6 @@ from uuid import UUID
 from typing import Optional, Union, Any
 
 from cognee.context_global_variables import set_database_global_context_variables
-from cognee.infrastructure.locks import dataset_lock
 from cognee.shared.logging_utils import get_logger
 from cognee.modules.observability import (
     new_span,
@@ -261,42 +260,37 @@ async def _forget_dataset_memory(dataset_ref: Union[str, UUID], user: Any) -> di
 
     dataset_id = await _resolve_dataset_id(dataset_ref, user)
 
-    # Same per-dataset lock as pipeline runs: wait for any in-flight pipeline
-    # on this dataset and exclude concurrent deletes.
-    async with dataset_lock(dataset_id):
-        # 1. Delete graph nodes/edges and vector embeddings
-        await delete_dataset_nodes_and_edges(dataset_id, user.id)
-
-        # 2. Reset pipeline_status on all data records in this dataset
-        db_engine = get_relational_engine()
-        async with db_engine.get_async_session() as session:
-            data_ids_query = select(DatasetData.data_id).where(DatasetData.dataset_id == dataset_id)
-            data_records = (
-                (await session.execute(select(Data).where(Data.id.in_(data_ids_query))))
-                .scalars()
-                .all()
-            )
+    # 1. Delete graph nodes/edges and vector embeddings
+    await delete_dataset_nodes_and_edges(dataset_id, user.id)
 
-            dataset_id_str = str(dataset_id)
-            for data_record in data_records:
-                if not data_record.pipeline_status:
-                    continue
-                updated = False
-                for pipeline_name in list(data_record.pipeline_status.keys()):
-                    if dataset_id_str in data_record.pipeline_status[pipeline_name]:
-                        del data_record.pipeline_status[pipeline_name][dataset_id_str]
-                        updated = True
-                if updated:
-                    orm_attributes.flag_modified(data_record, "pipeline_status")
+    # 2. Reset pipeline_status on all data records in this dataset
+    db_engine = get_relational_engine()
+    async with db_engine.get_async_session() as session:
+        data_ids_query = select(DatasetData.data_id).where(DatasetData.dataset_id == dataset_id)
+        data_records = (
+            (await session.execute(select(Data).where(Data.id.in_(data_ids_query)))).scalars().all()
+        )
 
-            await session.commit()
+        dataset_id_str = str(dataset_id)
+        for data_record in data_records:
+            if not data_record.pipeline_status:
+                continue
+            updated = False
+            for pipeline_name in list(data_record.pipeline_status.keys()):
+                if dataset_id_str in data_record.pipeline_status[pipeline_name]:
+                    del data_record.pipeline_status[pipeline_name][dataset_id_str]
+                    updated = True
+            if updated:
+                orm_attributes.flag_modified(data_record, "pipeline_status")
 
-        # 3. Reset dataset-level pipeline run status so cached cognify runs can execute again.
-        await reset_dataset_pipeline_run_status(
-            dataset_id=dataset_id,
-            user=user,
-            pipeline_names=["cognify_pipeline"],
-        )
+        await session.commit()
+
+    # 3. Reset dataset-level pipeline run status so cached cognify runs can execute again.
+    await reset_dataset_pipeline_run_status(
+        dataset_id=dataset_id,
+        user=user,
+        pipeline_names=["cognify_pipeline"],
+    )
 
     logger.info(
         "forget: cleared memory for dataset=%s, user=%s (%d data records reset)",
@@ -335,33 +329,30 @@ async def _forget_data_memory(data_id: UUID, dataset_ref: Union[str, UUID], user
 
     dataset_id = await _resolve_dataset_id(dataset_ref, user)
 
-    # Same per-dataset lock as pipeline runs: wait for any in-flight pipeline
-    # on this dataset and exclude concurrent deletes.
-    async with dataset_lock(dataset_id):
-        # 1. Delete graph nodes/edges and vector embeddings for this data item
-        await delete_data_nodes_and_edges(dataset_id, data_id, user.id)
-
-        # 2. Reset pipeline_status for this data record
-        db_engine = get_relational_engine()
-        async with db_engine.get_async_session() as session:
-            data_record = (
-                (await session.execute(select(Data).where(Data.id == data_id))).scalars().first()
-            )
+    # 1. Delete graph nodes/edges and vector embeddings for this data item
+    await delete_data_nodes_and_edges(dataset_id, data_id, user.id)
 
-            if data_record and data_record.pipeline_status:
-                dataset_id_str = str(dataset_id)
-                updated = False
-                # Memory-only forget removes cognify artifacts (graph/vector), so only
-                # cognify_pipeline status should be reset. Keep add status intact.
-                if (
-                    "cognify_pipeline" in data_record.pipeline_status
-                    and dataset_id_str in data_record.pipeline_status["cognify_pipeline"]
-                ):
-                    del data_record.pipeline_status["cognify_pipeline"][dataset_id_str]
-                    updated = True
-                if updated:
-                    orm_attributes.flag_modified(data_record, "pipeline_status")
-                await session.commit()
+    # 2. Reset pipeline_status for this data record
+    db_engine = get_relational_engine()
+    async with db_engine.get_async_session() as session:
+        data_record = (
+            (await session.execute(select(Data).where(Data.id == data_id))).scalars().first()
+        )
+
+        if data_record and data_record.pipeline_status:
+            dataset_id_str = str(dataset_id)
+            updated = False
+            # Memory-only forget removes cognify artifacts (graph/vector), so only
+            # cognify_pipeline status should be reset. Keep add status intact.
+            if (
+                "cognify_pipeline" in data_record.pipeline_status
+                and dataset_id_str in data_record.pipeline_status["cognify_pipeline"]
+            ):
+                del data_record.pipeline_status["cognify_pipeline"][dataset_id_str]
+                updated = True
+            if updated:
+                orm_attributes.flag_modified(data_record, "pipeline_status")
+            await session.commit()
 
     logger.info(
         "forget: cleared memory for data_id=%s in dataset=%s, user=%s",

{cognee-1.4.0.dev1 → cognee-1.4.0.dev3}/cognee/api/v1/improve/improve.py

@@ -334,13 +334,16 @@ async def _bridge_sessions(
         persist_sessions_in_knowledge_graph_pipeline,
     )
 
-    await persist_sessions_in_knowledge_graph_pipeline(
-        user=user,
-        session_ids=session_ids,
-        dataset=dataset_name,
-        run_in_background=run_in_background,
-    )
-    logger.info("improve: session Q&A persisted from %d session(s)", len(session_ids))
+    try:
+        await persist_sessions_in_knowledge_graph_pipeline(
+            user=user,
+            session_ids=session_ids,
+            dataset=dataset_name,
+            run_in_background=run_in_background,
+        )
+        logger.info("improve: session Q&A persisted from %d session(s)", len(session_ids))
+    except Exception as e:
+        logger.warning("improve: session persistence failed (non-fatal): %s", e)
 
 
 async def _extract_agent_context(