add auto token for vikingdb memory

Author: yaozheng-fang

veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import json
+import os
 import re
 import time
 import uuid
@@ -22,7 +23,7 @@
 from typing_extensions import override
 
 import veadk.config  # noqa E401
-from veadk.config import getenv
+from veadk.auth.veauth.utils import get_credential_from_vefaas_iam
 from veadk.integrations.ve_viking_db_memory.ve_viking_db_memory import (
     VikingDBMemoryClient,
 )
@@ -35,14 +36,16 @@
 
 
 class VikingDBLTMBackend(BaseLongTermMemoryBackend):
-    volcengine_access_key: str = Field(
-        default_factory=lambda: getenv("VOLCENGINE_ACCESS_KEY")
+    volcengine_access_key: str | None = Field(
+        default_factory=lambda: os.getenv("VOLCENGINE_ACCESS_KEY")
     )
 
-    volcengine_secret_key: str = Field(
-        default_factory=lambda: getenv("VOLCENGINE_SECRET_KEY")
+    volcengine_secret_key: str | None = Field(
+        default_factory=lambda: os.getenv("VOLCENGINE_SECRET_KEY")
     )
 
+    session_token: str = ""
+
     region: str = "cn-beijing"
     """VikingDB memory region"""
 
@@ -57,31 +60,41 @@ def precheck_index_naming(self):
             )
 
     def model_post_init(self, __context: Any) -> None:
-        self._client = VikingDBMemoryClient(
-            ak=self.volcengine_access_key,
-            sk=self.volcengine_secret_key,
-            region=self.region,
-        )
-
         # check whether collection exist, if not, create it
         if not self._collection_exist():
             self._create_collection()
 
     def _collection_exist(self) -> bool:
         try:
-            self._client.get_collection(collection_name=self.index)
+            client = self._get_client()
+            client.get_collection(collection_name=self.index)
             return True
         except Exception:
             return False
 
     def _create_collection(self) -> None:
-        response = self._client.create_collection(
+        client = self._get_client()
+        response = client.create_collection(
             collection_name=self.index,
             description="Created by Volcengine Agent Development Kit VeADK",
             builtin_event_types=["sys_event_v1"],
         )
         return response
 
+    def _get_client(self) -> VikingDBMemoryClient:
+        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
+
+        return VikingDBMemoryClient(
+            ak=self.volcengine_access_key,
+            sk=self.volcengine_secret_key,
+            sts_token=self.session_token,
+            region=self.region,
+        )
+
     @override
     def save_memory(self, event_strings: list[str], **kwargs) -> bool:
         user_id = kwargs.get("user_id")
@@ -101,7 +114,9 @@ def save_memory(self, event_strings: list[str], **kwargs) -> bool:
             "default_assistant_id": "assistant",
             "time": int(time.time() * 1000),
         }
-        response = self._client.add_messages(
+
+        client = self._get_client()
+        response = client.add_messages(
             collection_name=self.index,
             messages=messages,
             metadata=metadata,
@@ -122,7 +137,9 @@ def search_memory(self, query: str, top_k: int, **kwargs) -> list[str]:
             "user_id": user_id,
             "memory_type": ["sys_event_v1"],
         }
-        response = self._client.search_memory(
+
+        client = self._get_client()
+        response = client.search_memory(
             collection_name=self.index, query=query, filter=filter, limit=top_k
         )