feat(kb): enhance knowledgebase functionality with metadata support and improve search handling

Author: Fengrui-Liu

.gitignore

@@ -167,6 +167,13 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 .idea/
 
+# Visual Studio Code
+#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
+#  that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#  and can be added to the global gitignore or merged into this file. However, if you prefer,
+#  you could uncomment the following to ignore the entire vscode folder
+.vscode/
+
 # Ruff stuff:
 .ruff_cache/
 

veadk/integrations/ve_tos/ve_tos.py

@@ -178,8 +178,21 @@ def _set_cors_rules(self, bucket_name: str) -> bool:
             logger.error(f"Failed to set CORS rules for bucket {bucket_name}: {str(e)}")
             return False
 
-    def _build_object_key_for_file(self, data_path: str) -> str:
-        """generate TOS object key"""
+    def build_tos_url(
+        self, user_id: str, app_name: str, session_id: str, data_path: str
+    ) -> tuple[str, str]:
+        """Builds the TOS object key and URL for the given parameters.
+
+        Args:
+            user_id (str): User ID
+            app_name (str): App name
+            session_id (str): Session ID
+            data_path (str): Data path
+
+        Returns:
+            tuple[str, str]: Object key and TOS URL.
+        """
+
         parsed_url = urlparse(data_path)
 
         # Generate object key
@@ -245,14 +258,24 @@ def build_tos_url(self, object_key: str, bucket_name: str = "") -> str:
 
     # deprecated
     def upload(
-        self, data: Union[str, bytes], bucket_name: str = "", object_key: str = ""
+        self, object_key: str, data: Union[str, bytes], metadata: dict | None = None
     ):
+        """Uploads data to TOS.
+
+        Args:
+            object_key (str): The object key for the uploaded data.
+            data (Union[str, bytes]): The data to upload, either as a file path or raw bytes.
+            metadata (dict | None, optional): Metadata to associate with the object. Defaults to None.
+
+        Raises:
+            ValueError: If the data type is unsupported.
+        """
         if isinstance(data, str):
             # data is a file path
-            return asyncio.to_thread(self.upload_file, data, bucket_name, object_key)
+            return asyncio.to_thread(self.upload_file, data, bucket_name, object_key, metadata)
         elif isinstance(data, bytes):
             # data is bytes content
-            return asyncio.to_thread(self.upload_bytes, data, bucket_name, object_key)
+            return asyncio.to_thread(self.upload_bytes, data, bucket_name, object_key, metadata)
         else:
             error_msg = f"Upload failed: data type error. Only str (file path) and bytes are supported, got {type(data)}"
             logger.error(error_msg)

veadk/knowledgebase/backends/base_backend.py

@@ -29,19 +29,19 @@ def precheck_index_naming(self) -> None:
         """
 
     @abstractmethod
-    def add_from_directory(self, directory: str, **kwargs) -> bool:
+    def add_from_directory(self, directory: str, *args, **kwargs) -> bool:
         """Add knowledge from file path to knowledgebase"""
 
     @abstractmethod
-    def add_from_files(self, files: list[str], **kwargs) -> bool:
+    def add_from_files(self, files: list[str], *args, **kwargs) -> bool:
         """Add knowledge (e.g, documents, strings, ...) to knowledgebase"""
 
     @abstractmethod
-    def add_from_text(self, text: str | list[str], **kwargs) -> bool:
+    def add_from_text(self, text: str | list[str], *args, **kwargs) -> bool:
         """Add knowledge from text to knowledgebase"""
 
     @abstractmethod
-    def search(self, **kwargs) -> list:
+    def search(self, *args, **kwargs) -> list:
         """Search knowledge from knowledgebase"""
 
     # Optional methods for future use:

veadk/knowledgebase/backends/vikingdb_knowledge_backend.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import asyncio
+import json
 import re
 from pathlib import Path
 from typing import Any, Literal
@@ -26,6 +27,7 @@
 from veadk.configs.database_configs import NormalTOSConfig, TOSConfig
 from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
 from veadk.knowledgebase.backends.utils import build_vikingdb_knowledgebase_request
+from veadk.knowledgebase.entry import KnowledgebaseEntry
 from veadk.utils.logger import get_logger
 from veadk.utils.misc import formatted_timestamp
 
@@ -106,6 +108,7 @@ def add_from_directory(
         directory: str,
         tos_bucket_name: str | None = None,
         tos_bucket_path: str = "knowledgebase",
+        metadata: dict | None = None,
         **kwargs,
     ) -> bool:
         """Add knowledge from a directory to the knowledgebase.
@@ -114,6 +117,8 @@ def add_from_directory(
             directory (str): The directory to add to knowledgebase.
             tos_bucket_name (str | None, optional): The bucket name of TOS. Defaults to None.
             tos_bucket_path (str, optional): The path of TOS bucket. Defaults to "knowledgebase".
+            metadata (dict | None, optional): The metadata of the files. Defaults to None.
+            **kwargs: Additional keyword arguments.
 
         Returns:
             bool: True if successful, False otherwise.
@@ -126,6 +131,7 @@ def add_from_directory(
                 content,
                 tos_bucket_name=tos_bucket_name,
                 object_key=f"{tos_bucket_path}/{file_name}",
+                metadata=metadata,
             )
             self._add_doc(tos_url=tos_url)
         return True
@@ -136,6 +142,7 @@ def add_from_files(
         files: list[str],
         tos_bucket_name: str | None = None,
         tos_bucket_path: str = "knowledgebase",
+        metadata: dict | None = None,
         **kwargs,
     ) -> bool:
         """Add knowledge from a directory to the knowledgebase.
@@ -144,6 +151,8 @@ def add_from_files(
             files (list[str]): The files to add to knowledgebase.
             tos_bucket_name (str | None, optional): The bucket name of TOS. Defaults to None.
             tos_bucket_path (str, optional): The path of TOS bucket. Defaults to "knowledgebase".
+            metadata (dict | None, optional): The metadata of the files. Defaults to None.
+            **kwargs: Additional keyword arguments.
 
         Returns:
             bool: True if successful, False otherwise.
@@ -155,6 +164,7 @@ def add_from_files(
                 content,
                 tos_bucket_name=tos_bucket_name,
                 object_key=f"{tos_bucket_path}/{file_name}",
+                metadata=metadata,
             )
             self._add_doc(tos_url=tos_url)
         return True
@@ -165,6 +175,7 @@ def add_from_text(
         text: str | list[str],
         tos_bucket_name: str | None = None,
         tos_bucket_path: str = "knowledgebase",
+        metadata: dict | None = None,
         **kwargs,
     ) -> bool:
         """Add knowledge from text to the knowledgebase.
@@ -189,7 +200,7 @@ def add_from_text(
             for _text, _object_key in zip(text, object_keys):
                 _content = _text.encode("utf-8")
                 tos_url = self._upload_bytes_to_tos(
-                    _content, tos_bucket_name, _object_key
+                    _content, tos_bucket_name, _object_key, metadata=metadata
                 )
                 self._add_doc(tos_url=tos_url)
             return True
@@ -198,7 +209,9 @@ def add_from_text(
             object_key = kwargs.get(
                 "object_key", f"veadk/knowledgebase/{formatted_timestamp()}.txt"
             )
-            tos_url = self._upload_bytes_to_tos(content, tos_bucket_name, object_key)
+            tos_url = self._upload_bytes_to_tos(
+                content, tos_bucket_name, object_key, metadata=metadata
+            )
             self._add_doc(tos_url=tos_url)
         else:
             raise ValueError("text must be str or list[str]")
@@ -210,6 +223,7 @@ def add_from_bytes(
         file_name: str,
         tos_bucket_name: str | None = None,
         tos_bucket_path: str = "knowledgebase",
+        metadata: dict | None = None,
         **kwargs,
     ) -> bool:
         """Add knowledge from bytes to the knowledgebase.
@@ -219,6 +233,8 @@ def add_from_bytes(
             file_name (str): The file name of the content.
             tos_bucket_name (str | None, optional): The bucket name of TOS. Defaults to None.
             tos_bucket_path (str, optional): The path of TOS bucket. Defaults to "knowledgebase".
+            metadata (dict | None, optional): The metadata of the files. Defaults to None.
+            **kwargs: Additional keyword arguments.
 
         Returns:
             bool: True if successful, False otherwise.
@@ -228,15 +244,24 @@ def add_from_bytes(
             content,
             tos_bucket_name=tos_bucket_name,
             object_key=f"{tos_bucket_path}/{file_name}",
+            metadata=metadata,
         )
         response = self._add_doc(tos_url=tos_url)
         if response["code"] == 0:
             return True
         return False
 
     @override
-    def search(self, query: str, top_k: int = 5) -> list:
-        return self._search_knowledge(query=query, top_k=top_k)
+    def search(
+        self,
+        query: str,
+        top_k: int = 5,
+        metadata: dict | None = None,
+        rerank: bool = True,
+    ) -> list:
+        return self._search_knowledge(
+            query=query, top_k=top_k, metadata=metadata, rerank=rerank
+        )
 
     def delete_collection(self) -> bool:
         DELETE_COLLECTION_PATH = "/api/knowledge/collection/delete"
@@ -359,7 +384,7 @@ def create_collection(self) -> None:
         response = self._do_request(
             body={
                 "name": self.index,
-                "project": "default",
+                "project": self.volcengine_project,
                 "description": "Created by Volcengine Agent Development Kit (VeADK).",
             },
             path=CREATE_COLLECTION_PATH,
@@ -372,10 +397,17 @@ def create_collection(self) -> None:
             )
 
     def _upload_bytes_to_tos(
-        self, content: bytes, tos_bucket_name: str, object_key: str
+        self,
+        content: bytes,
+        tos_bucket_name: str,
+        object_key: str,
+        metadata: dict | None = None,
     ) -> str:
+        # Here, we set the metadata via the TOS object, ref: https://www.volcengine.com/docs/84313/1254624
         self._tos_client.bucket_name = tos_bucket_name
-        coro = self._tos_client.upload(object_key=object_key, data=content)
+        coro = self._tos_client.upload(
+            object_key=object_key, data=content, metadata=metadata
+        )
         try:
             loop = asyncio.get_running_loop()
             loop.run_until_complete(
@@ -391,7 +423,7 @@ def _add_doc(self, tos_url: str) -> Any:
         response = self._do_request(
             body={
                 "collection_name": self.index,
-                "project": "default",
+                "project": self.volcengine_project,
                 "add_type": "tos",
                 "tos_path": tos_url,
             },
@@ -400,14 +432,43 @@ def _add_doc(self, tos_url: str) -> Any:
         )
         return response
 
-    def _search_knowledge(self, query: str, top_k: int = 5) -> list[str]:
+    def _search_knowledge(
+        self,
+        query: str,
+        top_k: int = 5,
+        metadata: dict | None = None,
+        rerank: bool = True,
+        chunk_diffusion_count: int | None = 3,
+    ) -> list[KnowledgebaseEntry]:
         SEARCH_KNOWLEDGE_PATH = "/api/knowledge/collection/search_knowledge"
 
+        query_param = (
+            {
+                "doc_filter": {
+                    "op": "and",
+                    "conds": [
+                        {"op": "must", "field": str(k), "conds": [str(v)]}
+                        for k, v in metadata.items()
+                    ],
+                }
+            }
+            if metadata
+            else None
+        )
+
+        post_precessing = {
+            "rerank_swich": rerank,
+            "chunk_diffusion_count": chunk_diffusion_count,
+        }
+
         response = self._do_request(
             body={
                 "name": self.index,
+                "project": self.volcengine_project,
                 "query": query,
                 "limit": top_k,
+                "query_param": query_param,
+                "post_processing": post_precessing,
             },
             path=SEARCH_KNOWLEDGE_PATH,
             method="POST",
@@ -418,11 +479,19 @@ def _search_knowledge(self, query: str, top_k: int = 5) -> list[str]:
                 f"Error during knowledge search: {response.get('code')}, message: {response.get('message')}"
             )
 
-        search_result_list = response.get("data", {}).get("result_list", [])
+        entries = []
+        for result in response.get("data", {}).get("result_list", []):
+            doc_meta_raw_str = result.get("doc_info", {}).get("doc_meta")
+            doc_meta_list = json.loads(doc_meta_raw_str) if doc_meta_raw_str else []
+            metadata = {}
+            for meta in doc_meta_list:
+                metadata[meta["field_name"]] = meta["field_value"]
 
-        return [
-            search_result.get("content", "") for search_result in search_result_list
-        ]
+            entries.append(
+                KnowledgebaseEntry(content=result.get("content", ""), metadata=metadata)
+            )
+
+        return entries
 
     def _do_request(
         self,
@@ -445,4 +514,8 @@ def _do_request(
             headers=request.headers,
             data=request.body,
         )
+        if not response.ok:
+            logger.error(
+                f"VikingDBKnowledgeBackend error during request: {response.json()}"
+            )
         return response.json()

veadk/knowledgebase/entry.py

@@ -0,0 +1,25 @@
+# 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.
+
+from pydantic import BaseModel
+
+
+class KnowledgebaseEntry(BaseModel):
+    """Represents a single entry in the knowledgebase."""
+
+    # The main content of the knowledgebase entry.
+    content: str
+
+    # Optional metadata associated with the entry.
+    metadata: dict | None = None