fix: fix knowledgebase credential get (#444)

* fix: fix knowledgebase cridential get

* fix self.sk

* fix none ak

Author: yaozheng-fang

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "veadk-python"
-version = "0.5.5"
+version = "0.5.6"
 description = "Volcengine agent development kit, integrations with Volcengine cloud services."
 readme = "README.md"
 requires-python = ">=3.10"

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 (
@@ -136,14 +139,6 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
     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
@@ -165,20 +160,19 @@ 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
+        ak = None
+        sk = None
+        sts_token = None
+        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,
         )
@@ -552,6 +546,23 @@ def _search_knowledge(
             "chunk_diffusion_count": chunk_diffusion_count,
         }
 
+        ak = None
+        sk = None
+        sts_token = None
+        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=ak or self.volcengine_access_key,
+            sk=sk or self.volcengine_secret_key,
+            sts_token=sts_token or self.session_token,
+            scheme=self.schema,
+        )
+
         response = self._viking_sdk_client.search_knowledge(
             collection_name=self.index,
             project=self.volcengine_project,
@@ -575,7 +586,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,
@@ -592,11 +603,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,
@@ -606,11 +614,20 @@ def _do_request(
     ) -> dict:
         full_path = f"{self.base_url}{path}"
 
+        ak = None
+        sk = None
+        sts_token = None
+        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,
         )

veadk/version.py

@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-VERSION = "0.5.5"
+VERSION = "0.5.6"