veadk-python 0.2.27__py3-none-any.whl

veadk/agents/sequential_agent.py

@@ -0,0 +1,64 @@
+# 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 __future__ import annotations
+
+from google.adk.agents import SequentialAgent as GoogleADKSequentialAgent
+from google.adk.agents.base_agent import BaseAgent
+from pydantic import ConfigDict, Field
+from typing_extensions import Any
+
+from veadk.prompts.agent_default_prompt import DEFAULT_DESCRIPTION, DEFAULT_INSTRUCTION
+from veadk.tracing.base_tracer import BaseTracer
+from veadk.utils.logger import get_logger
+from veadk.utils.patches import patch_asyncio
+
+patch_asyncio()
+logger = get_logger(__name__)
+
+
+class SequentialAgent(GoogleADKSequentialAgent):
+    """Sequential Agent that executes sub-agents sequentially.
+
+    This agent is designed to execute multiple sub-agents in a predefined sequential order.
+    It ensures that each sub-agent is executed one after the other, making it suitable for
+    workflows where the output of one sub-agent is needed as input for the next. The agent
+    is well-suited for tasks that require a linear progression of steps or operations, ensuring
+    that the execution flow is maintained.
+
+    Attributes:
+        model_config (ConfigDict): Configuration dictionary for the model.
+        name (str): The name of the agent, default is "veSequentialAgent".
+        description (str): A description of the agent, useful in A2A scenarios.
+        instruction (str): Instructions or principles for function calling and agent execution.
+        sub_agents (list[BaseAgent]): A list of sub-agents managed by the sequential agent.
+            Each sub-agent is executed in the order they are listed.
+        tracers (list[BaseTracer]): A list of tracers used for monitoring the agent's performance
+            and behavior during sequential execution.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
+
+    name: str = "veSequentialAgent"
+    description: str = DEFAULT_DESCRIPTION
+    instruction: str = DEFAULT_INSTRUCTION
+
+    sub_agents: list[BaseAgent] = Field(default_factory=list, exclude=True)
+
+    tracers: list[BaseTracer] = []
+
+    def model_post_init(self, __context: Any) -> None:
+        super().model_post_init(None)  # for sub_agents init
+
+        logger.info(f"{self.__class__.__name__} `{self.name}` init done.")

veadk/auth/__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/auth/base_auth.py

@@ -0,0 +1,22 @@
+# 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.
+
+
+class BaseAuth:
+    def __init__(self) -> None: ...
+
+    def _fetch_token(self) -> str | dict: ...
+
+    @property
+    def token(self) -> str | dict: ...

veadk/auth/ve_credential_service.py

@@ -0,0 +1,203 @@
+# 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.
+
+"""A2A Credential Service for VeADK.
+
+This module provides a credential service that supports direct credential
+management by app_name and user_id, extending the ADK BaseCredentialService.
+"""
+
+import logging
+from typing import Optional
+
+from typing_extensions import override
+
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_tool import AuthConfig
+from google.adk.auth.credential_service.base_credential_service import (
+    BaseCredentialService,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class VeCredentialService(BaseCredentialService):
+    """In-memory credential service with direct app_name/user_id access.
+
+    This service extends BaseCredentialService to support both:
+    1. Standard ADK credential operations (load_credential, save_credential)
+    2. Direct credential access by app_name and user_id
+
+    The credential store is organized as:
+    {
+        app_name: {
+            user_id: {
+                credential_key: AuthCredential
+            }
+        }
+    }
+
+    Examples:
+        ```python
+        # Create service
+        service = VeCredentialService()
+
+        # Direct set/get
+        await service.set_credential(
+            app_name="my_app",
+            user_id="user123",
+            credential_key="bearer_token",
+            credential=AuthCredential(...)
+        )
+
+        credential = await service.get_credential(
+            app_name="my_app",
+            user_id="user123",
+            credential_key="bearer_token"
+        )
+
+        # Standard ADK operations
+        await service.save_credential(auth_config, callback_context)
+        credential = await service.load_credential(auth_config, callback_context)
+        ```
+    """
+
+    def __init__(self):
+        """Initialize the credential service with empty storage."""
+        super().__init__()
+        self._credentials: dict[str, dict[str, dict[str, AuthCredential]]] = {}
+
+    @override
+    async def load_credential(
+        self,
+        auth_config: AuthConfig,
+        callback_context: CallbackContext,
+    ) -> Optional[AuthCredential]:
+        """Load credential from the store using auth config and callback context.
+
+        Args:
+            auth_config: Auth configuration containing credential_key
+            callback_context: Callback context containing app_name and user_id
+
+        Returns:
+            The stored AuthCredential, or None if not found
+        """
+        app_name = callback_context._invocation_context.app_name
+        user_id = callback_context._invocation_context.user_id
+
+        return await self.get_credential(
+            app_name=app_name,
+            user_id=user_id,
+            credential_key=auth_config.credential_key,
+        )
+
+    @override
+    async def save_credential(
+        self,
+        auth_config: AuthConfig,
+        callback_context: CallbackContext,
+    ) -> None:
+        """Save credential to the store using auth config and callback context.
+
+        Args:
+            auth_config: Auth configuration containing credential_key and exchanged_auth_credential
+            callback_context: Callback context containing app_name and user_id
+        """
+        app_name = callback_context._invocation_context.app_name
+        user_id = callback_context._invocation_context.user_id
+
+        await self.set_credential(
+            app_name=app_name,
+            user_id=user_id,
+            credential_key=auth_config.credential_key,
+            credential=auth_config.exchanged_auth_credential,
+        )
+
+    async def set_credential(
+        self,
+        app_name: str,
+        user_id: str,
+        credential_key: str,
+        credential: AuthCredential,
+    ) -> None:
+        """Directly set a credential by app_name, user_id, and credential_key (async).
+
+        This method allows setting credentials without requiring a CallbackContext,
+        useful for middleware and request interceptors.
+
+        Args:
+            app_name: Application name
+            user_id: User identifier
+            credential_key: Key to identify the credential
+            credential: The AuthCredential to store
+
+        Examples:
+            ```python
+            from google.adk.auth.auth_credential import AuthCredential, AuthCredentialTypes
+
+            service = VeCredentialService()
+            await service.set_credential(
+                app_name="my_app",
+                user_id="user123",
+                credential_key="bearer_token",
+                credential=AuthCredential(
+                    auth_type=AuthCredentialTypes.BEARER_TOKEN,
+                    bearer_token="eyJhbGc..."
+                )
+            )
+            ```
+        """
+        if app_name not in self._credentials:
+            self._credentials[app_name] = {}
+        if user_id not in self._credentials[app_name]:
+            self._credentials[app_name][user_id] = {}
+
+        self._credentials[app_name][user_id][credential_key] = credential
+        logger.debug(
+            f"Set credential for app={app_name}, user={user_id}, key={credential_key}"
+        )
+
+    async def get_credential(
+        self,
+        app_name: str,
+        user_id: str,
+        credential_key: str,
+    ) -> Optional[AuthCredential]:
+        """Directly get a credential by app_name, user_id, and credential_key (async).
+
+        This method allows retrieving credentials without requiring a CallbackContext,
+        useful for middleware and request interceptors.
+
+        Args:
+            app_name: Application name
+            user_id: User identifier
+            credential_key: Key to identify the credential
+
+        Returns:
+            The stored AuthCredential, or None if not found
+
+        Examples:
+            ```python
+            service = VeCredentialService()
+            credential = await service.get_credential(
+                app_name="my_app",
+                user_id="user123",
+                credential_key="bearer_token"
+            )
+            if credential:
+                print(f"Found token: {credential.bearer_token}")
+            ```
+        """
+        return self._credentials.get(app_name, {}).get(user_id, {}).get(credential_key)

veadk/auth/veauth/__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/auth/veauth/apmplus_veauth.py

@@ -0,0 +1,58 @@
+# 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 os
+
+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__)
+
+
+def get_apmplus_token(region: str = "cn-beijing") -> str:
+    logger.info("Fetching APMPlus token...")
+
+    access_key = os.getenv("VOLCENGINE_ACCESS_KEY")
+    secret_key = os.getenv("VOLCENGINE_SECRET_KEY")
+    session_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={},
+        header={
+            "X-Security-Token": session_token,
+            # APMPlus frontend required
+            "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

@@ -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.
+
+import os
+
+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__)
+
+
+def get_ark_token(region: str = "cn-beijing") -> str:
+    logger.info("Fetching ARK token...")
+
+    access_key = os.getenv("VOLCENGINE_ACCESS_KEY")
+    secret_key = os.getenv("VOLCENGINE_SECRET_KEY")
+    session_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={"ProjectName": "default", "Filter": {"AllowAll": True}},
+        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"]
+        logger.warning("By default, VeADK fetches the first API Key in the list.")
+        logger.info(
+            f"Try to fetch ARK API Key with id={first_api_key_id}, name={res['Result']['Items'][0]['Name']}"
+        )
+    except KeyError:
+        raise ValueError(f"Failed to get ARK api key list: {res}")
+
+    # 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:
+        api_key = res["Result"]["ApiKey"]
+        logger.info(f"Successfully fetched ARK API Key (starts with {api_key[:8]}).")
+        return api_key
+    except KeyError:
+        raise ValueError(f"Failed to get ARK api key: {res}")

veadk/auth/veauth/base_veauth.py

@@ -0,0 +1,50 @@
+# 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 os
+from abc import ABC, abstractmethod
+
+from veadk.auth.base_auth import BaseAuth
+
+
+class BaseVeAuth(ABC, BaseAuth):
+    volcengine_access_key: str
+    """Volcengine Access Key"""
+
+    volcengine_secret_key: str
+    """Volcengine Secret Key"""
+
+    def __init__(
+        self,
+        access_key: str | None = None,
+        secret_key: str | None = None,
+    ) -> None:
+        super().__init__()
+
+        final_ak = access_key or os.getenv("VOLCENGINE_ACCESS_KEY")
+        final_sk = secret_key or os.getenv("VOLCENGINE_SECRET_KEY")
+
+        assert final_ak, "Volcengine access key cannot be empty."
+        assert final_sk, "Volcengine secret key cannot be empty."
+
+        self.access_key = final_ak
+        self.secret_key = final_sk
+
+        self._token: str = ""
+
+    @abstractmethod
+    def _fetch_token(self) -> None: ...
+
+    @property
+    def token(self) -> str: ...

veadk/auth/veauth/cozeloop_veauth.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/auth/veauth/opensearch_veauth.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.
+
+# 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 os
+
+from typing_extensions import override
+
+from veadk.auth.veauth.base_veauth import BaseVeAuth
+from veadk.utils.logger import get_logger
+
+# from veadk.utils.volcengine_sign import ve_request
+
+logger = get_logger(__name__)
+
+
+class OpensearchVeAuth(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)
+
+        self._token: str = ""
+
+    @override
+    def _fetch_token(self) -> None:
+        logger.info("Fetching Opensearch STS token...")
+
+        # res = ve_request(
+        #     request_body={},
+        #     action="GetOrCreatePromptPilotAPIKeys",
+        #     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"]["APIKeys"][0]["APIKey"]
+        # except KeyError:
+        #     raise ValueError(f"Failed to get Prompt Pilot token: {res}")
+
+    @property
+    def token(self) -> str:
+        if self._token:
+            return self._token
+        self._fetch_token()
+        return self._token

veadk/auth/veauth/postgresql_veauth.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.
+
+# 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 os
+
+from typing_extensions import override
+
+from veadk.auth.veauth.base_veauth import BaseVeAuth
+from veadk.utils.logger import get_logger
+
+# from veadk.utils.volcengine_sign import ve_request
+
+logger = get_logger(__name__)
+
+
+class PostgreSqlVeAuth(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)
+
+        self._token: str = ""
+
+    @override
+    def _fetch_token(self) -> None:
+        logger.info("Fetching PostgreSQL STS token...")
+
+        # res = ve_request(
+        #     request_body={},
+        #     action="GetOrCreatePromptPilotAPIKeys",
+        #     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"]["APIKeys"][0]["APIKey"]
+        # except KeyError:
+        #     raise ValueError(f"Failed to get Prompt Pilot token: {res}")
+
+    @property
+    def token(self) -> str:
+        if self._token:
+            return self._token
+        self._fetch_token()
+        return self._token