veadk-python 0.2.27__py3-none-any.whl

veadk/a2a/ve_a2a_server.py

@@ -0,0 +1,93 @@
+# 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.server.apps.jsonrpc.fastapi_app import A2AFastAPIApplication
+from a2a.server.request_handlers.default_request_handler import DefaultRequestHandler
+from a2a.server.tasks.inmemory_task_store import InMemoryTaskStore
+from fastapi import FastAPI
+
+from veadk import Agent
+from veadk.a2a.agent_card import get_agent_card
+from veadk.runner import Runner
+from veadk.memory.short_term_memory import ShortTermMemory
+
+from google.adk.a2a.executor.a2a_agent_executor import A2aAgentExecutor
+from google.adk.auth.credential_service.base_credential_service import (
+    BaseCredentialService,
+)
+
+
+class VeA2AServer:
+    def __init__(
+        self,
+        agent: Agent,
+        url: str,
+        app_name: str,
+        short_term_memory: ShortTermMemory,
+        credential_service: BaseCredentialService | None = None,
+    ):
+        self.agent_card = get_agent_card(agent, url)
+
+        self.agent_executor = A2aAgentExecutor(
+            runner=Runner(
+                agent=agent,
+                app_name=app_name,
+                short_term_memory=short_term_memory,
+                credential_service=credential_service,
+            )
+        )
+
+        self.task_store = InMemoryTaskStore()
+
+        self.request_handler = DefaultRequestHandler(
+            agent_executor=self.agent_executor, task_store=self.task_store
+        )
+
+    def build(self) -> FastAPI:
+        app_application = A2AFastAPIApplication(
+            agent_card=self.agent_card,
+            http_handler=self.request_handler,
+        )
+        app = app_application.build()  # build routes
+
+        return app
+
+
+def init_app(
+    server_url: str,
+    app_name: str,
+    agent: Agent,
+    short_term_memory: ShortTermMemory,
+    credential_service: BaseCredentialService | None = None,
+) -> FastAPI:
+    """Init the fastapi application in terms of VeADK agent.
+
+    Args:
+        server_url: str, the url of the server
+        app_name: str, the name of the app
+        agent: Agent, the agent of the app
+        short_term_memory: ShortTermMemory, the short term memory of the app
+
+    Returns:
+        FastAPI, the fastapi app
+    """
+
+    server = VeA2AServer(
+        agent=agent,
+        url=server_url,
+        app_name=app_name,
+        short_term_memory=short_term_memory,
+        credential_service=credential_service,
+    )
+    return server.build()

veadk/a2a/ve_agent_executor.py

@@ -0,0 +1,78 @@
+# 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.server.agent_execution.agent_executor import AgentExecutor
+from a2a.server.agent_execution.context import RequestContext
+from a2a.server.events.event_queue import EventQueue
+from a2a.utils import new_agent_text_message
+from typing_extensions import override
+
+from veadk import Agent
+from veadk.memory.short_term_memory import ShortTermMemory
+from veadk.runner import Runner
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class VeAgentExecutor(AgentExecutor):
+    def __init__(self, app_name: str, agent: Agent, short_term_memory: ShortTermMemory):
+        super().__init__()
+        self.app_name = app_name
+        self.agent = agent
+        self.short_term_memory = short_term_memory
+
+        self.runner = Runner(
+            agent=self.agent,
+            short_term_memory=self.short_term_memory,
+            app_name=app_name,
+            user_id="",  # waiting for `execute` change
+        )
+
+    @override
+    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
+        # extract metadata
+        user_id = (
+            context.metadata["user_id"]
+            if "user_id" in context.metadata
+            else "unkonwn_user"
+        )
+        self.runner.user_id = user_id
+
+        session_id = (
+            context.metadata["session_id"]
+            if "session_id" in context.metadata
+            else "unkonwn_session"
+        )
+
+        # process user input
+        user_input = context.get_user_input()
+
+        logger.debug(
+            f"Request: user_id: {user_id}, session_id: {session_id}, user_input: {user_input}"
+        )
+
+        # running
+        final_output = await self.runner.run(
+            messages=user_input,
+            session_id=session_id,
+        )
+
+        logger.debug(f"Final output: {final_output}")
+
+        await event_queue.enqueue_event(new_agent_text_message(final_output))
+
+    @override
+    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
+        return None

veadk/a2a/ve_middlewares.py

@@ -0,0 +1,313 @@
+# 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 Authentication Middleware for FastAPI/Starlette.
+
+This module provides middleware for extracting authentication credentials
+from HTTP requests and storing them in the credential service.
+"""
+
+import logging
+from typing import Callable, Literal, Optional
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from starlette.responses import Response
+from volcenginesdkcore.rest import ApiException
+
+
+from veadk.auth.ve_credential_service import VeCredentialService
+from veadk.utils.auth import (
+    extract_delegation_chain_from_jwt,
+    build_auth_config,
+    VE_TIP_TOKEN_HEADER,
+    VE_TIP_TOKEN_CREDENTIAL_KEY,
+)
+from veadk.integrations.ve_identity import (
+    WorkloadToken,
+    IdentityClient,
+    get_default_identity_client,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class A2AAuthMiddleware(BaseHTTPMiddleware):
+    """Middleware to extract and store authentication credentials from requests.
+
+    This middleware:
+    1. Extracts auth tokens from Authorization header or query string
+    2. Parses JWT tokens to extract user_id and delegation chain
+    3. Builds AuthConfig based on the authentication method
+    4. Stores credentials in the credential service
+    5. Extracts TIP token from X-Ve-TIP-Token header for trust propagation
+    6. Exchanges TIP token for workload access token using IdentityClient
+    7. Sets workload token in request.scope["auth"] for downstream use
+
+    Examples:
+        ```python
+        from fastapi import FastAPI
+        from veadk.a2a.ve_middlewares import build_a2a_auth_middleware
+        from veadk.auth.ve_credential_service  import VeCredentialService
+
+        app = FastAPI()
+        credential_service = VeCredentialService()
+
+        # Add middleware with Authorization header support
+        app.add_middleware(
+            build_a2a_auth_middleware(
+                app_name="my_app",
+                credential_service=credential_service,
+                auth_method="header"
+            )
+        )
+
+        # Or with query string support
+        app.add_middleware(
+            build_a2a_auth_middleware(
+                app_name="my_app",
+                credential_service=credential_service,
+                auth_method="querystring",
+                token_param="token"
+            )
+        )
+        ```
+    """
+
+    def __init__(
+        self,
+        app,
+        app_name: str,
+        credential_service: VeCredentialService,
+        auth_method: Literal["header", "querystring"] = "header",
+        token_param: str = "token",
+        credential_key: str = "inbound_auth",
+        identity_client: Optional[IdentityClient] = None,
+    ):
+        """Initialize the middleware.
+
+        Args:
+            app: The ASGI application
+            app_name: Application name for credential storage
+            credential_service: Credential service to store credentials
+            auth_method: Authentication method - "header" or "querystring"
+            token_param: Query parameter name for token (when auth_method="querystring")
+            credential_key: Key to identify the credential in the store
+            identity_client: Optional IdentityClient for TIP token exchange.
+                If not provided, uses the global IdentityClient from VeIdentityConfig.
+        """
+        super().__init__(app)
+        self.app_name = app_name
+        self.credential_service = credential_service
+        self.auth_method = auth_method
+        self.token_param = token_param
+        self.credential_key = credential_key
+        self.identity_client = identity_client or get_default_identity_client()
+
+    async def dispatch(self, request: Request, call_next: Callable) -> Response:
+        """Process the request and extract authentication credentials.
+
+        This method:
+        1. Extracts token from Authorization header or query string
+        2. Parses JWT to get user_id and stores credentials
+        3. Sets request.scope["user"] with SimpleUser instance
+        4. Extracts TIP token from X-Ve-TIP-Token header
+        5. Exchanges TIP token for workload token via IdentityClient
+        6. Sets request.scope["auth"] with WorkloadToken object
+
+        Args:
+            request: The incoming HTTP request
+            call_next: The next middleware or route handler
+
+        Returns:
+            The HTTP response
+        """
+        from starlette.authentication import SimpleUser
+
+        token, has_prefix = self._extract_token(request)
+        user_id = None
+
+        if token:
+            user_id, _ = extract_delegation_chain_from_jwt(token)
+            if user_id:
+                # Build auth config based on authentication method
+                auth_config = build_auth_config(
+                    token=token,
+                    auth_method=self.auth_method,
+                    credential_key=self.credential_key,
+                    header_scheme="bearer" if has_prefix else None,
+                    query_param_name=self.token_param,
+                )
+
+                await self.credential_service.set_credential(
+                    app_name=self.app_name,
+                    user_id=user_id,
+                    credential_key=self.credential_key,
+                    credential=auth_config.exchanged_auth_credential,
+                )
+
+                logger.debug(
+                    f"Stored credential for app={self.app_name}, user={user_id}, "
+                    f"method={self.auth_method}"
+                )
+
+                request.scope["user"] = SimpleUser(user_id)
+            else:
+                logger.warning("Failed to extract user_id from JWT token")
+
+        # Extract TIP token from X-Ve-TIP-Token header for trust propagation
+        tip_token = request.headers.get(VE_TIP_TOKEN_HEADER)
+
+        try:
+            workload_token: WorkloadToken = (
+                self.identity_client.get_workload_access_token(
+                    user_token=tip_token, user_id=user_id
+                )
+            )
+            workload_auth_config = build_auth_config(
+                token=workload_token.workload_access_token,
+                auth_method="apikey",
+                credential_key=VE_TIP_TOKEN_CREDENTIAL_KEY,
+                header_name=VE_TIP_TOKEN_HEADER,
+            )
+
+            await self.credential_service.set_credential(
+                app_name=self.app_name,
+                user_id=user_id,
+                credential_key=VE_TIP_TOKEN_CREDENTIAL_KEY,
+                credential=workload_auth_config.exchanged_auth_credential,
+            )
+        except ApiException as e:
+            logger.warning(f"Failed to get workload token: {e.reason}")
+            workload_token = None
+        request.scope["auth"] = workload_token
+        # Continue processing the request
+        response = await call_next(request)
+        return response
+
+    def _extract_token(self, request: Request) -> tuple[Optional[str], bool]:
+        """Extract authentication token from the request.
+
+        Args:
+            request: The HTTP request
+
+        Returns:
+            The extracted token, or None if not found
+        """
+        has_prefix = False
+        token = None
+
+        if self.auth_method == "header":
+            # Extract from Authorization header
+            auth_header = request.headers.get("Authorization") or request.headers.get(
+                "authorization"
+            )
+            if auth_header:
+                # Strip "Bearer " prefix if present
+                if auth_header.lower().startswith("bearer "):
+                    has_prefix = True
+                    token = auth_header[7:]
+                else:
+                    token = auth_header
+            else:
+                token = None
+        elif self.auth_method == "querystring":
+            # Extract from query string
+            token = request.query_params.get(self.token_param)
+
+        return token, has_prefix
+
+
+def build_a2a_auth_middleware(
+    app_name: str,
+    credential_service: VeCredentialService,
+    auth_method: Literal["header", "querystring"] = "header",
+    token_param: str = "token",
+    credential_key: str = "inbound_auth",
+    identity_client: Optional[IdentityClient] = None,
+) -> type[A2AAuthMiddleware]:
+    """Build an A2A authentication middleware class.
+
+    This is a factory function that creates a middleware class with the
+    specified configuration. Use this with FastAPI's add_middleware().
+
+    The middleware extracts authentication tokens from incoming requests,
+    parses JWT delegation chains, stores credentials, and sets user information
+    in the request state for downstream handlers.
+
+    TIP Token Support:
+        The middleware will:
+        1. Extract TIP token from X-Ve-TIP-Token header
+        2. Exchange TIP token for workload token using IdentityClient
+        3. Set WorkloadToken object in request.scope["auth"] for downstream use
+
+        If identity_client is not provided, uses the global IdentityClient
+        from VeIdentityConfig.
+
+    Args:
+        app_name: Application name for credential storage
+        credential_service: Credential service to store credentials
+        auth_method: Authentication method - "header" or "querystring"
+        token_param: Query parameter name for token (when auth_method="querystring")
+        credential_key: Key to identify the credential in the store
+        identity_client: Optional IdentityClient for TIP token exchange.
+            If not provided, uses the global IdentityClient from VeIdentityConfig.
+
+    Returns:
+        A configured middleware class
+
+    Request Attributes:
+        After successful authentication, the following attributes are set:
+        - request.scope["user"]: SimpleUser instance with the user_id from JWT
+        - request.scope["auth"]: WorkloadToken object containing workload_access_token
+
+    Examples:
+        ```python
+        from fastapi import FastAPI, Request
+        from veadk.a2a.ve_middlewares import build_a2a_auth_middleware
+        from veadk.auth.ve_credential_service  import VeCredentialService
+        from veadk.integrations.ve_identity import IdentityClient
+
+        app = FastAPI()
+        credential_service = VeCredentialService()
+
+        # Optional: Create identity client for TIP token support
+        # If not provided, uses global client from VeIdentityConfig
+        identity_client = IdentityClient(region="cn-beijing")
+
+        # Add middleware with TIP token support
+        app.add_middleware(
+            build_a2a_auth_middleware(
+                app_name="my_app",
+                credential_service=credential_service,
+                auth_method="header",
+                identity_client=identity_client,  # Optional, uses global if not provided
+            )
+        )
+        ```
+    """
+
+    class ConfiguredA2AAuthMiddleware(A2AAuthMiddleware):
+        def __init__(self, app):
+            super().__init__(
+                app=app,
+                app_name=app_name,
+                credential_service=credential_service,
+                auth_method=auth_method,
+                token_param=token_param,
+                credential_key=credential_key,
+                identity_client=identity_client,
+            )
+
+    return ConfiguredA2AAuthMiddleware

veadk/a2a/ve_task_store.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 a2a.server.tasks import TaskStore
+from a2a.types import Task
+from typing_extensions import override
+
+
+class VeTaskStore(TaskStore):
+    def __init__(self):
+        super().__init__()
+
+    @override
+    async def save(self, task: Task) -> None:
+        """Saves or updates a task in the store."""
+        return None
+
+    @override
+    async def get(self, task_id: str) -> Task | None:
+        """Retrieves a task from the store by ID."""
+        return None
+
+    @override
+    async def delete(self, task_id: str) -> None:
+        """Deletes a task from the store by ID."""
+        return None