diff --git a/pyproject.toml b/pyproject.toml index f325280f..ec51b7f1 100644 --- a/pyproject.toml +++ b/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" diff --git a/veadk/knowledgebase/backends/vikingdb_knowledge_backend.py b/veadk/knowledgebase/backends/vikingdb_knowledge_backend.py index 11bdb96b..d360435a 100644 --- a/veadk/knowledgebase/backends/vikingdb_knowledge_backend.py +++ b/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, ) diff --git a/veadk/version.py b/veadk/version.py index ddf32e6b..8c1547e8 100644 --- a/veadk/version.py +++ b/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"