Feature/cog 1369 quick fix dynamic collection handling in search

cognee/infrastructure/databases/exceptions/exceptions.py

@@ -2,7 +2,7 @@
 from fastapi import status
 
 
-class EntityNotFoundError(CogneeApiError):
+class EntityNotFoundError(Exception):
     """Database returns nothing"""
 
     def __init__(

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

@@ -157,7 +157,7 @@ async def get_table(self, collection_name: str) -> Table:
             if collection_name in metadata.tables:
                 return metadata.tables[collection_name]
             else:
-                raise EntityNotFoundError(message=f"Table '{collection_name}' not found.")
+                raise EntityNotFoundError()
 
     async def retrieve(self, collection_name: str, data_point_ids: List[str]):
         # Get PGVectorDataPoint Table from database

cognee/modules/graph/cognee_graph/CogneeGraph.py

@@ -126,8 +126,6 @@ async def map_vector_distances_to_graph_nodes(self, node_distances) -> None:
                 node = self.get_node(node_id)
                 if node:
                     node.add_attribute("vector_distance", score)
-                else:
-                    print(f"Node with id {node_id} not found in the graph.")
 
     async def map_vector_distances_to_graph_edges(self, vector_engine, query) -> None:
         try:

cognee/tasks/completion/graph_query_completion.py

@@ -1,4 +1,5 @@
 from cognee.infrastructure.engine import ExtendableDataPoint
+from cognee.infrastructure.engine.models.DataPoint import DataPoint
 from cognee.modules.graph.utils.convert_node_to_data_point import get_all_subclasses
 from cognee.tasks.completion.exceptions import NoRelevantDataFound
 from cognee.infrastructure.llm.get_llm_client import get_llm_client
@@ -38,7 +39,7 @@ async def graph_query_completion(query: str, context_resolver: Callable = None)
     - Ensure that the LLM client and graph database are properly configured and accessible.
     """
 
-    subclasses = get_all_subclasses(ExtendableDataPoint)
+    subclasses = get_all_subclasses(DataPoint)
 
     vector_index_collections = []