Add type to DataPoint metadata

cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py

@@ -1,21 +1,26 @@
 import asyncio
 # from datetime import datetime
 import json
-from uuid import UUID
 from textwrap import dedent
+from uuid import UUID
+
 from falkordb import FalkorDB
 
 from cognee.exceptions import InvalidValueError
-from cognee.infrastructure.engine import DataPoint
-from cognee.infrastructure.databases.graph.graph_db_interface import GraphDBInterface
+from cognee.infrastructure.databases.graph.graph_db_interface import \
+    GraphDBInterface
 from cognee.infrastructure.databases.vector.embeddings import EmbeddingEngine
-from cognee.infrastructure.databases.vector.vector_db_interface import VectorDBInterface
+from cognee.infrastructure.databases.vector.vector_db_interface import \
+    VectorDBInterface
+from cognee.infrastructure.engine import DataPoint
+
 
 class IndexSchema(DataPoint):
     text: str
 
     _metadata: dict = {
-        "index_fields": ["text"]
+        "index_fields": ["text"],
+        "type": "IndexSchema"
     }
 
 class FalkorDBAdapter(VectorDBInterface, GraphDBInterface):

cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py

@@ -1,25 +1,29 @@
-from typing import List, Optional, get_type_hints, Generic, TypeVar
 import asyncio
+from typing import Generic, List, Optional, TypeVar, get_type_hints
 from uuid import UUID
+
 import lancedb
+from lancedb.pydantic import LanceModel, Vector
 from pydantic import BaseModel
-from lancedb.pydantic import Vector, LanceModel
 
 from cognee.exceptions import InvalidValueError
 from cognee.infrastructure.engine import DataPoint
 from cognee.infrastructure.files.storage import LocalStorage
 from cognee.modules.storage.utils import copy_model, get_own_properties
+
+from ..embeddings.EmbeddingEngine import EmbeddingEngine
 from ..models.ScoredResult import ScoredResult
-from ..vector_db_interface import VectorDBInterface
 from ..utils import normalize_distances
-from ..embeddings.EmbeddingEngine import EmbeddingEngine
+from ..vector_db_interface import VectorDBInterface
+
 
 class IndexSchema(DataPoint):
     id: str
     text: str
 
     _metadata: dict = {
-        "index_fields": ["text"]
+        "index_fields": ["text"],
+        "type": "IndexSchema"
     }
 
 class LanceDBAdapter(VectorDBInterface):

cognee/infrastructure/databases/vector/milvus/MilvusAdapter.py

@@ -4,10 +4,12 @@
 import logging
 from typing import List, Optional
 from uuid import UUID
+
 from cognee.infrastructure.engine import DataPoint
-from ..vector_db_interface import VectorDBInterface
-from ..models.ScoredResult import ScoredResult
+
 from ..embeddings.EmbeddingEngine import EmbeddingEngine
+from ..models.ScoredResult import ScoredResult
+from ..vector_db_interface import VectorDBInterface
 
 logger = logging.getLogger("MilvusAdapter")
 
@@ -16,7 +18,8 @@ class IndexSchema(DataPoint):
     text: str
 
     _metadata: dict = {
-        "index_fields": ["text"]
+        "index_fields": ["text"],
+        "type": "IndexSchema"
     }
 
 

cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py

@@ -1,27 +1,30 @@
 import asyncio
-from uuid import UUID
 from typing import List, Optional, get_type_hints
+from uuid import UUID
+
+from sqlalchemy import JSON, Column, Table, delete, select
+from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
 from sqlalchemy.orm import Mapped, mapped_column
-from sqlalchemy import JSON, Column, Table, select, delete
-from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker
 
 from cognee.exceptions import InvalidValueError
 from cognee.infrastructure.databases.exceptions import EntityNotFoundError
 from cognee.infrastructure.engine import DataPoint
 
-from .serialize_data import serialize_data
+from ...relational.ModelBase import Base
+from ...relational.sqlalchemy.SqlAlchemyAdapter import SQLAlchemyAdapter
+from ..embeddings.EmbeddingEngine import EmbeddingEngine
 from ..models.ScoredResult import ScoredResult
-from ..vector_db_interface import VectorDBInterface
 from ..utils import normalize_distances
-from ..embeddings.EmbeddingEngine import EmbeddingEngine
-from ...relational.sqlalchemy.SqlAlchemyAdapter import SQLAlchemyAdapter
-from ...relational.ModelBase import Base
+from ..vector_db_interface import VectorDBInterface
+from .serialize_data import serialize_data
+
 
 class IndexSchema(DataPoint):
     text: str
 
     _metadata: dict = {
-        "index_fields": ["text"]
+        "index_fields": ["text"],
+        "type": "IndexSchema"
     }
 
 class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface):
@@ -87,6 +90,7 @@ def __init__(self, id, payload, vector):
     async def create_data_points(
         self, collection_name: str, data_points: List[DataPoint]
     ):
+        data_point_types = get_type_hints(DataPoint)
         if not await self.has_collection(collection_name):
             await self.create_collection(
                 collection_name = collection_name,
@@ -106,7 +110,7 @@ class PGVectorDataPoint(Base):
             primary_key: Mapped[int] = mapped_column(
                 primary_key=True, autoincrement=True
             )
-            id: Mapped[type(data_points[0].id)]
+            id: Mapped[data_point_types["id"]]
             payload = Column(JSON)
             vector = Column(self.Vector(vector_size))
 

cognee/infrastructure/databases/vector/qdrant/QDrantAdapter.py

@@ -1,21 +1,25 @@
 import logging
+from typing import Dict, List, Optional
 from uuid import UUID
-from typing import List, Dict, Optional
+
 from qdrant_client import AsyncQdrantClient, models
 
 from cognee.exceptions import InvalidValueError
-from cognee.infrastructure.databases.vector.models.ScoredResult import ScoredResult
+from cognee.infrastructure.databases.vector.models.ScoredResult import \
+    ScoredResult
 from cognee.infrastructure.engine import DataPoint
-from ..vector_db_interface import VectorDBInterface
+
 from ..embeddings.EmbeddingEngine import EmbeddingEngine
+from ..vector_db_interface import VectorDBInterface
 
 logger = logging.getLogger("QDrantAdapter")
 
 class IndexSchema(DataPoint):
     text: str
 
     _metadata: dict = {
-        "index_fields": ["text"]
+        "index_fields": ["text"],
+        "type": "IndexSchema"
     }
 
 # class CollectionConfig(BaseModel, extra = "forbid"):

cognee/infrastructure/databases/vector/weaviate_db/WeaviateAdapter.py

@@ -5,17 +5,19 @@
 
 from cognee.exceptions import InvalidValueError
 from cognee.infrastructure.engine import DataPoint
-from ..vector_db_interface import VectorDBInterface
-from ..models.ScoredResult import ScoredResult
+
 from ..embeddings.EmbeddingEngine import EmbeddingEngine
+from ..models.ScoredResult import ScoredResult
+from ..vector_db_interface import VectorDBInterface
 
 logger = logging.getLogger("WeaviateAdapter")
 
 class IndexSchema(DataPoint):
     text: str
 
     _metadata: dict = {
-        "index_fields": ["text"]
+        "index_fields": ["text"],
+        "type": "IndexSchema"
     }
 
 class WeaviateAdapter(VectorDBInterface):

cognee/infrastructure/engine/models/DataPoint.py

@@ -1,8 +1,10 @@
-from typing_extensions import TypedDict
-from uuid import UUID, uuid4
-from typing import Optional
 from datetime import datetime, timezone
+from typing import Optional
+from uuid import UUID, uuid4
+
 from pydantic import BaseModel, Field
+from typing_extensions import TypedDict
+
 
 class MetaData(TypedDict):
     index_fields: list[str]
@@ -13,7 +15,8 @@ class DataPoint(BaseModel):
     updated_at: Optional[datetime] = datetime.now(timezone.utc)
     topological_rank: Optional[int] = 0
     _metadata: Optional[MetaData] = {
-        "index_fields": []
+        "index_fields": [],
+        "type": "DataPoint"
     }
 
     # class Config:
@@ -39,4 +42,4 @@ def get_embeddable_properties(self, data_point):
 
     @classmethod
     def get_embeddable_property_names(self, data_point):
-        return data_point._metadata["index_fields"] or []
+        return data_point._metadata["index_fields"] or []

cognee/modules/chunking/models/DocumentChunk.py

@@ -1,8 +1,10 @@
 from typing import List, Optional
+
 from cognee.infrastructure.engine import DataPoint
 from cognee.modules.data.processing.document_types import Document
 from cognee.modules.engine.models import Entity
 
+
 class DocumentChunk(DataPoint):
     __tablename__ = "document_chunk"
     text: str
@@ -14,4 +16,5 @@ class DocumentChunk(DataPoint):
 
     _metadata: Optional[dict] = {
         "index_fields": ["text"],
+        "type": "DocumentChunk"
     }

cognee/modules/data/processing/document_types/Document.py

@@ -1,12 +1,17 @@
-from cognee.infrastructure.engine import DataPoint
 from uuid import UUID
 
+from cognee.infrastructure.engine import DataPoint
+
+
 class Document(DataPoint):
-    type: str
     name: str
     raw_data_location: str
     metadata_id: UUID
     mime_type: str
+    _metadata: dict = {
+        "index_fields": ["name"],
+        "type": "Document"
+    }
 
     def read(self, chunk_size: int) -> str:
-        pass
+        pass

cognee/modules/engine/models/Entity.py

@@ -10,4 +10,5 @@ class Entity(DataPoint):
 
     _metadata: dict = {
         "index_fields": ["name"],
+        "type": "Entity"
     }

cognee/modules/engine/models/EntityType.py

@@ -1,11 +1,12 @@
 from cognee.infrastructure.engine import DataPoint
 
+
 class EntityType(DataPoint):
     __tablename__ = "entity_type"
     name: str
-    type: str
     description: str
 
     _metadata: dict = {
         "index_fields": ["name"],
+        "type": "EntityType"
     }

cognee/modules/graph/models/EdgeType.py

@@ -1,11 +1,14 @@
 from typing import Optional
+
 from cognee.infrastructure.engine import DataPoint
 
+
 class EdgeType(DataPoint):
     __tablename__ = "edge_type"
     relationship_name: str
     number_of_edges: int
 
     _metadata: Optional[dict] = {
         "index_fields": ["relationship_name"],
+        "type": "EdgeType"
     }

cognee/modules/graph/utils/convert_node_to_data_point.py

@@ -2,7 +2,7 @@
 
 
 def convert_node_to_data_point(node_data: dict) -> DataPoint:
-    subclass = find_subclass_by_name(DataPoint, node_data["type"])
+    subclass = find_subclass_by_name(DataPoint, node_data._metadata["type"])
 
     return subclass(**node_data)
 

cognee/shared/CodeGraphEntities.py

@@ -1,40 +1,47 @@
 from typing import List, Optional
+
 from cognee.infrastructure.engine import DataPoint
 
+
 class Repository(DataPoint):
     __tablename__ = "Repository"
     path: str
-    type: Optional[str] = "Repository"
+    _metadata: dict = {
+        "index_fields": ["source_code"],
+        "type": "Repository"
+    }
 
 class CodeFile(DataPoint):
     __tablename__ = "codefile"
     extracted_id: str  # actually file path
-    type: Optional[str] = "CodeFile"
     source_code: Optional[str] = None
     part_of: Optional[Repository] = None
     depends_on: Optional[List["CodeFile"]] = None
     depends_directly_on: Optional[List["CodeFile"]] = None
     contains: Optional[List["CodePart"]] = None
 
     _metadata: dict = {
-        "index_fields": ["source_code"]
+        "index_fields": ["source_code"],
+        "type": "CodeFile"
     }
 
 class CodePart(DataPoint):
     __tablename__ = "codepart"
     # part_of: Optional[CodeFile]
     source_code: str
-    type: Optional[str] = "CodePart"
-
+
     _metadata: dict = {
-        "index_fields": ["source_code"]
+        "index_fields": ["source_code"],
+        "type": "CodePart"
     }
 
 class CodeRelationship(DataPoint):
     source_id: str
     target_id: str
-    type: str  # between files
     relation: str  # depends on or depends directly
+    _metadata: dict = {
+        "type": "CodeRelationship"
+    }
 
 CodeFile.model_rebuild()
 CodePart.model_rebuild()