fix self.sk

Author: yaozheng-fang

veadk/knowledgebase/backends/vikingdb_knowledge_backend.py

@@ -25,7 +25,10 @@
 from volcengine.viking_knowledgebase import VikingKnowledgeBaseService
 
 import veadk.config  # noqa E401
-from veadk.auth.veauth.utils import get_credential_from_vefaas_iam
+from veadk.auth.veauth.utils import (
+    VeIAMCredential,
+    get_credential_from_vefaas_iam,
+)
 from veadk.configs.database_configs import NormalTOSConfig, TOSConfig
 from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
 from veadk.knowledgebase.backends.utils import (
@@ -157,20 +160,16 @@ def precheck_index_naming(self):
             )
 
     def _get_tos_client(self, tos_bucket_name: str) -> VeTOS:
-        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):
-            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
+        if not (self.volcengine_access_key and self.volcengine_secret_key):
+            cred = self._set_service_info()
+            ak = cred.access_key_id
+            sk = cred.secret_access_key
+            sts_token = cred.session_token
 
         return VeTOS(
-            ak=volcengine_access_key,
-            sk=volcengine_secret_key,
-            session_token=session_token,
+            ak=ak or self.volcengine_access_key,
+            sk=sk or self.volcengine_secret_key,
+            session_token=sts_token or self.session_token,
             region=self.tos_config.region,
             bucket_name=tos_bucket_name or self.tos_config.bucket,
         )
@@ -544,13 +543,17 @@ def _search_knowledge(
             "chunk_diffusion_count": chunk_diffusion_count,
         }
 
-        self._set_service_info()
+        if not (self.volcengine_access_key and self.volcengine_secret_key):
+            cred = self._set_service_info()
+            ak = cred.access_key_id
+            sk = cred.secret_access_key
+            sts_token = cred.session_token
 
         self._viking_sdk_client = VikingKnowledgeBaseService(
             host=self.host,
-            ak=self.volcengine_access_key,
-            sk=self.volcengine_secret_key,
-            sts_token=self.session_token,
+            ak=ak or self.volcengine_access_key,
+            sk=sk or self.volcengine_secret_key,
+            sts_token=sts_token or self.session_token,
             scheme=self.schema,
         )
 
@@ -577,7 +580,7 @@ def _search_knowledge(
 
         return entries
 
-    def _set_service_info(self):
+    def _set_service_info(self) -> VeIAMCredential:
         env_host = getenv(
             "DATABASE_VIKING_BASE_URL",
             default_value=None,
@@ -594,11 +597,8 @@ def _set_service_info(self):
                     "DATABASE_VIKING_BASE_URL must start with http:// or https://"
                 )
 
-        if not (self.volcengine_access_key and self.volcengine_secret_key):
-            cred = get_credential_from_vefaas_iam()
-            self.volcengine_access_key = cred.access_key_id
-            self.volcengine_secret_key = cred.secret_access_key
-            self.session_token = cred.session_token
+        cred = get_credential_from_vefaas_iam()
+        return cred
 
     def _do_request(
         self,
@@ -608,13 +608,17 @@ def _do_request(
     ) -> dict:
         full_path = f"{self.base_url}{path}"
 
-        self._set_service_info()
+        if not (self.volcengine_access_key and self.volcengine_secret_key):
+            cred = self._set_service_info()
+            ak = cred.access_key_id
+            sk = cred.secret_access_key
+            sts_token = cred.session_token
 
         request = build_vikingdb_knowledgebase_request(
             path=path,
-            volcengine_access_key=self.volcengine_access_key,
-            volcengine_secret_key=self.volcengine_secret_key,
-            session_token=self.session_token,
+            volcengine_access_key=ak or self.volcengine_access_key,
+            volcengine_secret_key=sk or self.volcengine_secret_key,
+            session_token=sts_token or self.session_token,
             method=method,
             data=body,
         )