feat: adds pydantic types to graph layer models

cognee/modules/chunking/models/DocumentChunk.py

@@ -12,6 +12,7 @@ class DocumentChunk(DataPoint):
     chunk_index: int
     cut_type: str
     is_part_of: Document
+    pydantic_type: str = "DocumentChunk"
     contains: List[Entity] = None
 
     _metadata: dict = {"index_fields": ["text"], "type": "DocumentChunk"}

cognee/modules/engine/models/Entity.py

@@ -7,5 +7,6 @@ class Entity(DataPoint):
     name: str
     is_a: EntityType
     description: str
+    pydantic_type: str = "Entity"
 
     _metadata: dict = {"index_fields": ["name"], "type": "Entity"}

cognee/modules/engine/models/EntityType.py

@@ -5,5 +5,6 @@ class EntityType(DataPoint):
     __tablename__ = "entity_type"
     name: str
     description: str
+    pydantic_type: str = "EntityType"
 
     _metadata: dict = {"index_fields": ["name"], "type": "EntityType"}

cognee/shared/CodeGraphEntities.py

@@ -5,12 +5,14 @@ from cognee.infrastructure.engine import DataPoint
 class Repository(DataPoint):
     __tablename__ = "Repository"
     path: str
+    pydantic_type: str = "Repository"
     _metadata: dict = {"index_fields": [], "type": "Repository"}
 
 
 class CodeFile(DataPoint):
     __tablename__ = "codefile"
     extracted_id: str  # actually file path
+    pydantic_type: str = "CodeFile"
     source_code: Optional[str] = None
     part_of: Optional[Repository] = None
     depends_on: Optional[List["CodeFile"]] = None
@@ -22,6 +24,7 @@ class CodeFile(DataPoint):
 class CodePart(DataPoint):
     __tablename__ = "codepart"
     # part_of: Optional[CodeFile] = None
+    pydantic_type: str = "CodePart"
     source_code: Optional[str] = None
     _metadata: dict = {"index_fields": [], "type": "CodePart"}
 
@@ -30,6 +33,7 @@ class SourceCodeChunk(DataPoint):
     __tablename__ = "sourcecodechunk"
     code_chunk_of: Optional[CodePart] = None
     source_code: Optional[str] = None
+    pydantic_type: str = "SourceCodeChunk"
     previous_chunk: Optional["SourceCodeChunk"] = None
 
     _metadata: dict = {"index_fields": ["source_code"], "type": "SourceCodeChunk"}

cognee/tasks/summarization/models.py

@@ -17,5 +17,6 @@ class CodeSummary(DataPoint):
     __tablename__ = "code_summary"
     text: str
     summarizes: Union[CodeFile, CodePart, SourceCodeChunk]
+    pydantic_type: str = "CodeSummary"
 
     _metadata: dict = {"index_fields": ["text"], "type": "CodeSummary"}