feat: GraphDBBlock

camel/memories/__init__.py

@@ -14,11 +14,13 @@
 
 from .agent_memories import (
     ChatHistoryMemory,
+    GraphDBMemory,
     LongtermAgentMemory,
     VectorDBMemory,
 )
 from .base import AgentMemory, BaseContextCreator, MemoryBlock
 from .blocks.chat_history_block import ChatHistoryBlock
+from .blocks.graphdb_block import GraphDBBlock
 from .blocks.vectordb_block import VectorDBBlock
 from .context_creators.score_based import ScoreBasedContextCreator
 from .records import ContextRecord, MemoryRecord
@@ -27,12 +29,14 @@
     'MemoryRecord',
     'ContextRecord',
     'MemoryBlock',
-    "AgentMemory",
+    'AgentMemory',
     'BaseContextCreator',
     'ScoreBasedContextCreator',
     'ChatHistoryMemory',
     'VectorDBMemory',
     'ChatHistoryBlock',
     'VectorDBBlock',
     'LongtermAgentMemory',
+    'GraphDBBlock',
+    'GraphDBMemory',
 ]

camel/memories/agent_memories.py

@@ -16,11 +16,12 @@
 from typing import List, Optional
 
 from camel.memories.base import AgentMemory, BaseContextCreator
-from camel.memories.blocks import ChatHistoryBlock, VectorDBBlock
+from camel.memories.blocks import ChatHistoryBlock, GraphDBBlock, VectorDBBlock
 from camel.memories.records import ContextRecord, MemoryRecord
 from camel.storages.key_value_storages.base import BaseKeyValueStorage
 from camel.storages.vectordb_storages.base import BaseVectorStorage
 from camel.types import OpenAIBackendRole
+from camel.utils import retry_on_error
 
 
 class ChatHistoryMemory(AgentMemory):
@@ -75,6 +76,7 @@ def retrieve(self) -> List[ContextRecord]:
             )
         return records
 
+    @retry_on_error()
     def write_records(self, records: List[MemoryRecord]) -> None:
         for record in records:
             # assign the agent_id to the record
@@ -153,6 +155,82 @@ def clear(self) -> None:
         self._vectordb_block.clear()
 
 
+class GraphDBMemory(AgentMemory):
+    r"""An agent memory wrapper of :obj:`GraphDBBlock`. This memory queries
+    messages stored in the graph database using similarity search or retrieves
+    recent messages if no query is provided.
+
+    Args:
+        context_creator (BaseContextCreator): A model context creator.
+        graph_db_block (GraphDBBlock): A graph database block instance.
+        retrieve_limit (int, optional): The maximum number of messages
+            to retrieve. Defaults to 25.
+        agent_id (str, optional): The ID of the agent associated with the
+            messages stored in the graph database.
+    """
+
+    def __init__(
+        self,
+        context_creator: BaseContextCreator,
+        graph_db_block: GraphDBBlock,
+        retrieve_limit: int = 25,
+        agent_id: Optional[str] = None,
+    ) -> None:
+        self._context_creator = context_creator
+        self._graph_db_block = graph_db_block
+        self._retrieve_limit = retrieve_limit
+        self._agent_id = agent_id
+        self._current_query: Optional[str] = None
+
+    @property
+    def agent_id(self) -> Optional[str]:
+        return self._agent_id
+
+    @agent_id.setter
+    def agent_id(self, val: Optional[str]) -> None:
+        self._agent_id = val
+
+    def retrieve(self) -> List[ContextRecord]:
+        r"""Retrieves context records from the graph database.
+
+        If a current query is set (based on the last user message), it performs
+        a similarity search. Otherwise, it retrieves recent records.
+
+        Returns:
+            List[ContextRecord]: A list of context records.
+        """
+        return self._graph_db_block.retrieve(
+            query=self._current_query,
+            number_of_nearest_neighbours=self._retrieve_limit,
+        )
+
+    def write_records(self, records: List[MemoryRecord]) -> None:
+        r"""Writes records to the graph database and updates the current query
+        based on the last user message.
+
+        Args:
+            records (List[MemoryRecord]): List of memory records to write.
+        """
+        for record in records:
+            if record.role_at_backend == OpenAIBackendRole.USER:
+                self._current_query = record.message.content
+            if record.agent_id == "" and self.agent_id is not None:
+                record.agent_id = self.agent_id
+        self._graph_db_block.write_records(records)
+
+    def get_context_creator(self) -> BaseContextCreator:
+        r"""Returns the context creator used by the memory.
+
+        Returns:
+            BaseContextCreator: The context creator.
+        """
+        return self._context_creator
+
+    def clear(self) -> None:
+        r"""Clears all records from the graph database."""
+        self._graph_db_block.clear()
+
+
 class LongtermAgentMemory(AgentMemory):
     r"""An implementation of the :obj:`AgentMemory` abstract base class for
     augmenting ChatHistoryMemory with VectorDBMemory.

camel/memories/blocks/__init__.py

@@ -13,9 +13,11 @@
 # ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
 
 from .chat_history_block import ChatHistoryBlock
+from .graphdb_block import GraphDBBlock
 from .vectordb_block import VectorDBBlock
 
 __all__ = [
     'ChatHistoryBlock',
     'VectorDBBlock',
+    'GraphDBBlock',
 ]