update auto auth

Author: yaozheng-fang

tests/test_agent.py

@@ -199,17 +199,6 @@ def test_agent_with_tracers():
     assert tracer2 in agent.tracers
 
 
-@patch.dict(
-    "os.environ",
-    {"MODEL_AGENT_NAME": "env_model_name", "MODEL_AGENT_API_KEY": "mock_api_key"},
-    clear=True,
-)
-def test_agent_environment_variables():
-    agent = Agent()
-    print(agent)
-    assert agent.model_name == "env_model_name"
-
-
 @patch.dict("os.environ", {"MODEL_AGENT_API_KEY": "mock_api_key"})
 def test_agent_custom_name_and_description():
     custom_name = "CustomAgent"

tests/test_tos.py

@@ -26,12 +26,9 @@
 from pydantic import ConfigDict, Field
 from typing_extensions import Any
 
-from veadk.config import getenv
+from veadk.config import settings
 from veadk.consts import (
     DEFAULT_AGENT_NAME,
-    DEFAULT_MODEL_AGENT_API_BASE,
-    DEFAULT_MODEL_AGENT_NAME,
-    DEFAULT_MODEL_AGENT_PROVIDER,
     DEFAULT_MODEL_EXTRA_CONFIG,
 )
 from veadk.evaluation import EvalSetRecorder
@@ -63,26 +60,16 @@ class Agent(LlmAgent):
     instruction: str = DEFAULT_INSTRUCTION
     """The instruction for the agent, such as principles of function calling."""
 
-    model_name: str = Field(
-        default_factory=lambda: getenv("MODEL_AGENT_NAME", DEFAULT_MODEL_AGENT_NAME)
-    )
+    model_name: str = Field(default_factory=lambda: settings.model.name)
     """The name of the model for agent running."""
 
-    model_provider: str = Field(
-        default_factory=lambda: getenv(
-            "MODEL_AGENT_PROVIDER", DEFAULT_MODEL_AGENT_PROVIDER
-        )
-    )
+    model_provider: str = Field(default_factory=lambda: settings.model.provider)
     """The provider of the model for agent running."""
 
-    model_api_base: str = Field(
-        default_factory=lambda: getenv(
-            "MODEL_AGENT_API_BASE", DEFAULT_MODEL_AGENT_API_BASE
-        )
-    )
+    model_api_base: str = Field(default_factory=lambda: settings.model.api_base)
     """The api base of the model for agent running."""
 
-    model_api_key: str = Field(default_factory=lambda: getenv("MODEL_AGENT_API_KEY"))
+    model_api_key: str = Field(default_factory=lambda: settings.model.api_key)
     """The api key of the model for agent running."""
 
     model_extra_config: dict = Field(default_factory=dict)

veadk/agent.py

@@ -28,9 +28,12 @@ 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
+
         self._token: str = ""
 
     @override
@@ -44,8 +47,10 @@ def _fetch_token(self) -> None:
             sk=self.secret_key,
             service="apmplus_server",
             version="2024-07-30",
-            region="cn-beijing",
+            region=self.region,
             host="open.volcengineapi.com",
+            # APMPlus frontend required
+            header={"X-Apmplus-Region": self.region.replace("-", "_")},
         )
         try:
             self._token = res["data"]["app_key"]

veadk/auth/veauth/apmplus_veauth.py

@@ -35,6 +35,7 @@ def __init__(
 
     @override
     def _fetch_token(self) -> None:
+        logger.info("Fetching ARK token...")
         # list api keys
         first_api_key_id = ""
         res = ve_request(

veadk/auth/veauth/ark_veauth.py

@@ -14,7 +14,6 @@
 
 import os
 from abc import ABC, abstractmethod
-from typing import Type
 
 from veadk.auth.base_auth import BaseAuth
 
@@ -49,36 +48,3 @@ def _fetch_token(self) -> None: ...
 
     @property
     def token(self) -> str: ...
-
-
-def veauth(auth_token_name: str, auth_cls: Type[BaseVeAuth]):
-    def decorator(cls: Type):
-        # api_key -> _api_key
-        # for cache
-        private_auth_token_name = f"_{auth_token_name}"
-        setattr(cls, private_auth_token_name, "")
-
-        # init a auth cls for fetching token
-        auth_cls_instance = "_auth_cls_instance"
-        setattr(cls, auth_cls_instance, auth_cls())
-
-        def getattribute(self, name: str):
-            if name != auth_token_name:
-                return object.__getattribute__(self, name)
-            if name == auth_token_name:
-                token = object.__getattribute__(self, name)
-
-                if token:
-                    return token
-                elif not token and not getattr(cls, private_auth_token_name):
-                    token = getattr(cls, auth_cls_instance).token
-                    setattr(cls, private_auth_token_name, token)
-                    return token
-                elif not token and getattr(cls, private_auth_token_name):
-                    return getattr(cls, private_auth_token_name)
-            return token
-
-        setattr(cls, "__getattribute__", getattribute)
-        return cls
-
-    return decorator

veadk/auth/veauth/base_veauth.py

@@ -31,7 +31,7 @@ def __init__(
     ) -> None:
         super().__init__(access_key, secret_key)
 
-        self._token: str | dict = ""
+        self._token: str = ""
 
     @override
     def _fetch_token(self) -> None:
@@ -53,7 +53,7 @@ def _fetch_token(self) -> None:
             raise ValueError(f"Failed to get Prompt Pilot token: {res}")
 
     @property
-    def token(self) -> str | dict:
+    def token(self) -> str:
         if self._token:
             return self._token
         self._fetch_token()

veadk/auth/veauth/prompt_pilot_veauth.py

@@ -14,6 +14,7 @@
 
 
 import click
+
 from veadk.version import VERSION
 
 TEMP_PATH = "/tmp"

veadk/cli/cli_deploy.py

@@ -31,7 +31,7 @@ def prompt(path: str, feedback: str, api_key: str, model_name: str) -> None:
     from pathlib import Path
 
     from veadk.agent import Agent
-    from veadk.config import getenv
+    from veadk.config import settings
     from veadk.integrations.ve_prompt_pilot.ve_prompt_pilot import VePromptPilot
     from veadk.utils.misc import load_module_from_file
 
@@ -55,7 +55,7 @@ def prompt(path: str, feedback: str, api_key: str, model_name: str) -> None:
         click.echo(f"Found {len(agents)} agents in {module_abs_path}")
 
         if not api_key:
-            api_key = getenv("PROMPT_PILOT_API_KEY")
+            api_key = settings.prompt_pilot.api_key
         ve_prompt_pilot = VePromptPilot(api_key)
         ve_prompt_pilot.optimize(
             agents=agents, feedback=feedback, model_name=model_name