fix: fix todo

Author: zakahan

veadk/database/database_adapter.py

@@ -11,7 +11,7 @@
 # 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 re
 import time
 from typing import BinaryIO, TextIO
 
@@ -120,8 +120,19 @@ class VectorDatabaseAdapter(BaseModel):
     client: OpenSearchVectorDatabase
 
     def _validate_index(self, index: str):
-        # TODO
-        pass
+        """
+        Verify whether the string conforms to the naming rules of index_name in OpenSearch.
+        https://docs.opensearch.org/2.8/api-reference/index-apis/create-index/
+        """
+        if not (
+            isinstance(index, str)
+            and not index.startswith(("_", "-"))
+            and index.islower()
+            and re.match(r"^[a-z0-9_\-.]+$", index)
+        ):
+            raise ValueError(
+                "The index name does not conform to the naming rules of OpenSearch"
+            )
 
     def add(self, data: list[str], index: str):
         self._validate_index(index)
@@ -133,9 +144,6 @@ 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(
             f"Querying vector database: collection_name={index} query={query} top_k={top_k}"
         )
@@ -153,19 +161,34 @@ class VikingDatabaseAdapter(BaseModel):
     client: VikingDatabase
 
     def _validate_index(self, index: str):
-        # TODO
-        pass
+        """
+        Only English letters, numbers, and underscores (_) are allowed.
+        It must start with an English letter and cannot be empty. Length requirement: [1, 128].
+        For details, please see: https://www.volcengine.com/docs/84313/1254542?lang=zh
+        """
+        if not (
+            isinstance(index, str)
+            and 0 < len(index) <= 128
+            and re.fullmatch(r"^[a-zA-Z][a-zA-Z0-9_]*$", index)
+        ):
+            raise ValueError(
+                "The index name does not conform to the rules: it must start with an English letter, contain only letters, numbers, and underscores, and have a length of 1-128."
+            )
 
     def get_or_create_collection(self, collection_name: str):
         if not self.client.collection_exists(collection_name):
+            logger.warning(
+                f"Collection {collection_name} does not exist, creating a new collection."
+            )
             self.client.create_collection(collection_name)
 
-        # FIXME
+        # After creation, it is necessary to wait for a while.
         count = 0
         while not self.client.collection_exists(collection_name):
+            print("here")
             time.sleep(1)
             count += 1
-            if count > 50:
+            if count > 60:
                 raise TimeoutError(
                     f"Collection {collection_name} not created after 50 seconds"
                 )
@@ -185,9 +208,8 @@ def query(self, query: str, index: str, top_k: int) -> list[str]:
 
         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 []
 
         return self.client.query(query, collection_name=index, top_k=top_k)
 
@@ -198,8 +220,14 @@ class VikingMemoryDatabaseAdapter(BaseModel):
     client: VikingMemoryDatabase
 
     def _validate_index(self, index: str):
-        # TODO
-        pass
+        if not (
+            isinstance(index, str)
+            and 1 <= len(index) <= 128
+            and re.fullmatch(r"^[a-zA-Z][a-zA-Z0-9_]*$", index)
+        ):
+            raise ValueError(
+                "The index name does not conform to the rules: it must start with an English letter, contain only letters, numbers, and underscores, and have a length of 1-128."
+            )
 
     def add(self, data: list[str], index: str, **kwargs):
         self._validate_index(index)
@@ -208,17 +236,16 @@ def add(self, data: list[str], index: str, **kwargs):
             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)
+        self.client.add(data, collection_name=index, **kwargs)
 
-    def query(self, query: str, index: str, top_k: int):
+    def query(self, query: str, index: str, top_k: int, **kwargs):
         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)
+        result = self.client.query(query, collection_name=index, top_k=top_k, **kwargs)
         return result
 
 
@@ -245,8 +272,8 @@ def query(self, query: str, **kwargs):
 
 
 def get_knowledgebase_database_adapter(database_client: BaseDatabase):
-    return MAPPING[type(database_client)](database_client=database_client)
+    return MAPPING[type(database_client)](client=database_client)
 
 
 def get_long_term_memory_database_adapter(database_client: BaseDatabase):
-    return MAPPING[type(database_client)](database_client=database_client)
+    return MAPPING[type(database_client)](client=database_client)

veadk/database/database_factory.py

@@ -69,12 +69,12 @@ def create(backend: str, config=None) -> BaseDatabase:
             return VikingDatabase() if config is None else VikingDatabase(config=config)
 
         if backend == DatabaseBackend.VIKING_MEM:
-            from .viking.viking_memory_db import VikingDatabaseMemory
+            from .viking.viking_memory_db import VikingMemoryDatabase
 
             return (
-                VikingDatabaseMemory()
+                VikingMemoryDatabase()
                 if config is None
-                else VikingDatabaseMemory(config=config)
+                else VikingMemoryDatabase(config=config)
             )
         else:
             raise ValueError(f"Unsupported database type: {backend}")

veadk/database/viking/viking_database.py

@@ -40,6 +40,7 @@
 get_collections_path = "/api/knowledge/collection/info"
 doc_add_path = "/api/knowledge/doc/add"
 doc_info_path = "/api/knowledge/doc/info"
+doc_del_path = "/api/collection/drop"
 
 
 class VolcengineTOSConfig(BaseModel):
@@ -246,9 +247,23 @@ def add(
         }
 
     def delete(self, **kwargs: Any):
-        # collection_name = kwargs.get("collection_name")
-        # todo: delete vikingdb
-        ...
+        collection_name = kwargs.get("collection_name")
+        resource_id = kwargs.get("resource_id")
+        request_param = {"collection_name": collection_name, "resource_id": resource_id}
+        doc_del_req = prepare_request(
+            method="POST", path=doc_del_path, config=self.config, data=request_param
+        )
+        rsp = requests.request(
+            method=doc_del_req.method,
+            url="http://{}{}".format(g_knowledge_base_domain, doc_del_req.path),
+            headers=doc_del_req.headers,
+            data=doc_del_req.body,
+        )
+        result = rsp.json()
+        if result["code"] != 0:
+            logger.error(f"Error in add_doc: {result['message']}")
+            return {"error": result["message"]}
+        return {}
 
     def query(self, query: str, **kwargs: Any) -> list[str]:
         """

veadk/database/viking/viking_memory_db.py

@@ -34,7 +34,6 @@
 logger = get_logger(__name__)
 
 
-# FIXME
 class VikingMemConfig(BaseModel):
     volcengine_ak: Optional[str] = Field(
         default=getenv("VOLCENGINE_ACCESS_KEY"),
@@ -54,8 +53,8 @@ class VikingMemConfig(BaseModel):
     )
 
 
-# ======= adapted from ... =======
-class VikingDBMemoryException(Exception):
+# ======= adapted from https://github.com/volcengine/mcp-server/blob/main/server/mcp_server_vikingdb_memory/src/mcp_server_vikingdb_memory/common/memory_client.py =======
+class VikingMemoryException(Exception):
     def __init__(self, code, request_id, message=None):
         self.code = code
         self.request_id = request_id
@@ -67,15 +66,15 @@ def __str__(self):
         return self.message
 
 
-class VikingDBMemoryService(Service):
+class VikingMemoryService(Service):
     _instance_lock = threading.Lock()
 
     def __new__(cls, *args, **kwargs):
-        if not hasattr(VikingDBMemoryService, "_instance"):
-            with VikingDBMemoryService._instance_lock:
-                if not hasattr(VikingDBMemoryService, "_instance"):
-                    VikingDBMemoryService._instance = object.__new__(cls)
-        return VikingDBMemoryService._instance
+        if not hasattr(VikingMemoryService, "_instance"):
+            with VikingMemoryService._instance_lock:
+                if not hasattr(VikingMemoryService, "_instance"):
+                    VikingMemoryService._instance = object.__new__(cls)
+        return VikingMemoryService._instance
 
     def __init__(
         self,
@@ -88,11 +87,11 @@ def __init__(
         connection_timeout=30,
         socket_timeout=30,
     ):
-        self.service_info = VikingDBMemoryService.get_service_info(
+        self.service_info = VikingMemoryService.get_service_info(
             host, region, scheme, connection_timeout, socket_timeout
         )
-        self.api_info = VikingDBMemoryService.get_api_info()
-        super(VikingDBMemoryService, self).__init__(self.service_info, self.api_info)
+        self.api_info = VikingMemoryService.get_api_info()
+        super(VikingMemoryService, self).__init__(self.service_info, self.api_info)
         if ak:
             self.set_ak(ak)
         if sk:
@@ -102,12 +101,12 @@ def __init__(
         try:
             self.get_body("Ping", {}, json.dumps({}))
         except Exception as e:
-            raise VikingDBMemoryException(
+            raise VikingMemoryException(
                 1000028, "missed", "host or region is incorrect: {}".format(str(e))
             ) from None
 
     def setHeader(self, header):
-        api_info = VikingDBMemoryService.get_api_info()
+        api_info = VikingMemoryService.get_api_info()
         for key in api_info:
             for item in header:
                 api_info[key].header[item] = header[item]
@@ -213,17 +212,17 @@ def get_body_exception(self, api, params, body):
             try:
                 res_json = json.loads(e.args[0].decode("utf-8"))
             except Exception:
-                raise VikingDBMemoryException(
+                raise VikingMemoryException(
                     1000028, "missed", "json load res error, res:{}".format(str(e))
                 ) from None
             code = res_json.get("code", 1000028)
             request_id = res_json.get("request_id", 1000028)
             message = res_json.get("message", None)
 
-            raise VikingDBMemoryException(code, request_id, message)
+            raise VikingMemoryException(code, request_id, message)
 
         if res == "":
-            raise VikingDBMemoryException(
+            raise VikingMemoryException(
                 1000028,
                 "missed",
                 "empty response due to unknown error, please contact customer service",
@@ -237,15 +236,15 @@ def get_exception(self, api, params):
             try:
                 res_json = json.loads(e.args[0].decode("utf-8"))
             except Exception:
-                raise VikingDBMemoryException(
+                raise VikingMemoryException(
                     1000028, "missed", "json load res error, res:{}".format(str(e))
                 ) from None
             code = res_json.get("code", 1000028)
             request_id = res_json.get("request_id", 1000028)
             message = res_json.get("message", None)
-            raise VikingDBMemoryException(code, request_id, message)
+            raise VikingMemoryException(code, request_id, message)
         if res == "":
-            raise VikingDBMemoryException(
+            raise VikingMemoryException(
                 1000028,
                 "missed",
                 "empty response due to unknown error, please contact customer service",
@@ -365,7 +364,7 @@ def format_milliseconds(timestamp_ms):
     return dt.strftime("%Y%m%d %H:%M:%S")
 
 
-# ======= adapted from ... =======
+# ======= adapted from https://github.com/volcengine/mcp-server/blob/main/server/mcp_server_vikingdb_memory/src/mcp_server_vikingdb_memory/common/memory_client.py =======
 
 
 class VikingMemoryDatabase(BaseModel, BaseDatabase):
@@ -375,7 +374,7 @@ class VikingMemoryDatabase(BaseModel, BaseDatabase):
     )
 
     def model_post_init(self, context: Any, /) -> None:
-        self._vm = VikingDBMemoryService(
+        self._vm = VikingMemoryService(
             ak=self.config.volcengine_ak, sk=self.config.volcengine_sk
         )
 
@@ -516,8 +515,8 @@ def query(self, query: str, **kwargs: Any) -> list[str]:
         assert collection_name is not None, "collection_name is required"
         user_id = kwargs.get("user_id")
         assert user_id is not None, "user_id is required"
-
-        resp = self.search_memory(collection_name, query, user_id=user_id)
+        top_k = kwargs.get("top_k", 5)
+        resp = self.search_memory(collection_name, query, user_id=user_id, top_k=top_k)
         return resp
 
     def delete(self, **kwargs: Any):

veadk/knowledgebase/knowledgebase.py

@@ -58,8 +58,12 @@ def add(
         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'.
         """
-        # TODO: add check for data type
-        ...
+        if self.backend != "viking" and not (
+            isinstance(data, str) or isinstance(data, list)
+        ):
+            raise ValueError(
+                "Only vikingdb supports uploading files or file characters."
+            )
 
         index = build_knowledgebase_index(app_name)
 
@@ -73,7 +77,8 @@ def search(self, query: str, app_name: str, top_k: int = None) -> list[str]:
         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)
+        index = build_knowledgebase_index(app_name)
+        result = self.adapter.query(query=query, index=index, top_k=top_k)
         if len(result) == 0:
             logger.warning(f"No documents found in knowledgebase. Query: {query}")
         return result

veadk/memory/long_term_memory.py

@@ -97,7 +97,10 @@ async def add_session_to_memory(
         )
 
         # check if viking memory database, should give a user id： if/else
-        self.adapter.add(data=event_strings, index=index)
+        if self.backend == "viking_mem":
+            self.adapter.add(data=event_strings, index=index, user_id=session.user_id)
+        else:
+            self.adapter.add(data=event_strings, index=index)
 
         logger.info(
             f"Added {len(event_strings)} events to long term memory: index={index}"
@@ -112,15 +115,14 @@ async def search_memory(self, *, app_name: str, user_id: str, query: str):
         )
 
         # 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:
-        #     logger.info(f"Found no memory chunks for query: {query} index={index}")
-        #     return SearchMemoryResponse()
-
-        # logger.info(
-        #     f"Found {len(memory_chunks)} memory chunks for query: {query} index={index}"
-        # )
+        if self.backend == "viking_mem":
+            memory_chunks = self.adapter.query(
+                query=query, index=index, top_k=self.top_k, user_id=user_id
+            )
+        else:
+            memory_chunks = self.adapter.query(
+                query=query, index=index, top_k=self.top_k
+            )
 
         memory_events = []
         for memory in memory_chunks: