refine knowledgebase and memory logs

Author: yaozheng-fang

veadk/knowledgebase/knowledgebase.py

@@ -26,10 +26,13 @@ class KnowledgeBase:
     def __init__(
         self,
         backend: Literal["local", "opensearch", "viking", "redis", "mysql"] = "local",
-        top_k: int = 5,
+        top_k: int = 10,
         db_config=None,
     ):
-        logger.debug(f"Create knowledgebase, backend is {backend}")
+        logger.info(
+            f"Initializing knowledgebase: backend={backend} top_k={top_k} db_config={db_config}"
+        )
+
         self.backend = backend
         self.top_k = top_k
 
@@ -38,6 +41,10 @@ def __init__(
             database_client=self.db_client
         )
 
+        logger.info(
+            f"Initialized knowledgebase: db_client={self.db_client} adapter={self.adapter}"
+        )
+
     def add(
         self,
         data: str | list[str] | TextIO | BinaryIO | bytes,
@@ -51,9 +58,8 @@ def add(
             for example, if you read the file stream of a pdf, then you need to pass an additional parameter file_ext = '.pdf'.
         """
         kwargs.pop("session_id", None)  # remove session_id
-        self.adapter.add(
-            data, app_name, user_id="user_id", session_id="session_id", **kwargs
-        )
+        logger.info(f"Adding documents to knowledgebase: app_name={app_name}")
+        self.adapter.add(data, app_name, **kwargs)
 
     def search(self, query: str, app_name: str, top_k: int = None) -> list[str]:
         """Retrieve documents similar to the query text in the vector database.
@@ -66,9 +72,10 @@ def search(self, query: str, app_name: str, top_k: int = None) -> list[str]:
         """
         top_k = self.top_k if top_k is None else top_k
 
-        result = self.adapter.query(
-            query=query, app_name=app_name, user_id="user_id", top_k=top_k
+        logger.info(
+            f"Searching knowledgebase: app_name={app_name} query={query} top_k={top_k}"
         )
+        result = self.adapter.query(query=query, app_name=app_name, top_k=top_k)
         if len(result) == 0:
             logger.warning(f"No documents found in knowledgebase. Query: {query}")
         return result

veadk/knowledgebase/knowledgebase_database_adapter.py

@@ -35,6 +35,14 @@ def format_collection_name(collection_name: str) -> str:
     return re.sub(r"[^a-z0-9_]", "", replaced_str).lower()
 
 
+def build_index(**kwargs):
+    """
+    Build the index name for the knowledgebase.
+    """
+    # TODO
+    ...
+
+
 def get_knowledgebase_adapter(backend: str):
     if backend == DatabaseBackend.REDIS:
         return KnowledgebaseKVDatabaseAdapter
@@ -55,30 +63,22 @@ class KnowledgebaseKVDatabaseAdapter(BaseModel):
 
     database_client: BaseDatabase
 
-    def add(self, content: list[str], app_name: str, user_id: str, session_id: str):
-        """Add texts to Redis.
-
-        Key: app_name
-        Field: app_name:user_id
-        Value: text in List
-        """
-        # key = f"{app_name}:{user_id}"
+    def add(self, content: list[str], app_name: str, **kwargs):
         key = f"{app_name}"
+        logger.debug(f"Adding documents to Redis database: key={key}")
 
         try:
             for _content in content:
                 self.database_client.add(key, _content)
-            logger.debug(
-                f"Successfully added {len(content)} texts to Redis list key `{key}`."
-            )
+            logger.debug(f"Added {len(content)} texts to Redis database: key={key}")
         except Exception as e:
-            logger.error(f"Failed to add texts to Redis list key `{key}`: {e}")
+            logger.error(f"Failed to add texts to Redis database key `{key}`: {e}")
             raise e
 
-    def query(self, query: str, app_name: str, user_id: str, **kwargs):
-        # key = f"{app_name}:{user_id}"
+    def query(self, query: str, app_name: str, **kwargs):
         key = f"{app_name}"
         top_k = 10
+        logger.debug(f"Querying Redis database: key={key} query={query}")
 
         try:
             result = self.database_client.query(key, query)
@@ -87,9 +87,8 @@ def query(self, query: str, app_name: str, user_id: str, **kwargs):
             logger.error(f"Failed to search from Redis list key '{key}': {e}")
             raise e
 
-    def delete(self, app_name: str, user_id: str, session_id: str):
+    def delete(self, app_name: str, **kwargs):
         try:
-            # key = f"{app_name}:{user_id}:{session_id}"
             key = f"{app_name}"
             self.database_client.delete(key=key)
             logger.info(f"Successfully deleted data for app {app_name}")
@@ -104,6 +103,8 @@ class KnowledgebaseRelationalDatabaseAdapter(BaseModel):
     database_client: BaseDatabase
 
     def create_table(self, table_name: str):
+        logger.debug(f"Creating table for SQL database: table_name={table_name}")
+
         sql = f"""
             CREATE TABLE `{table_name}` (
                 `id` BIGINT AUTO_INCREMENT PRIMARY KEY,
@@ -113,8 +114,11 @@ def create_table(self, table_name: str):
         """
         self.database_client.add(sql)
 
-    def add(self, content: list[str], app_name: str, user_id: str, session_id: str):
+    def add(self, content: list[str], app_name: str, **kwargs):
         table = app_name
+        logger.debug(
+            f"Adding documents to SQL database: table_name={table} content_len={len(content)}"
+        )
 
         if not self.database_client.table_exists(table):
             logger.warning(f"Table {table} does not exist, creating...")
@@ -126,16 +130,19 @@ def add(self, content: list[str], app_name: str, user_id: str, session_id: str):
                 VALUES (%s);
                 """
             self.database_client.add(sql, params=(_content,))
-        logger.info(f"Successfully added {len(content)} texts to table {table}.")
+        logger.debug(f"Added {len(content)} texts to table {table}.")
 
-    def query(self, query: str, app_name: str, user_id: str, **kwargs):
+    def query(self, query: str, app_name: str, **kwargs):
         """Search content from table app_name."""
         table = app_name
         top_k = 10
+        logger.debug(
+            f"Querying SQL database: table_name={table} query={query} top_k={top_k}"
+        )
 
         if not self.database_client.table_exists(table):
             logger.warning(
-                f"querying {query}, but table `{table}` does not exist, returning empty list."
+                f"Querying SQL database, but table `{table}` does not exist, returning empty list."
             )
             return []
 
@@ -145,7 +152,7 @@ def query(self, query: str, app_name: str, user_id: str, **kwargs):
         results = self.database_client.query(sql)
         return [item["data"] for item in results]
 
-    def delete(self, app_name: str, user_id: str, session_id: str):
+    def delete(self, app_name: str, **kwargs):
         table = app_name
         try:
             self.database_client.delete(table=table)
@@ -160,23 +167,23 @@ class KnowledgebaseVectorDatabaseAdapter(BaseModel):
 
     database_client: BaseDatabase
 
-    def add(self, content: list[str], app_name: str, user_id: str, session_id: str):
-        # collection_name = format_collection_name(f"{app_name}_{user_id}")
-        # knowledgebase  is application specific
+    def add(self, content: list[str], app_name: str, **kwargs):
+        logger.debug(
+            f"Adding documents to vector database: app_name={app_name} content_len={len(content)}"
+        )
         collection_name = format_collection_name(f"{app_name}")
         self.database_client.add(content, collection_name=collection_name)
 
-    def query(self, query: str, app_name: str, user_id: str, **kwargs):
-        # collection_name = format_collection_name(f"{app_name}_{user_id}")
-        # knowledgebase  is application specific
+    def query(self, query: str, app_name: str, **kwargs):
+        logger.debug(
+            f"Querying vector database: collection_name={app_name} query={query}"
+        )
         collection_name = format_collection_name(f"{app_name}")
         return self.database_client.query(
             query, collection_name=collection_name, **kwargs
         )
 
-    def delete(self, app_name: str, user_id: str, session_id: str):
-        # collection_name = format_collection_name(f"{app_name}_{user_id}")
-        # knowledgebase  is application specific
+    def delete(self, app_name: str, **kwargs):
         collection_name = format_collection_name(f"{app_name}")
         try:
             self.database_client.delete(collection_name=collection_name)
@@ -193,13 +200,13 @@ class KnowledgebaseLocalDatabaseAdapter(BaseModel):
 
     database_client: BaseDatabase
 
-    def add(self, content: list[str], app_name: str, user_id: str, session_id: str):
+    def add(self, content: list[str], **kwargs):
         self.database_client.add(content)
 
     def query(self, query: str, app_name: str, user_id: str, **kwargs):
         return self.database_client.query(query, **kwargs)
 
-    def delete(self, app_name: str, user_id: str, session_id: str):
+    def delete(self, app_name: str, **kwargs):
         try:
             self.database_client.delete()
             logger.info(f"Successfully cleared local database for app {app_name}")
@@ -229,24 +236,30 @@ def add(
         self,
         content: str | list[str] | TextIO | BinaryIO | bytes,
         app_name: str,
-        user_id: str,
-        session_id: str,
         **kwargs,
     ):
-        # collection_name = format_collection_name(f"{app_name}_{user_id}")
         collection_name = format_collection_name(f"{app_name}")
+        logger.debug(
+            f"Adding documents to Viking database: collection_name={collection_name}"
+        )
         self.get_or_create_collection(collection_name)
         self.database_client.add(content, collection_name=collection_name, **kwargs)
 
-    def query(self, query: str, app_name: str, user_id: str, **kwargs):
+    def query(self, query: str, app_name: str, **kwargs):
         collection_name = format_collection_name(f"{app_name}")
+        logger.debug(
+            f"Querying Viking database: collection_name={collection_name} query={query}"
+        )
+
+        # FIXME(): fix here
         if not self.database_client.collection_exists(collection_name):
             raise ValueError(f"Collection {collection_name} does not exist")
+
         return self.database_client.query(
             query, collection_name=collection_name, **kwargs
         )
 
-    def delete(self, app_name: str, user_id: str, session_id: str):
+    def delete(self, app_name: str, **kwargs):
         # collection_name = format_collection_name(f"{app_name}_{user_id}")
         collection_name = format_collection_name(f"{app_name}")
         try:

veadk/memory/long_term_memory.py

@@ -25,7 +25,6 @@
 from google.genai import types
 from typing_extensions import override
 
-from veadk.config import getenv
 from veadk.database import DatabaseFactory
 from veadk.utils.logger import get_logger
 
@@ -40,20 +39,27 @@ def __init__(
         backend: Literal[
             "local", "opensearch", "redis", "mysql", "viking"
         ] = "opensearch",
-        top_k: int = getenv("LONGTERM_MEMORY_TOP_K", 3),
+        top_k: int = 5,
     ):
         if backend == "viking":
             backend = "viking_mem"
         self.top_k = top_k
         self.backend = backend
 
+        logger.info(
+            f"Initializing long term memory: backend={self.backend} top_k={self.top_k}"
+        )
+
         self.db_client = DatabaseFactory.create(
             backend=backend,
         )
-        logger.info(f"Long term memory backend is `{backend}`.")
 
         self.adapter = get_memory_adapter(backend)(database_client=self.db_client)
 
+        logger.info(
+            f"Initialized long term memory: db_client={self.db_client} adapter={self.adapter}"
+        )
+
     @override
     async def add_session_to_memory(
         self,
@@ -63,6 +69,8 @@ async def add_session_to_memory(
         for event in session.events:
             if not event.content or not event.content.parts:
                 continue
+            if not event.author == "user":  # we only add user event to memory
+                continue
 
             message = event.content.model_dump(exclude_none=True, mode="json")
             if (
@@ -77,37 +85,58 @@ async def add_session_to_memory(
             session_id=session.id,
         )
 
+        logger.info(
+            f"Added {len(event_list)} events to long term memory: app_name={session.app_name} user_id={session.user_id} session_id={session.id}"
+        )
+
     @override
     async def search_memory(self, *, app_name: str, user_id: str, query: str):
+        logger.info(
+            f"Searching long term memory: query={query} app_name={app_name} user_id={user_id}"
+        )
         memory_chunks = self.adapter.query(
             query=query,
             app_name=app_name,
             user_id=user_id,
         )
         if len(memory_chunks) == 0:
+            logger.info(
+                f"Found no memory chunks for query: {query} app_name={app_name} user_id={user_id}"
+            )
             return SearchMemoryResponse()
 
+        logger.info(
+            f"Found {len(memory_chunks)} memory chunks for query: {query} app_name={app_name} user_id={user_id}"
+        )
+
         memory_events = []
         for memory in memory_chunks:
             try:
                 memory_dict = json.loads(memory)
                 try:
                     text = memory_dict["parts"][0]["text"]
                     role = memory_dict["role"]
-                except KeyError as e:
-                    logger.error(
-                        f"Memory content: {memory_dict}. Error parsing memory: {e}"
+                except KeyError as _:
+                    # prevent not a standard text-based event
+                    logger.warning(
+                        f"Memory content: {memory_dict}. Skip return this memory."
                     )
                     continue
             except json.JSONDecodeError:
+                # prevent the memory string is not dumped by `event`
                 text = memory
                 role = "user"
+
             memory_events.append(
                 MemoryEntry(
                     author="user",
                     content=types.Content(parts=[types.Part(text=text)], role=role),
                 )
             )
+
+        logger.info(
+            f"Return {len(memory_events)} memory events for query: {query} app_name={app_name} user_id={user_id}"
+        )
         return SearchMemoryResponse(memories=memory_events)
 
     @override

veadk/memory/memory_database_adapter.py

@@ -33,6 +33,14 @@ def format_collection_name(collection_name: str) -> str:
     return re.sub(r"[^a-z0-9_]", "", replaced_str).lower()
 
 
+def build_index(**kwargs):
+    """
+    Build the index name for the long-term memory.
+    """
+    # TODO
+    ...
+
+
 def get_memory_adapter(backend: str):
     if backend == DatabaseBackend.REDIS:
         return MemoryKVDatabaseAdapter