veadk-python 0.2.9__py3-none-any.whl → 0.2.11__py3-none-any.whl

veadk/a2a/remote_ve_agent.py

@@ -13,21 +13,78 @@
 # limitations under the License.
 
 import json
+from typing import Literal
 
+import httpx
 import requests
 from a2a.types import AgentCard
 from google.adk.agents.remote_a2a_agent import RemoteA2aAgent
 
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
 AGENT_CARD_WELL_KNOWN_PATH = "/.well-known/agent-card.json"
 
 
+def _convert_agent_card_dict_to_obj(agent_card_dict: dict) -> AgentCard:
+    agent_card_json_str = json.dumps(agent_card_dict, ensure_ascii=False, indent=2)
+    agent_card_object = AgentCard.model_validate_json(str(agent_card_json_str))
+    return agent_card_object
+
+
 class RemoteVeAgent(RemoteA2aAgent):
-    def __init__(self, name: str, url: str):
-        agent_card_dict = requests.get(url + AGENT_CARD_WELL_KNOWN_PATH).json()
-        agent_card_dict["url"] = url
+    """Connect to remote agent on VeFaaS platform."""
+
+    def __init__(
+        self,
+        name: str,
+        url: str,
+        auth_token: str | None = None,
+        auth_method: Literal["header", "querystring"] | None = None,
+    ):
+        if not auth_token:
+            agent_card_dict = requests.get(url + AGENT_CARD_WELL_KNOWN_PATH).json()
+            # replace agent_card_url with actual host
+            agent_card_dict["url"] = url
+
+            agent_card_object = _convert_agent_card_dict_to_obj(agent_card_dict)
+
+            logger.debug(f"Agent card of {name}: {agent_card_object}")
+            super().__init__(name=name, agent_card=agent_card_object)
+        else:
+            if auth_method == "header":
+                headers = {"Authorization": f"Bearer {auth_token}"}
+                agent_card_dict = requests.get(
+                    url + AGENT_CARD_WELL_KNOWN_PATH, headers=headers
+                ).json()
+                agent_card_dict["url"] = url
+
+                agent_card_object = _convert_agent_card_dict_to_obj(agent_card_dict)
+                httpx_client = httpx.AsyncClient(
+                    base_url=url, headers=headers, timeout=600
+                )
 
-        agent_card_json_str = json.dumps(agent_card_dict, ensure_ascii=False, indent=2)
+                logger.debug(f"Agent card of {name}: {agent_card_object}")
+                super().__init__(
+                    name=name, agent_card=agent_card_object, httpx_client=httpx_client
+                )
+            elif auth_method == "querystring":
+                agent_card_dict = requests.get(
+                    url + AGENT_CARD_WELL_KNOWN_PATH + f"?token={auth_token}"
+                ).json()
+                agent_card_dict["url"] = url
 
-        agent_card_object = AgentCard.model_validate_json(str(agent_card_json_str))
+                agent_card_object = _convert_agent_card_dict_to_obj(agent_card_dict)
+                httpx_client = httpx.AsyncClient(
+                    base_url=url, params={"token": auth_token}, timeout=600
+                )
 
-        super().__init__(name=name, agent_card=agent_card_object)
+                logger.debug(f"Agent card of {name}: {agent_card_object}")
+                super().__init__(
+                    name=name, agent_card=agent_card_object, httpx_client=httpx_client
+                )
+            else:
+                raise ValueError(
+                    f"Unsupported auth method {auth_method}, use `header` or `querystring` instead."
+                )

veadk/agent.py

@@ -133,14 +133,21 @@ class Agent(LlmAgent):
             )
 
         if self.knowledgebase:
-            from veadk.tools import load_knowledgebase_tool
+            from veadk.tools.builtin_tools.load_knowledgebase import (
+                LoadKnowledgebaseTool,
+            )
 
-            load_knowledgebase_tool.knowledgebase = self.knowledgebase
-            self.tools.append(load_knowledgebase_tool.load_knowledgebase_tool)
+            load_knowledgebase_tool = LoadKnowledgebaseTool(
+                knowledgebase=self.knowledgebase
+            )
+            self.tools.append(load_knowledgebase_tool)
 
         if self.long_term_memory is not None:
             from google.adk.tools import load_memory
 
+            if not load_memory.custom_metadata:
+                load_memory.custom_metadata = {}
+            load_memory.custom_metadata["backend"] = self.long_term_memory.backend
             self.tools.append(load_memory)
 
         logger.info(f"VeADK version: {VERSION}")

veadk/agent_builder.py

@@ -52,9 +52,8 @@ class AgentBuilder:
         tools = []
         if agent_config.get("tools", []):
             for tool in agent_config["tools"]:
-                module_name = tool["module"]
-                func_name = tool["func"]
-
+                name = tool["name"]
+                module_name, func_name = name.rsplit(".", 1)
                 module = importlib.import_module(module_name)
                 func = getattr(module, func_name)
 

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/cli/cli.py

@@ -18,6 +18,7 @@ import click
 from veadk.cli.cli_deploy import deploy
 from veadk.cli.cli_eval import eval
 from veadk.cli.cli_init import init
+from veadk.cli.cli_kb import kb
 from veadk.cli.cli_pipeline import pipeline
 from veadk.cli.cli_prompt import prompt
 from veadk.cli.cli_web import web
@@ -39,6 +40,7 @@ veadk.add_command(prompt)
 veadk.add_command(web)
 veadk.add_command(pipeline)
 veadk.add_command(eval)
+veadk.add_command(kb)
 
 if __name__ == "__main__":
     veadk()

veadk/cli/cli_kb.py

@@ -0,0 +1,75 @@
+# 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.
+
+from pathlib import Path
+from typing import Literal
+
+import click
+
+
+@click.command()
+@click.option(
+    "--backend",
+    type=click.Choice(
+        ["local", "opensearch", "viking", "redis"],
+        case_sensitive=False,
+    ),
+    required=True,
+)
+@click.option(
+    "--app_name",
+    default="",
+    help="`app_name` for init your knowledgebase",
+)
+@click.option(
+    "--index",
+    default="",
+    help="Knowledgebase index",
+)
+@click.option(
+    "--path",
+    required=True,
+    help="Knowledge file or directory path",
+)
+def add(
+    backend: Literal["local", "opensearch", "viking", "redis"],
+    app_name: str,
+    index: str,
+    path: str,
+):
+    """Add files to knowledgebase"""
+    _path = Path(path)
+    assert _path.exists(), f"Path {path} not exists. Please check your input."
+
+    from veadk.knowledgebase import KnowledgeBase
+
+    knowledgebase = KnowledgeBase(backend=backend, app_name=app_name, index=index)
+
+    if _path.is_file():
+        knowledgebase.add_from_files(files=[path])
+    elif _path.is_dir():
+        knowledgebase.add_from_directory(directory=path)
+    else:
+        raise RuntimeError(
+            "Unsupported knowledgebase file type, only support a single file and a directory."
+        )
+
+
+@click.group()
+def kb():
+    """VeADK Knowledgebase management"""
+    pass
+
+
+kb.add_command(add)

veadk/cli/cli_web.py

@@ -154,6 +154,7 @@ def web(host: str) -> None:
         eval_set_results_manager: Any,
         agents_dir: str,
         extra_plugins: Optional[list[str]] = None,
+        **kwargs: Any,
     ):
         self.agent_loader = agent_loader
         self.artifact_service = artifact_service
@@ -166,6 +167,9 @@ def web(host: str) -> None:
         self.runner_dict = {}
         self.extra_plugins = extra_plugins or []
 
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
         # parse VeADK memories
         short_term_memory, long_term_memory = _get_memory(module_path=agents_dir)
         self.session_service = short_term_memory.session_service

veadk/configs/model_configs.py

@@ -17,7 +17,7 @@ from functools import cached_property
 
 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

@@ -63,6 +63,15 @@ DEFAULT_TOS_BUCKET_NAME = "ark-tutorial"
 DEFAULT_COZELOOP_SPACE_NAME = "VeADK Space"
 
 DEFAULT_TEXT_TO_IMAGE_MODEL_NAME = "doubao-seedream-3-0-t2i-250415"
+DEFAULT_TEXT_TO_IMAGE_MODEL_API_BASE = "https://ark.cn-beijing.volces.com/api/v3/"
+
 DEFAULT_IMAGE_EDIT_MODEL_NAME = "doubao-seededit-3-0-i2i-250628"
+DEFAULT_IMAGE_EDIT_MODEL_API_BASE = "https://ark.cn-beijing.volces.com/api/v3/"
+
 DEFAULT_VIDEO_MODEL_NAME = "doubao-seedance-1-0-pro-250528"
+DEFAULT_VIDEO_MODEL_API_BASE = "https://ark.cn-beijing.volces.com/api/v3/"
+
 DEFAULT_IMAGE_GENERATE_MODEL_NAME = "doubao-seedream-4-0-250828"
+DEFAULT_IMAGE_GENERATE_MODEL_API_BASE = "https://ark.cn-beijing.volces.com/api/v3/"
+
+VEFAAS_IAM_CRIDENTIAL_PATH = "/var/run/secrets/iam/credential"

veadk/integrations/__init__.py

@@ -0,0 +1,13 @@
+# 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.

veadk/integrations/ve_viking_db_memory/__init__.py

@@ -0,0 +1,13 @@
+# 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.