feat: add delete_doc and list_docs

Author: zakahan

veadk/database/database_adapter.py

@@ -54,6 +54,12 @@ def query(self, query: str, index: str, top_k: int = 0) -> list:
             logger.error(f"Failed to search from Redis: index={index} error={e}")
             raise e
 
+    def delete_docs(self, index: str, ids: list[int]): ...
+
+    def list_docs(
+        self, index: str, offset: int = 0, limit: int = 100
+    ) -> list[dict]: ...
+
 
 class RelationalDatabaseAdapter:
     def __init__(self, client):
@@ -108,6 +114,12 @@ def query(self, query: str, index: str, top_k: int) -> list[str]:
 
         return [item["data"] for item in results]
 
+    def delete_docs(self, index: str, ids: list[int]): ...
+
+    def list_docs(
+        self, index: str, offset: int = 0, limit: int = 100
+    ) -> list[dict]: ...
+
 
 class VectorDatabaseAdapter:
     def __init__(self, client):
@@ -152,6 +164,23 @@ def query(self, query: str, index: str, top_k: int) -> list[str]:
             top_k=top_k,
         )
 
+    def delete_doc(self, index: str, id: str) -> bool:
+        self._validate_index(index)
+        logger.debug(f"Deleting documents from vector database: index={index} id={id}")
+        try:
+            self.client.delete_by_id(collection_name=index, id=id)
+            return True
+        except Exception as e:
+            logger.error(
+                f"Failed to delete document from vector database: index={index} id={id} error={e}"
+            )
+            return False
+
+    def list_docs(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)
+
 
 class VikingDatabaseAdapter:
     def __init__(self, client):
@@ -212,6 +241,16 @@ def query(self, query: str, index: str, top_k: int) -> list[str]:
 
         return self.client.query(query, collection_name=index, top_k=top_k)
 
+    def delete_doc(self, index: str, id: str) -> bool:
+        self._validate_index(index)
+        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_docs(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_docs(collection_name=index, offset=offset, limit=limit)
+
 
 class VikingMemoryDatabaseAdapter:
     def __init__(self, client):
@@ -248,6 +287,12 @@ def query(self, query: str, index: str, top_k: int, **kwargs):
         result = self.client.query(query, collection_name=index, top_k=top_k, **kwargs)
         return result
 
+    def delete_docs(self, index: str, ids: list[int]):
+        raise NotImplementedError("VikingMemoryDatabase does not support delete_docs")
+
+    def list_docs(self, index: str):
+        raise NotImplementedError("VikingMemoryDatabase does not support list_docs")
+
 
 class LocalDatabaseAdapter:
     def __init__(self, client):
@@ -261,6 +306,12 @@ def add(self, data: list[str], **kwargs):
     def query(self, query: str, **kwargs):
         return self.client.query(query, **kwargs)
 
+    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]:
+        return self.client.list_docs(offset=offset, limit=limit)
+
 
 MAPPING = {
     "RedisDatabase": KVDatabaseAdapter,

veadk/database/local_database.py

@@ -24,20 +24,35 @@ class LocalDataBase(BaseDatabase):
 
     def __init__(self, **kwargs):
         super().__init__()
-        self.data = []
+        self.data = {}  # 改为字典
         self._type = "local"
+        self._next_id = 0  # 用于生成唯一ID
 
     def add_texts(self, texts: list[str], **kwargs):
-        self.data.extend(texts)
+        for text in texts:
+            self.data[str(self._next_id)] = text
+            self._next_id += 1
 
     def is_empty(self):
         return len(self.data) == 0
 
     def query(self, query: str, **kwargs: Any) -> list[str]:
-        return self.data
+        return list(self.data.values())
 
     def delete(self, **kwargs: Any):
-        self.data = []
+        self.data = {}
 
     def add(self, texts: list[str], **kwargs: Any):
         return self.add_texts(texts)
+
+    def list_docs(self, **kwargs: Any) -> list[dict]:
+        return [{"id": id, "content": content} for id, content in self.data.items()]
+
+    def delete_doc(self, id: str, **kwargs: Any):
+        if id not in self.data:
+            raise ValueError(f"id {id} not found")
+        try:
+            del self.data[id]
+            return True
+        except Exception:
+            return False

veadk/database/vector/opensearch_vector_database.py

@@ -219,20 +219,22 @@ def list_all_collection(self) -> list:
         response = self._opensearch_client.indices.get_alias()
         return list(response.keys())
 
-    def get_all_docs(self, collection_name: str, size: int = 10000) -> list[dict]:
+    def list_docs(
+        self, collection_name: str, offset: int = 0, limit: int = 10000
+    ) -> list[dict]:
         """Match all docs in one index of OpenSearch"""
         if not self.collection_exists(collection_name):
             logger.warning(
                 f"Get all docs, but collection {collection_name} does not exist. return a empty list."
             )
             return []
 
-        query = {"size": size, "query": {"match_all": {}}}
+        query = {"size": limit, "from": offset, "query": {"match_all": {}}}
         response = self._opensearch_client.search(index=collection_name, body=query)
         return [
             {
                 "id": hit["_id"],
-                "page_content": hit["_source"]["page_content"],
+                "content": hit["_source"]["page_content"],
             }
             for hit in response["hits"]["hits"]
         ]

veadk/database/viking/viking_database.py

@@ -41,6 +41,8 @@
 doc_add_path = "/api/knowledge/doc/add"
 doc_info_path = "/api/knowledge/doc/info"
 doc_del_path = "/api/collection/drop"
+list_docs_path = "/api/knowledge/point/list"
+delete_docs_path = "/api/knowledge/point/delete"
 
 
 class VolcengineTOSConfig(BaseModel):
@@ -400,3 +402,66 @@ def collection_exists(self, collection_name: str) -> bool:
             return True
         else:
             return False
+
+    def list_docs(
+        self, collection_name: str, offset: int = 0, limit: int = -1
+    ) -> list[dict]:
+        request_params = {
+            "collection_name": collection_name,
+            "project": self.config.project,
+            "offset": offset,
+            "limit": limit,
+        }
+
+        create_collection_req = prepare_request(
+            method="POST",
+            path=list_docs_path,
+            config=self.config,
+            data=request_params,
+        )
+        resp = requests.request(
+            method=create_collection_req.method,
+            url="https://{}{}".format(
+                g_knowledge_base_domain, create_collection_req.path
+            ),
+            headers=create_collection_req.headers,
+            data=create_collection_req.body,
+        )
+
+        result = resp.json()
+        if result["code"] != 0:
+            logger.error(f"Error in list_docs: {result['message']}")
+            raise ValueError(f"Error in list_docs: {result['message']}")
+
+        data = [
+            {"id": res["point_id"], "content": res["content"]}
+            for res in result["data"]["point_list"]
+        ]
+        return data
+
+    def delete_by_id(self, collection_name: str, id: str) -> bool:
+        request_params = {
+            "collection_name": collection_name,
+            "project": self.config.project,
+            "point_id": id,
+        }
+
+        create_collection_req = prepare_request(
+            method="POST",
+            path=delete_docs_path,
+            config=self.config,
+            data=request_params,
+        )
+        resp = requests.request(
+            method=create_collection_req.method,
+            url="https://{}{}".format(
+                g_knowledge_base_domain, create_collection_req.path
+            ),
+            headers=create_collection_req.headers,
+            data=create_collection_req.body,
+        )
+
+        result = resp.json()
+        if result["code"] != 0:
+            return False
+        return True

veadk/knowledgebase/knowledgebase.py

@@ -80,3 +80,11 @@ def search(self, query: str, app_name: str, top_k: int | None = None) -> list[st
         if len(result) == 0:
             logger.warning(f"No documents found in knowledgebase. Query: {query}")
         return result
+
+    def delete_doc(self, app_name: str, id: str) -> bool:
+        index = build_knowledgebase_index(app_name)
+        return self.adapter.delete_doc(index=index, id=id)
+
+    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)