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

cognee/modules/migration/snapshot.py

@@ -19,7 +19,7 @@ so no property is lost.
 """
 
 import json
-from typing import Any, Dict, List, Optional, Type, Union
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, Union
 
 from pydantic import BaseModel, ConfigDict, Field, SerializeAsAny, create_model, field_validator
 
@@ -59,13 +59,41 @@ class _DynamicDataPoint(DataPoint):
     model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
 
 
-_dynamic_models: Dict[str, Type[DataPoint]] = {}
+_dynamic_models: Dict[Tuple[str, Tuple[str, ...], Tuple[str, ...]], Type[DataPoint]] = {}
 
 
-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 _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 _clean_properties(properties: Dict[str, Any]) -> Dict[str, Any]:
@@ -101,14 +129,20 @@ 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)
+        return _dynamic_model(type_name, props, index_fields)(**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)
+    return _dynamic_model(type_name, base_safe, index_fields)(**base_safe)
 
 
 class GraphEdge(BaseModel):

cognee/tests/test_cogx_roundtrip.py

@@ -155,6 +155,18 @@ async def main():
         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.")
 
 

cognee/tests/unit/migration/test_snapshot.py

@@ -1,5 +1,6 @@
 """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
@@ -68,6 +69,58 @@ 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.dev1.dist-info → cognee-1.2.2.dev2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognee
-Version: 1.2.2.dev1
+Version: 1.2.2.dev2
 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.dev1.dist-info → cognee-1.2.2.dev2.dist-info}/RECORD

@@ -787,7 +787,7 @@ cognee/modules/migration/export.py,sha256=v-GPc1oT41cmXpuU_vqOCX0Us_9sWQZFoVwMTN
 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=wYOSuX_N9R3AF5NaHUyX4YbgXil2y7bhJWS3A0srAU4,9084
+cognee/modules/migration/snapshot.py,sha256=z81JRsbLr7Z4jawueOn48OfRkL9fHGe8lRVHFJfm7R8,10773
 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
@@ -1261,7 +1261,7 @@ 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=QgISGfRT97TGs1zH4Z-A0-27QuMw5UjEk2xWDPQB_AY,6883
+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
@@ -1547,7 +1547,7 @@ cognee/tests/unit/migration/test_migration.py,sha256=5tqGzk01dsfhGxDDRzyumQ2fotb
 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=tdzSUAHq187ZaN8gMBlF2WBe9ABv6kBw6VYu6Os4sHU,4898
+cognee/tests/unit/migration/test_snapshot.py,sha256=dfXsSqT6X9DIRFS9Qac5emimLyjOR08j5YvJyc9KPOk,7383
 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
@@ -1713,9 +1713,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.dev1.dist-info/METADATA,sha256=Rl32xhDD3k-Se-Ur0hybsHY5vSZim72PFlVlBZ5yed0,26799
-cognee-1.2.2.dev1.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
-cognee-1.2.2.dev1.dist-info/entry_points.txt,sha256=GCCTsNg8gzOJkolq7dR7OK1VlIAO202dGDnMI8nm8oQ,55
-cognee-1.2.2.dev1.dist-info/licenses/LICENSE,sha256=pHHjSQj1DD8SDppW88MMs04TPk7eAanL1c5xj8NY7NQ,11344
-cognee-1.2.2.dev1.dist-info/licenses/NOTICE.md,sha256=6L3saP3kSpcingOxDh-SGjMS8GY79Rlh2dBNLaO0o5c,339
-cognee-1.2.2.dev1.dist-info/RECORD,,
+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,,