feat(auth): support apmplus auto auth (#270)

Author: yaozheng-fang

veadk/auth/veauth/apmplus_veauth.py

@@ -14,52 +14,45 @@
 
 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 APMPlusVeAuth(BaseVeAuth):
-    def __init__(
-        self,
-        access_key: str = os.getenv("VOLCENGINE_ACCESS_KEY", ""),
-        secret_key: str = os.getenv("VOLCENGINE_SECRET_KEY", ""),
-        region: str = "cn-beijing",
-    ) -> None:
-        super().__init__(access_key, secret_key)
-
-        self.region = region
+def get_apmplus_token(region: str = "cn-beijing") -> str:
+    logger.info("Fetching APMPlus 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 APMPlus token...")
+    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
 
-        res = ve_request(
-            request_body={},
-            action="GetAppKey",
-            ak=self.access_key,
-            sk=self.secret_key,
-            service="apmplus_server",
-            version="2024-07-30",
-            region=self.region,
-            host="open.volcengineapi.com",
+    res = ve_request(
+        request_body={},
+        header={
+            "X-Security-Token": session_token,
             # APMPlus frontend required
-            header={"X-Apmplus-Region": self.region.replace("-", "_")},
-        )
-        try:
-            self._token = res["data"]["app_key"]
-        except KeyError:
-            raise ValueError(f"Failed to get APMPlus token: {res}")
-
-    @property
-    def token(self) -> str:
-        if self._token:
-            return self._token
-        self._fetch_token()
-        return self._token
+            "X-Apmplus-Region": region.replace("-", "_"),
+        },
+        action="GetAppKey",
+        ak=access_key,
+        sk=secret_key,
+        service="apmplus_server",
+        version="2024-07-30",
+        region=region,
+        host="open.volcengineapi.com",
+    )
+    try:
+        api_key = res["data"]["app_key"]
+        logger.info("Successfully fetching APMPlus API Key.")
+        return api_key
+    except KeyError:
+        raise ValueError(f"Failed to get APMPlus token: {res}")

veadk/auth/veauth/ark_veauth.py

@@ -64,6 +64,8 @@ def get_ark_token(region: str = "cn-beijing") -> str:
         host="open.volcengineapi.com",
     )
     try:
-        return res["Result"]["ApiKey"]
+        api_key = res["Result"]["ApiKey"]
+        logger.info("Successfully fetching ARK API Key.")
+        return api_key
     except KeyError:
         raise ValueError(f"Failed to get ARK api key: {res}")

veadk/auth/veauth/utils.py

@@ -50,6 +50,9 @@ def get_credential_from_vefaas_iam() -> VeIAMCredential:
         access_key = cred_dict["access_key_id"]
         secret_key = cred_dict["secret_access_key"]
         session_token = cred_dict["session_token"]
+
+        logger.info("Get credential from IAM file successfully.")
+
         return VeIAMCredential(
             access_key_id=access_key,
             secret_access_key=secret_key,

veadk/cli/cli_web.py

@@ -132,6 +132,12 @@ async def wrapper(*args, **kwargs) -> ADKRunner:
 
     from google.adk.cli.cli_tools_click import cli_web
 
-    extra_args = ctx.args
-    logger.debug(f"User args: {ctx.args}")
+    extra_args: list = ctx.args
+    logger.debug(f"User args: {extra_args}")
+
+    # set a default log level to avoid unnecessary outputs
+    # from Google ADK and Litellm
+    if "--log_level" not in extra_args:
+        extra_args.extend(["--log_level", "ERROR"])
+
     cli_web.main(args=extra_args, standalone_mode=False)

veadk/configs/model_configs.py

@@ -56,7 +56,11 @@ class EmbeddingModelConfig(BaseSettings):
 
     @cached_property
     def api_key(self) -> str:
-        return os.getenv("MODEL_EMBEDDING_API_KEY") or get_ark_token()
+        return (
+            os.getenv("MODEL_EMBEDDING_API_KEY")
+            or os.getenv("MODEL_AGENT_API_KEY")  # try to use agent's model api key
+            or get_ark_token()
+        )
 
 
 class NormalEmbeddingModelConfig(BaseSettings):

veadk/configs/tracing_configs.py

@@ -18,7 +18,7 @@
 from pydantic import Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
-from veadk.auth.veauth.apmplus_veauth import APMPlusVeAuth
+from veadk.auth.veauth.apmplus_veauth import get_apmplus_token
 from veadk.consts import (
     DEFAULT_APMPLUS_OTEL_EXPORTER_ENDPOINT,
     DEFAULT_APMPLUS_OTEL_EXPORTER_SERVICE_NAME,
@@ -46,7 +46,7 @@ class APMPlusConfig(BaseSettings):
     def otel_exporter_api_key(self) -> str:
         return (
             os.getenv("OBSERVABILITY_OPENTELEMETRY_APMPLUS_API_KEY")
-            or APMPlusVeAuth().token
+            or get_apmplus_token()
         )