reconstruct database adapter

Author: yaozheng-fang

veadk/database/database_adapter.py

@@ -0,0 +1,248 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import time
+from typing import BinaryIO, TextIO
+
+from pydantic import BaseModel, ConfigDict
+
+from veadk.database.base_database import BaseDatabase
+from veadk.database.kv.redis_database import RedisDatabase
+from veadk.database.local_database import LocalDataBase
+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.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class KVDatabaseAdapter(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    client: RedisDatabase
+
+    def add(self, data: list[str], index: str):
+        logger.debug(f"Adding documents to Redis database: index={index}")
+
+        try:
+            for _data in data:
+                self.client.add(key=index, data=_data)
+            logger.debug(f"Added {len(data)} texts to Redis database: index={index}")
+        except Exception as e:
+            logger.error(
+                f"Failed to add data to Redis database: index={index} error={e}"
+            )
+            raise e
+
+    def query(self, query: str, index: str, top_k: int = 0) -> list[str]:
+        logger.debug(f"Querying Redis database: index={index} query={query}")
+
+        # ignore top_k, as KV search only return one result
+        _ = top_k
+
+        try:
+            result = self.client.query(key=index, query=query)
+            return result
+        except Exception as e:
+            logger.error(f"Failed to search from Redis: index={index} error={e}")
+            raise e
+
+
+class RelationalDatabaseAdapter(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    client: MysqlDatabase
+
+    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,
+                `data` TEXT NOT NULL,
+                `created_at` TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+            ) ENGINE=InnoDB DEFAULT CHARSET={self.client.config.charset};
+        """
+        self.client.add(sql)
+
+    def add(self, data: list[str], index: str):
+        logger.debug(
+            f"Adding documents to SQL database: table_name={index} data_len={len(data)}"
+        )
+
+        if not self.client.table_exists(index):
+            logger.warning(f"Table {index} does not exist, creating...")
+            self.create_table(index)
+
+        for _data in data:
+            sql = f"""
+                INSERT INTO `{index}` (`data`)
+                VALUES (%s);
+                """
+            self.client.add(sql, params=(_data,))
+        logger.debug(f"Added {len(data)} texts to table {index}.")
+
+    def query(self, query: str, index: str, top_k: int) -> list[str]:
+        logger.debug(
+            f"Querying SQL database: table_name={index} query={query} top_k={top_k}"
+        )
+
+        if not self.client.table_exists(index):
+            logger.warning(
+                f"Querying SQL database, but table `{index}` does not exist, returning empty list."
+            )
+            return []
+
+        sql = f"""
+            SELECT `data` FROM `{index}` ORDER BY `created_at` DESC LIMIT {top_k};
+            """
+        results = self.client.query(sql)
+
+        return [item["data"] for item in results]
+
+
+class VectorDatabaseAdapter(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    client: OpenSearchVectorDatabase
+
+    def _validate_index(self, index: str):
+        # TODO
+        pass
+
+    def add(self, data: list[str], index: str):
+        self._validate_index(index)
+
+        logger.debug(
+            f"Adding documents to vector database: index={index} data_len={len(data)}"
+        )
+
+        self.client.add(data, collection_name=index)
+
+    def query(self, query: str, index: str, top_k: int) -> list[str]:
+        self._validate_index(index)
+
+        logger.debug(
+            f"Querying vector database: collection_name={index} query={query} top_k={top_k}"
+        )
+
+        return self.client.query(
+            query=query,
+            collection_name=index,
+            top_k=top_k,
+        )
+
+
+class VikingDatabaseAdapter(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    client: VikingDatabase
+
+    def _validate_index(self, index: str):
+        # TODO
+        pass
+
+    def get_or_create_collection(self, collection_name: str):
+        if not self.client.collection_exists(collection_name):
+            self.client.create_collection(collection_name)
+
+        count = 0
+        while not self.client.collection_exists(collection_name):
+            time.sleep(1)
+            count += 1
+            if count > 50:
+                raise TimeoutError(
+                    f"Collection {collection_name} not created after 50 seconds"
+                )
+
+    def add(
+        self, data: str | list[str] | TextIO | BinaryIO | bytes, index: str, **kwargs
+    ):
+        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)
+
+    def query(self, query: str, index: str, top_k: int) -> list[str]:
+        self._validate_index(index)
+
+        logger.debug(f"Querying Viking database: collection_name={index} query={query}")
+
+        # FIXME(): maybe do not raise, but just return []
+        if not self.client.collection_exists(index):
+            raise ValueError(f"Collection {index} does not exist")
+
+        return self.client.query(query, collection_name=index, top_k=top_k)
+
+
+class VikingDatabaseMemoryAdapter(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    client: VikingDatabaseMemory
+
+    def _validate_index(self, index: str):
+        # TODO
+        pass
+
+    def add(self, data: list[str], index: str):
+        self._validate_index(index)
+
+        logger.debug(
+            f"Adding documents to Viking database memory: collection_name={index} data_len={len(data)}"
+        )
+
+        self.client.add(data, collection_name=index)
+
+    def query(self, query: str, index: str, top_k: int):
+        self._validate_index(index)
+
+        logger.debug(
+            f"Querying Viking database memory: collection_name={index} query={query} top_k={top_k}"
+        )
+
+        result = self.client.query(query, collection_name=index, top_k=top_k)
+        return result
+
+
+class LocalDatabaseAdapter(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    client: LocalDataBase
+
+    def add(self, data: list[str], **kwargs):
+        self.client.add(data)
+
+    def query(self, query: str, **kwargs):
+        return self.client.query(query, **kwargs)
+
+
+MAPPING = {
+    RedisDatabase: KVDatabaseAdapter,
+    MysqlDatabase: RelationalDatabaseAdapter,
+    LocalDataBase: LocalDatabaseAdapter,
+    VikingDatabase: VikingDatabaseAdapter,
+    OpenSearchVectorDatabase: VectorDatabaseAdapter,
+    VikingDatabaseMemory: VikingDatabaseMemoryAdapter,
+}
+
+
+def get_knowledgebase_database_adapter(database_client: BaseDatabase):
+    return MAPPING[type(database_client)](database_client=database_client)
+
+
+def get_long_term_memory_database_adapter(database_client: BaseDatabase):
+    return MAPPING[type(database_client)](database_client=database_client)

veadk/database/viking/viking_database.py

@@ -215,7 +215,12 @@ def _add_doc(self, collection_name: str, tos_url: str, doc_id: str, **kwargs: An
 
         return doc_id
 
-    def add(self, data: str | list[str] | TextIO | BinaryIO | bytes, **kwargs: Any):
+    def add(
+        self,
+        data: str | list[str] | TextIO | BinaryIO | bytes,
+        collection_name: str,
+        **kwargs,
+    ):
         """
         Args:
             data: str, file path or file stream:  Both file or file.read() are acceptable.
@@ -226,8 +231,6 @@ def add(self, data: str | list[str] | TextIO | BinaryIO | bytes, **kwargs: Any):
                 "doc_id": "<doc_id>",
             }
         """
-        collection_name = kwargs.get("collection_name")
-        assert collection_name is not None, "collection_name is required"
 
         status, tos_url = self._upload_to_tos(data=data, **kwargs)
         if status != 200:

veadk/knowledgebase/knowledgebase.py

@@ -14,14 +14,17 @@
 
 from typing import BinaryIO, Literal, TextIO
 
+from veadk.database.database_adapter import get_knowledgebase_database_adapter
 from veadk.database.database_factory import DatabaseFactory
 from veadk.utils.logger import get_logger
 
-from .knowledgebase_database_adapter import get_knowledgebase_adapter
-
 logger = get_logger(__name__)
 
 
+def build_knowledgebase_index(app_name: str):
+    return f"{app_name}"
+
+
 class KnowledgeBase:
     def __init__(
         self,
@@ -37,9 +40,7 @@ def __init__(
         self.top_k = top_k
 
         self.db_client = DatabaseFactory.create(backend=backend, config=db_config)
-        self.adapter = get_knowledgebase_adapter(backend)(
-            database_client=self.db_client
-        )
+        self.adapter = get_knowledgebase_database_adapter(self.db_client)
 
         logger.info(
             f"Initialized knowledgebase: db_client={self.db_client} adapter={self.adapter}"
@@ -53,23 +54,18 @@ def add(
     ):
         """
         Add documents to the vector database.
-        You can only upload files or file characters when the adapter type you use is vikingdb.
+        You can only upload files or file characters when the adapter type used is vikingdb.
         In addition, if you upload data of the bytes type,
             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
-        logger.info(f"Adding documents to knowledgebase: app_name={app_name}")
-        self.adapter.add(data, app_name=app_name, **kwargs)
+        # TODO: add check for data type
+        ...
 
-    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.
+        index = build_knowledgebase_index(app_name)
+        logger.info(f"Adding documents to knowledgebase: index={index}")
+        self.adapter.add(data=data, index=index)
 
-        Args:
-            query (str): The query text to be retrieved (e.g., "Who proposed the Turing machine model?")
-
-        Returns:
-            list[str]: A list of the top most similar document contents retrieved (sorted by vector similarity)
-        """
+    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
 
         logger.info(