cognee 1.2.2.dev2__py3-none-any.whl → 1.2.2.dev3__py3-none-any.whl

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/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/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/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/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/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/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/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/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/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.dist-info → cognee-1.2.2.dev3.dist-info}/METADATA

@@ -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.dist-info → cognee-1.2.2.dev3.dist-info}/RECORD

@@ -602,8 +602,8 @@ cognee/memify_pipelines/consolidate_entity_descriptions.py,sha256=E3lNaVSQBZG9gm
 cognee/memify_pipelines/create_triplet_embeddings.py,sha256=oniU5JUClgpmxgR860nOX0Wu7TjlXNXE9zH2OnSlD_E,1730
 cognee/memify_pipelines/global_context_index.py,sha256=-fi--mYVazD_Kzjbek_IE-q7GRquP77pmnaqbRpMMEk,2700
 cognee/memify_pipelines/memify_default_tasks.py,sha256=SmXqkQCufI8zYdx239-ncjNQqZE-E10AIGbajPMPln8,908
-cognee/memify_pipelines/persist_agent_trace_feedbacks_in_knowledge_graph.py,sha256=gycJh3PvuHZjrscbg7Dp2esN4U34FT-uM27shb_aZLc,3601
-cognee/memify_pipelines/persist_sessions_in_knowledge_graph.py,sha256=jRV1mwTlNZi08vGkYW6IRMNkwstDFxUC4dqjAB7bHFg,2469
+cognee/memify_pipelines/persist_agent_trace_feedbacks_in_knowledge_graph.py,sha256=5i35i9KTpOCMYuxncqXDId6rDcNnmv1BTSN7hipTHNQ,3620
+cognee/memify_pipelines/persist_sessions_in_knowledge_graph.py,sha256=gtpqVR6qb4FKl6gG4C0Q3av3h3e1GjYlEjkdnfX_q2c,2503
 cognee/memory/__init__.py,sha256=LOsCYZ0aMKAb0wdpe1aUBdu4RoTjKy_4_XLNCULW_Iw,298
 cognee/memory/entries.py,sha256=obAHjOybB3Hqr-y37_9ZKa9KMycBLqjieIfsdzWKIfA,5872
 cognee/migration/__init__.py,sha256=KIgvAJwxU9AnthJmwIxw_qvz-KDHteHpWf7ZXyv9i_4,1946
@@ -785,9 +785,9 @@ cognee/modules/migration/archive.py,sha256=ptBMXeJiWQwas2n0Yf-tD2cdqVfU_WaBxw9sa
 cognee/modules/migration/cogx.py,sha256=Z_phNpzX3p9upeKxevtRE_RjLht1Ly9jYw6koYkfPXU,11079
 cognee/modules/migration/export.py,sha256=v-GPc1oT41cmXpuU_vqOCX0Us_9sWQZFoVwMTNyLOCQ,8210
 cognee/modules/migration/formats.py,sha256=-o2hqVwSHXz2O7d_bsxdACZjTd-h8goKoLQ_UqFBcSw,5785
-cognee/modules/migration/import_source.py,sha256=ch7hSReqAKdt2jzbuY-FyL-hXgDJhrzQWyzE1vgNFx4,11475
-cognee/modules/migration/loader.py,sha256=1fWim7Zblt-DggcIJ56EjZj0IFP3Up0AxmKztCHNoSU,23935
-cognee/modules/migration/snapshot.py,sha256=z81JRsbLr7Z4jawueOn48OfRkL9fHGe8lRVHFJfm7R8,10773
+cognee/modules/migration/import_source.py,sha256=Ndy_wJsyP4ZMHtdAKkjD1sIP4GMlkMbz4dlsld1MOHU,10907
+cognee/modules/migration/loader.py,sha256=IXyH0svZm8VoURymfZCQ-3WZdqPGNPzBQ-y6UjNenpk,23699
+cognee/modules/migration/snapshot.py,sha256=wYOSuX_N9R3AF5NaHUyX4YbgXil2y7bhJWS3A0srAU4,9084
 cognee/modules/migration/sources/__init__.py,sha256=HofNJVBO45VaaEWHnC_I7lV5XFtsMLqbm_bpiD3-0Ck,348
 cognee/modules/migration/sources/base.py,sha256=WXhgaRX9M7GCvLmyzvWkREgGCvTliJwk_DGK7LjicVE,2135
 cognee/modules/migration/sources/cogx_archive.py,sha256=x4wz5E-Gim7_vLkB9JFJ9NiZB_JIFKt5ISA1VZNToGs,1569
@@ -1179,7 +1179,7 @@ cognee/tasks/memify/__init__.py,sha256=Bpkm0TdrCnhY76b1orosMVfA3H7a5Dn-mIxIVodNb
 cognee/tasks/memify/apply_feedback_weights.py,sha256=vlAQUb0Xb4egS3Bf-EyWnxwfhsZP-WPAO14nBwNT_Gw,7717
 cognee/tasks/memify/apply_frequency_weights.py,sha256=CZb5ITSIBkgRoQuaYq8yletm6vEsDLKiUGZ9Ng3y8O8,6570
 cognee/tasks/memify/cognify_agent_trace_feedback.py,sha256=fS7upXpo_-N_l5x0nl9mKiDElMww5yvH1iO5fPh75OA,2005
-cognee/tasks/memify/cognify_session.py,sha256=m-t3U-IrWJMPM-csTay854qf3rOzPxfdVcoZprJb8X0,1687
+cognee/tasks/memify/cognify_session.py,sha256=z4zj88E0mEV7LCjbFDXh3AdLnaEazmgaWSXHoHVl0RM,1857
 cognee/tasks/memify/extract_agent_trace_feedbacks.py,sha256=V7_ZscgwlCZh07cthAXdOy1EEG4gsk0xDAOIpEOEG14,5226
 cognee/tasks/memify/extract_feedback_qas.py,sha256=VANI24LkOfFxvES8B9bHXuWiUD8OaAW00JHFBIdI9Vc,3218
 cognee/tasks/memify/extract_subgraph.py,sha256=Al1A5d3bhmd-rd9T9is_ppTfK9JsTb3Es9mBWTiJjB8,220
@@ -1261,7 +1261,6 @@ cognee/tests/test_advanced_pdf_loader.py,sha256=0Jdg4ZZ67_3dg5Z0bCYzubHSLYDQtNgH
 cognee/tests/test_agent_memory_e2e.py,sha256=I1d52atkk0r0dycBG_FTu7_gOyQL-tYt067dW-GTQF0,6681
 cognee/tests/test_chunk_associations.py,sha256=OnsX3a9pLV_Yaj7Xyb297mOciihEz3TFK5uwhbjh_80,8182
 cognee/tests/test_cognee_server_start.py,sha256=nld8SkG4WHUPiHVz6872ImWmcERcUnNN_-6Gf8fUj0E,7572
-cognee/tests/test_cogx_roundtrip.py,sha256=K0kiGxZ1COFsytTS5zfhymMwClWCUTOQK7VTsWOZxgk,7522
 cognee/tests/test_concurrent_cognify_e2e.py,sha256=tdmp0QA6uatFd5YL0bBahk8GWq-GG8LKrhiDfrP4jWQ,5682
 cognee/tests/test_concurrent_subprocess_access.py,sha256=xh72FclRY1LnwUbpmIQTzvK7Yayoov6oewWmHivPTxE,3041
 cognee/tests/test_conversation_history.py,sha256=80be_RWiT4acwZ6agSQY2ch69hmJ8MN-RSeL2qIx0pM,18186
@@ -1540,14 +1539,13 @@ cognee/tests/unit/interfaces/graph/get_graph_from_model_unit_test.py,sha256=IZQ7
 cognee/tests/unit/interfaces/graph/test_expand_with_nodes_and_edges.py,sha256=h3ZRdvtfdJA77aC_Cz7UCXMWXuGXxeLtVwk4CnVD1uc,6429
 cognee/tests/unit/interfaces/graph/test_weighted_edges.py,sha256=Xi-iVyhVJ4YzAgG-eGvx9kZLtwCMhq2Ca43xNkaO_K4,13315
 cognee/tests/unit/memify_pipelines/test_apply_feedback_weights_pipeline.py,sha256=B96e_wv7TBbLANknkgR7Wu26-TwxQrYDgWg6bFWQNeI,3248
-cognee/tests/unit/memify_pipelines/test_persist_agent_trace_feedbacks_pipeline.py,sha256=5qfut4sRZvlqghjK_o-d1uceSqitrMRFDHknsSZIsiw,2589
+cognee/tests/unit/memify_pipelines/test_persist_agent_trace_feedbacks_pipeline.py,sha256=T01IppNOKbRL9apCxb5eQEjdzt3Zy0Py8ppOGqF9iNc,2630
 cognee/tests/unit/migration/test_export_roundtrip.py,sha256=CSB9G_KtHMEbDWa8UR6oU0hi3jy5G2Qv3WGc1cs5Vtw,18167
-cognee/tests/unit/migration/test_import_source.py,sha256=gkcX2ZvMA8M5crvVQrtha7PT4i5Cs1CsJkb22V6mRow,17514
+cognee/tests/unit/migration/test_import_source.py,sha256=ZyuY4AjHugKEa6T4vWu4Vh_RtwcHH7JwnFxHTrtUPxw,15932
 cognee/tests/unit/migration/test_migration.py,sha256=5tqGzk01dsfhGxDDRzyumQ2fotbiKW7XWfjonrMUy3M,19833
-cognee/tests/unit/migration/test_provenance_ctx.py,sha256=GHiAA8NLs2JWaeztyFJuaeeuo9NScgzUNtQGX5wguuA,3057
 cognee/tests/unit/migration/test_push.py,sha256=GYRZXhAKHsyWSNLYa23PDCM9qX5BoXlRY9EWRe-8TwY,16905
 cognee/tests/unit/migration/test_remember_router_cogx.py,sha256=Q72DFoaW9gjdGZI07EXSgU_pQlhnr0-5LlDx1TY_h-A,12156
-cognee/tests/unit/migration/test_snapshot.py,sha256=dfXsSqT6X9DIRFS9Qac5emimLyjOR08j5YvJyc9KPOk,7383
+cognee/tests/unit/migration/test_snapshot.py,sha256=tdzSUAHq187ZaN8gMBlF2WBe9ABv6kBw6VYu6Os4sHU,4898
 cognee/tests/unit/modules/agent_memory/test_agent_memory.py,sha256=AGgsjsS1BTvdJbgX6M4peNgkbMBbRSHGH2BRD2sLuXc,40791
 cognee/tests/unit/modules/agents/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cognee/tests/unit/modules/agents/conftest.py,sha256=5QcS2mitl-wl1ytZVm1_T7SAofrPhgBbs4P3I95nLJE,1027
@@ -1570,7 +1568,7 @@ cognee/tests/unit/modules/graph/test_resolve_edges_to_text.py,sha256=Pz4CTFHEE8R
 cognee/tests/unit/modules/graph/test_retrieve_existing_edges.py,sha256=VIolEF-RdQ_9kqOed8XQ5Irl6NH2MreIXPxA1uG7SeI,3004
 cognee/tests/unit/modules/memify_tasks/test_apply_feedback_weights.py,sha256=tISeQ7T_6sokaVpvWb6aEec3Hgg6bHPsMYSArqlYpow,9874
 cognee/tests/unit/modules/memify_tasks/test_cognify_agent_trace_feedback.py,sha256=pF1ELFjcHgCIxHnV4hVdUEy9lvWTICap6v93zuMV6DI,2753
-cognee/tests/unit/modules/memify_tasks/test_cognify_session.py,sha256=PUduJ3D0x0hqW7eOuFwuz0ifMpR6Fa67nxyQfa2MwjA,3837
+cognee/tests/unit/modules/memify_tasks/test_cognify_session.py,sha256=rp_uRURlfujpY6e2Mo61X2kj-VFAfqVUDrHdAAXNLzI,4783
 cognee/tests/unit/modules/memify_tasks/test_extract_agent_trace_feedbacks.py,sha256=XOvmx3dSCgSfDldFOjGlds9mPPk074QcYTXmj5m_HxU,13729
 cognee/tests/unit/modules/memify_tasks/test_extract_feedback_qas.py,sha256=CZdduhZCXczui7XNk2Y0R2a9gDSX-SvfnKOdMTzkUpo,6320
 cognee/tests/unit/modules/memify_tasks/test_extract_user_sessions.py,sha256=6WLyquWcVWeRVCoUk2i6uwpFk-xl-0s_WMruQYijzas,7542
@@ -1713,9 +1711,9 @@ distributed/workers/data_point_saving_worker.py,sha256=kmaQy2A2J7W3k9Gd5lyoiT0XY
 distributed/workers/graph_saving_worker.py,sha256=KVds58rBh-m3TnOyouE4cqSDnzIHNRqw-Td06abOGS0,4961
 kuzu/__init__.py,sha256=NqYsYyycqMZOUte0_3r_eO5VJGe_BCNdOdLPpJTvgdk,224
 kuzu/database.py,sha256=JfHFK1YAFKrnSId0p_5sqFwhX2IWfKo0MFNEPPNeQ9k,154
-cognee-1.2.2.dev2.dist-info/METADATA,sha256=iNvFhyIiE25-hlxU5l-vrhJn-uosLnwDT8JnFnxGGss,26799
-cognee-1.2.2.dev2.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
-cognee-1.2.2.dev2.dist-info/entry_points.txt,sha256=GCCTsNg8gzOJkolq7dR7OK1VlIAO202dGDnMI8nm8oQ,55
-cognee-1.2.2.dev2.dist-info/licenses/LICENSE,sha256=pHHjSQj1DD8SDppW88MMs04TPk7eAanL1c5xj8NY7NQ,11344
-cognee-1.2.2.dev2.dist-info/licenses/NOTICE.md,sha256=6L3saP3kSpcingOxDh-SGjMS8GY79Rlh2dBNLaO0o5c,339
-cognee-1.2.2.dev2.dist-info/RECORD,,
+cognee-1.2.2.dev3.dist-info/METADATA,sha256=XMMcyeQT0DAv0pnHBXUlg-TKN6SGk4-fK6j31Cf2T9E,26799
+cognee-1.2.2.dev3.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
+cognee-1.2.2.dev3.dist-info/entry_points.txt,sha256=GCCTsNg8gzOJkolq7dR7OK1VlIAO202dGDnMI8nm8oQ,55
+cognee-1.2.2.dev3.dist-info/licenses/LICENSE,sha256=pHHjSQj1DD8SDppW88MMs04TPk7eAanL1c5xj8NY7NQ,11344
+cognee-1.2.2.dev3.dist-info/licenses/NOTICE.md,sha256=6L3saP3kSpcingOxDh-SGjMS8GY79Rlh2dBNLaO0o5c,339
+cognee-1.2.2.dev3.dist-info/RECORD,,

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/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