update naming

Author: yaozheng-fang

veadk/database/database_adapter.py

@@ -23,7 +23,7 @@
 from veadk.database.relational.mysql_database import MysqlDatabase
 from veadk.database.vector.opensearch_vector_database import OpenSearchVectorDatabase
 from veadk.database.viking.viking_database import VikingDatabase
-from veadk.database.viking.viking_memory_db import VikingDatabaseMemory
+from veadk.database.viking.viking_memory_db import VikingMemoryDatabase
 from veadk.utils.logger import get_logger
 
 logger = get_logger(__name__)
@@ -39,7 +39,7 @@ def add(self, data: list[str], index: str):
 
         try:
             for _data in data:
-                self.client.add(key=index, data=_data)
+                self.client.add(key=index, value=_data)
             logger.debug(f"Added {len(data)} texts to Redis database: index={index}")
         except Exception as e:
             logger.error(
@@ -84,7 +84,7 @@ def add(self, data: list[str], index: str):
         )
 
         if not self.client.table_exists(index):
-            logger.warning(f"Table {index} does not exist, creating...")
+            logger.warning(f"Table {index} does not exist, creating a new table.")
             self.create_table(index)
 
         for _data in data:
@@ -133,6 +133,7 @@ def add(self, data: list[str], index: str):
         self.client.add(data, collection_name=index)
 
     def query(self, query: str, index: str, top_k: int) -> list[str]:
+        # FIXME: confirm
         self._validate_index(index)
 
         logger.debug(
@@ -159,6 +160,7 @@ def get_or_create_collection(self, collection_name: str):
         if not self.client.collection_exists(collection_name):
             self.client.create_collection(collection_name)
 
+        # FIXME
         count = 0
         while not self.client.collection_exists(collection_name):
             time.sleep(1)
@@ -174,6 +176,7 @@ def add(
         self._validate_index(index)
 
         logger.debug(f"Adding documents to Viking database: collection_name={index}")
+
         self.get_or_create_collection(index)
         self.client.add(data, collection_name=index, **kwargs)
 
@@ -189,22 +192,23 @@ def query(self, query: str, index: str, top_k: int) -> list[str]:
         return self.client.query(query, collection_name=index, top_k=top_k)
 
 
-class VikingDatabaseMemoryAdapter(BaseModel):
+class VikingMemoryDatabaseAdapter(BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    client: VikingDatabaseMemory
+    client: VikingMemoryDatabase
 
     def _validate_index(self, index: str):
         # TODO
         pass
 
-    def add(self, data: list[str], index: str):
+    def add(self, data: list[str], index: str, **kwargs):
         self._validate_index(index)
 
         logger.debug(
             f"Adding documents to Viking database memory: collection_name={index} data_len={len(data)}"
         )
 
+        # TODO: parse user_id
         self.client.add(data, collection_name=index)
 
     def query(self, query: str, index: str, top_k: int):
@@ -236,7 +240,7 @@ def query(self, query: str, **kwargs):
     LocalDataBase: LocalDatabaseAdapter,
     VikingDatabase: VikingDatabaseAdapter,
     OpenSearchVectorDatabase: VectorDatabaseAdapter,
-    VikingDatabaseMemory: VikingDatabaseMemoryAdapter,
+    VikingMemoryDatabase: VikingMemoryDatabaseAdapter,
 }
 
 

veadk/database/viking/viking_memory_db.py

@@ -34,6 +34,7 @@
 logger = get_logger(__name__)
 
 
+# FIXME
 class VikingMemConfig(BaseModel):
     volcengine_ak: Optional[str] = Field(
         default=getenv("VOLCENGINE_ACCESS_KEY"),
@@ -53,6 +54,7 @@ class VikingMemConfig(BaseModel):
     )
 
 
+# ======= adapted from ... =======
 class VikingDBMemoryException(Exception):
     def __init__(self, code, request_id, message=None):
         self.code = code
@@ -363,7 +365,10 @@ def format_milliseconds(timestamp_ms):
     return dt.strftime("%Y%m%d %H:%M:%S")
 
 
-class VikingDatabaseMemory(BaseModel, BaseDatabase):
+# ======= adapted from ... =======
+
+
+class VikingMemoryDatabase(BaseModel, BaseDatabase):
     config: VikingMemConfig = Field(
         default_factory=VikingMemConfig,
         description="VikingDB configuration",

veadk/knowledgebase/knowledgebase.py

@@ -22,7 +22,7 @@
 
 
 def build_knowledgebase_index(app_name: str):
-    return f"{app_name}"
+    return f"veadk_kb_{app_name}"
 
 
 class KnowledgeBase:
@@ -62,7 +62,9 @@ def add(
         ...
 
         index = build_knowledgebase_index(app_name)
+
         logger.info(f"Adding documents to knowledgebase: index={index}")
+
         self.adapter.add(data=data, index=index)
 
     def search(self, query: str, app_name: str, top_k: int = None) -> list[str]:
@@ -75,12 +77,3 @@ def search(self, query: str, app_name: str, top_k: int = None) -> list[str]:
         if len(result) == 0:
             logger.warning(f"No documents found in knowledgebase. Query: {query}")
         return result
-
-    def delete(self, app_name: str, user_id: str, session_id: str):
-        """Delete documents in the vector database.
-        Args:
-            app_name (str): The name of the application
-            user_id (str): The user ID
-            session_id (str): The session ID
-        """
-        self.adapter.delete(app_name=app_name, user_id=user_id, session_id=session_id)

veadk/memory/long_term_memory.py

@@ -74,12 +74,13 @@ def _filter_and_convert_events(self, events: list[Event]) -> list[str]:
             if not event.author == "user":
                 continue
 
-            # filter: discard function call and function_response
+            # filter: discard function call and function response
             if not event.content.parts[0].text:
                 continue
 
             # convert: to string-format for storage
             message = event.content.model_dump(exclude_none=True, mode="json")
+
             final_events.append(json.dumps(message))
         return final_events
 
@@ -95,6 +96,7 @@ async def add_session_to_memory(
             f"Adding {len(event_strings)} events to long term memory: index={index}"
         )
 
+        # check if viking memory database, should give a user id： if/else
         self.adapter.add(data=event_strings, index=index)
 
         logger.info(
@@ -109,6 +111,7 @@ async def search_memory(self, *, app_name: str, user_id: str, query: str):
             f"Searching long term memory: query={query} index={index} top_k={self.top_k}"
         )
 
+        # user id if viking memory db
         memory_chunks = self.adapter.query(query=query, index=index, top_k=self.top_k)
 
         # if len(memory_chunks) == 0: