google-adk-extras 0.1.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

google_adk_extras/credentials/x_oauth2_credential_service.py

@@ -0,0 +1,240 @@
+"""X (Twitter) OAuth2 credential service implementation."""
+
+from typing import Optional, List
+import logging
+
+from google.adk.auth.credential_service.session_state_credential_service import SessionStateCredentialService
+from google.adk.auth.credential_service.base_credential_service import CallbackContext
+from google.adk.auth import AuthConfig, AuthCredential, AuthCredentialTypes
+from google.adk.auth.auth_credential import OAuth2Auth
+from fastapi.openapi.models import OAuth2
+from fastapi.openapi.models import OAuthFlowAuthorizationCode, OAuthFlows
+
+from .base_custom_credential_service import BaseCustomCredentialService
+
+logger = logging.getLogger(__name__)
+
+
+class XOAuth2CredentialService(BaseCustomCredentialService):
+    """X (Twitter) OAuth2 credential service for handling X API authentication flows.
+    
+    This service provides pre-configured OAuth2 flows for X API v2 including
+    reading tweets, posting content, and managing user data.
+    
+    Args:
+        client_id: The X OAuth2 client ID from X Developer Portal.
+        client_secret: The X OAuth2 client secret from X Developer Portal.
+        scopes: List of OAuth2 scopes to request. Common scopes include:
+            - "tweet.read" - Read tweets
+            - "tweet.write" - Write tweets
+            - "tweet.moderate.write" - Moderate tweets
+            - "users.read" - Read user information
+            - "follows.read" - Read follows information
+            - "follows.write" - Manage follows
+            - "offline.access" - Maintain access when user is offline
+            - "space.read" - Read Spaces information
+            - "mute.read" - Read muted accounts
+            - "mute.write" - Manage muted accounts
+            - "like.read" - Read likes information
+            - "like.write" - Manage likes
+            - "list.read" - Read list information
+            - "list.write" - Manage lists
+            - "block.read" - Read blocked accounts
+            - "block.write" - Manage blocked accounts
+        use_session_state: If True, stores credentials in session state. If False,
+            uses in-memory storage. Default is True for persistence.
+    
+    Example:
+        ```python
+        credential_service = XOAuth2CredentialService(
+            client_id="your-x-client-id",
+            client_secret="your-x-client-secret",
+            scopes=["tweet.read", "tweet.write", "users.read", "offline.access"]
+        )
+        await credential_service.initialize()
+        
+        # Use with Runner
+        runner = Runner(
+            agent=agent,
+            session_service=session_service,
+            credential_service=credential_service,
+            app_name="my_app"
+        )
+        ```
+    """
+
+    # X OAuth2 endpoints
+    X_AUTH_URL = "https://twitter.com/i/oauth2/authorize"
+    X_TOKEN_URL = "https://api.twitter.com/2/oauth2/token"
+    
+    # Common X OAuth2 scopes
+    COMMON_SCOPES = {
+        # Tweet scopes
+        "tweet.read": "Read tweets",
+        "tweet.write": "Write tweets",
+        "tweet.moderate.write": "Hide and unhide replies to your tweets",
+        
+        # User scopes
+        "users.read": "Read user profile information",
+        
+        # Follows scopes
+        "follows.read": "Read who a user is following or who is following a user",
+        "follows.write": "Follow and unfollow other users",
+        
+        # Offline access
+        "offline.access": "Maintain access to accounts when users are offline",
+        
+        # Space scopes
+        "space.read": "Read Spaces information",
+        
+        # Mute scopes
+        "mute.read": "Read muted accounts",
+        "mute.write": "Mute and unmute accounts",
+        
+        # Like scopes
+        "like.read": "Read liked tweets",
+        "like.write": "Like and unlike tweets",
+        
+        # List scopes
+        "list.read": "Read list information",
+        "list.write": "Create and manage lists",
+        
+        # Block scopes
+        "block.read": "Read blocked accounts",
+        "block.write": "Block and unblock accounts",
+        
+        # Bookmark scopes
+        "bookmark.read": "Read bookmarked tweets",
+        "bookmark.write": "Bookmark and unbookmark tweets",
+        
+        # Direct message scopes
+        "dm.read": "Read direct messages",
+        "dm.write": "Send direct messages",
+    }
+
+    def __init__(
+        self,
+        client_id: str,
+        client_secret: str,
+        scopes: Optional[List[str]] = None,
+        use_session_state: bool = True
+    ):
+        """Initialize the X OAuth2 credential service.
+        
+        Args:
+            client_id: X OAuth2 client ID.
+            client_secret: X OAuth2 client secret.
+            scopes: List of OAuth2 scopes to request.
+            use_session_state: Whether to use session state for credential storage.
+        """
+        super().__init__()
+        self.client_id = client_id
+        self.client_secret = client_secret
+        self.scopes = scopes or ["tweet.read", "users.read", "offline.access"]
+        self.use_session_state = use_session_state
+        
+        # Underlying credential service for storage
+        if use_session_state:
+            self._storage_service = SessionStateCredentialService()
+        else:
+            from google.adk.auth.credential_service.in_memory_credential_service import InMemoryCredentialService
+            self._storage_service = InMemoryCredentialService()
+
+    async def _initialize_impl(self) -> None:
+        """Initialize the X OAuth2 credential service.
+        
+        Validates the client credentials and sets up the OAuth2 auth scheme.
+        
+        Raises:
+            ValueError: If client_id or client_secret is missing.
+        """
+        if not self.client_id:
+            raise ValueError("X OAuth2 client_id is required")
+        if not self.client_secret:
+            raise ValueError("X OAuth2 client_secret is required")
+        if not self.scopes:
+            raise ValueError("At least one OAuth2 scope is required")
+            
+        # Validate scopes against known X scopes
+        unknown_scopes = set(self.scopes) - set(self.COMMON_SCOPES.keys())
+        if unknown_scopes:
+            logger.warning(f"Unknown X OAuth2 scopes: {unknown_scopes}")
+            
+        logger.info(f"Initialized X OAuth2 credential service with scopes: {self.scopes}")
+
+    def create_auth_config(self) -> AuthConfig:
+        """Create an AuthConfig for X OAuth2 authentication.
+        
+        Returns:
+            AuthConfig: Configured auth config for X OAuth2 flow.
+        """
+        self._check_initialized()
+        
+        # Create OAuth2 auth scheme
+        auth_scheme = OAuth2(
+            flows=OAuthFlows(
+                authorizationCode=OAuthFlowAuthorizationCode(
+                    authorizationUrl=self.X_AUTH_URL,
+                    tokenUrl=self.X_TOKEN_URL,
+                    scopes={
+                        scope: self.COMMON_SCOPES.get(scope, f"X API scope: {scope}")
+                        for scope in self.scopes
+                    }
+                )
+            )
+        )
+        
+        # Create OAuth2 credential
+        auth_credential = AuthCredential(
+            auth_type=AuthCredentialTypes.OAUTH2,
+            oauth2=OAuth2Auth(
+                client_id=self.client_id,
+                client_secret=self.client_secret
+            )
+        )
+        
+        return AuthConfig(
+            auth_scheme=auth_scheme,
+            raw_auth_credential=auth_credential
+        )
+
+    async def load_credential(
+        self,
+        auth_config: AuthConfig,
+        callback_context: CallbackContext,
+    ) -> Optional[AuthCredential]:
+        """Load X OAuth2 credential from storage.
+        
+        Args:
+            auth_config: The auth config containing credential key information.
+            callback_context: The current callback context.
+            
+        Returns:
+            Optional[AuthCredential]: The stored credential or None if not found.
+        """
+        self._check_initialized()
+        return await self._storage_service.load_credential(auth_config, callback_context)
+
+    async def save_credential(
+        self,
+        auth_config: AuthConfig,
+        callback_context: CallbackContext,
+    ) -> None:
+        """Save X OAuth2 credential to storage.
+        
+        Args:
+            auth_config: The auth config containing the credential to save.
+            callback_context: The current callback context.
+        """
+        self._check_initialized()
+        await self._storage_service.save_credential(auth_config, callback_context)
+        
+        logger.info(f"Saved X OAuth2 credential for user {callback_context._invocation_context.user_id}")
+
+    def get_supported_scopes(self) -> dict:
+        """Get dictionary of supported X OAuth2 scopes and their descriptions.
+        
+        Returns:
+            dict: Mapping of scope names to descriptions.
+        """
+        return self.COMMON_SCOPES.copy()

google_adk_extras/custom_agent_loader.py

@@ -0,0 +1,156 @@
+"""CustomAgentLoader - Enhanced agent loader for programmatic agent management.
+
+This module provides CustomAgentLoader which extends Google ADK's agent loading
+capabilities to support programmatically registered agent instances with
+thread-safe registry management.
+"""
+
+import logging
+import threading
+from typing import Dict, List, Optional
+
+from google.adk.agents.base_agent import BaseAgent
+from google.adk.cli.utils.base_agent_loader import BaseAgentLoader
+
+logger = logging.getLogger(__name__)
+
+
+class CustomAgentLoader(BaseAgentLoader):
+    """Enhanced agent loader for programmatic agent management.
+    
+    This loader allows you to:
+    1. Register agent instances directly for programmatic control
+    2. Dynamically add/remove agents at runtime
+    3. Thread-safe access to agent registry
+    
+    Examples:
+        # Register and use agents
+        loader = CustomAgentLoader()
+        loader.register_agent("my_agent", my_agent_instance)
+        agent = loader.load_agent("my_agent")
+        
+        # List all registered agents
+        agents = loader.list_agents()  # ['my_agent']
+    """
+    
+    def __init__(self):
+        """Initialize CustomAgentLoader."""
+        self._registered_agents: Dict[str, BaseAgent] = {}
+        self._lock = threading.RLock()  # Thread-safe access to registry
+        
+        logger.debug("CustomAgentLoader initialized")
+    
+    def register_agent(self, name: str, agent: BaseAgent) -> None:
+        """Register an agent instance by name.
+        
+        Args:
+            name: Agent name for discovery and loading.
+            agent: BaseAgent instance to register.
+            
+        Raises:
+            ValueError: If name is empty or agent is not a BaseAgent instance.
+        """
+        if not name or not name.strip():
+            raise ValueError("Agent name cannot be empty")
+            
+        if not isinstance(agent, BaseAgent):
+            raise ValueError(f"Agent must be BaseAgent instance, got {type(agent)}")
+        
+        with self._lock:
+            if name in self._registered_agents:
+                logger.info("Replacing existing registered agent: %s", name)
+            else:
+                logger.info("Registering new agent instance: %s", name)
+            
+            self._registered_agents[name] = agent
+    
+    def unregister_agent(self, name: str) -> bool:
+        """Unregister an agent instance by name.
+        
+        Args:
+            name: Name of agent to unregister.
+            
+        Returns:
+            bool: True if agent was found and removed, False otherwise.
+        """
+        with self._lock:
+            if name in self._registered_agents:
+                del self._registered_agents[name]
+                logger.info("Unregistered agent instance: %s", name)
+                return True
+            else:
+                logger.debug("Agent not found in registry: %s", name)
+                return False
+    
+    def is_registered(self, name: str) -> bool:
+        """Check if an agent is registered by name.
+        
+        Args:
+            name: Agent name to check.
+            
+        Returns:
+            bool: True if agent is registered, False otherwise.
+        """
+        with self._lock:
+            return name in self._registered_agents
+    
+    def get_registered_agents(self) -> Dict[str, BaseAgent]:
+        """Get a copy of all registered agents.
+        
+        Returns:
+            Dict[str, BaseAgent]: Copy of registered agents mapping.
+        """
+        with self._lock:
+            return self._registered_agents.copy()
+    
+    def clear_registry(self) -> None:
+        """Clear all registered agents from the registry."""
+        with self._lock:
+            count = len(self._registered_agents)
+            self._registered_agents.clear()
+            logger.info("Cleared %d registered agents", count)
+    
+    def load_agent(self, name: str) -> BaseAgent:
+        """Load an agent by name.
+        
+        Args:
+            name: Name of agent to load.
+            
+        Returns:
+            BaseAgent: The loaded agent instance.
+            
+        Raises:
+            ValueError: If agent is not found in registry.
+        """
+        with self._lock:
+            if name in self._registered_agents:
+                logger.debug("Loading registered agent: %s", name)
+                return self._registered_agents[name]
+        
+        # Agent not found
+        available_agents = self.list_agents()
+        raise ValueError(
+            f"Agent '{name}' not found. "
+            f"Available agents: {available_agents if available_agents else 'None'}"
+        )
+    
+    def list_agents(self) -> List[str]:
+        """List all available agents from registry.
+        
+        Returns:
+            List[str]: Sorted list of all registered agent names.
+        """
+        with self._lock:
+            agent_names = list(self._registered_agents.keys())
+        
+        sorted_agents = sorted(agent_names)
+        logger.debug("Total registered agents: %d", len(sorted_agents))
+        return sorted_agents
+    
+    
+    def __repr__(self) -> str:
+        """String representation of the loader."""
+        with self._lock:
+            registered_count = len(self._registered_agents)
+        
+        return f"CustomAgentLoader(registered={registered_count})"

google_adk_extras/enhanced_adk_web_server.py

@@ -0,0 +1,137 @@
+"""Enhanced ADK Web Server that uses EnhancedRunner.
+
+This module provides the EnhancedAdkWebServer class which extends Google ADK's
+AdkWebServer to use our EnhancedRunner with advanced features.
+"""
+
+import os
+from typing import Optional
+
+from google.adk.cli.adk_web_server import AdkWebServer
+from google.adk.auth.credential_service.in_memory_credential_service import InMemoryCredentialService
+from google.adk.cli.utils import cleanup
+from google.adk.cli.utils import envs
+from google.adk.runners import Runner
+
+from .enhanced_runner import EnhancedRunner
+
+
+class EnhancedAdkWebServer(AdkWebServer):
+    """Enhanced ADK Web Server that creates EnhancedRunner instances.
+    
+    This class extends Google's AdkWebServer to use our EnhancedRunner with:
+    - Advanced tool execution strategies (MCP, OpenAPI, Function tools)
+    - Circuit breakers and retry policies for resilience
+    - YAML system context and enhanced configuration
+    - Performance monitoring and debugging capabilities
+    - Credential service integration (inherited)
+    
+    The EnhancedAdkWebServer is a drop-in replacement for AdkWebServer that
+    provides significantly enhanced capabilities while maintaining full
+    backward compatibility with all existing APIs.
+    
+    Examples:
+        Basic usage (drop-in replacement):
+        ```python
+        enhanced_server = EnhancedAdkWebServer(
+            agent_loader=agent_loader,
+            session_service=session_service,
+            artifact_service=artifact_service,
+            memory_service=memory_service,
+            credential_service=credential_service,
+            eval_sets_manager=eval_sets_manager,
+            eval_set_results_manager=eval_set_results_manager,
+            agents_dir="./agents"
+        )
+        ```
+        
+        With enhanced features:
+        ```python
+        enhanced_config = EnhancedRunConfig.from_yaml_dict({
+            'max_llm_calls': 200,
+            'tool_timeouts': {'mcp_tools': 30.0},
+            'circuit_breaker': {'failure_threshold': 3},
+            'debug': {'enabled': True}
+        })
+        
+        strategy_manager = ToolExecutionStrategyManager()
+        strategy_manager.register_strategy('mcp', McpToolExecutionStrategy(timeout=45.0))
+        
+        enhanced_server = EnhancedAdkWebServer(
+            agent_loader=agent_loader,
+            session_service=session_service,
+            # ... other services ...
+            enhanced_config=enhanced_config,
+            yaml_context=YamlSystemContext(
+                system_name="my-agent-system",
+                config_path="/path/to/config.yaml"
+            ),
+            tool_strategy_manager=strategy_manager
+        )
+        ```
+    """
+    
+    def __init__(self, **kwargs):
+        """Initialize EnhancedAdkWebServer.
+        
+        Args:
+            enhanced_config: Enhanced configuration for runners (optional)
+            yaml_context: YAML system context for error handling (optional)
+            tool_strategy_manager: Tool execution strategy manager (optional)
+            **kwargs: All other parameters passed to AdkWebServer
+        """
+        # Ensure a credential service exists; default to InMemory if not provided
+        if 'credential_service' not in kwargs or kwargs.get('credential_service') is None:
+            kwargs['credential_service'] = InMemoryCredentialService()
+
+        # Initialize base AdkWebServer with all standard parameters
+        super().__init__(**kwargs)
+        
+        # No enhanced configuration retained in simplified scope
+
+    async def get_runner_async(self, app_name: str) -> EnhancedRunner:
+        """Returns the enhanced runner for the given app.
+        
+        This method overrides AdkWebServer.get_runner_async to create
+        EnhancedRunner instances instead of standard Runner instances.
+        
+        The logic is identical to the parent class except:
+        1. Creates EnhancedRunner instead of Runner
+        2. Passes enhanced configuration parameters
+        3. Maintains full compatibility with cleanup and caching
+        
+        Args:
+            app_name: The name of the application/agent to get runner for
+            
+        Returns:
+            An EnhancedRunner instance for the specified app
+        """
+        # EXACT copy of parent logic for cleanup and caching
+        if app_name in self.runners_to_clean:
+            self.runners_to_clean.remove(app_name)
+            runner = self.runner_dict.pop(app_name, None)
+            await cleanup.close_runners(list([runner]))
+
+        # Load environment for the agent
+        envs.load_dotenv_for_agent(os.path.basename(app_name), self.agents_dir)
+        
+        # Return cached runner if available
+        if app_name in self.runner_dict:
+            return self.runner_dict[app_name]
+            
+        # Load agent and create new EnhancedRunner
+        root_agent = self.agent_loader.load_agent(app_name)
+        
+        # Create EnhancedRunner (thin wrapper over ADK Runner)
+        runner = EnhancedRunner(
+            app_name=app_name,
+            agent=root_agent,
+            artifact_service=self.artifact_service,
+            session_service=self.session_service,
+            memory_service=self.memory_service,
+            credential_service=self.credential_service,
+        )
+        
+        # Cache and return runner (same as parent)
+        self.runner_dict[app_name] = runner
+        return runner