fix: knowledgebase add method & list_chunks

Author: zakahan

veadk/database/database_adapter.py

@@ -28,7 +28,7 @@ def __init__(self, client):
 
         self.client: RedisDatabase = client
 
-    def add(self, data: list[str], index: str):
+    def add(self, data: list[str], index: str, **kwargs):
         logger.debug(f"Adding documents to Redis database: index={index}")
 
         try:
@@ -78,7 +78,7 @@ def delete_doc(self, index: str, id: str) -> bool:
             )
             return False
 
-    def list_docs(self, index: str, offset: int = 0, limit: int = 100) -> list[dict]:
+    def list_chunks(self, index: str, offset: int = 0, limit: int = 100) -> list[dict]:
         logger.debug(f"Listing documents from Redis database: index={index}")
         try:
             # Get all documents from Redis
@@ -111,7 +111,7 @@ def create_table(self, table_name: str):
         """
         self.client.add(sql)
 
-    def add(self, data: list[str], index: str):
+    def add(self, data: list[str], index: str, **kwargs):
         logger.debug(
             f"Adding documents to SQL database: table_name={index} data_len={len(data)}"
         )
@@ -203,7 +203,7 @@ def _validate_index(self, index: str):
                 "The index name does not conform to the naming rules of OpenSearch"
             )
 
-    def add(self, data: list[str], index: str):
+    def add(self, data: list[str], index: str, **kwargs):
         self._validate_index(index)
 
         logger.debug(
@@ -247,7 +247,7 @@ def delete_doc(self, index: str, id: str) -> bool:
             )
             return False
 
-    def list_docs(self, index: str, offset: int = 0, limit: int = 1000) -> list[dict]:
+    def list_chunks(self, index: str, offset: int = 0, limit: int = 1000) -> list[dict]:
         self._validate_index(index)
         logger.debug(f"Listing documents from vector database: index={index}")
         return self.client.list_docs(collection_name=index, offset=offset, limit=limit)
@@ -322,6 +322,13 @@ def delete_doc(self, index: str, id: str) -> bool:
         logger.debug(f"Deleting documents from vector database: index={index} id={id}")
         return self.client.delete_by_id(collection_name=index, id=id)
 
+    def list_chunks(self, index: str, offset: int, limit: int) -> list[dict]:
+        self._validate_index(index)
+        logger.debug(f"Listing documents from vector database: index={index}")
+        return self.client.list_chunks(
+            collection_name=index, offset=offset, limit=limit
+        )
+
     def list_docs(self, index: str, offset: int, limit: int) -> list[dict]:
         self._validate_index(index)
         logger.debug(f"Listing documents from vector database: index={index}")
@@ -371,7 +378,7 @@ def delete(self, index: str) -> bool:
     def delete_docs(self, index: str, ids: list[int]):
         raise NotImplementedError("VikingMemoryDatabase does not support delete_docs")
 
-    def list_docs(self, index: str):
+    def list_chunks(self, index: str):
         raise NotImplementedError("VikingMemoryDatabase does not support list_docs")
 
 
@@ -393,7 +400,7 @@ def delete(self, index: str) -> bool:
     def delete_doc(self, index: str, id: str) -> bool:
         return self.client.delete_doc(id)
 
-    def list_docs(self, index: str, offset: int = 0, limit: int = 100) -> list[dict]:
+    def list_chunks(self, index: str, offset: int = 0, limit: int = 100) -> list[dict]:
         return self.client.list_docs(offset=offset, limit=limit)
 
 

veadk/database/viking/viking_database.py

@@ -403,7 +403,7 @@ def collection_exists(self, collection_name: str) -> bool:
         else:
             return False
 
-    def list_docs(
+    def list_chunks(
         self, collection_name: str, offset: int = 0, limit: int = -1
     ) -> list[dict]:
         request_params = {
@@ -431,6 +431,9 @@ def list_docs(
             logger.error(f"Error in list_docs: {result['message']}")
             raise ValueError(f"Error in list_docs: {result['message']}")
 
+        if not result["data"]["point_list"]:
+            return []
+
         data = [
             {
                 "id": res["point_id"],

veadk/knowledgebase/knowledgebase.py

@@ -11,13 +11,15 @@
 # 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 io
+import os.path
 from typing import Any, BinaryIO, Literal, TextIO
 
 from pydantic import BaseModel
 
 from veadk.database.database_adapter import get_knowledgebase_database_adapter
 from veadk.database.database_factory import DatabaseFactory
+from veadk.utils.misc import formatted_timestamp
 from veadk.utils.logger import get_logger
 
 logger = get_logger(__name__)
@@ -66,10 +68,65 @@ def add(
             )
 
         index = build_knowledgebase_index(app_name)
-
         logger.info(f"Adding documents to knowledgebase: index={index}")
 
-        self._adapter.add(data=data, index=index)
+        if self.backend == "viking":
+            # Case 1: Handling file paths or lists of file paths (str)
+            if isinstance(data, str) and os.path.isfile(data):
+                # 单个文件路径，直接调用client.add
+                # 获取文件名（包括后缀名）
+                if "file_name" not in kwargs or not kwargs["file_name"]:
+                    kwargs["file_name"] = os.path.basename(data)
+                return self._adapter.add(data=data, index=index, **kwargs)
+            # Case 2: Handling when list[str] is a full path  (list[str])
+            if isinstance(data, list):
+                if all(isinstance(item, str) for item in data):
+                    all_paths = all(os.path.isfile(item) for item in data)
+                    all_not_paths = all(not os.path.isfile(item) for item in data)
+                    if all_paths:
+                        if "file_name" not in kwargs or not kwargs["file_name"]:
+                            kwargs["file_name"] = [
+                                os.path.basename(item) for item in data
+                            ]
+                        return self._adapter.add(data=data, index=index, **kwargs)
+                    elif (
+                        not all_not_paths
+                    ):  # Prevent the occurrence of non-existent paths
+                        # There is a mixture of paths and non-paths
+                        raise ValueError(
+                            "Mixed file paths and content strings in list are not allowed"
+                        )
+            # Case 3: Handling strings or string arrays (content)  (str or list[str])
+            if isinstance(data, str) or (
+                isinstance(data, list) and all(isinstance(item, str) for item in data)
+            ):
+                if "file_name" not in kwargs or not kwargs["file_name"]:
+                    if isinstance(data, str):
+                        kwargs["file_name"] = f"{formatted_timestamp()}.txt"
+                    else:  # list[str] without file_names
+                        prefix_file_name = formatted_timestamp()
+                        kwargs["file_name"] = [
+                            f"{prefix_file_name}_{i}.txt" for i in range(len(data))
+                        ]
+                return self._adapter.add(data=data, index=index, **kwargs)
+
+            # Case 4: Handling binary data (bytes)
+            if isinstance(data, bytes):
+                # user must give file_name
+                if "file_name" not in kwargs:
+                    raise ValueError("file_name must be provided for binary data")
+                return self._adapter.add(data=data, index=index, **kwargs)
+
+            # Case 5: Handling file objects TextIO or BinaryIO
+            if isinstance(data, (io.TextIOWrapper, io.BufferedReader)):
+                if not kwargs.get("file_name") and hasattr(data, "name"):
+                    kwargs["file_name"] = os.path.basename(data.name)
+                return self._adapter.add(data=data, index=index, **kwargs)
+            # Case6: Unsupported data type
+            raise TypeError(f"Unsupported data type: {type(data)}")
+
+        # not viking
+        return self._adapter.add(data=data, index=index, **kwargs)
 
     def search(self, query: str, app_name: str, top_k: int | None = None) -> list[str]:
         top_k = self.top_k if top_k is None else top_k
@@ -93,4 +150,4 @@ def delete_doc(self, app_name: str, id: str) -> bool:
 
     def list_docs(self, app_name: str, offset: int = 0, limit: int = 100) -> list[dict]:
         index = build_knowledgebase_index(app_name)
-        return self._adapter.list_docs(index=index, offset=offset, limit=limit)
+        return self._adapter.list_chunks(index=index, offset=offset, limit=limit)