From 3d4312577e30de84d8440a2be5c3037b2eed70c4 Mon Sep 17 00:00:00 2001
From: lxobr <122801072+lxobr@users.noreply.github.com>
Date: Thu, 27 Feb 2025 19:05:09 +0100
Subject: [PATCH] fix: Use DataPoint instead of ExtendableDataPoint in
 get_all_subclasses (#588)

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

## Description
<!-- Provide a clear description of the changes in this PR -->
- Use DataPoint instead of ExtendableDataPoint when calling
get_all_subclasses in the get_triplets function of the
GraphCompletionRetriever
## 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

- **Refactor**
- Updated the internal data handling for retrieving information,
ensuring a more consistent and reliable output for end-users.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---
 cognee/modules/retrieval/graph_completion_retriever.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cognee/modules/retrieval/graph_completion_retriever.py b/cognee/modules/retrieval/graph_completion_retriever.py
index 034c1d40d..709415fa7 100644
--- a/cognee/modules/retrieval/graph_completion_retriever.py
+++ b/cognee/modules/retrieval/graph_completion_retriever.py
@@ -1,6 +1,6 @@
 from typing import Any, Optional
 
-from cognee.infrastructure.engine import ExtendableDataPoint
+from cognee.infrastructure.engine import DataPoint
 from cognee.modules.graph.utils.convert_node_to_data_point import get_all_subclasses
 from cognee.modules.retrieval.base_retriever import BaseRetriever
 from cognee.modules.retrieval.utils.brute_force_triplet_search import brute_force_triplet_search
@@ -35,7 +35,7 @@ class GraphCompletionRetriever(BaseRetriever):
 
     async def get_triplets(self, query: str) -> list:
         """Retrieves relevant graph triplets."""
-        subclasses = get_all_subclasses(ExtendableDataPoint)
+        subclasses = get_all_subclasses(DataPoint)
         vector_index_collections = []
 
         for subclass in subclasses: