veadk-python 0.2.27__py3-none-any.whl

veadk/__init__.py

@@ -0,0 +1,37 @@
+# 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 typing import TYPE_CHECKING
+
+from veadk.version import VERSION
+
+if TYPE_CHECKING:
+    from veadk.agent import Agent
+    from veadk.runner import Runner
+
+
+# Lazy loading for `Agent` class
+def __getattr__(name):
+    if name == "Agent":
+        from veadk.agent import Agent
+
+        return Agent
+    if name == "Runner":
+        from veadk.runner import Runner
+
+        return Runner
+    raise AttributeError(f"module 'veadk' has no attribute '{name}'")
+
+
+__all__ = ["Agent", "Runner", "VERSION"]

veadk/a2a/__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/a2a/agent_card.py

@@ -0,0 +1,45 @@
+# 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 a2a.types import AgentCapabilities, AgentCard, AgentProvider, AgentSkill
+
+from veadk import Agent
+from veadk.version import VERSION
+
+
+def get_agent_card(
+    agent: Agent, url: str, version: str = VERSION, provider: str = "veadk"
+) -> AgentCard:
+    agent_provider = AgentProvider(organization=provider, url="")
+    agent_capabilities = AgentCapabilities()
+    agent_skills = [
+        AgentSkill(
+            id="0",
+            name="chat",
+            description="Basically chat with user.",
+            tags=["chat", "talk"],
+        )
+    ]
+    agent_card = AgentCard(
+        capabilities=agent_capabilities,
+        description=agent.description,
+        name=agent.name,
+        defaultInputModes=["text"],
+        defaultOutputModes=["text"],
+        provider=agent_provider,
+        skills=agent_skills,
+        url=url,
+        version=version,
+    )
+    return agent_card

veadk/a2a/remote_ve_agent.py

@@ -0,0 +1,390 @@
+# 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
+import functools
+from typing import AsyncGenerator, Literal, Optional
+
+from a2a.client.base_client import BaseClient
+import httpx
+import requests
+from a2a.types import AgentCard
+from google.adk.agents.remote_a2a_agent import RemoteA2aAgent
+
+from veadk.integrations.ve_identity.utils import generate_headers
+from veadk.utils.auth import VE_TIP_TOKEN_CREDENTIAL_KEY, VE_TIP_TOKEN_HEADER
+from veadk.utils.logger import get_logger
+from google.adk.utils.context_utils import Aclosing
+from google.adk.events.event import Event
+from google.adk.agents.invocation_context import InvocationContext
+
+
+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):
+    """Connect to a remote agent on the VeFaaS platform.
+
+    This class provides an interface to remotely connect with an agent deployed on the
+    VeFaaS platform. It automatically fetches the agent card (metadata) and configures
+    an HTTP client for secure communication.
+
+    The class extends `RemoteA2aAgent` to provide compatibility with the A2A
+    (Agent-to-Agent) communication layer.
+
+    This constructor handles agent discovery and HTTP client setup. It determines the
+    agent's URL, fetches its metadata (`agent_card`), and prepares an
+    `httpx.AsyncClient` for subsequent communication. You can either provide a URL
+    directly, or pass a pre-configured `httpx.AsyncClient` with a `base_url`.
+
+    Authentication can be handled via a bearer token in the HTTP header or via a
+    query string parameter. If a custom `httpx_client` is provided, authentication
+    details will be added to it.
+
+    Attributes:
+        name (str):
+            A unique name identifying this remote agent instance.
+        url (Optional[str]):
+            The base URL of the remote agent. This is optional if an `httpx_client`
+            with a configured `base_url` is provided. If both are given, they must
+            not conflict.
+        auth_token (Optional[str]):
+            Optional authentication token used for secure access during initialization.
+            If not provided, the agent will be accessed without authentication.
+            Note: For runtime authentication, use the credential service in InvocationContext.
+        auth_method (Literal["header", "querystring"] | None):
+            The method of attaching the authentication token at runtime.
+            - `"header"`: Token is retrieved from credential service and passed via HTTP `Authorization` header.
+            - `"querystring"`: Token is retrieved from credential service and passed as a query parameter.
+            - `None`: No runtime authentication injection (default).
+            The credential is loaded from `InvocationContext.credential_service` using the
+            app_name and user_id from the context.
+        httpx_client (Optional[httpx.AsyncClient]):
+            An optional, pre-configured `httpx.AsyncClient` to use for communication.
+            This allows for client sharing and advanced configurations (e.g., proxies).
+            If its `base_url` is set, it will be used as the agent's location.
+
+    Raises:
+        ValueError:
+            - If `url` and `httpx_client.base_url` are both provided and conflict.
+            - If neither `url` nor an `httpx_client` with a `base_url` is provided.
+            - If an unsupported `auth_method` is provided when `auth_token` is set.
+        requests.RequestException:
+            If fetching the agent card from the remote URL fails.
+
+    Examples:
+        ```python
+        # Example 1: Connect using a URL (no authentication)
+        agent = RemoteVeAgent(
+            name="public_agent",
+            url="https://vefaas.example.com/agents/public"
+        )
+
+        # Example 2: Using static Bearer token in header for initialization
+        agent = RemoteVeAgent(
+            name="secured_agent",
+            url="https://vefaas.example.com/agents/secure",
+            auth_token="my_secret_token",
+            auth_method="header"
+        )
+
+        # Example 3: Using runtime authentication with credential service
+        # The auth token will be automatically injected from InvocationContext.credential_service
+        agent = RemoteVeAgent(
+            name="dynamic_auth_agent",
+            url="https://vefaas.example.com/agents/secure",
+            auth_method="header"  # Will load credential at runtime
+        )
+
+        # Example 4: Using a pre-configured httpx_client
+        import httpx
+        client = httpx.AsyncClient(
+            base_url="https://vefaas.example.com/agents/query",
+            timeout=600
+        )
+        agent = RemoteVeAgent(
+            name="query_agent",
+            auth_method="querystring",  # Will load credential at runtime
+            httpx_client=client
+        )
+        ```
+    """
+
+    auth_method: Literal["header", "querystring"] | None = None
+
+    def __init__(
+        self,
+        name: str,
+        url: Optional[str] = None,
+        auth_token: Optional[str] = None,
+        auth_method: Literal["header", "querystring"] | None = None,
+        httpx_client: Optional[httpx.AsyncClient] = None,
+    ):
+        # Determine the effective URL for the agent and handle conflicts.
+        effective_url = url
+        if httpx_client and httpx_client.base_url:
+            client_url_str = str(httpx_client.base_url).rstrip("/")
+            if url and url.rstrip("/") != client_url_str:
+                raise ValueError(
+                    f"The `url` parameter ('{url}') conflicts with the `base_url` of the provided "
+                    f"httpx_client ('{client_url_str}'). Please provide only one or ensure they match."
+                )
+            effective_url = client_url_str
+
+        if not effective_url:
+            raise ValueError(
+                "Could not determine agent URL. Please provide the `url` parameter or an `httpx_client` with a configured `base_url`."
+            )
+
+        req_headers = {}
+        req_params = {}
+
+        if auth_token:
+            if auth_method == "header":
+                req_headers = {"Authorization": f"Bearer {auth_token}"}
+            elif auth_method == "querystring":
+                req_params = {"token": auth_token}
+            elif auth_method:
+                raise ValueError(
+                    f"Unsupported auth method {auth_method}, use `header` or `querystring` instead."
+                )
+
+        agent_card_dict = requests.get(
+            effective_url + AGENT_CARD_WELL_KNOWN_PATH,
+            headers=req_headers,
+            params=req_params,
+        ).json()
+        # replace agent_card_url with actual host
+        agent_card_dict["url"] = effective_url
+
+        agent_card_object = _convert_agent_card_dict_to_obj(agent_card_dict)
+
+        logger.debug(f"Agent card of {name}: {agent_card_object}")
+
+        client_was_provided = httpx_client is not None
+        client_to_use = httpx_client
+
+        if client_was_provided:
+            # If a client was provided, update it with auth info
+            if auth_token:
+                if auth_method == "header":
+                    client_to_use.headers.update(req_headers)
+                elif auth_method == "querystring":
+                    new_params = dict(client_to_use.params)
+                    new_params.update(req_params)
+                    client_to_use.params = new_params
+        else:
+            # If no client was provided, create a new one with auth info
+            if auth_token:
+                if auth_method == "header":
+                    client_to_use = httpx.AsyncClient(
+                        base_url=effective_url, headers=req_headers, timeout=600
+                    )
+                elif auth_method == "querystring":
+                    client_to_use = httpx.AsyncClient(
+                        base_url=effective_url, params=req_params, timeout=600
+                    )
+            else:  # No auth, no client provided
+                client_to_use = httpx.AsyncClient(base_url=effective_url, timeout=600)
+
+        super().__init__(
+            name=name, agent_card=agent_card_object, httpx_client=client_to_use
+        )
+
+        # The parent class sets _httpx_client_needs_cleanup based on whether
+        # the httpx_client it received was None. Since we always pass a
+        # client (either the user's or one we create), it will always set
+        # it to False. We must override this to ensure clients we create
+        # are properly cleaned up.
+        if not client_was_provided:
+            self._httpx_client_needs_cleanup = True
+
+        # Set auth_method if provided
+        if auth_method:
+            self.auth_method = auth_method
+
+        # Wrap _run_async_impl with pre-run hook to ensure initialization
+        # and authentication logic always executes, even if users override _run_async_impl
+        self._wrap_run_async_impl()
+
+    def _wrap_run_async_impl(self) -> None:
+        """Wrap _run_async_impl with a decorator that ensures pre-run logic executes.
+
+        This method wraps the _run_async_impl method with a decorator that:
+        1. Executes _pre_run before the actual implementation
+        2. Handles errors from _pre_run and yields error events
+        3. Ensures the wrapper works even if users override _run_async_impl
+
+        The wrapper is applied by replacing the bound method on the instance.
+        """
+        # Store the original _run_async_impl method
+        original_run_async_impl = self._run_async_impl
+
+        @functools.wraps(original_run_async_impl)
+        async def wrapped_run_async_impl(
+            ctx: InvocationContext,
+        ) -> AsyncGenerator[Event, None]:
+            """Wrapped version of _run_async_impl with pre-run hook."""
+            # Execute pre-run initialization
+            try:
+                await self._pre_run(ctx)
+            except Exception as e:
+                yield Event(
+                    author=self.name,
+                    error_message=f"Failed to initialize remote A2A agent: {e}",
+                    invocation_id=ctx.invocation_id,
+                    branch=ctx.branch,
+                )
+                return
+
+            # Call the original (or overridden) _run_async_impl
+            async with Aclosing(original_run_async_impl(ctx)) as agen:
+                async for event in agen:
+                    yield event
+
+        # Replace the instance method with the wrapped version
+        self._run_async_impl = wrapped_run_async_impl
+
+    async def _pre_run(self, ctx: InvocationContext) -> None:
+        """Pre-run initialization and authentication setup.
+
+        This method is called before the actual agent execution to:
+        1. Ensure the agent is resolved (agent card fetched, client initialized)
+        2. Inject authentication token from credential service if available
+
+        This method is separated from _run_async_impl to ensure these critical
+        initialization steps are always executed, even if users override _run_async_impl.
+
+        Args:
+            ctx: Invocation context containing session and user information
+
+        Raises:
+            Exception: If agent initialization fails
+        """
+        # Ensure agent is resolved
+        await self._ensure_resolved()
+
+        # Inject auth token if credential service is available
+        await self._inject_auth_token(ctx)
+
+    async def _inject_auth_token(self, ctx: InvocationContext) -> None:
+        """Inject authentication token from credential service into the HTTP client.
+
+        This method retrieves the authentication token from the credential service
+        in the InvocationContext and updates the HTTP client headers or query params
+        based on the configured auth_method.
+
+        Args:
+            ctx: Invocation context containing credential service and user information
+        """
+        # Skip if no credential service in context
+        if not ctx.credential_service:
+            logger.debug(
+                "No credential service in InvocationContext, skipping auth token injection"
+            )
+            return
+
+        # Skip if client is not initialized or not a BaseClient
+        if not hasattr(self, "_a2a_client") or not isinstance(
+            self._a2a_client, BaseClient
+        ):
+            logger.debug(
+                "A2A client not initialized or not a BaseClient, skipping auth token injection"
+            )
+            return
+
+        # Skip if transport is not available
+        if not hasattr(self._a2a_client, "_transport"):
+            logger.debug(
+                "A2A client transport not available, skipping auth token injection"
+            )
+            return
+
+        # Skip if httpx_client is not available
+        if not hasattr(self._a2a_client._transport, "httpx_client"):
+            logger.debug(
+                "A2A client httpx_client not available, skipping auth token injection"
+            )
+            return
+
+        try:
+            from veadk.utils.auth import build_auth_config
+            from google.adk.agents.callback_context import CallbackContext
+
+            # Inject TIP token via header
+            workload_auth_config = build_auth_config(
+                auth_method="apikey",
+                credential_key=VE_TIP_TOKEN_CREDENTIAL_KEY,
+                header_name=VE_TIP_TOKEN_HEADER,
+            )
+
+            tip_credential = await ctx.credential_service.load_credential(
+                auth_config=workload_auth_config,
+                callback_context=CallbackContext(ctx),
+            )
+
+            if tip_credential:
+                self._a2a_client._transport.httpx_client.headers.update(
+                    {VE_TIP_TOKEN_HEADER: tip_credential.api_key}
+                )
+                logger.debug(
+                    f"Injected TIP token via header for app={ctx.app_name}, user={ctx.user_id}"
+                )
+
+            # Build auth config based on auth_method
+            auth_config = build_auth_config(
+                credential_key="inbound_auth",
+                auth_method=self.auth_method or "header",
+                header_scheme="bearer",
+            )
+
+            # Load credential from credential service
+            credential = await ctx.credential_service.load_credential(
+                auth_config=auth_config,
+                callback_context=CallbackContext(ctx),
+            )
+
+            if not credential:
+                logger.debug(
+                    f"No credential loaded, skipping auth token injection for app={ctx.app_name}, user={ctx.user_id}"
+                )
+                return
+
+            # Inject credential based on auth_method
+            if self.auth_method == "querystring":
+                # Extract API key
+                api_key = credential.api_key
+                new_params = dict(self._a2a_client._transport.httpx_client.params)
+                new_params.update({"token": api_key})
+                self._a2a_client._transport.httpx_client.params = new_params
+                logger.debug(
+                    f"Injected auth token via querystring for app={ctx.app_name}, user={ctx.user_id}"
+                )
+            else:
+                if headers := generate_headers(credential):
+                    self._a2a_client._transport.httpx_client.headers.update(headers)
+                    logger.debug(
+                        f"Injected auth token via header for app={ctx.app_name}, user={ctx.user_id}"
+                    )
+
+        except Exception as e:
+            logger.warning(f"Failed to inject auth token: {e}", exc_info=True)

veadk/a2a/utils/__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/a2a/utils/agent_to_a2a.py

@@ -0,0 +1,170 @@
+# 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 typing import Literal, Optional, Union
+
+from a2a.types import AgentCard
+from google.adk.agents import BaseAgent
+from starlette.applications import Starlette
+from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.memory.in_memory_memory_service import InMemoryMemoryService
+from google.adk.a2a.utils.agent_to_a2a import to_a2a as google_adk_to_a2a
+from veadk import Runner
+from veadk.a2a.ve_middlewares import build_a2a_auth_middleware
+from veadk.auth.ve_credential_service import VeCredentialService
+from veadk.consts import DEFAULT_AGENT_NAME
+
+
+def to_a2a(
+    agent: BaseAgent,
+    *,
+    host: str = "localhost",
+    port: int = 8000,
+    protocol: str = "http",
+    agent_card: Optional[Union[AgentCard, str]] = None,
+    runner: Optional[Runner] = None,
+    enable_auth: bool = False,
+    auth_method: Literal["header", "querystring"] = "header",
+) -> Starlette:
+    """Convert an ADK agent to a A2A Starlette application with optional VeADK enhancements.
+
+    This function wraps Google ADK's to_a2a utility and optionally adds:
+    - VeCredentialService for authentication management
+    - A2A authentication middleware for token validation
+
+    Args:
+        agent: The ADK agent to convert to A2A server
+        host: The host for the A2A RPC URL (default: "localhost")
+        port: The port for the A2A RPC URL (default: 8000)
+        protocol: The protocol for the A2A RPC URL (default: "http")
+        agent_card: Optional pre-built AgentCard object or path to agent card
+                    JSON. If not provided, will be built automatically from the
+                    agent.
+        runner: Optional pre-built Runner object. If not provided, a default
+                runner will be created using in-memory services.
+                When enable_auth=True:
+                - If runner is provided and has a credential_service, it must be
+                  a VeCredentialService instance (raises TypeError otherwise)
+                - If runner is provided without credential_service, a new
+                  VeCredentialService will be created and set
+                - If runner is not provided, a new runner with VeCredentialService
+                  will be created
+        auth_method: Authentication method for A2A requests (only used when
+                    enable_auth=True). Options:
+                    - "header": Extract token from Authorization header (default)
+                    - "querystring": Extract token from query parameter
+        enable_auth: Whether to enable VeADK authentication features.
+                          When True, enables credential service and auth middleware.
+                          When False, uses standard Google ADK behavior.
+                          Default: False
+
+    Returns:
+        A Starlette application that can be run with uvicorn
+
+    Raises:
+        TypeError: If enable_auth=True and runner has a credential_service
+                  that is not a VeCredentialService instance
+
+    Example:
+        Basic usage (without VeADK auth):
+        ```python
+        from veadk import Agent
+        from veadk.a2a.utils.agent_to_a2a import to_a2a
+
+        agent = Agent(name="my_agent", tools=[...])
+        app = to_a2a(agent, host="localhost", port=8000)
+        # Run with: uvicorn module:app --host localhost --port 8000
+        ```
+
+        With VeADK authentication enabled:
+        ```python
+        app = to_a2a(agent, enable_auth=True)
+        ```
+
+        With custom runner and VeADK auth:
+        ```python
+        from veadk import Agent, Runner
+        from veadk.memory.short_term_memory import ShortTermMemory
+        from veadk.auth.ve_credential_service import VeCredentialService
+
+        agent = Agent(name="my_agent")
+        runner = Runner(
+            agent=agent,
+            short_term_memory=ShortTermMemory(),
+            app_name="my_app",
+            credential_service=VeCredentialService()  # Optional
+        )
+        app = to_a2a(agent, runner=runner, enable_auth=True)
+        ```
+
+        With querystring authentication:
+        ```python
+        app = to_a2a(agent, enable_auth=True, auth_method="querystring")
+        ```
+    """
+    app_name = agent.name or DEFAULT_AGENT_NAME
+    middleware = None  # May need support multiple middlewares in the future
+
+    # Handle VeADK authentication setup
+    if enable_auth:
+        # Create credential service if not provided
+        credential_service = VeCredentialService()
+        if runner is not None:
+            # Check if runner has credential_service
+            if runner.credential_service is not None:
+                # Validate that it's a VeCredentialService
+                if not isinstance(runner.credential_service, VeCredentialService):
+                    raise TypeError(
+                        f"When enable_auth=True, runner.credential_service must be "
+                        f"a VeCredentialService instance, got {type(runner.credential_service).__name__}"
+                    )
+                # Use existing credential service
+                credential_service = runner.credential_service
+            else:
+                # Add credential_service to runner
+                runner.credential_service = credential_service
+        else:
+            # Create runner with credential_service
+            runner = Runner(
+                app_name=app_name,
+                agent=agent,
+                artifact_service=InMemoryArtifactService(),
+                session_service=InMemorySessionService(),
+                memory_service=InMemoryMemoryService(),
+                credential_service=credential_service,
+            )
+
+        middleware = build_a2a_auth_middleware(
+            app_name=app_name,
+            credential_service=credential_service,
+            auth_method=auth_method,
+        )
+
+    # Convert agent to A2A Starlette app using Google ADK utility
+    app: Starlette = google_adk_to_a2a(
+        agent=agent,
+        host=host,
+        port=port,
+        protocol=protocol,
+        agent_card=agent_card,
+        runner=runner,
+    )
+
+    # Add VeADK authentication middleware only if enabled
+    if middleware:
+        app.add_middleware(middleware)
+
+    return app