cognee 1.2.2.dev1__tar.gz → 1.2.2.dev2__tar.gz

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

@@ -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 → cognee-1.2.2.dev2}/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-1.2.2.dev1 → cognee-1.2.2.dev2}/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-1.2.2.dev1 → cognee-1.2.2.dev2}/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 → cognee-1.2.2.dev2}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "cognee"
 
-version = "1.2.2.dev1"
+version = "1.2.2.dev2"
 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.dev1 → cognee-1.2.2.dev2}/uv.lock

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