xpander-sdk 1.60.4__py3-none-any.whl → 2.0.155__py3-none-any.whl

xpander_sdk/modules/backend/utils/mcp_oauth.py

@@ -0,0 +1,95 @@
+import asyncio
+from datetime import datetime, timezone
+from loguru import logger
+from xpander_sdk.consts.api_routes import APIRoute
+from xpander_sdk.core.xpander_api_client import APIClient
+from xpander_sdk.models.events import TaskUpdateEventType
+from xpander_sdk.modules.tasks.sub_modules.task import Task, TaskUpdateEvent
+from xpander_sdk.modules.tools_repository.models.mcp import MCPOAuthGetTokenGenericResponse, MCPOAuthGetTokenResponse, MCPOAuthResponseType, MCPServerDetails
+
+POLLING_INTERVAL = 1 # every 1s
+MAX_WAIT_FOR_LOGIN = 600 # 10 mintutes
+
+async def push_event(task: Task, event: TaskUpdateEvent):
+    client = APIClient(configuration=task.configuration)
+    await client.make_request(
+        path=APIRoute.PushExecutionEventToQueue.format(task_id=task.id),
+        method="POST",
+        payload=[
+            TaskUpdateEvent(
+                task_id=task.id,
+                organization_id=task.organization_id,
+                time=datetime.now(timezone.utc).isoformat(),
+                type=TaskUpdateEventType.AuthEvent,
+                data=event
+                ).model_dump_safe()
+            ]
+    )
+
+async def get_token(mcp_server: MCPServerDetails, task: Task, user_id: str) -> MCPOAuthGetTokenResponse:
+    client = APIClient(configuration=task.configuration)
+    result: MCPOAuthGetTokenResponse = await client.make_request(
+        path=APIRoute.GetUserMCPAuthToken.format(agent_id=task.agent_id, user_id=user_id),
+        method="POST",
+        payload=mcp_server.model_dump(),
+        model=MCPOAuthGetTokenResponse
+    )
+    if result.type == MCPOAuthResponseType.TOKEN_READY:
+        return result
+    
+    return None
+
+async def authenticate_mcp_server(mcp_server: MCPServerDetails, task: Task, user_id: str) -> MCPOAuthGetTokenResponse:
+    try:
+        logger.info(f"Authenticating MCP Server {mcp_server.url}")
+        
+        user_identifier = user_id if mcp_server.share_user_token_across_other_agents else f"{task.agent_id}_{user_id}"
+        
+        client = APIClient(configuration=task.configuration)
+        result: MCPOAuthGetTokenResponse = await client.make_request(
+            path=APIRoute.GetUserMCPAuthToken.format(agent_id=task.agent_id, user_id=user_identifier),
+            method="POST",
+            payload=mcp_server.model_dump(),
+            model=MCPOAuthGetTokenResponse
+        )
+        
+        if not result:
+            raise Exception("Invalid response")
+        
+        if result.type == MCPOAuthResponseType.LOGIN_REQUIRED:
+            logger.info(f"Initiating login for MCP Server {mcp_server.url}")
+            # Notify user about login requirement
+            await push_event(task=task, event=result)
+            
+            # Poll for token with timeout
+            elapsed_time = 0
+            while elapsed_time < MAX_WAIT_FOR_LOGIN:
+                await asyncio.sleep(POLLING_INTERVAL)
+                elapsed_time += POLLING_INTERVAL
+                
+                # Check for token
+                token_result = await get_token(mcp_server=mcp_server, task=task, user_id=user_identifier)
+                if token_result and token_result.type == MCPOAuthResponseType.TOKEN_READY:
+                    logger.info(f"Successful login for MCP Server {mcp_server.url}")
+                    redacted_token_result = MCPOAuthGetTokenResponse(**token_result.model_dump_safe())
+                    redacted_token_result.data.access_token = "REDACTED"
+                    await push_event(task=task, event=redacted_token_result)
+                    return token_result
+            
+            # Timeout reached
+            raise TimeoutError(f"Authentication timeout for MCP Server {mcp_server.url} after {MAX_WAIT_FOR_LOGIN} seconds")
+        
+        if result.type == MCPOAuthResponseType.TOKEN_READY:
+            logger.info(f"Token ready for MCP Server {mcp_server.url}")
+            redacted_token_result = MCPOAuthGetTokenResponse(**result.model_dump_safe())
+            redacted_token_result.data.access_token = "REDACTED"
+            await push_event(task=task, event=redacted_token_result)
+        
+        return result
+    except Exception as e:
+        error = f"Failed to process MCP Auth - {str(e)}"
+        logger.error(error)
+        return MCPOAuthGetTokenResponse(
+            type=MCPOAuthResponseType.TOKEN_ISSUE,
+            data=MCPOAuthGetTokenGenericResponse(message=error)
+        )

xpander_sdk/modules/events/decorators/on_boot.py

@@ -0,0 +1,94 @@
+"""
+xpander_sdk.decorators.on_boot
+
+This module provides the `@on_boot` decorator, which allows developers to define
+boot functions that execute before event listeners are set up in xpander.ai's event system.
+
+The decorator ensures that the registered function:
+- Can be either synchronous or asynchronous
+- Executes before any event listeners are initialized
+- Handles initialization logic, setup tasks, or pre-processing
+
+Execution Notes:
+- Boot handlers execute once during the application startup phase
+- They run before event listeners are registered and task processing begins
+- If an exception is raised in the function, the application will fail to start
+- Use boot handlers for initialization tasks like database connections, configuration loading, etc.
+
+Example usage:
+--------------
+>>> @on_boot
+... async def initialize_database():
+...     database = Database()
+...     await database.connect()
+...     logger.info("Database connected successfully")
+
+>>> @on_boot  
+... def load_configuration():
+...     config = load_config_file()
+...     validate_config(config)
+...     logger.info("Configuration loaded and validated")
+"""
+
+import asyncio
+from functools import wraps
+from inspect import iscoroutinefunction
+from typing import Optional, Callable
+
+from xpander_sdk.models.configuration import Configuration
+from xpander_sdk.modules.events.events_module import Events
+
+
+def on_boot(
+    _func: Optional[Callable] = None,
+    *,
+    configuration: Optional[Configuration] = None
+):
+    """
+    Decorator to register a handler as a boot event executor.
+
+    The decorated function will be executed during the application boot phase, before
+    any event listeners are set up. The function:
+    - Can be either synchronous or asynchronous
+    - Should handle initialization logic or setup tasks
+    - Will be executed only once during application startup
+
+    Args:
+        _func (Optional[Callable]):
+            The function to decorate (for direct usage like `@on_boot`).
+        configuration (Optional[Configuration]):
+            An optional configuration object used to initialize the Events module.
+
+    Example:
+        >>> @on_boot
+        ... async def initialize_services():
+        ...     await setup_database()
+        ...     await connect_external_apis()
+        ...     logger.info("All services initialized")
+        
+        >>> @on_boot
+        ... def load_environment():
+        ...     load_dotenv()
+        ...     validate_environment_variables()
+    """
+
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            return await func(*args, **kwargs)
+
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            return func(*args, **kwargs)
+
+        wrapped = async_wrapper if iscoroutinefunction(func) else sync_wrapper
+
+        # Register handler directly on Events class without instantiating
+        Events.register_boot_handler(wrapped)
+
+        return wrapped
+
+    if _func and callable(_func):
+        return decorator(_func)
+
+    return decorator

xpander_sdk/modules/events/decorators/on_shutdown.py

@@ -0,0 +1,94 @@
+"""
+xpander_sdk.decorators.on_shutdown
+
+This module provides the `@on_shutdown` decorator, which allows developers to define
+shutdown functions that execute when the application is exiting from xpander.ai's event system.
+
+The decorator ensures that the registered function:
+- Can be either synchronous or asynchronous
+- Executes during application shutdown phase
+- Handles cleanup logic, resource disposal, or finalization tasks
+
+Execution Notes:
+- Shutdown handlers execute when the application is terminating
+- They run after event listeners are stopped but before final cleanup
+- If an exception is raised in the function, it will be logged but won't prevent shutdown
+- Use shutdown handlers for cleanup tasks like closing database connections, saving state, etc.
+
+Example usage:
+--------------
+>>> @on_shutdown
+... async def cleanup_database():
+...     database = get_database_connection()
+...     await database.close()
+...     logger.info("Database connection closed")
+
+>>> @on_shutdown  
+... def save_application_state():
+...     state = get_current_state()
+...     save_state_to_file(state)
+...     logger.info("Application state saved")
+"""
+
+import asyncio
+from functools import wraps
+from inspect import iscoroutinefunction
+from typing import Optional, Callable
+
+from xpander_sdk.models.configuration import Configuration
+from xpander_sdk.modules.events.events_module import Events
+
+
+def on_shutdown(
+    _func: Optional[Callable] = None,
+    *,
+    configuration: Optional[Configuration] = None
+):
+    """
+    Decorator to register a handler as a shutdown event executor.
+
+    The decorated function will be executed during the application shutdown phase, after
+    event listeners are stopped. The function:
+    - Can be either synchronous or asynchronous
+    - Should handle cleanup logic or resource disposal
+    - Will be executed once during application shutdown
+
+    Args:
+        _func (Optional[Callable]):
+            The function to decorate (for direct usage like `@on_shutdown`).
+        configuration (Optional[Configuration]):
+            An optional configuration object used to initialize the Events module.
+
+    Example:
+        >>> @on_shutdown
+        ... async def cleanup_resources():
+        ...     await close_database_connections()
+        ...     await cleanup_temp_files()
+        ...     logger.info("All resources cleaned up")
+        
+        >>> @on_shutdown
+        ... def save_final_state():
+        ...     save_application_metrics()
+        ...     logger.info("Final state saved")
+    """
+
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            return await func(*args, **kwargs)
+
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            return func(*args, **kwargs)
+
+        wrapped = async_wrapper if iscoroutinefunction(func) else sync_wrapper
+
+        # Register handler directly on Events class without instantiating
+        Events.register_shutdown_handler(wrapped)
+
+        return wrapped
+
+    if _func and callable(_func):
+        return decorator(_func)
+
+    return decorator

xpander_sdk/modules/events/decorators/on_task.py

@@ -0,0 +1,203 @@
+"""
+xpander_sdk.decorators.on_task
+
+This module provides the `@on_task` decorator, which allows developers to define
+task execution functions that integrate with xpander.ai’s event system.
+
+The decorator ensures that the registered function:
+- Accepts a single parameter named `task` of type `Task`
+- Returns the modified `Task` after processing
+- Optionally works with both sync and async function styles
+
+Execution Notes:
+- The task executor automatically handles status transitions (e.g., success or failure).
+- You must set `task.result = ...` to assign your output.
+- If an exception is raised in the function, the task will be marked as failed.
+- You do not need to manually manage status or lifecycle hooks.
+
+Example usage:
+--------------
+>>> @on_task
+... async def handle_agent_task(task: Task):
+...     backend = Backend(task.configuration)
+...     agno_args = await backend.aget_args(
+...         agent_id=task.agent_id,
+...         agent_version=task.agent_version,
+...         task=task
+...     )
+...     agno_agent = Agent(**agno_args)
+...     result = await agno_agent.arun(message=task.to_message())
+...     task.result = result.content
+...     return task
+"""
+
+import asyncio
+import argparse
+import json
+import sys
+from functools import wraps
+from inspect import iscoroutinefunction, signature
+from typing import Optional, Callable
+
+from xpander_sdk.models.configuration import Configuration
+from xpander_sdk.modules.events.events_module import Events
+from xpander_sdk.modules.tasks.models.task import LocalTaskTest, AgentExecutionInput
+from xpander_sdk.models.shared import OutputFormat
+from xpander_sdk.modules.tasks.sub_modules.task import Task
+
+
+def on_task(
+    _func: Optional[Callable] = None,
+    *,
+    configuration: Optional[Configuration] = None,
+    test_task: Optional[LocalTaskTest] = None
+):
+    """
+    Decorator to register a handler as an event-driven task executor.
+
+    The decorated function will be executed when the task is triggered through
+    xpander.ai’s event system. The function:
+    - Must accept a single parameter named `task` of type `Task`
+    - Must return the `Task` instance after applying changes
+    - Should assign the final output using `task.result`
+    - Does not need to set the task status; it's handled automatically
+
+    - Supports both synchronous and asynchronous functions
+    - Automatically transitions task statuses based on function outcome
+    Args:
+        _func (Optional[Callable]):
+            The function to decorate (for direct usage like `@on_task`).
+        configuration (Optional[Configuration]):
+            An optional configuration object used to initialize the Events module.
+        test_task (Optional[LocalTaskTest]):
+            Optional simulated task used for local development or testing.
+
+    Raises:
+        TypeError: If the decorated function does not accept a `task` parameter.
+        TypeError: If `test_task` is provided but is not a `LocalTaskTest`.
+        TypeError: If the input `task` is not an instance of `Task`.
+        TypeError: If the return value is not an instance of `Task`.
+
+    Example:
+        >>> @on_task
+        ... async def handle_agent_task(task: Task):
+        ...     backend = Backend(task.configuration)
+        ...     agno_args = await backend.aget_args(
+        ...         agent_id=task.agent_id,
+        ...         agent_version=task.agent_version,
+        ...         task=task
+        ...     )
+        ...     agno_agent = Agent(**agno_args)
+        ...     result = await agno_agent.arun(message=task.to_message())
+        ...     task.result = result.content
+        ...     return task
+    """
+
+    def decorator(func: Callable) -> Callable:
+        sig = signature(func)
+
+        if 'task' not in sig.parameters:
+            raise TypeError(f"Function '{func.__name__}' must have a 'task' parameter.")
+
+        # Initialize effective_test_task with the passed parameter
+        effective_test_task = test_task
+        
+        if effective_test_task and not isinstance(effective_test_task, LocalTaskTest):
+            raise TypeError(
+                f"'test_task' must be an instance of LocalTaskTest, got {type(effective_test_task).__name__}"
+            )
+
+        param_names = list(sig.parameters)
+
+        def get_task_from_call(args, kwargs) -> Task:
+            if 'task' in kwargs:
+                task_value = kwargs['task']
+            else:
+                task_index = param_names.index('task')
+                task_value = args[task_index] if len(args) > task_index else None
+            if not isinstance(task_value, Task):
+                raise TypeError(
+                    f"'task' must be an instance of Task, got {type(task_value).__name__}"
+                )
+            return task_value
+
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            get_task_from_call(args, kwargs)
+            result = await func(*args, **kwargs)
+            if not isinstance(result, Task):
+                raise TypeError(
+                    f"Function '{func.__name__}' must return an instance of Task, got {type(result).__name__}"
+                )
+            return result
+
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            get_task_from_call(args, kwargs)
+            result = func(*args, **kwargs)
+            if not isinstance(result, Task):
+                raise TypeError(
+                    f"Function '{func.__name__}' must return an instance of Task, got {type(result).__name__}"
+                )
+            return result
+
+        wrapped = async_wrapper if iscoroutinefunction(func) else sync_wrapper
+
+        # Check for direct invocation with --invoke and --prompt arguments
+        if any('--invoke' in arg for arg in sys.argv):
+            # Parse command line arguments
+            parser = argparse.ArgumentParser(description='Run task handler locally')
+            parser.add_argument('--invoke', action='store_true', help='Invoke the task handler directly')
+            parser.add_argument('--prompt', type=str, required='--invoke' in sys.argv, help='Input text for the task')
+            parser.add_argument('--output_format', type=str, choices=['json', 'markdown', 'text'], 
+                               help='Output format (json, markdown, text)')
+            parser.add_argument('--output_schema', type=str, help='JSON schema as a string for output validation')
+            
+            # Parse only known args to avoid conflicts with other args that might be present
+            args, _ = parser.parse_known_args()
+            
+            if args.invoke:
+                # Create a LocalTaskTest instance from command line arguments (like existing test_task logic)
+                output_format_map = {
+                    'json': OutputFormat.Json,
+                    'markdown': OutputFormat.Markdown,
+                    'text': OutputFormat.Text
+                }
+                
+                # Parse output schema if provided
+                output_schema = None
+                if args.output_schema:
+                    try:
+                        output_schema = json.loads(args.output_schema)
+                    except json.JSONDecodeError as e:
+                        print(f"Error parsing output_schema JSON: {e}")
+                        sys.exit(1)
+                
+                # Create LocalTaskTest (same as existing test_task functionality)
+                # Only include output_format and output_schema if explicitly provided
+                task_kwargs = {
+                    'input': AgentExecutionInput(text=args.prompt)
+                }
+                
+                if args.output_format:
+                    task_kwargs['output_format'] = output_format_map.get(args.output_format.lower())
+                    
+                if output_schema is not None:
+                    task_kwargs['output_schema'] = output_schema
+                
+                cli_test_task = LocalTaskTest(**task_kwargs)
+                # Override effective_test_task with CLI-provided one
+                effective_test_task = cli_test_task
+                print(f"Running task handler directly with prompt: {args.prompt}")
+                if output_schema:
+                    print(f"Using output schema: {output_schema}")
+        
+        events_module = Events(configuration=configuration)
+        events_module.register(on_task=wrapped, test_task=effective_test_task)
+        
+        return wrapped
+
+    if _func and callable(_func):
+        return decorator(_func)
+
+    return decorator