From 9997e5f734832bcfcc00b10e14f63a8cbea17d86 Mon Sep 17 00:00:00 2001
From: baoyumeng <baoyumeng@bytedance.com>
Date: Fri, 23 Jan 2026 00:21:49 +0800
Subject: [PATCH] fix: fix VikingDB memory credential get

---
 pyproject.toml                                |  2 +-
 .../vikingdb_memory_backend.py                | 63 +++++++++++--------
 veadk/version.py                              |  2 +-
 3 files changed, 40 insertions(+), 27 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index a8e43d4c..82b37726 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "veadk-python"
-version = "0.5.11"
+version = "0.5.12"
 description = "Volcengine agent development kit, integrations with Volcengine cloud services."
 readme = "README.md"
 requires-python = ">=3.10"
diff --git a/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py b/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py
index 9c9d415d..443ff989 100644
--- a/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py
+++ b/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py
@@ -116,30 +116,55 @@ def _create_collection(self) -> None:
         logger.debug(f"Create collection with response {response}")
         return response
 
-    def _get_client(self) -> VikingDBMemoryClient:
-        if not (self.volcengine_access_key and self.volcengine_secret_key):
+    def _get_ak_sk_sts(self) -> tuple[str, str, str]:
+        ak = ""
+        sk = ""
+        sts_token = ""
+
+        if self.volcengine_access_key and self.volcengine_secret_key:
+            ak = self.volcengine_access_key
+            sk = self.volcengine_secret_key
+            sts_token = self.session_token
+            logger.debug(
+                "Get volcengine credential from Environments for VikingMEMBackend"
+            )
+
+        else:
             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
+            ak = cred.access_key_id
+            sk = cred.secret_access_key
+            sts_token = cred.session_token
+            logger.debug(
+                "Get volcengine credential from VeFaaS IAM file for VikingMEMBackend."
+            )
 
+        return ak, sk, sts_token
+
+    def _get_client(self) -> VikingDBMemoryClient:
+        ak, sk, sts_token = self._get_ak_sk_sts()
         return VikingDBMemoryClient(
-            ak=self.volcengine_access_key,
-            sk=self.volcengine_secret_key,
-            sts_token=self.session_token,
+            ak=ak,
+            sk=sk,
+            sts_token=sts_token,
             region=self.region,
         )
 
     def _get_sdk_client(self) -> VikingMem:
-        client = self._get_client()
+        ak, sk, sts_token = self._get_ak_sk_sts()
+        client = VikingDBMemoryClient(
+            ak=ak,
+            sk=sk,
+            sts_token=sts_token,
+            region=self.region,
+        )
         return VikingMem(
             host=client.get_host(),
             region=self.region,
             auth=IAM(
-                ak=self.volcengine_access_key,
-                sk=self.volcengine_secret_key,
+                ak=ak,
+                sk=sk,
             ),
-            sts_token=self.session_token,
+            sts_token=sts_token,
         )
 
     @override
@@ -237,19 +262,7 @@ def get_user_profile(self, user_id: str) -> str:
             f"Get user profile for user_id={user_id} from Viking Memory with mem_id={mem_id}"
         )
 
-        ak = ""
-        sk = ""
-        sts_token = ""
-        if self.volcengine_access_key and self.volcengine_secret_key:
-            ak = self.volcengine_access_key
-            sk = self.volcengine_secret_key
-            sts_token = self.session_token
-        else:
-            cred = get_credential_from_vefaas_iam()
-            ak = cred.access_key_id
-            sk = cred.secret_access_key
-            sts_token = cred.session_token
-
+        ak, sk, sts_token = self._get_ak_sk_sts()
         response = ve_request(
             request_body={
                 "filter": {
diff --git a/veadk/version.py b/veadk/version.py
index eb4379a1..428f4b2e 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.11"
+VERSION = "0.5.12"