feat(auth): support sts token auth (#236)

* feat(auth): support sts token auth
* add auto token for vikingdb memory

Author: yaozheng-fang

veadk/auth/veauth/ark_veauth.py

@@ -14,64 +14,56 @@
 
 import os
 
-from typing_extensions import override
-
-from veadk.auth.veauth.base_veauth import BaseVeAuth
+from veadk.auth.veauth.utils import get_credential_from_vefaas_iam
 from veadk.utils.logger import get_logger
 from veadk.utils.volcengine_sign import ve_request
 
 logger = get_logger(__name__)
 
 
-class ARKVeAuth(BaseVeAuth):
-    def __init__(
-        self,
-        access_key: str = os.getenv("VOLCENGINE_ACCESS_KEY", ""),
-        secret_key: str = os.getenv("VOLCENGINE_SECRET_KEY", ""),
-    ) -> None:
-        super().__init__(access_key, secret_key)
+def get_ark_token(region: str = "cn-beijing") -> str:
+    logger.info("Fetching ARK token...")
 
-        self._token: str = ""
+    access_key = os.getenv("VOLCENGINE_ACCESS_KEY")
+    secret_key = os.getenv("VOLCENGINE_SECRET_KEY")
+    session_token = ""
 
-    @override
-    def _fetch_token(self) -> None:
-        logger.info("Fetching ARK token...")
-        # list api keys
-        first_api_key_id = ""
-        res = ve_request(
-            request_body={"ProjectName": "default", "Filter": {}},
-            action="ListApiKeys",
-            ak=self.access_key,
-            sk=self.secret_key,
-            service="ark",
-            version="2024-01-01",
-            region="cn-beijing",
-            host="open.volcengineapi.com",
-        )
-        try:
-            first_api_key_id = res["Result"]["Items"][0]["Id"]
-        except KeyError:
-            raise ValueError(f"Failed to get ARK api key list: {res}")
+    if not (access_key and secret_key):
+        # try to get from vefaas iam
+        cred = get_credential_from_vefaas_iam()
+        access_key = cred.access_key_id
+        secret_key = cred.secret_access_key
+        session_token = cred.session_token
 
-        # get raw api key
-        res = ve_request(
-            request_body={"Id": first_api_key_id},
-            action="GetRawApiKey",
-            ak=self.access_key,
-            sk=self.secret_key,
-            service="ark",
-            version="2024-01-01",
-            region="cn-beijing",
-            host="open.volcengineapi.com",
-        )
-        try:
-            self._token = res["Result"]["ApiKey"]
-        except KeyError:
-            raise ValueError(f"Failed to get ARK api key: {res}")
+    res = ve_request(
+        request_body={"ProjectName": "default", "Filter": {}},
+        header={"X-Security-Token": session_token},
+        action="ListApiKeys",
+        ak=access_key,
+        sk=secret_key,
+        service="ark",
+        version="2024-01-01",
+        region=region,
+        host="open.volcengineapi.com",
+    )
+    try:
+        first_api_key_id = res["Result"]["Items"][0]["Id"]
+    except KeyError:
+        raise ValueError(f"Failed to get ARK api key list: {res}")
 
-    @property
-    def token(self) -> str:
-        if self._token:
-            return self._token
-        self._fetch_token()
-        return self._token
+    # get raw api key
+    res = ve_request(
+        request_body={"Id": first_api_key_id},
+        header={"X-Security-Token": session_token},
+        action="GetRawApiKey",
+        ak=access_key,
+        sk=secret_key,
+        service="ark",
+        version="2024-01-01",
+        region=region,
+        host="open.volcengineapi.com",
+    )
+    try:
+        return res["Result"]["ApiKey"]
+    except KeyError:
+        raise ValueError(f"Failed to get ARK api key: {res}")

veadk/auth/veauth/utils.py

@@ -0,0 +1,57 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+from pathlib import Path
+
+from pydantic import BaseModel
+
+from veadk.consts import VEFAAS_IAM_CRIDENTIAL_PATH
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class VeIAMCredential(BaseModel):
+    access_key_id: str
+    secret_access_key: str
+    session_token: str
+
+
+def get_credential_from_vefaas_iam() -> VeIAMCredential:
+    """Get credential from VeFaaS IAM file"""
+    logger.info(
+        f"Get Volcegnine access key or secret key from environment variables failed, try to get from VeFaaS IAM file (path={VEFAAS_IAM_CRIDENTIAL_PATH})."
+    )
+
+    path = Path(VEFAAS_IAM_CRIDENTIAL_PATH)
+
+    if not path.exists():
+        logger.error(
+            f"Get Volcegnine access key or secret key from environment variables failed, and VeFaaS IAM file (path={VEFAAS_IAM_CRIDENTIAL_PATH}) not exists. Please check your configuration."
+        )
+        raise FileNotFoundError(
+            f"Get Volcegnine access key or secret key from environment variables failed, and VeFaaS IAM file (path={VEFAAS_IAM_CRIDENTIAL_PATH}) not exists. Please check your configuration."
+        )
+
+    with open(VEFAAS_IAM_CRIDENTIAL_PATH, "r") as f:
+        cred_dict = json.load(f)
+        access_key = cred_dict["access_key_id"]
+        secret_key = cred_dict["secret_access_key"]
+        session_token = cred_dict["session_token"]
+        return VeIAMCredential(
+            access_key_id=access_key,
+            secret_access_key=secret_key,
+            session_token=session_token,
+        )

veadk/configs/model_configs.py

@@ -17,7 +17,7 @@
 
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
-from veadk.auth.veauth.ark_veauth import ARKVeAuth
+from veadk.auth.veauth.ark_veauth import get_ark_token
 from veadk.consts import (
     DEFAULT_MODEL_AGENT_API_BASE,
     DEFAULT_MODEL_AGENT_NAME,
@@ -39,7 +39,7 @@ class ModelConfig(BaseSettings):
 
     @cached_property
     def api_key(self) -> str:
-        return os.getenv("MODEL_AGENT_API_KEY") or ARKVeAuth().token
+        return os.getenv("MODEL_AGENT_API_KEY") or get_ark_token()
 
 
 class EmbeddingModelConfig(BaseSettings):
@@ -56,7 +56,7 @@ class EmbeddingModelConfig(BaseSettings):
 
     @cached_property
     def api_key(self) -> str:
-        return os.getenv("MODEL_EMBEDDING_API_KEY") or ARKVeAuth().token
+        return os.getenv("MODEL_EMBEDDING_API_KEY") or get_ark_token()
 
 
 class NormalEmbeddingModelConfig(BaseSettings):

veadk/consts.py

@@ -66,3 +66,5 @@
 DEFAULT_IMAGE_EDIT_MODEL_NAME = "doubao-seededit-3-0-i2i-250628"
 DEFAULT_VIDEO_MODEL_NAME = "doubao-seedance-1-0-pro-250528"
 DEFAULT_IMAGE_GENERATE_MODEL_NAME = "doubao-seedream-4-0-250828"
+
+VEFAAS_IAM_CRIDENTIAL_PATH = "/var/run/secrets/iam/credential"

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,24 +36,20 @@
 
 
 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"""
 
     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()
@@ -69,19 +66,35 @@ def precheck_index_naming(self):
 
     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, user_id: str, event_strings: list[str], **kwargs) -> bool:
         session_id = str(uuid.uuid1())
@@ -103,7 +116,8 @@ def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool:
             f"Request for add {len(messages)} memory to VikingDB: collection_name={self.index}, metadata={metadata}, session_id={session_id}"
         )
 
-        response = self._client.add_messages(
+        client = self._get_client()
+        response = client.add_messages(
             collection_name=self.index,
             messages=messages,
             metadata=metadata,
@@ -130,7 +144,8 @@ def search_memory(
             f"Request for search memory in VikingDB: filter={filter}, collection_name={self.index}, query={query}, limit={top_k}"
         )
 
-        response = self._client.search_memory(
+        client = self._get_client()
+        response = client.search_memory(
             collection_name=self.index, query=query, filter=filter, limit=top_k
         )
 

veadk/utils/volcengine_sign.py

@@ -144,6 +144,8 @@ def request(method, date, query, header, ak, sk, action, body):
         )
     )
     header = {**header, **sign_result}
+    if "X-Security-Token" in header and header["X-Security-Token"] == "":
+        del header["X-Security-Token"]
     # header = {**header, **{"X-Security-Token": SessionToken}}
     # 第六步：将 Signature 签名写入 HTTP Header 中，并发送 HTTP 请求。
     r = requests.request(