feat: replace search_knowledge API with Volcengine Viking SDK (#354)

* feat: replace search_knowledge API with Volcengine Viking SDK

* fix: add sts token in viking knowledge sdk client

Author: ethfoo

veadk/knowledgebase/backends/vikingdb_knowledge_backend.py

@@ -31,6 +31,7 @@
 from veadk.knowledgebase.entry import KnowledgebaseEntry
 from veadk.utils.logger import get_logger
 from veadk.utils.misc import formatted_timestamp
+from volcengine.viking_knowledgebase import VikingKnowledgeBaseService
 
 try:
     from veadk.integrations.ve_tos.ve_tos import VeTOS
@@ -113,10 +114,25 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
     volcengine_project: str = "default"
 
     region: str = "cn-beijing"
+    base_url: str = "https://api-knowledgebase.mlp.cn-beijing.volces.com"
+    host: str = "api-knowledgebase.mlp.cn-beijing.volces.com"
+    schema: str = "https"
 
     tos_config: TOSConfig | NormalTOSConfig = Field(default_factory=TOSConfig)
 
+    _viking_sdk_client = None
+
     def model_post_init(self, __context: Any) -> None:
+        self._set_service_info()
+
+        self._viking_sdk_client = VikingKnowledgeBaseService(
+            host=self.host,
+            ak=self.volcengine_access_key,
+            sk=self.volcengine_secret_key,
+            sts_token=self.session_token,
+            scheme=self.schema,
+        )
+
         self.precheck_index_naming()
 
         # check whether collection exist, if not, create it
@@ -502,8 +518,6 @@ def _search_knowledge(
         rerank: bool = True,
         chunk_diffusion_count: int | None = 0,
     ) -> list[KnowledgebaseEntry]:
-        SEARCH_KNOWLEDGE_PATH = "/api/knowledge/collection/search_knowledge"
-
         query_param = (
             {
                 "doc_filter": {
@@ -523,26 +537,17 @@ def _search_knowledge(
             "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",
+        response = self._viking_sdk_client.search_knowledge(
+            collection_name=self.index,
+            project=self.volcengine_project,
+            query=query,
+            limit=top_k,
+            query_param=query_param,
+            post_processing=post_precessing,
         )
 
-        if response.get("code") != 0:
-            raise ValueError(
-                f"Error during knowledge search: {response.get('code')}, message: {response.get('message')}"
-            )
-
         entries = []
-        for result in response.get("data", {}).get("result_list", []):
+        for result in response.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 = {}
@@ -555,43 +560,40 @@ def _search_knowledge(
 
         return entries
 
-    def _do_request(
-        self,
-        body: dict,
-        path: str,
-        method: Literal["GET", "POST", "PUT", "DELETE"] = "POST",
-    ) -> dict:
-        VIKINGDB_KNOWLEDGEBASE_BASE_URL = (
-            "https://api-knowledgebase.mlp.cn-beijing.volces.com"
-        )
-        full_path = f"{VIKINGDB_KNOWLEDGEBASE_BASE_URL}{path}"
-
+    def _set_service_info(self):
         env_host = getenv(
             "DATABASE_VIKING_BASE_URL", default_value=None, allow_false_values=True
         )
         if env_host:
             if env_host.startswith("http://") or env_host.startswith("https://"):
-                full_path = f"{env_host}{path}"
+                self.base_url = env_host
+                split_url = env_host.split("://")
+                self.host = split_url[-1]
+                self.schema = split_url[0]
             else:
                 raise ValueError(
                     "DATABASE_VIKING_BASE_URL must start with http:// or https://"
                 )
 
-        volcengine_access_key = self.volcengine_access_key
-        volcengine_secret_key = self.volcengine_secret_key
-        session_token = self.session_token
-
-        if not (volcengine_access_key and volcengine_secret_key):
+        if not (self.volcengine_access_key and self.volcengine_secret_key):
             cred = get_credential_from_vefaas_iam()
-            volcengine_access_key = cred.access_key_id
-            volcengine_secret_key = cred.secret_access_key
-            session_token = cred.session_token
+            self.volcengine_access_key = cred.access_key_id
+            self.volcengine_secret_key = cred.secret_access_key
+            self.session_token = cred.session_token
+
+    def _do_request(
+        self,
+        body: dict,
+        path: str,
+        method: Literal["GET", "POST", "PUT", "DELETE"] = "POST",
+    ) -> dict:
+        full_path = f"{self.base_url}{path}"
 
         request = build_vikingdb_knowledgebase_request(
             path=path,
-            volcengine_access_key=volcengine_access_key,
-            volcengine_secret_key=volcengine_secret_key,
-            session_token=session_token,
+            volcengine_access_key=self.volcengine_access_key,
+            volcengine_secret_key=self.volcengine_secret_key,
+            session_token=self.session_token,
             method=method,
             data=body,
         )