Merge branch 'main' into feat/COG-341-log-interactions

cognee/infrastructure/databases/graph/graph_db_interface.py

@@ -62,3 +62,8 @@ async def add_edges(
     async def delete_graph(
         self,
     ): raise NotImplementedError
+
+    @abstractmethod
+    async def get_graph_data(
+        self
+    ): raise NotImplementedError

cognee/infrastructure/databases/graph/neo4j_driver/adapter.py

@@ -27,9 +27,6 @@ def __init__(
             max_connection_lifetime = 120
         )
 
-    async def close(self) -> None:
-        await self.driver.close()
-
     @asynccontextmanager
     async def get_session(self) -> AsyncSession:
         async with self.driver.session() as session:

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

@@ -112,10 +112,18 @@ def create_lance_data_point(data_point: DataPoint, vector: list[float]) -> Lance
                 for (data_point_index, data_point) in enumerate(data_points)
         ]
 
-        await collection.merge_insert("id") \
-            .when_matched_update_all() \
-            .when_not_matched_insert_all() \
-            .execute(lance_data_points)
+        # TODO: This enables us to work with pydantic version but shouldn't
+        #  stay like this, existing rows should be updated
+
+        await collection.delete("id IS NOT NULL")
+
+        original_size = await collection.count_rows()
+        await collection.add(lance_data_points)
+        new_size = await collection.count_rows()
+
+        if new_size <= original_size:
+            raise ValueError(
+                "LanceDB create_datapoints error: data points did not get added.")
 
 
     async def retrieve(self, collection_name: str, data_point_ids: list[str]):

cognee/modules/chunking/TextChunker.py

@@ -17,6 +17,7 @@ def __init__(self, document, get_text: callable, chunk_size: int = 1024):
         self.get_text = get_text
 
     def read(self):
+        self.paragraph_chunks = []
         for content_text in self.get_text():
             for chunk_data in chunk_by_paragraph(
                 content_text,

cognee/modules/graph/cognee_graph/CogneeAbstractGraph.py

@@ -0,0 +1,35 @@
+from abc import ABC, abstractmethod
+from typing import List, Dict, Union
+from cognee.modules.graph.cognee_graph.CogneeGraphElements import Node, Edge
+from cognee.infrastructure.databases.graph.graph_db_interface import GraphDBInterface
+
+class CogneeAbstractGraph(ABC):
+    """
+    Abstract base class for representing a graph structure.
+
+    """
+
+    @abstractmethod
+    def add_node(self, node: Node) -> None:
+        """Add a node to the graph."""
+        pass
+
+    @abstractmethod
+    def add_edge(self, edge: Edge) -> None:
+        """Add an edge to the graph."""
+        pass
+
+    @abstractmethod
+    def get_node(self, node_id: str) -> Node:
+        """Retrieve a node by its ID."""
+        pass
+
+    @abstractmethod
+    def get_edges(self, node_id: str) -> List[Edge]:
+        """Retrieve edges connected to a specific node."""
+        pass
+
+    @abstractmethod
+    async def project_graph_from_db(self, adapter: GraphDBInterface, directed: bool, dimension: int) -> None:
+        """Project the graph structure from a database using the provided adapter."""
+        pass