From 1b630366c98ca978c1e1ca281cc18d6ac87bc070 Mon Sep 17 00:00:00 2001
From: hajdul88 <52442977+hajdul88@users.noreply.github.com>
Date: Tue, 11 Feb 2025 19:23:19 +0100
Subject: [PATCH] Adds types property to pydantic Datapoint inherited classes
 (#523)

<!-- .github/pull_request_template.md -->

## Description
This PR adds types to DataPoint pydantic class + fixes visualization
colors

## DCO Affirmation
I affirm that all code in every commit of this pull request conforms to
the terms of the Topoteretes Developer Certificate of Origin


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Added a `type` field to the `DataPoint` model for clearer data
classification.
- Enhanced color mapping in visualizations by assigning a distinct color
to "TextSummary" nodes.

- **Refactor**
- Improved default settings for version control and ordering to ensure
consistent data behavior.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---
 cognee/infrastructure/engine/models/DataPoint.py             | 5 +++++
 cognee/modules/visualization/cognee_network_visualization.py | 1 +
 2 files changed, 6 insertions(+)

diff --git a/cognee/infrastructure/engine/models/DataPoint.py b/cognee/infrastructure/engine/models/DataPoint.py
index 9ab8cf5ca..9e504be9a 100644
--- a/cognee/infrastructure/engine/models/DataPoint.py
+++ b/cognee/infrastructure/engine/models/DataPoint.py
@@ -25,6 +25,11 @@ class DataPoint(BaseModel):
     version: int = 1  # Default version
     topological_rank: Optional[int] = 0
     metadata: Optional[MetaData] = {"index_fields": []}
+    type: str = Field(default_factory=lambda: DataPoint.__name__)
+
+    def __init__(self, **data):
+        super().__init__(**data)
+        object.__setattr__(self, "type", self.__class__.__name__)
 
     @classmethod
     def get_embeddable_data(self, data_point: "DataPoint"):
diff --git a/cognee/modules/visualization/cognee_network_visualization.py b/cognee/modules/visualization/cognee_network_visualization.py
index 24c8e1f66..cf4162044 100644
--- a/cognee/modules/visualization/cognee_network_visualization.py
+++ b/cognee/modules/visualization/cognee_network_visualization.py
@@ -19,6 +19,7 @@ async def cognee_network_visualization(graph_data, destination_file_path: str =
         "Entity": "#f47710",
         "EntityType": "#6510f4",
         "DocumentChunk": "#801212",
+        "TextSummary": "#1077f4",
         "default": "#D3D3D3",
     }