veadk-python 0.2.27__py3-none-any.whl

veadk/integrations/ve_identity/models.py

@@ -0,0 +1,228 @@
+# 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.
+
+"""Data models for veIdentity integration."""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Optional, TYPE_CHECKING, List
+
+from pydantic import BaseModel, model_validator, field_validator
+from google.adk.auth.auth_credential import OAuth2Auth
+
+if TYPE_CHECKING:
+    from veadk.integrations.ve_identity.identity_client import IdentityClient
+else:
+    # For runtime, use Any to avoid circular import issues
+    IdentityClient = Any
+
+
+# Forward declaration for type hints
+class OAuth2AuthPoller:
+    """Abstract base class for OAuth2 authentication polling implementations.
+
+    OAuth2 auth pollers are used to retrieve complete OAuth2 authentication data
+    after user authorization. Implementations should poll the identity service
+    until the authentication becomes available or a timeout occurs.
+    """
+
+    async def poll_for_auth(self) -> OAuth2Auth:
+        """Poll for OAuth2 authentication data and return it when available.
+
+        Returns:
+            The complete OAuth2Auth object containing tokens and metadata.
+
+        Raises:
+            asyncio.TimeoutError: If polling times out before auth data is available.
+        """
+        raise NotImplementedError("Subclasses must implement poll_for_auth")
+
+
+class AuthRequestConfig(BaseModel):
+    """Configuration for authentication request processing.
+
+    Attributes:
+        on_auth_url: Optional callback function invoked when an authorization URL is generated.
+                     Can be sync or async. Receives the auth URL as a parameter.
+        oauth2_auth_poller: Optional custom token poller implementation for retrieving tokens
+                      after user authorization. If None, a default poller will be used.
+    """
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    on_auth_url: Optional[Callable[[str], Any]] = None
+    # Currently we only use auth_uri to initialize poller, may extend to support other fields like exchanged_auth_credential.
+    oauth2_auth_poller: Optional[Callable[[Any], OAuth2AuthPoller]] = None
+    max_auth_cycles: Optional[int] = None
+    identity_client: Optional[IdentityClient] = None
+    region: Optional[str] = None
+
+
+class OAuth2TokenResponse(BaseModel):
+    """Response from OAuth2 token request.
+
+    Attributes:
+        response_type: Type of response - either "token" or "auth_url".
+        access_token: The OAuth2 access token (present when response_type is "token").
+        authorization_url: The authorization URL for user consent (present when response_type is "auth_url").
+        resource_ref: When response_type is "auth_url", this field contains the serialized request parameters
+                      needed to poll for the final OAuth2 tokens after user authorization.
+
+    """
+
+    response_type: str
+    access_token: Optional[str] = None
+    authorization_url: Optional[str] = None
+    resource_ref: Optional[str] = None
+
+    @field_validator("response_type")
+    @classmethod
+    def validate_response_type(cls, v: str) -> str:
+        """Validate that response_type is either 'token' or 'auth_url'."""
+        if v not in ("token", "auth_url"):
+            raise ValueError("response_type must be either 'token' or 'auth_url'")
+        return v
+
+    @model_validator(mode="after")
+    def validate_response_fields(self):
+        """Validate that required fields are present based on response_type."""
+        if self.response_type == "token":
+            if not self.access_token:
+                raise ValueError(
+                    "access_token is required when response_type is 'token'"
+                )
+        elif self.response_type == "auth_url":
+            if not self.authorization_url:
+                raise ValueError(
+                    "authorization_url is required when response_type is 'auth_url'"
+                )
+        return self
+
+
+class DCRRegistrationRequest(BaseModel):
+    """Dynamic Client Registration (DCR) request model.
+
+    Based on RFC 7591 - OAuth 2.0 Dynamic Client Registration Protocol.
+    """
+
+    client_name: str = "VeADK Framework"
+    redirect_uris: Optional[List[str]] = None
+    scope: Optional[str] = None
+    grant_types: Optional[List[str]] = None
+    response_types: Optional[List[str]] = None
+    token_endpoint_auth_method: Optional[str] = None
+
+    @field_validator("client_name")
+    @classmethod
+    def validate_client_name_not_empty(cls, v: str) -> str:
+        """Validate that client_name is not empty."""
+        if not v or not v.strip():
+            raise ValueError("client_name cannot be empty")
+        return v.strip()
+
+
+class DCRRegistrationResponse(BaseModel):
+    """Dynamic Client Registration (DCR) response model.
+
+    Based on RFC 7591 - OAuth 2.0 Dynamic Client Registration Protocol.
+    """
+
+    client_id: str
+    client_secret: Optional[str] = None
+    client_id_issued_at: Optional[int] = None
+    client_secret_expires_at: Optional[int] = None
+    redirect_uris: Optional[List[str]] = None
+    grant_types: Optional[List[str]] = None
+    response_types: Optional[List[str]] = None
+    scope: Optional[str] = None
+    token_endpoint_auth_method: Optional[str] = None
+
+    @field_validator("client_id")
+    @classmethod
+    def validate_client_id_not_empty(cls, v: str) -> str:
+        """Validate that client_id is not empty."""
+        if not v or not v.strip():
+            raise ValueError("client_id cannot be empty")
+        return v.strip()
+
+
+class AuthorizationServerMetadata(BaseModel):
+    """Extended Authorization Server Metadata with DCR support.
+
+    Based on RFC 8414 - OAuth 2.0 Authorization Server Metadata
+    and RFC 7591 - OAuth 2.0 Dynamic Client Registration Protocol.
+    """
+
+    authorization_endpoint: str
+    token_endpoint: str
+    issuer: str
+    register_endpoint: Optional[str] = None  # DCR endpoint
+    response_types: Optional[List[str]] = None
+
+    @field_validator("authorization_endpoint", "token_endpoint", "issuer")
+    @classmethod
+    def validate_url_not_empty(cls, v: str) -> str:
+        """Validate that URL fields are not empty."""
+        if not v or not v.strip():
+            raise ValueError("URL fields cannot be empty")
+        return v.strip()
+
+
+class OAuth2Discovery(BaseModel):
+    """OAuth2 Discovery configuration with DCR support."""
+
+    authorization_server_metadata: AuthorizationServerMetadata
+    discovery_url: Optional[str] = None
+
+    @field_validator("discovery_url")
+    @classmethod
+    def validate_discovery_url(cls, v: Optional[str]) -> Optional[str]:
+        """Validate discovery URL if provided."""
+        if v is not None and (not v or not v.strip()):
+            raise ValueError("discovery_url cannot be empty string")
+        return v.strip() if v else None
+
+
+class WorkloadToken(BaseModel):
+    """Workload access token and expiration time.
+
+    Attributes:
+        workload_access_token: The workload access token.
+        expires_at: Unix timestamp (in seconds) when the token expires.
+    """
+
+    workload_access_token: str
+    expires_at: int
+
+    @field_validator("workload_access_token")
+    @classmethod
+    def validate_token_not_empty(cls, v: str) -> str:
+        """Validate that the workload access token is not empty."""
+        if not v or not v.strip():
+            raise ValueError("workload_access_token cannot be empty")
+        return v.strip()
+
+    @field_validator("expires_at")
+    @classmethod
+    def validate_expires_at_positive(cls, v: int) -> int:
+        """Validate that expires_at is a positive timestamp."""
+        if v <= 0:
+            raise ValueError("expires_at must be a positive Unix timestamp")
+        return v
+
+
+class AssumeRoleCredential(BaseModel):
+    access_key_id: str
+    secret_access_key: str
+    session_token: str

veadk/integrations/ve_identity/token_manager.py

@@ -0,0 +1,188 @@
+# 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.
+
+"""Token manager for agent identity tokens with caching and expiration handling."""
+
+from __future__ import annotations
+
+import time
+from typing import Optional, Union
+
+from google.adk.tools.tool_context import ToolContext
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.agents.readonly_context import ReadonlyContext
+
+from veadk.integrations.ve_identity.auth_config import (
+    get_default_identity_client,
+)
+from veadk.utils.logger import get_logger
+
+from veadk.integrations.ve_identity.identity_client import IdentityClient
+from veadk.integrations.ve_identity.models import WorkloadToken
+
+logger = get_logger(__name__)
+
+# Token expiration buffer in seconds - tokens will be refreshed this many seconds before actual expiration
+TOKEN_EXPIRATION_BUFFER_SECONDS = 60
+
+
+class WorkloadTokenManager:
+    """Manager for workload access tokens with automatic caching and expiration handling.
+
+    This class manages the lifecycle of workload access tokens, including:
+    - Caching tokens in session state
+    - Automatic token refresh when expired
+    - Support for different authentication modes (JWT, user ID, workload-only)
+
+    Attributes:
+        identity_client: The IdentityClient instance for making API requests.
+        region: VolcEngine region for the identity client. Defaults to "cn-beijing".
+    """
+
+    def __init__(
+        self,
+        identity_client: Optional[IdentityClient] = None,
+        region: Optional[str] = None,
+    ):
+        """Initialize the token manager.
+
+        Args:
+            identity_client: Optional IdentityClient instance to use for token requests.
+                If not provided and use_global_client is True, uses the global client
+                from VeIdentityConfig.
+            region: Optional region for creating a new IdentityClient.
+                Only used if identity_client is not provided and use_global_client is False.
+        """
+
+        self._identity_client = identity_client or get_default_identity_client(
+            region=region
+        )
+
+    def _build_cache_key(
+        self, tool_context: Union[ToolContext | CallbackContext | ReadonlyContext]
+    ) -> str:
+        """Build a unique cache key for storing the workload token.
+
+        The cache key is composed of the agent name and user ID to ensure
+        tokens are properly scoped per agent and user.
+
+        Args:
+            tool_context: The tool context containing agent and user information.
+
+        Returns:
+            A unique cache key string in the format "workload_token:{agent}:{user}".
+        """
+        return f"workload_token:{tool_context.agent_name}:{tool_context._invocation_context.user_id}"
+
+    def _is_token_expired(self, expires_at: int) -> bool:
+        """Check if a token has expired or will expire soon.
+
+        Args:
+            expires_at: The expiration timestamp in seconds since Unix epoch.
+
+        Returns:
+            True if the token has expired or will expire within the buffer period,
+            False otherwise.
+        """
+        current_time = int(time.time())
+        return current_time >= (expires_at - TOKEN_EXPIRATION_BUFFER_SECONDS)
+
+    async def get_workload_token(
+        self,
+        tool_context: Union[ToolContext | CallbackContext | ReadonlyContext],
+        workload_name: Optional[str] = None,
+        user_token: Optional[str] = None,
+    ) -> str:
+        """Get or refresh the workload access token.
+
+        This method implements intelligent token caching:
+        1. Checks if a valid cached token exists in session state
+        2. Returns cached token if not expired
+        3. Fetches and caches a new token if expired or not found
+
+        Args:
+            tool_context: The tool context containing session state and user information.
+            workload_name: Optional workload name. If not provided, uses tool_context.agent_name.
+            user_token: Optional JWT token for user-scoped authentication.
+
+        Returns:
+            The workload access token string.
+
+        Raises:
+            ValueError: If the identity service response is missing required fields.
+        """
+        cache_key = self._build_cache_key(tool_context)
+
+        # Attempt to retrieve cached token from session state
+        cached_data: Optional[WorkloadToken | None] = (
+            tool_context._invocation_context.session.state.get(cache_key)
+        )
+
+        # Validate and return cached token if still valid, and type check
+        if cached_data and isinstance(cached_data, WorkloadToken):
+            if cached_data.workload_access_token and cached_data.expires_at:
+                if not self._is_token_expired(cached_data.expires_at):
+                    return cached_data.workload_access_token
+                else:
+                    logger.info(
+                        f"Cached workload token expired for agent '{tool_context.agent_name}', refreshing..."
+                    )
+
+        # Determine user_id based on authentication mode
+        user_id = None if user_token else tool_context._invocation_context.user_id
+
+        # Request new token from identity service
+        workload_token: WorkloadToken = self._identity_client.get_workload_access_token(
+            workload_name=workload_name,
+            user_token=user_token,
+            user_id=user_id,
+        )
+
+        tool_context._invocation_context.session.state[cache_key] = workload_token
+
+        return workload_token.workload_access_token
+
+
+async def get_workload_token(
+    tool_context: Union[ToolContext | CallbackContext | ReadonlyContext],
+    identity_client: Optional[IdentityClient] = None,
+    workload_name: Optional[str] = None,
+    user_token: Optional[str] = None,
+    region: str = "cn-beijing",
+) -> str:
+    """Convenience function to get a workload access token.
+
+    This function creates a token manager and retrieves the token with automatic
+    caching and expiration handling. It's a simplified interface for common use cases.
+
+    Args:
+        tool_context: The tool context containing session state and user information.
+        identity_client: Optional IdentityClient instance. If not provided, creates a new one.
+        workload_name: Optional workload name. If not provided, uses tool_context.agent_name.
+        user_token: Optional JWT token for user-scoped authentication.
+        region: The VolcEngine region for the identity client. Defaults to "cn-beijing".
+
+    Returns:
+        The workload access token string.
+
+    Raises:
+        ValueError: If the identity service response is missing required fields.
+    """
+    return await WorkloadTokenManager(
+        identity_client=identity_client, region=region
+    ).get_workload_token(
+        tool_context=tool_context,
+        workload_name=workload_name,
+        user_token=user_token,
+    )

veadk/integrations/ve_identity/utils.py

@@ -0,0 +1,151 @@
+# 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.
+
+"""Utility functions for handling authentication events in the veADK framework."""
+
+from __future__ import annotations
+import base64
+from typing import Optional
+
+from google.adk.events import Event
+from google.adk.auth import AuthConfig
+from google.adk.auth.auth_credential import AuthCredential
+
+
+def is_pending_auth_event(event: Event) -> bool:
+    """Check if an ADK event represents a pending authentication request.
+
+    The ADK framework emits a special function call ('adk_request_credential')
+    when a tool requires user authentication that hasn't been satisfied yet.
+
+    Args:
+        event: The ADK Event object to inspect.
+
+    Returns:
+        True if the event is an 'adk_request_credential' function call, False otherwise.
+    """
+    return (
+        event.content
+        and event.content.parts
+        and event.content.parts[0]
+        and event.content.parts[0].function_call
+        and event.content.parts[0].function_call.name == "adk_request_credential"
+    )
+
+
+def get_function_call_id(event: Event) -> str:
+    """Extract the unique function call ID from an ADK event.
+
+    This ID is used to correlate function responses back to the specific
+    function call that initiated the authentication request.
+
+    Args:
+        event: The ADK Event object containing the function call.
+
+    Returns:
+        The unique identifier string of the function call.
+
+    Raises:
+        ValueError: If the function call ID cannot be found in the event structure.
+    """
+    if (
+        event
+        and event.content
+        and event.content.parts
+        and event.content.parts[0]
+        and event.content.parts[0].function_call
+        and event.content.parts[0].function_call.id
+    ):
+        return event.content.parts[0].function_call.id
+
+    raise ValueError(f"Cannot extract function call ID from event: {event}")
+
+
+def get_function_call_auth_config(event: Event) -> AuthConfig:
+    """Extract authentication configuration from an 'adk_request_credential' event.
+
+    The client should use this AuthConfig to provide necessary authentication details
+    (like OAuth codes and state) and send it back to the ADK to continue the OAuth
+    token exchange process.
+
+    Args:
+        event: The ADK Event object containing the 'adk_request_credential' call.
+
+    Returns:
+        An AuthConfig object populated with details from the function call arguments.
+
+    Raises:
+        ValueError: If the 'authConfig' argument cannot be found in the event.
+    """
+    if (
+        event
+        and event.content
+        and event.content.parts
+        and event.content.parts[0]
+        and event.content.parts[0].function_call
+        and event.content.parts[0].function_call.args
+        and event.content.parts[0].function_call.args.get("authConfig")
+    ):
+        auth_config_dict = event.content.parts[0].function_call.args.get("authConfig")
+        return AuthConfig(**auth_config_dict)
+
+    raise ValueError(f"Cannot extract auth config from event: {event}")
+
+
+def generate_headers(credential: AuthCredential) -> Optional[dict[str, str]]:
+    """Extracts authentication headers from credentials.
+
+    Args:
+        credential: The authentication credential to process.
+
+    Returns:
+        Dictionary of headers to add to the request, or None if no auth.
+    """
+    headers: Optional[dict[str, str]] = None
+    if credential:
+        if credential.oauth2:
+            headers = {"Authorization": f"Bearer {credential.oauth2.access_token}"}
+        elif credential.http:
+            # Handle HTTP authentication schemes
+            if (
+                credential.http.scheme.lower() == "bearer"
+                and credential.http.credentials.token
+            ):
+                headers = {
+                    "Authorization": f"Bearer {credential.http.credentials.token}"
+                }
+            elif credential.http.scheme.lower() == "basic":
+                # Handle basic auth
+                if (
+                    credential.http.credentials.username
+                    and credential.http.credentials.password
+                ):
+                    credentials = f"{credential.http.credentials.username}:{credential.http.credentials.password}"
+                    encoded_credentials = base64.b64encode(
+                        credentials.encode()
+                    ).decode()
+                    headers = {"Authorization": f"Basic {encoded_credentials}"}
+            elif credential.http.credentials.token:
+                # Handle other HTTP schemes with token
+                headers = {
+                    "Authorization": (
+                        f"{credential.http.scheme} {credential.http.credentials.token}"
+                    )
+                }
+        elif credential.api_key:
+            headers = {"Authorization": credential.api_key}
+        elif credential.service_account:
+            pass
+
+    return headers

veadk/integrations/ve_prompt_pilot/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

veadk/integrations/ve_prompt_pilot/ve_prompt_pilot.py

@@ -0,0 +1,85 @@
+# 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 agent_pilot as ap
+from agent_pilot.models import TaskType
+from veadk import Agent
+from veadk.prompts import prompt_optimization
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class VePromptPilot:
+    def __init__(
+        self,
+        api_key: str,
+        workspace_id: str,
+        path: str = "",
+        task_id: str | None = None,
+    ) -> None:
+        self.api_key = api_key
+        self.workspace_id = workspace_id
+
+        self.path = path
+
+    def optimize(
+        self,
+        agents: list[Agent],
+        feedback: str = "",
+        model_name: str = "doubao-1.5-pro-32k-250115",
+    ) -> str:
+        for idx, agent in enumerate(agents):
+            optimized_prompt = ""
+            if not feedback:
+                logger.info("Optimizing prompt without feedback.")
+                task_description = prompt_optimization.render_prompt_with_jinja2(agent)
+            else:
+                logger.info(f"Optimizing prompt with feedback: {feedback}")
+                task_description = (
+                    prompt_optimization.render_prompt_feedback_with_jinja2(
+                        agent, feedback
+                    )
+                )
+
+            logger.info(
+                f"Optimizing prompt for agent {agent.name} by {model_name} [{idx + 1}/{len(agents)}]"
+            )
+
+            usage = None
+            for chunk in ap.generate_prompt_stream(
+                task_description=task_description,
+                current_prompt=str(agent.instruction),
+                model_name=model_name,
+                task_type=TaskType.DIALOG,
+                temperature=1.0,
+                top_p=0.7,
+                api_key=self.api_key,
+                workspace_id=self.workspace_id,
+            ):  # stream chunks of optimized prompt
+                # Process each chunk as it arrives
+                optimized_prompt += chunk.data.content if chunk.data else ""
+                # print(chunk.data.content, end="", flush=True)
+                if chunk.event == "usage":
+                    usage = chunk.data.usage if chunk.data else 0
+            optimized_prompt = optimized_prompt.replace("\\n", "\n")
+            print(f"Optimized prompt for agent {agent.name}:\n{optimized_prompt}")
+            if usage:
+                logger.info(f"Token usage: {usage['total_tokens']}")
+            else:
+                logger.warning("No usage data.")
+
+        return optimized_prompt

veadk/integrations/ve_tls/__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.