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

xpander_sdk/modules/events/events_module.py

@@ -0,0 +1,629 @@
+"""
+Events module for handling background tasks and event streaming in the xpander.ai platform.
+
+This module provides functionality for managing Server Sent Events (SSE) and executing tasks
+based on events within the xpander.ai platform. It supports asynchronous execution and retry logic.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import json as py_json
+import os
+import signal
+import sys
+from os import getenv
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Awaitable, Callable, Optional, Set, Union, List
+
+import httpx
+from httpx_sse import aconnect_sse
+from loguru import logger
+from pydantic import BaseModel
+
+from xpander_sdk.core.module_base import ModuleBase
+from xpander_sdk.exceptions.module_exception import ModuleException
+from xpander_sdk.models.configuration import Configuration
+from xpander_sdk.models.shared import OutputFormat
+from xpander_sdk.modules.agents.models.agent import SourceNodeType
+from xpander_sdk.modules.tasks.tasks_module import Tasks
+
+from .utils.git_init import configure_git_credentials
+from .utils.generic import backoff_delay, get_events_base, get_events_headers
+from .models.deployments import DeployedAsset
+from .models.events import (
+    EventType,
+    WorkerEnvironmentConflict,
+    WorkerFinishedEvent,
+    WorkerHeartbeat,
+)
+from ..tasks.sub_modules.task import Task
+from ..tasks.models.task import AgentExecutionStatus, LocalTaskTest
+
+
+_MAX_RETRIES = 5  # total attempts (1 initial + 4 retries)
+
+ExecutionRequestHandler = Union[
+    Callable[[Task], Task],
+    Callable[[Task], Awaitable[Task]],
+]
+
+BootHandler = Union[
+    Callable[[], None],
+    Callable[[], Awaitable[None]],
+]
+
+ShutdownHandler = Union[
+    Callable[[], None],
+    Callable[[], Awaitable[None]],
+]
+
+
+class Events(ModuleBase):
+    """
+    Events module for managing SSE connections and task execution.
+
+    This class manages Server Sent Events (SSE) for real-time task execution requests
+    and integrates with agents deployed on xpander.ai. It handles event streaming,
+    retry logic, and background task management. The worker is directly attached
+    to the agent without a parent worker hierarchy.
+
+    Attributes:
+        worker (Optional[DeployedAsset]): Represents the deployed asset/agent worker.
+        test_task (Optional[LocalTaskTest]): Task to be tested within the local environment.
+        configuration (Configuration): SDK configuration with credentials and endpoints.
+
+    Example:
+        >>> events = Events()
+        >>> events.register(on_task=handle_task)
+    """
+
+    worker: Optional[DeployedAsset] = None
+    test_task: Optional[LocalTaskTest] = None
+    
+    # Class-level registries for boot and shutdown handlers
+    _boot_handlers: List[BootHandler] = []
+    _shutdown_handlers: List[ShutdownHandler] = []
+
+    def __init__(
+        self,
+        configuration: Optional[Configuration] = None,
+        max_sync_workers: Optional[int] = 4,
+        max_retries: Optional[int] = _MAX_RETRIES,
+    ):
+        """
+        Initialize the Events module with configuration and worker settings.
+
+        Configures event streaming parameters and validates essential environment setup.
+
+        Args:
+            configuration (Optional[Configuration]): SDK configuration with credentials and endpoints. Defaults to environment configuration.
+            max_sync_workers (Optional[int]): Maximum number of synchronous worker threads. Defaults to 4.
+            max_retries (Optional[int]): Maximum retry attempts for network calls. Defaults to 5.
+
+        Raises:
+            ModuleException: When required environment variables are missing or configuration is incorrect.
+        """
+        super().__init__(configuration)
+        configure_git_credentials()
+
+        self.is_xpander_cloud = getenv("IS_XPANDER_CLOUD", "false") == "true"
+        self.agent_id = self.configuration.agent_id or getenv("XPANDER_AGENT_ID", None)
+
+        if not self.agent_id:
+            raise ModuleException(
+                400, "XPANDER_AGENT_ID is missing from your environment variables."
+            )
+        if not self.configuration.organization_id:
+            raise ModuleException(
+                400,
+                "XPANDER_ORGANIZATION_ID is missing from your environment variables.",
+            )
+        if not self.configuration.api_key:
+            raise ModuleException(
+                400, "XPANDER_API_KEY is missing from your environment variables."
+            )
+
+        self.max_retries = max_retries
+
+        # Internal resources
+        self._pool: ThreadPoolExecutor = ThreadPoolExecutor(
+            max_workers=max_sync_workers,
+            thread_name_prefix="xpander-handler",
+        )
+        self._bg: Set[asyncio.Task] = set()
+
+        logger.debug(
+            f"Events initialised (base_url={self.configuration.base_url}, "
+            f"org_id={self.configuration.organization_id}, retries={self.max_retries})"
+        )
+
+    # lifecycle
+    async def start(
+        self,
+        on_execution_request: ExecutionRequestHandler,
+    ) -> None:
+        """
+        Start the event listener and register handlers for task execution events.
+
+        This method sets up signal handling for graceful shutdown, registers the
+        agent worker directly, and begins listening to task execution requests over SSE.
+        Use the @on_task decorator instead of calling this method directly.
+
+        Args:
+            on_execution_request (ExecutionRequestHandler): Callback handler
+                for processing task execution requests. Can be synchronous or asynchronous.
+        """
+        # Execute boot handlers first, before any event listeners are set up
+        await self._execute_boot_handlers()
+        
+        loop = asyncio.get_running_loop()
+        for sig in (signal.SIGINT, signal.SIGTERM):
+            loop.add_signal_handler(
+                sig, lambda s=sig: asyncio.create_task(self.stop(s))
+            )
+
+        # Register agent worker directly
+        self.track(
+            asyncio.create_task(
+                self.register_agent_worker(self.agent_id, on_execution_request)
+            )
+        )
+
+        logger.info("Listener started; waiting for events…")
+        await asyncio.gather(*self._bg)
+
+    async def stop(self, sig: signal.Signals | None = None) -> None:
+        """
+        Stop the event listener and cleanup background tasks.
+
+        Args:
+            sig (signal.Signals | None): Signal that triggered the stop request.
+
+        Example:
+            >>> await events.stop()
+        """
+        if sig:
+            logger.info(f"Received {sig.name} – shutting down…")
+
+        for t in self._bg:
+            t.cancel()
+        if self._bg:
+            await asyncio.gather(*self._bg, return_exceptions=True)
+
+        self._pool.shutdown(wait=False, cancel_futures=True)
+        self._bg.clear()
+        
+        # Execute shutdown handlers after stopping event listeners but before final cleanup
+        await self._execute_shutdown_handlers()
+        
+        logger.info("Listener stopped.")
+
+    async def __aenter__(self) -> "Events":
+        return self
+
+    async def __aexit__(self, *_exc) -> bool:  # noqa: D401
+        await self.stop()
+        return False
+
+    # ---------------------- HTTP helpers with retry ---------------------- #
+
+    async def _request_with_retries(
+        self,
+        method: str,
+        url: str,
+        *,
+        headers: dict[str, str],
+        json: Any | None = None,
+        timeout: float | None = 10.0,
+    ) -> httpx.Response:
+        """
+        Perform an HTTP request with automatic retries on failure.
+
+        Args:
+            method (str): HTTP method to use for the request (e.g., 'POST', 'GET').
+            url (str): The URL to which the request is sent.
+            headers (dict[str, str]): HTTP headers to include in the request.
+            json (Any | None, optional): JSON payload to send with the request.
+            timeout (float | None, optional): Timeout for the request.
+
+        Returns:
+            httpx.Response: The response object received from the request.
+
+        Raises:
+            Exception: If the request fails after the maximum retry attempts.
+        """
+        last_exc: Exception | None = None
+        for attempt in range(1, self.max_retries + 1):
+            try:
+                async with httpx.AsyncClient(timeout=timeout) as client:
+                    response = await client.request(
+                        method,
+                        url,
+                        headers=headers,
+                        json=json,
+                        follow_redirects=True,
+                    )
+                return response
+            except Exception as exc:  # noqa: BLE001 broad (includes timeouts)
+                last_exc = exc
+                if attempt < self.max_retries:
+                    delay = backoff_delay(attempt)
+                    await asyncio.sleep(delay)
+                else:
+                    logger.error(
+                        f"{method} {url} failed after {self.max_retries} attempts - exiting. ({exc})"
+                    )
+                    sys.exit(1)
+        assert last_exc is not None
+        raise last_exc  # for static checkers
+
+    async def _release_worker(self, worker_id: str) -> None:
+        """
+        Release the worker resource after task execution completion.
+
+        Args:
+            worker_id (str): The unique identifier of the worker to release.
+        """
+        url = f"{get_events_base(configuration=self.configuration)}/{worker_id}?type=worker&agent_id={self.agent_id}"
+        await self._request_with_retries(
+            "POST",
+            url,
+            headers=get_events_headers(configuration=self.configuration),
+            json=WorkerFinishedEvent(data={}).model_dump_safe(),
+        )
+
+    async def _make_heartbeat(self, worker_id: str) -> None:
+        """
+        Send a heartbeat signal to maintain the worker's active status.
+
+        Args:
+            worker_id (str): The unique identifier of the worker to update.
+        """
+        url = f"{get_events_base(configuration=self.configuration)}/{worker_id}?type=worker&agent_id={self.agent_id}"
+        await self._request_with_retries(
+            "POST",
+            url,
+            headers=get_events_headers(configuration=self.configuration),
+            json=WorkerHeartbeat().model_dump_safe(),
+        )
+
+
+    # ----------------------- SSE helpers with retry ---------------------- #
+
+    async def _sse_events_with_retries(self, url: str):
+        """Yield Server-Sent Events with reconnect/back‑off logic using httpx‑sse."""
+        attempt = 1
+        while True:
+            try:
+                async with httpx.AsyncClient(
+                    timeout=None, follow_redirects=True
+                ) as client:
+                    if not url.endswith('/'):
+                        url += "/"
+                    async with aconnect_sse(
+                        client,
+                        "GET",
+                        url,
+                        headers=get_events_headers(configuration=self.configuration),
+                        follow_redirects=True
+                    ) as event_source:
+                        async for sse in event_source.aiter_sse():
+                            yield sse
+
+                # Server closed the stream gracefully – reconnect
+                attempt = 1
+                await asyncio.sleep(backoff_delay(1))
+
+            except Exception as exc:  # noqa: BLE001 broad
+                if attempt >= self.max_retries:
+                    logger.error(
+                        f"SSE connection to {url} failed after {self.max_retries} attempts – exiting. ({exc})"
+                    )
+                    sys.exit(1)
+                await asyncio.sleep(backoff_delay(attempt))
+                attempt += 1
+
+    async def handle_task_execution_request(
+        self,
+        agent_worker: DeployedAsset,
+        task: Task,
+        on_execution_request: ExecutionRequestHandler,
+    ) -> None:
+        """
+        Handle an incoming task execution request.
+
+        Args:
+            agent_worker (DeployedAsset): The deployed asset (agent) to handle the task.
+            task (Task): The task object containing execution details.
+            on_execution_request (ExecutionRequestHandler): The handler function to process the task.
+        """
+        error = None
+        try:
+            logger.info(f"Handling task {task.id}")
+            await task.aset_status(status=AgentExecutionStatus.Executing)
+            if asyncio.iscoroutinefunction(on_execution_request):
+                task = await on_execution_request(task)
+            else:
+                task = await asyncio.get_running_loop().run_in_executor(
+                    self._pool,
+                    on_execution_request,
+                    task,
+                )
+        except Exception as e:
+            logger.exception(f"Execution handler failed - {str(e)}")
+            error = str(e)
+        finally:
+            task_used_tokens = task.tokens
+            task_used_tools = task.used_tools
+            await self._release_worker(agent_worker.id)
+
+            if error:
+                task.result = error
+                task.status = AgentExecutionStatus.Error
+            elif (
+                task.status == AgentExecutionStatus.Executing
+            ):  # let the handler set the status, if not set - mark as completed
+                task.status = AgentExecutionStatus.Completed
+
+            # in case of structured output, return as stringified json
+            try:
+                if task.output_format == OutputFormat.Json:
+                    if isinstance(task.result, BaseModel):
+                        task.result = task.result.model_dump_json()
+                    if isinstance(task.result, dict) or isinstance(task.result, list):
+                        task.result = py_json.dumps(task.result)
+            except Exception:
+                pass
+            
+            await task.asave()
+            task.tokens = task_used_tokens
+            task.used_tools = task_used_tools
+            
+            if task.tokens:
+                await task.areport_metrics()
+
+            logger.info(f"Finished handling task {task.id}")
+
+            # local test task, finish? kill the worker
+            if self.test_task:
+                logger.info("Local task handled, exiting")
+                
+                # Print the task result for CLI
+                if task.result:
+                    logger.info("\n" + "="*50)
+                    logger.info("TASK RESULT:")
+                    logger.info("="*50)
+                    if isinstance(task.result, (dict, list)):
+                        import json
+                        logger.info(json.dumps(task.result, indent=2))
+                    else:
+                        logger.info(task.result)
+                    logger.info("="*50 + "\n")
+                else:
+                    logger.info("\n" + "="*50)
+                    logger.info("TASK COMPLETED (No result set)")
+                    logger.info("="*50 + "\n")
+                
+                # Use os._exit to avoid exception traceback from asyncio
+                os._exit(0)
+
+    async def register_agent_worker(
+        self,
+        agent_id: str,
+        on_execution_request: ExecutionRequestHandler,
+    ) -> None:
+        """
+        Register a worker agent and start listening for task events.
+
+        Args:
+            agent_id (str): The unique identifier of the agent to register.
+            on_execution_request (ExecutionRequestHandler): The callback function to process task execution requests.
+        """
+        environment = "xpander" if self.is_xpander_cloud else "local"
+
+        url = f"{get_events_base(configuration=self.configuration)}/{agent_id}?environment={environment}"
+
+        async for event in self._sse_events_with_retries(url):
+            if event.event == EventType.EnvironmentConflict:
+                conflict = WorkerEnvironmentConflict(**json.loads(event.data))
+                logger.error(f"Conflict! - {conflict.error}")
+                return
+            if event.event == EventType.WorkerRegistration:
+                self.worker = agent_worker = DeployedAsset(**json.loads(event.data))
+                logger.info(f"Worker registered – id={agent_worker.id}")
+
+                # convenience URLs
+                agent_meta = agent_worker.metadata or {}
+                if agent_meta:
+                    is_stg = "stg." in get_events_base(
+                        configuration=self.configuration
+                    ) or "localhost" in get_events_base(
+                        configuration=self.configuration
+                    )
+                    chat_url = (
+                        f"https://{agent_meta.get('unique_name', agent_id)}.agents"
+                    )
+                    chat_url += ".stg" if is_stg else ""
+                    chat_url += ".xpander.ai"
+
+                    builder_url = (
+                        "https://"
+                        + ("stg." if is_stg else "")
+                        + f"app.xpander.ai/agents/{agent_id}"
+                    )
+                    logger.info(
+                        f"Agent '{agent_meta.get('name', agent_id)}' chat: {chat_url} | builder: {builder_url}"
+                    )
+
+                if self.test_task:
+                    logger.info(f"Invoking agent {self.test_task.model_dump_json()}")
+                    created_task = await Tasks(configuration=self.configuration).acreate(
+                        agent_id=self.agent_id,
+                        prompt=self.test_task.input.text,
+                        file_urls=self.test_task.input.files,
+                        user_details=self.test_task.input.user,
+                        agent_version=self.test_task.agent_version,
+                        worker_id=self.worker.id,
+                        output_format=self.test_task.output_format,
+                        output_schema=self.test_task.output_schema,
+                        run_locally=True,
+                        source=SourceNodeType.SDK.value
+                    )
+                    self.track(
+                    asyncio.create_task(
+                        self.handle_task_execution_request(
+                            agent_worker, created_task, on_execution_request
+                        )
+                    )
+                )
+
+                self.track(asyncio.create_task(self.heartbeat_loop(agent_worker.id)))
+
+            elif event.event == EventType.AgentExecution:
+                task = Task(**json.loads(event.data), configuration=self.configuration)
+                self.track(
+                    asyncio.create_task(
+                        self.handle_task_execution_request(
+                            agent_worker, task, on_execution_request
+                        )
+                    )
+                )
+
+
+    # --------------------------------------------------------------------- #
+    # Misc helpers                                                          #
+    # --------------------------------------------------------------------- #
+
+    def track(self, task: asyncio.Task) -> None:
+        """
+        Add a task to the background task set for auto-removal on completion.
+
+        Args:
+            task (asyncio.Task): The asynchronous task to track.
+        """
+        self._bg.add(task)
+        task.add_done_callback(self._bg.discard)
+
+    async def heartbeat_loop(self, worker_id: str) -> None:
+        """
+        Continuously send heartbeat signals to maintain worker's active status.
+
+        Args:
+            worker_id (str): The unique identifier of the worker.
+        """
+        while True:
+            try:
+                await self._make_heartbeat(worker_id)
+            except Exception:
+                # _request_with_retries handles fatal exit
+                pass
+            await asyncio.sleep(2)
+
+    def register(
+        self,
+        on_task: ExecutionRequestHandler,
+        test_task: Optional[LocalTaskTest] = None,
+    ) -> None:
+        """
+        Register the event listener with optional test task in synchronous or asynchronous environments.
+
+        Args:
+            on_task (ExecutionRequestHandler): Callback handler for task execution.
+            test_task (Optional[LocalTaskTest]): Optional local test task for diagnostics and testing.
+
+        Example:
+            >>> def handle_task(task):
+            ...     # process task execution
+
+            >>> events = Events()
+            >>> events.register(on_task=handle_task)
+        """
+        try:
+            self.test_task = test_task
+            loop = asyncio.get_running_loop()
+            if loop.is_running():
+                loop.create_task(self.start(on_task))
+            else:
+                asyncio.run(self.start(on_task))
+        except RuntimeError:
+            # No running loop, safe to run
+            asyncio.run(self.start(on_task))
+    
+    # --------------------------------------------------------------------- #
+    # Boot and Shutdown Handler Management                                  #
+    # --------------------------------------------------------------------- #
+    
+    @classmethod
+    def register_boot_handler(cls, handler: BootHandler) -> None:
+        """
+        Register a boot handler to be executed before event listeners are set up.
+        
+        Args:
+            handler (BootHandler): The boot handler function to register.
+        """
+        cls._boot_handlers.append(handler)
+        logger.debug(f"Boot handler registered: {handler.__name__ if hasattr(handler, '__name__') else 'anonymous'}")
+    
+    @classmethod
+    def register_shutdown_handler(cls, handler: ShutdownHandler) -> None:
+        """
+        Register a shutdown handler to be executed during application shutdown.
+        
+        Args:
+            handler (ShutdownHandler): The shutdown handler function to register.
+        """
+        cls._shutdown_handlers.append(handler)
+        logger.debug(f"Shutdown handler registered: {handler.__name__ if hasattr(handler, '__name__') else 'anonymous'}")
+    
+    @classmethod
+    async def _execute_boot_handlers(cls) -> None:
+        """
+        Execute all registered boot handlers.
+        
+        Raises:
+            Exception: If any boot handler fails, the application will not start.
+        """
+        if not cls._boot_handlers:
+            return
+            
+        logger.info(f"Executing {len(cls._boot_handlers)} boot handler(s)...")
+        
+        for handler in cls._boot_handlers:
+            try:
+                if asyncio.iscoroutinefunction(handler):
+                    await handler()
+                else:
+                    handler()
+                logger.debug(f"Boot handler executed successfully: {handler.__name__ if hasattr(handler, '__name__') else 'anonymous'}")
+            except Exception as e:
+                logger.error(f"Boot handler failed: {handler.__name__ if hasattr(handler, '__name__') else 'anonymous'} - {e}")
+                raise
+                
+        logger.info("All boot handlers executed successfully")
+    
+    @classmethod
+    async def _execute_shutdown_handlers(cls) -> None:
+        """
+        Execute all registered shutdown handlers.
+        
+        Note: Exceptions in shutdown handlers are logged but do not prevent shutdown.
+        """
+        if not cls._shutdown_handlers:
+            return
+            
+        logger.info(f"Executing {len(cls._shutdown_handlers)} shutdown handler(s)...")
+        
+        for handler in cls._shutdown_handlers:
+            try:
+                if asyncio.iscoroutinefunction(handler):
+                    await handler()
+                else:
+                    handler()
+                logger.debug(f"Shutdown handler executed successfully: {handler.__name__ if hasattr(handler, '__name__') else 'anonymous'}")
+            except Exception as e:
+                logger.error(f"Shutdown handler failed: {handler.__name__ if hasattr(handler, '__name__') else 'anonymous'} - {e}")
+                # Continue with other shutdown handlers even if one fails
+                
+        logger.info("All shutdown handlers executed")

xpander_sdk/modules/events/models/deployments.py

@@ -0,0 +1,25 @@
+from datetime import datetime
+from enum import Enum
+from typing import Dict, Optional
+
+from pydantic import BaseModel, Field
+
+class DeploymentType(str, Enum):
+    Gateway = "gateway"
+    Controller = "controller"
+    Worker = "worker"
+    Redis = "redis"
+
+class DeployedAsset(BaseModel):
+    id: str = Field(...,description="The asset's unique identifier")
+    name: str = Field(...,description="The asset's generated name")
+    organization_id: str = Field(...,description="The asset's organization")
+    type: DeploymentType = Field(...,description="The asset's type")
+    created_at: datetime = Field(...,description="The asset's creation date")
+    created_by: Optional[str] = Field(...,description="The asset's creator - used in local env workers")
+    last_heartbeat: datetime = Field(...,description="The asset's creation date")
+    configuration: Optional[Dict] = Field(default=None,description="The asset's configuration")
+    dedicated_agent_id: Optional[str] = Field(...,description="The asset agent id if used as a dedicated asset")
+    parent_asset_id: Optional[str] = Field(...,description="The asset parent id if used as a dedicated asset worker")
+    is_busy: Optional[bool] = Field(False,description="The asset busyness indication")
+    metadata: Optional[Dict] = Field(default=None,description="The asset's metadata")