cognee 1.2.2.dev2__tar.gz → 1.2.2.dev3__tar.gz

{cognee-1.2.2.dev2 → cognee-1.2.2.dev3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognee
-Version: 1.2.2.dev2
+Version: 1.2.2.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.2.2.dev2 → cognee-1.2.2.dev3}/cognee/memify_pipelines/persist_agent_trace_feedbacks_in_knowledge_graph.py

@@ -81,6 +81,7 @@ async def persist_agent_trace_feedbacks_in_knowledge_graph_pipeline(
         extraction_tasks=extraction_tasks,
         enrichment_tasks=enrichment_tasks,
         dataset=dataset_to_write[0].id,
+        user=user,
         data=[{}],
         run_in_background=run_in_background,
     )

{cognee-1.2.2.dev2 → cognee-1.2.2.dev3}/cognee/memify_pipelines/persist_sessions_in_knowledge_graph.py

@@ -52,13 +52,14 @@ async def persist_sessions_in_knowledge_graph_pipeline(
         extraction_tasks = [Task(extract_user_sessions, session_ids=session_ids)]
 
         enrichment_tasks = [
-            Task(cognify_session, dataset_id=dataset_to_write[0].id),
+            Task(cognify_session, dataset_id=dataset_to_write[0].id, user=user),
         ]
 
         result = await memify(
             extraction_tasks=extraction_tasks,
             enrichment_tasks=enrichment_tasks,
             dataset=dataset_to_write[0].id,
+            user=user,
             data=[{}],
             run_in_background=run_in_background,
         )

{cognee-1.2.2.dev2 → cognee-1.2.2.dev3}/cognee/modules/migration/import_source.py

@@ -77,16 +77,6 @@ async def import_memory_source(
     external_id, node ids from entity names) make re-running an interrupted
     or repeated import safe.
     """
-    from cognee.modules.migrations.startup import run_migrations_and_block
-
-    # Imports are writes, so they take the same migration gate as
-    # remember()/cognify() (the remember() MemorySource dispatch happens
-    # before its own gate). This also records the data-migration revision —
-    # stamping a fresh store at head — BEFORE the imported rows arrive;
-    # without it the populated store has no recorded revision and the first
-    # migration-aware startup replays the entire data chain over it.
-    await run_migrations_and_block(dataset_name, user)
-
     node_set = node_set or [f"import:{source.source_system}"]
 
     if source.mode == "preserve" and getattr(source, "replayable", False):

{cognee-1.2.2.dev2 → cognee-1.2.2.dev3}/cognee/modules/migration/loader.py

@@ -11,7 +11,6 @@ Two targets, selected by the source's import mode:
   the source (``valid_at``/``invalid_at``) is preserved as edge properties.
 """
 
-import dataclasses
 from dataclasses import dataclass, field
 from datetime import datetime
 from types import SimpleNamespace
@@ -428,12 +427,6 @@ def _provenance_ctx(ctx):
     add_data_points stamps ledger provenance from ``ctx.data_item.id``, which
     exists on Data ORM records but not on raw ingestion DataItems. Substitute
     the DataItem's deterministic ``data_id`` so provenance still lands.
-
-    Only ``data_item`` is replaced; every other PipelineContext field is
-    carried over via ``dataclasses.replace`` so consumers that read ctx
-    attributes (e.g. ``add_data_points`` reads ``ctx.pipeline_run_id``) never
-    see a field-stripped context — a hand-copied field list here silently
-    drops any parameter later added to PipelineContext.
     """
     if ctx is None:
         return None
@@ -441,7 +434,12 @@ def _provenance_ctx(ctx):
     if data_item is None or hasattr(data_item, "id"):
         return ctx
     data_id = getattr(data_item, "data_id", None)
-    return dataclasses.replace(ctx, data_item=SimpleNamespace(id=data_id) if data_id else None)
+    return SimpleNamespace(
+        user=getattr(ctx, "user", None),
+        dataset=getattr(ctx, "dataset", None),
+        data_item=SimpleNamespace(id=data_id) if data_id else None,
+        extras=getattr(ctx, "extras", None),
+    )
 
 
 # Edge batches can run larger than node batches: edges are small tuples and

{cognee-1.2.2.dev2 → cognee-1.2.2.dev3}/cognee/modules/migration/snapshot.py

@@ -19,7 +19,7 @@ so no property is lost.
 """
 
 import json
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, Union
+from typing import Any, Dict, List, Optional, Type, Union
 
 from pydantic import BaseModel, ConfigDict, Field, SerializeAsAny, create_model, field_validator
 
@@ -59,41 +59,13 @@ class _DynamicDataPoint(DataPoint):
     model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
 
 
-_dynamic_models: Dict[Tuple[str, Tuple[str, ...], Tuple[str, ...]], Type[DataPoint]] = {}
+_dynamic_models: Dict[str, Type[DataPoint]] = {}
 
 
-def _dynamic_model(
-    type_name: str,
-    property_names: Optional[Iterable[str]] = None,
-    index_fields: Optional[List[str]] = None,
-) -> Type[DataPoint]:
-    """A DataPoint subclass named ``type_name`` that DECLARES the record's
-    properties as fields and bakes ``index_fields`` into the CLASS-level
-    metadata default.
-
-    Instance-level state is not enough on either count: downstream model
-    machinery only sees declared fields. ``get_graph_from_model`` iterates
-    declared fields when collecting node properties (``extra="allow"`` values
-    are silently dropped — imported DocumentChunk text never reached storage),
-    and rebuilds nodes through ``copy_model``, which re-applies class field
-    defaults (a plain DataPoint metadata default of ``index_fields: []``
-    dropped the embeddable-field markers, so nothing was vector-indexed).
-    """
-    extra_fields = tuple(
-        sorted(
-            name
-            for name in (property_names or ())
-            if name != "metadata" and name not in DataPoint.model_fields
-        )
-    )
-    key = (type_name, extra_fields, tuple(index_fields or ()))
-    if key not in _dynamic_models:
-        field_definitions: Dict[str, Any] = {name: (Any, None) for name in extra_fields}
-        field_definitions["metadata"] = (dict, {"index_fields": list(index_fields or [])})
-        _dynamic_models[key] = create_model(
-            type_name, __base__=_DynamicDataPoint, **field_definitions
-        )
-    return _dynamic_models[key]
+def _dynamic_model(type_name: str) -> Type[DataPoint]:
+    if type_name not in _dynamic_models:
+        _dynamic_models[type_name] = create_model(type_name, __base__=_DynamicDataPoint)
+    return _dynamic_models[type_name]
 
 
 def _clean_properties(properties: Dict[str, Any]) -> Dict[str, Any]:
@@ -129,20 +101,14 @@ def rehydrate_node(
         except Exception:  # noqa: BLE001 — fall back to a dynamic model
             pass
 
-    # Carry the record's properties and embeddable fields onto the dynamic
-    # CLASS so they survive downstream field iteration and copy_model
-    # rebuilds (see _dynamic_model).
-    metadata = props.get("metadata")
-    index_fields = metadata.get("index_fields") if isinstance(metadata, dict) else None
-
     try:
         # Created lazily: only when the registered class is absent or fails.
-        return _dynamic_model(type_name, props, index_fields)(**props)
+        return _dynamic_model(type_name)(**props)
     except Exception:  # noqa: BLE001 — final fallback below
         pass
 
     base_safe = {key: value for key, value in props.items() if key in _BASE_FIELDS}
-    return _dynamic_model(type_name, base_safe, index_fields)(**base_safe)
+    return _dynamic_model(type_name)(**base_safe)
 
 
 class GraphEdge(BaseModel):

{cognee-1.2.2.dev2 → cognee-1.2.2.dev3}/cognee/tasks/memify/cognify_session.py

@@ -5,11 +5,16 @@ import cognee
 
 from cognee.exceptions import CogneeValidationError, CogneeSystemError
 from cognee.shared.logging_utils import get_logger
+from cognee.modules.users.models import User
 
 logger = get_logger("cognify_session")
 
 
-async def cognify_session(data: str, dataset_id: Optional[UUID | str] = None) -> None:
+async def cognify_session(
+    data: str,
+    dataset_id: Optional[UUID | str] = None,
+    user: Optional[User] = None,
+) -> None:
     """
     Process and cognify session data into the knowledge graph.
 
@@ -32,9 +37,14 @@ async def cognify_session(data: str, dataset_id: Optional[UUID | str] = None) ->
 
         logger.info("Processing session data for cognification")
 
-        await cognee.add(data, dataset_id=dataset_id, node_set=["user_sessions_from_cache"])
+        await cognee.add(
+            data,
+            dataset_id=dataset_id,
+            node_set=["user_sessions_from_cache"],
+            user=user,
+        )
         logger.debug("Session data added to cognee with node_set: user_sessions")
-        await cognee.cognify(datasets=[dataset_id])
+        await cognee.cognify(datasets=[dataset_id], user=user)
         logger.info("Session data successfully cognified")
 
     except CogneeValidationError:

{cognee-1.2.2.dev2 → cognee-1.2.2.dev3}/cognee/tests/unit/memify_pipelines/test_persist_agent_trace_feedbacks_pipeline.py

@@ -48,6 +48,7 @@ async def test_persist_agent_trace_feedbacks_pipeline_wires_memify_tasks():
 
     memify_kwargs = memify_mock.call_args.kwargs
     assert memify_kwargs["dataset"] == "dataset-1"
+    assert memify_kwargs["user"] == user
     assert memify_kwargs["data"] == [{}]
     assert len(memify_kwargs["extraction_tasks"]) == 1
     assert len(memify_kwargs["enrichment_tasks"]) == 1

{cognee-1.2.2.dev2 → cognee-1.2.2.dev3}/cognee/tests/unit/migration/test_import_source.py

@@ -22,7 +22,6 @@ from cognee.modules.migration.sources.base import MemorySource
 # fails here because package __init__ files rebind submodule names to
 # same-named functions (e.g. ``cognee.api.v1.remember.remember``).
 add_module = importlib.import_module("cognee.api.v1.add")
-migrations_startup_module = importlib.import_module("cognee.modules.migrations.startup")
 pipeline_module = importlib.import_module("cognee.modules.run_custom_pipeline")
 remember_module = importlib.import_module("cognee.api.v1.remember.remember")
 serve_state = importlib.import_module("cognee.api.v1.serve.state")
@@ -76,13 +75,7 @@ def install_sinks(monkeypatch):
     ``add_data_points`` sink. It returns a run-info dict so pipeline_run_id
     propagation can be asserted.
     """
-    sinks = SimpleNamespace(
-        pipeline_calls=[],
-        add_calls=[],
-        remember_calls=[],
-        graph_flushes=[],
-        migration_gate_calls=[],
-    )
+    sinks = SimpleNamespace(pipeline_calls=[], add_calls=[], remember_calls=[], graph_flushes=[])
 
     async def fake_run_custom_pipeline(**kwargs):
         sinks.pipeline_calls.append(kwargs)
@@ -105,24 +98,10 @@ def install_sinks(monkeypatch):
         result.items_processed = len(data)
         return result
 
-    async def fake_run_migrations_and_block(datasets, user):
-        # writes_before_gate lets tests assert the gate fired before any
-        # imported data was stored (stamping must precede the rows).
-        sinks.migration_gate_calls.append(
-            {
-                "datasets": datasets,
-                "user": user,
-                "writes_before_gate": len(sinks.add_calls) + len(sinks.graph_flushes),
-            }
-        )
-
     monkeypatch.setattr(pipeline_module, "run_custom_pipeline", fake_run_custom_pipeline)
     monkeypatch.setattr(storage_module, "add_data_points", fake_add_data_points)
     monkeypatch.setattr(add_module, "add", fake_add)
     monkeypatch.setattr(remember_module, "remember", fake_remember)
-    monkeypatch.setattr(
-        migrations_startup_module, "run_migrations_and_block", fake_run_migrations_and_block
-    )
     return sinks
 
 
@@ -341,27 +320,6 @@ class TestSkippedFacts:
         assert _flushed_edges(sinks) == []
 
 
-class TestMigrationGate:
-    """import_memory_source must take the run_migrations_and_block gate.
-
-    The gate stamps a fresh store's migration revision at head BEFORE the
-    imported rows arrive; skipping it leaves the populated store with no
-    recorded revision, so the first migration-aware startup replays the whole
-    data-migration chain over the freshly imported data.
-    """
-
-    @pytest.mark.parametrize("mode", ["preserve", "re-derive", "hybrid"])
-    def test_gate_fires_once_before_any_write(self, monkeypatch, mode):
-        sinks = install_sinks(monkeypatch)
-        source = FakeSource(_sample_records(), mode=mode)
-
-        asyncio.run(import_memory_source(source, dataset_name="ds"))
-
-        (gate_call,) = sinks.migration_gate_calls
-        assert gate_call["datasets"] == "ds"
-        assert gate_call["writes_before_gate"] == 0
-
-
 class TestRememberDispatch:
     """remember() routing for MemorySource inputs (guards, kwargs forwarding)."""
 

{cognee-1.2.2.dev2 → cognee-1.2.2.dev3}/cognee/tests/unit/migration/test_snapshot.py

@@ -1,6 +1,5 @@
 """Unit tests for the Pydantic-native export (GraphSnapshot). Pure: no DB, no LLM."""
 
-import asyncio
 from uuid import uuid4
 
 from cognee.infrastructure.engine import DataPoint
@@ -69,58 +68,6 @@ class TestRehydrateNode:
         assert isinstance(node, Entity)
         assert node.metadata["index_fields"] == ["name"]
 
-    def test_dynamic_model_bakes_index_fields_into_class_default(self):
-        node = rehydrate_node(
-            {
-                "id": str(uuid4()),
-                "type": "AlienChunk",
-                "text": "Zorg lives on Mars.",
-                "metadata": {"index_fields": ["text"]},
-            }
-        )
-        # The CLASS default must carry index_fields, not just the instance:
-        # get_graph_from_model rebuilds nodes via copy_model from class field
-        # defaults, so instance-only metadata never reaches the vector index.
-        assert type(node).model_fields["metadata"].default["index_fields"] == ["text"]
-
-    def test_index_fields_and_text_survive_graph_extraction(self):
-        """Regression: imported DocumentChunk raw nodes fail real-class
-        validation (flat properties lack the required relation fields) and
-        fall back to a dynamic model. Downstream model machinery only sees
-        DECLARED fields, so the dynamic class must declare the record's
-        properties AND carry index_fields as its class metadata default —
-        otherwise get_graph_from_model drops the chunk's text (extra="allow"
-        values are not iterated) and/or its embeddable-field markers
-        (copy_model re-applies class defaults), the restored archive has no
-        DocumentChunk_text collection, and hybrid retrieval fails with
-        NoDataError."""
-        from cognee.modules.graph.utils import get_graph_from_model
-
-        node = rehydrate_node(
-            {
-                "id": str(uuid4()),
-                "type": "DocumentChunk",
-                "text": "Alice lives in Berlin.",
-                "chunk_index": 0,
-                "metadata": {"index_fields": ["text"]},
-            }
-        )
-        assert node.metadata["index_fields"] == ["text"]
-        assert node.text == "Alice lives in Berlin."
-
-        nodes, _ = asyncio.run(
-            get_graph_from_model(node, added_nodes={}, added_edges={}, visited_properties={})
-        )
-        (extracted,) = nodes
-        assert extracted.metadata["index_fields"] == ["text"], (
-            "index_fields lost during graph extraction — imported chunk text "
-            "will not be vector-indexed"
-        )
-        assert getattr(extracted, "text", None) == "Alice lives in Berlin.", (
-            "chunk text lost during graph extraction — undeclared extra fields "
-            "are dropped, so the imported chunk stores and indexes as empty"
-        )
-
 
 class TestGraphSnapshot:
     def _snapshot(self, link_relations=False):

{cognee-1.2.2.dev2 → cognee-1.2.2.dev3}/cognee/tests/unit/modules/memify_tasks/test_cognify_session.py

@@ -19,9 +19,12 @@ async def test_cognify_session_success():
         await cognify_session(session_data, dataset_id="123")
 
         mock_add.assert_called_once_with(
-            session_data, dataset_id="123", node_set=["user_sessions_from_cache"]
+            session_data,
+            dataset_id="123",
+            node_set=["user_sessions_from_cache"],
+            user=None,
         )
-        mock_cognify.assert_called_once()
+        mock_cognify.assert_called_once_with(datasets=["123"], user=None)
 
 
 @pytest.mark.asyncio
@@ -106,6 +109,32 @@ async def test_cognify_session_with_special_characters():
         await cognify_session(session_data, dataset_id="123")
 
         mock_add.assert_called_once_with(
-            session_data, dataset_id="123", node_set=["user_sessions_from_cache"]
+            session_data,
+            dataset_id="123",
+            node_set=["user_sessions_from_cache"],
+            user=None,
         )
-        mock_cognify.assert_called_once()
+        mock_cognify.assert_called_once_with(datasets=["123"], user=None)
+
+
+@pytest.mark.asyncio
+async def test_cognify_session_passes_user_to_add_and_cognify():
+    """Test user is forwarded to cognee.add/cognee.cognify."""
+    session_data = (
+        "Session ID: test_session\n\nQuestion: What is AI?\n\nAnswer: AI is artificial intelligence"
+    )
+    user = object()
+
+    with (
+        patch("cognee.add", new_callable=AsyncMock) as mock_add,
+        patch("cognee.cognify", new_callable=AsyncMock) as mock_cognify,
+    ):
+        await cognify_session(session_data, dataset_id="123", user=user)
+
+        mock_add.assert_called_once_with(
+            session_data,
+            dataset_id="123",
+            node_set=["user_sessions_from_cache"],
+            user=user,
+        )
+        mock_cognify.assert_called_once_with(datasets=["123"], user=user)

{cognee-1.2.2.dev2 → cognee-1.2.2.dev3}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "cognee"
 
-version = "1.2.2.dev2"
+version = "1.2.2.dev3"
 description = "Cognee - is a library for enriching LLM context with a semantic layer for better understanding and reasoning."
 authors = [
     { name = "Vasilije Markovic" },

{cognee-1.2.2.dev2 → cognee-1.2.2.dev3}/uv.lock

@@ -1208,7 +1208,7 @@ wheels = [
 
 [[package]]
 name = "cognee"
-version = "1.2.2.dev2"
+version = "1.2.2.dev3"
 source = { editable = "." }
 dependencies = [
     { name = "aiofiles" },

cognee-1.2.2.dev2/cognee/tests/test_cogx_roundtrip.py

@@ -1,174 +0,0 @@
-"""E2E: export a cognified dataset as a COGX archive, restore it into a fresh
-store, and search the restored knowledge graph.
-
-This is the Cognee-to-Cognee instance-migration path (the local half of
-push/pull): ``cognee.export(format="cogx")`` on the source machine,
-``cognee.remember(COGXArchiveSource(...))`` on the target. The test covers
-the full system contract:
-
-- the preserve-mode import survives the pipeline context adaptation
-  (regression: ``ctx.pipeline_run_id`` was dropped, crashing every import),
-- the import stamps the fresh store's migration revision at head BEFORE the
-  rows arrive (regression: an unstamped store replayed the entire
-  data-migration chain on the first server start after a restore),
-- the restored graph is populated and answerable via search.
-
-The test runs in multi-user mode (backend access control enabled — pinned
-below): each dataset_database row must be head-stamped when the import
-creates it.
-"""
-
-import asyncio
-import os
-import pathlib
-import tempfile
-
-import cognee
-from cognee.modules.migration.sources.cogx_archive import COGXArchiveSource
-from cognee.modules.search.types import SearchType
-from cognee.shared.logging_utils import get_logger
-
-logger = get_logger()
-
-TEXT = """The Aurora Skyline Hotel is a boutique hotel in Berlin founded by Clara Novak in 2015.
-The hotel is known for its rooftop observatory and partners with the Berlin Astronomy Club.
-Clara Novak previously managed the Lindenhof Resort in Munich before moving to Berlin.
-"""
-
-QUERY = "Who founded the Aurora Skyline Hotel and what is the hotel known for?"
-
-
-async def _wipe_store():
-    """Reset to an empty store, as on a machine that never ran Cognee."""
-    await cognee.prune.prune_data()
-    await cognee.prune.prune_system(metadata=True)
-
-
-async def _assert_revision_at_head():
-    """The restored store's data-migration bookkeeping must be stamped at head.
-
-    Without the stamp the store is indistinguishable from a pre-migration one,
-    and the first migration-aware startup replays the whole chain over the
-    imported data. In multi-user mode the stamp lives on the per-dataset
-    ``dataset_database`` rows, head-stamped when the import creates them.
-    """
-    from cognee.modules.data.methods.get_dataset_databases import get_dataset_databases
-    from cognee.modules.migrations.migration import head_revision
-    from cognee.modules.migrations.registry import MIGRATIONS
-
-    head = head_revision(MIGRATIONS)
-
-    rows = await get_dataset_databases()
-    assert rows, "Import created no dataset_database bookkeeping row for the fresh store."
-    for row in rows:
-        assert row.migration_revision == head, (
-            f"Imported dataset {row.dataset_id} is stamped at {row.migration_revision!r}, "
-            f"expected head {head!r} — the first startup would replay the migration chain."
-        )
-
-
-async def main():
-    # Disable session-turn gating so searches return direct retrieval results
-    # (same rationale as test_library.py).
-    os.environ["AUTO_FEEDBACK"] = "False"
-
-    # This test runs in multi-user mode: the stamp assertion reads the
-    # per-dataset bookkeeping rows that mode maintains. The flag is read
-    # dynamically, so setting it here wins over any ambient .env.
-    os.environ["ENABLE_BACKEND_ACCESS_CONTROL"] = "true"
-
-    data_directory_path = str(
-        pathlib.Path(
-            os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_cogx_roundtrip")
-        ).resolve()
-    )
-    cognee.config.data_root_directory(data_directory_path)
-    cognee_directory_path = str(
-        pathlib.Path(
-            os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_cogx_roundtrip")
-        ).resolve()
-    )
-    cognee.config.system_root_directory(cognee_directory_path)
-
-    dataset_name = "hotel_knowledge"
-
-    with tempfile.TemporaryDirectory() as temporary_directory:
-        # The archive must live OUTSIDE the data/system roots so wiping the
-        # store between phases cannot delete it (it is the "file shipped to
-        # the other machine").
-        archive_dir = pathlib.Path(temporary_directory) / "hotel_cogx"
-
-        # --- Phase 1: source instance — build a graph and export it. -------
-        await _wipe_store()
-
-        await cognee.add([TEXT], dataset_name)
-        await cognee.cognify([dataset_name])
-
-        source_results = await cognee.search(
-            query_type=SearchType.GRAPH_COMPLETION, query_text=QUERY, datasets=[dataset_name]
-        )
-        assert len(source_results) != 0, "Search on the source instance returned nothing."
-        logger.info("Source search results: %s", source_results)
-
-        export_result = await cognee.export(dataset_name, format="cogx", destination=archive_dir)
-        assert export_result.num_nodes > 0, "Export produced an empty archive (0 nodes)."
-        assert export_result.num_edges > 0, "Export produced an archive with 0 edges."
-        logger.info(
-            "Exported %d nodes, %d edges to %s",
-            export_result.num_nodes,
-            export_result.num_edges,
-            export_result.destination,
-        )
-
-        # --- Phase 2: target instance — restore into a fresh store. --------
-        await _wipe_store()
-
-        # run_migrations is once-per-process; on a real restore the import
-        # runs in a brand-new process on the target machine. Reset the guard
-        # so the import's migration gate treats the wiped store as exactly
-        # that: a fresh instance.
-        import cognee.modules.migrations.startup as migrations_startup
-
-        migrations_startup._startup_migrations_done = False
-
-        import_result = await cognee.remember(
-            COGXArchiveSource(archive_dir), dataset_name=dataset_name
-        )
-        assert import_result.status == "completed", (
-            f"Import did not complete: {import_result.status!r}"
-        )
-        assert import_result.items_processed > 0, "Import processed nothing."
-
-        (import_summary,) = [
-            item for item in import_result.items if item.get("kind") == "migration_import"
-        ]
-        assert import_summary["graph_nodes"] > 0, "Import restored no graph nodes."
-        assert import_summary["graph_edges"] > 0, "Import restored no graph edges."
-        logger.info("Import summary: %s", import_summary)
-
-        await _assert_revision_at_head()
-
-        # --- Phase 3: the restored graph answers searches. ------------------
-        restored_results = await cognee.search(
-            query_type=SearchType.GRAPH_COMPLETION, query_text=QUERY, datasets=[dataset_name]
-        )
-        assert len(restored_results) != 0, "Search on the restored instance returned nothing."
-        logger.info("Restored search results: %s", restored_results)
-
-        # Chunk-level retrieval needs the imported DocumentChunk text to be
-        # vector-indexed (regression: rehydrated raw nodes lost their
-        # index_fields, so restored archives had no DocumentChunk_text
-        # collection and chunk/hybrid retrieval failed with NoDataError).
-        chunk_results = await cognee.search(
-            query_type=SearchType.CHUNKS, query_text="rooftop observatory", datasets=[dataset_name]
-        )
-        assert len(chunk_results) != 0, (
-            "Chunk search on the restored instance returned nothing — imported "
-            "DocumentChunk text was not vector-indexed."
-        )
-
-        print("COGX roundtrip e2e passed: export -> fresh store -> import -> search.")
-
-
-if __name__ == "__main__":
-    asyncio.run(main())

cognee-1.2.2.dev2/cognee/tests/unit/migration/test_provenance_ctx.py

@@ -1,75 +0,0 @@
-"""Drift guard between PipelineContext and the COGX import ctx adaptation.
-
-The COGX import path adapts the pipeline context before handing it to
-``add_data_points`` (``_provenance_ctx`` in ``cognee/modules/migration/loader.py``).
-A hand-copied field list there silently drops any parameter later added to
-``PipelineContext`` — exactly how ``ctx.pipeline_run_id`` was lost, crashing
-every preserve-mode import with AttributeError because ``add_data_points``
-reads it unconditionally.
-
-These tests introspect ``dataclasses.fields(PipelineContext)``, so they fail
-for ANY field that does not survive the adaptation — adding a new parameter
-to PipelineContext without carrying it through the cogx import path breaks
-them by construction. Pure: no databases, no LLM calls, no network.
-"""
-
-import dataclasses
-from types import SimpleNamespace
-from uuid import uuid4
-
-from cognee.modules.migration.loader import _provenance_ctx
-from cognee.modules.pipelines.models.PipelineContext import PipelineContext
-
-
-def _full_context() -> PipelineContext:
-    """A PipelineContext with every field set to a distinct sentinel.
-
-    ``data_item`` is DataItem-shaped (has ``data_id``, no ``id``) so
-    ``_provenance_ctx`` takes its substitution path instead of returning the
-    context unchanged. Sentinels are per-field strings, so a swapped or
-    defaulted field cannot pass equality by accident.
-    """
-    ctx = PipelineContext()
-    for context_field in dataclasses.fields(PipelineContext):
-        setattr(ctx, context_field.name, f"sentinel:{context_field.name}")
-    ctx.data_item = SimpleNamespace(data_id=uuid4())
-    return ctx
-
-
-def test_every_pipeline_context_field_survives_adaptation():
-    ctx = _full_context()
-
-    adapted = _provenance_ctx(ctx)
-
-    assert adapted is not ctx, "expected the data_item substitution path, not the pass-through"
-    for context_field in dataclasses.fields(PipelineContext):
-        assert hasattr(adapted, context_field.name), (
-            f"PipelineContext.{context_field.name} does not survive the COGX import "
-            "ctx adaptation (_provenance_ctx in cognee/modules/migration/loader.py). "
-            "If you added a parameter to PipelineContext, the cogx import path must "
-            "carry it: task consumers such as add_data_points read ctx attributes "
-            "unconditionally (this is how pipeline_run_id regressed)."
-        )
-        if context_field.name == "data_item":
-            continue
-        assert getattr(adapted, context_field.name) == getattr(ctx, context_field.name), (
-            f"PipelineContext.{context_field.name} changed value across _provenance_ctx"
-        )
-
-
-def test_data_item_id_substituted_from_data_id():
-    ctx = _full_context()
-
-    adapted = _provenance_ctx(ctx)
-
-    assert adapted.data_item.id == ctx.data_item.data_id
-
-
-def test_passthrough_when_data_item_already_has_id():
-    ctx = PipelineContext(data_item=SimpleNamespace(id=uuid4()))
-
-    assert _provenance_ctx(ctx) is ctx
-
-
-def test_none_ctx_stays_none():
-    assert _provenance_ctx(None) is None