veadk-python 0.2.27__py3-none-any.whl

veadk/integrations/ve_identity/auth_processor.py

@@ -0,0 +1,385 @@
+# 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.
+
+"""Authentication request processor for handling OAuth2 flows in agent conversations."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import time
+from typing import Any, Awaitable, Callable, Optional, TYPE_CHECKING
+
+from google.adk.auth import AuthConfig
+from google.genai import types
+from google.adk.auth.auth_credential import OAuth2Auth
+
+from veadk.integrations.ve_identity.auth_config import get_default_identity_client
+from veadk.processors.base_run_processor import BaseRunProcessor
+from veadk.integrations.ve_identity.models import AuthRequestConfig, OAuth2AuthPoller
+from veadk.integrations.ve_identity.utils import (
+    get_function_call_auth_config,
+    get_function_call_id,
+    is_pending_auth_event,
+)
+from a2a.utils import new_agent_text_message
+
+from veadk.utils.logger import get_logger
+from a2a.server.tasks import TaskUpdater
+from a2a.types import TaskState
+
+
+if TYPE_CHECKING:
+    from veadk.runner import Runner
+
+
+logger = get_logger(__name__)
+
+# Default configuration for token polling
+DEFAULT_POLLING_INTERVAL_SECONDS = 5
+DEFAULT_POLLING_TIMEOUT_SECONDS = 600
+# Authentication loop will break after this many cycles
+DEFAULT_MAX_CYCLES = 10
+
+
+class MockOauth2AuthPoller(OAuth2AuthPoller):
+    """Mock OAuth2 auth poller for testing purposes."""
+
+    async def poll_for_auth(self, *args, **kwargs) -> OAuth2Auth:
+        """Return a mock oauth2 auth immediately.
+
+        Returns:
+            A mock token string.
+        """
+        return OAuth2Auth(access_token="mock_token")
+
+
+class WaitInputOauth2AuthPoller(OAuth2AuthPoller):
+    """Wait for user input to complete OAuth2 flow."""
+
+    def __init__(self, auth_uri: str, *args, **kwargs):
+        self.auth_uri = auth_uri
+
+    async def poll_for_auth(self) -> OAuth2Auth:
+        """Wait for user input to complete OAuth2 flow.
+
+        Returns:
+            The complete OAuth2Auth object containing auth_response_uri.
+        """
+
+        async def get_user_input(prompt: str) -> str:
+            loop = asyncio.get_event_loop()
+            # Run the blocking `input()` function in a separate thread managed by the executor.
+            return await loop.run_in_executor(None, input, prompt)
+
+        logger.info(
+            f"Please open this URL in your browser to authorize: {self.auth_uri}"
+        )
+        auth_response_uri = await get_user_input("Please enter the callback URL:\n> ")
+        auth_response_uri = auth_response_uri.strip()
+
+        if not auth_response_uri:
+            raise Exception("No auth response URI provided. Aborting.")
+
+        return OAuth2Auth(auth_response_uri=auth_response_uri)
+
+
+class _DefaultOauth2AuthPoller(OAuth2AuthPoller):
+    """Default implementation of OAuth2 auth polling.
+
+    This poller repeatedly calls a polling function until a OAuth2 auth becomes available
+    or a timeout occurs.
+    """
+
+    def __init__(
+        self,
+        auth_url: str,
+        polling_func: (
+            Callable[[], Optional[OAuth2Auth]]
+            | Callable[[], Awaitable[Optional[OAuth2Auth]]]
+        ),
+    ):
+        """Initialize the OAuth2 auth poller.
+
+        Args:
+            auth_url: The authorization URL for logging purposes.
+            polling_func: Function that returns the OAuth2 auth when ready, or None if not yet available.
+                         Can be either sync or async function.
+        """
+        self.auth_url = auth_url
+        self.polling_func = polling_func
+
+    async def poll_for_auth(self) -> OAuth2Auth:
+        """Poll for a oauth2 auth until it becomes available or timeout occurs.
+
+        Returns:
+            The complete OAuth2Auth object containing tokens or auth_response_uri.
+
+        Raises:
+            asyncio.TimeoutError: If polling times out before oauth2 auth is available.
+        """
+        start_time = time.time()
+
+        while time.time() - start_time < DEFAULT_POLLING_TIMEOUT_SECONDS:
+            await asyncio.sleep(DEFAULT_POLLING_INTERVAL_SECONDS)
+
+            logger.info(
+                f"Polling for oauth2 auth at authorization URL: {self.auth_url}"
+            )
+
+            # Check if polling_func is async or sync
+            import inspect
+
+            if inspect.iscoroutinefunction(self.polling_func):
+                oauth2auth = await self.polling_func()
+            else:
+                oauth2auth = self.polling_func()
+
+            if oauth2auth is not None:
+                logger.info("OAuth2 auth successfully retrieved")
+                return oauth2auth
+
+        raise asyncio.TimeoutError(
+            f"OAuth2 auth polling timed out after {DEFAULT_POLLING_TIMEOUT_SECONDS} seconds. "
+            "User may not have completed authorization."
+        )
+
+
+class AuthRequestProcessor(BaseRunProcessor):
+    """Processor for handling authentication requests in agent conversations.
+
+    This class manages the OAuth2 authentication flow when tools require user authorization.
+    It handles displaying authorization URLs to users and polling for tokens after authorization.
+
+    Attributes:
+        config: Configuration for authentication request handling.
+    """
+
+    def __init__(self, *, config: Optional[AuthRequestConfig] = None):
+        """Initialize the authentication request processor.
+
+        Args:
+            config: Authentication configuration. If None, uses default configuration
+                   that prints the authorization URL to console.
+        """
+        self.config = config or AuthRequestConfig(
+            on_auth_url=lambda url: print(
+                f"Please open this URL in your browser to authorize: {url}"
+            )
+        )
+
+        self._identity_client = (
+            self.config.identity_client
+            or get_default_identity_client(self.config.region)
+        )
+
+    async def process_auth_request(
+        self,
+        auth_request_event_id: str,
+        auth_config: AuthConfig,
+        task_updater: Optional[TaskUpdater] = None,
+    ) -> types.Content:
+        """Process a single authentication request.
+
+        This method handles the complete OAuth2 flow:
+        1. Displays the authorization URL to the user (via on_auth_url callback)
+        2. Polls for the access token after user authorization
+        3. Constructs and returns the authentication response
+
+        Args:
+            auth_request_event_id: Unique ID of the authentication request event.
+            auth_config: Authentication configuration containing the auth URI.
+
+        Returns:
+            Content object containing the authentication response to send back to the agent.
+        """
+        logger.info(f"Processing auth request: {auth_request_event_id}")
+
+        auth_uri = auth_config.exchanged_auth_credential.oauth2.auth_uri
+        request_dict = (
+            json.loads(resource_ref_str)
+            if (resource_ref_str := auth_config.exchanged_auth_credential.resource_ref)
+            else {}
+        )
+
+        # Invoke the auth URL callback (sync or async)
+        if on_auth_url := self.config.on_auth_url:
+            if asyncio.iscoroutinefunction(on_auth_url):
+                await on_auth_url(auth_uri)
+            else:
+                on_auth_url(auth_uri)
+
+        # Use custom poller or default poller
+        active_poller = (
+            self.config.oauth2_auth_poller(auth_uri, request_dict)
+            if self.config.oauth2_auth_poller
+            else _DefaultOauth2AuthPoller(
+                auth_uri,
+                lambda: (
+                    lambda response: (
+                        OAuth2Auth(access_token=response.access_token)
+                        if response.access_token and response.access_token.strip()
+                        else None
+                    )
+                )(self._identity_client.get_oauth2_token_or_auth_url(**request_dict)),
+            )
+        )
+
+        # Poll for the oauth2 auth
+        updated_oauth2_auth = await active_poller.poll_for_auth()
+        if task_updater:
+            await task_updater.update_status(
+                TaskState.working,
+                message=new_agent_text_message("Authorization received, continuing..."),
+            )
+        for k, v in updated_oauth2_auth.__dict__.items():
+            if (
+                v is not None
+                and k
+                in auth_config.exchanged_auth_credential.oauth2.__pydantic_fields__
+            ):
+                auth_config.exchanged_auth_credential.oauth2.__dict__[k] = v
+
+        # Construct the authentication response
+        auth_content = types.Content(
+            role="user",
+            parts=[
+                types.Part(
+                    function_response=types.FunctionResponse(
+                        id=auth_request_event_id,
+                        name="adk_request_credential",
+                        response=auth_config.model_dump(),
+                    )
+                )
+            ],
+        )
+
+        logger.info(f"Auth request {auth_request_event_id} processed successfully")
+        return auth_content
+
+    def process_run(
+        self,
+        runner: Runner,
+        message: types.Content,
+        **kwargs: Any,
+    ):
+        """Process the agent run by wrapping the event generator with authentication loop.
+
+        This method implements the BaseRunProcessor interface and adds authentication
+        loop handling to event generators.
+
+        This decorator intercepts runner.run_async calls and automatically handles
+        authentication loops. The event_generator code can remain completely unchanged!
+
+        Usage example:
+            @auth_processor.process_run(
+                runner=runner,
+                message=message,
+            )
+            async def event_generator():
+                async for event in runner.run_async(
+                    user_id=user_id,
+                    session_id=session_id,
+                    new_message=message,
+                    run_config=RunConfig(streaming_mode=stream_mode),
+                ):
+                    if event.get_function_calls():
+                        for function_call in event.get_function_calls():
+                            logger.debug(f"Function call: {function_call}")
+                    elif event.content is not None:
+                        yield event.content.parts[0].text
+
+            # Then call (no parameters needed):
+            async for chunk in event_generator():
+                print(chunk)
+
+        The decorator automatically:
+        1. Intercepts runner.run_async calls
+        2. Detects authentication events
+        3. Processes authentication requests
+        4. Loops until no more authentication events
+
+        Args:
+            runner: Runner instance (will be wrapped).
+            message: Initial message to send.
+            **kwargs: Additional keyword arguments. Supports:
+                     - task_updater: Optional TaskUpdater for status updates.
+
+        Returns:
+            Decorated generator function.
+        """
+        # Extract task_updater from kwargs
+        task_updater = kwargs.get("task_updater")
+
+        def decorator(event_generator_func):
+            async def wrapper():
+                current_message = message
+
+                for _ in range(self.config.max_auth_cycles or DEFAULT_MAX_CYCLES):
+                    auth_request_event_id = None
+                    auth_config = None
+
+                    # Buffer to collect chunks from this cycle
+                    cycle_buffer = []
+
+                    # Create a wrapped runner to intercept run_async calls
+                    original_run_async = runner.run_async
+
+                    async def wrapped_run_async(**run_kwargs):
+                        nonlocal auth_request_event_id, auth_config
+
+                        # Override the message with the current message
+                        run_kwargs["new_message"] = current_message
+
+                        async for event in original_run_async(**run_kwargs):
+                            # Detect authentication events
+                            if is_pending_auth_event(event):
+                                auth_request_event_id = get_function_call_id(event)
+                                auth_config = get_function_call_auth_config(event)
+                                logger.info(
+                                    f"Found auth request {auth_request_event_id}, breaking to process"
+                                )
+
+                                break
+
+                            # Pass events to the caller
+                            yield event
+
+                    # Temporarily replace runner.run_async
+                    runner.run_async = wrapped_run_async
+
+                    try:
+                        # Call the original event_generator and buffer the output
+                        async for chunk in event_generator_func():
+                            cycle_buffer.append(chunk)
+                    finally:
+                        # Restore the original run_async
+                        runner.run_async = original_run_async
+
+                    # Check if there was an authentication request
+                    if auth_request_event_id and auth_config:
+                        # Process the authentication request
+                        current_message = await self.process_auth_request(
+                            auth_request_event_id, auth_config, task_updater
+                        )
+                        # Continue to next cycle
+                    else:
+                        # No more auth events, yield the final chunk
+                        logger.info("No more auth events found, processing complete")
+                        yield cycle_buffer[-1]
+                        break
+
+            return wrapper
+
+        return decorator

veadk/integrations/ve_identity/function_tool.py

@@ -0,0 +1,158 @@
+# 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
+
+import inspect
+from typing import Any, Callable
+
+from typing_extensions import override
+
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.tools.function_tool import FunctionTool
+from google.adk.tools.tool_context import ToolContext
+
+from veadk.integrations.ve_identity.auth_config import (
+    VeIdentityAuthConfig,
+    ApiKeyAuthConfig,
+    OAuth2AuthConfig,
+    WorkloadAuthConfig,
+)
+from veadk.integrations.ve_identity.auth_mixins import (
+    VeIdentityAuthMixin,
+    AuthRequiredException,
+)
+
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class VeIdentityFunctionTool(VeIdentityAuthMixin, FunctionTool):
+    """Unified function tool with automatic VeIdentity authentication.
+
+    This tool wraps a function and automatically handles authentication based on the
+    provided auth configuration. It supports both API Key and OAuth2 authentication.
+
+    Examples:
+        # API Key authentication
+        api_key_tool = VeIdentityFunctionTool(
+            func=my_function,
+            auth_config=api_key_auth("my-provider"),
+            into="api_key"
+        )
+
+        # OAuth2 authentication
+        oauth2_tool = VeIdentityFunctionTool(
+            func=my_function,
+            auth_config=oauth2_auth(
+                provider_name="my-provider",
+                scopes=["read", "write"],
+                auth_flow="USER_FEDERATION"
+            ),
+            into="access_token"
+        )
+    """
+
+    def __init__(
+        self,
+        *,
+        func: Callable[..., Any],
+        auth_config: VeIdentityAuthConfig,
+        into: str = None,
+    ):
+        """Initialize the unified Identity function tool.
+
+        Args:
+            func: The function to wrap with Identity authentication.
+            auth_config: Authentication configuration (ApiKeyAuthConfig or OAuth2AuthConfig).
+            into: Parameter name to inject the credential into. If None, uses default
+                 based on auth type ("api_key" for API key, "access_token" for OAuth2).
+        """
+        # Determine default parameter name based on auth type
+        if into is None:
+            if isinstance(auth_config, ApiKeyAuthConfig):
+                into = "api_key"
+            elif isinstance(auth_config, OAuth2AuthConfig):
+                into = "access_token"
+            elif isinstance(auth_config, WorkloadAuthConfig):
+                into = "access_token"
+            else:
+                raise ValueError(f"Unsupported auth config type: {type(auth_config)}")
+
+        # Initialize mixins first
+        super().__init__(
+            func=func,
+            auth_config=auth_config,
+        )
+        self._ignore_params.append(into)
+        self._into = into
+
+    @override
+    async def run_async(
+        self, *, args: dict[str, Any], tool_context: ToolContext
+    ) -> Any:
+        """Execute the wrapped function with Identity authentication.
+
+        This method handles authentication based on the configured auth type.
+
+        Args:
+            args: Arguments to pass to the wrapped function.
+            tool_context: The tool context for accessing session state and auth.
+
+        Returns:
+            The result from the wrapped function, or an auth pending message for OAuth2.
+        """
+        try:
+            return await self.run_with_identity_auth(
+                args=args, tool_context=tool_context
+            )
+        except AuthRequiredException as e:
+            # Only OAuth2 can raise this exception
+            return e.message
+
+    async def _execute_with_credential(
+        self,
+        *,
+        args: dict[str, Any],
+        tool_context: ToolContext,
+        credential: AuthCredential,
+    ) -> Any:
+        """Execute the wrapped function with the provided credential.
+
+        Args:
+            args: Arguments for the wrapped function.
+            tool_context: The tool context.
+            credential: The authentication credential (API key or OAuth2).
+
+        Returns:
+            The result from the wrapped function.
+        """
+        args_to_call = args.copy()
+        signature = inspect.signature(self.func)
+
+        # Inject the appropriate credential based on type
+        if self._into in signature.parameters:
+            if isinstance(self._auth_config, ApiKeyAuthConfig):
+                args_to_call[self._into] = credential.api_key
+            elif isinstance(self._auth_config, OAuth2AuthConfig):
+                args_to_call[self._into] = credential.oauth2.access_token
+            elif isinstance(self._auth_config, WorkloadAuthConfig):
+                args_to_call[self._into] = credential.oauth2.access_token
+            else:
+                raise ValueError(
+                    f"Unsupported auth config type: {type(self._auth_config)}"
+                )
+
+        return await super().run_async(args=args_to_call, tool_context=tool_context)