genai-protocol-lite 1.0.0__py3-none-any.whl

AIConnector/core/job_manager.py

@@ -0,0 +1,219 @@
+import asyncio
+import uuid
+import logging
+from typing import Any, Callable, Awaitable, Dict, List, Optional
+
+from AIConnector.common.message import MessageTypes
+from .job import Job
+
+logger = logging.getLogger(__name__)
+
+
+class JobManager:
+    """
+    Manages asynchronous jobs by creating, monitoring, and reporting their status.
+    """
+
+    def __init__(self, send_message_callback: Callable[..., Awaitable[None]]) -> None:
+        """
+        Initialize the JobManager.
+
+        Args:
+            send_message_callback (Callable[..., Awaitable[None]]):
+                Asynchronous callback to send messages about job status.
+        """
+        self.send_message_callback: Callable[..., Awaitable[None]] = send_message_callback
+        self.active_jobs: Dict[str, Job] = {}  # Maps job IDs to Job instances
+        self.jobs: List[Callable[[Any], Awaitable[Any]]] = []  # Registered job callables
+        self.monitoring_task: Optional[asyncio.Task[Any]] = None  # Task for monitoring jobs
+        self.job_futures: Dict[str, asyncio.Future[Any]] = {}  # Futures for job results
+
+    async def create_job(
+        self,
+        data: Any,
+        job_callable: Callable[[Any], Awaitable[Any]],
+        peer_id: str,
+        queue_id: str,
+    ) -> str:
+        """
+        Create a new job, start its asynchronous task, and notify via system messages.
+
+        Args:
+            data (Any): Data associated with the job.
+            job_callable (Callable[[Any], Awaitable[Any]]): Asynchronous callable to process the job.
+            peer_id (str): Identifier of the peer associated with the job.
+            queue_id (str): Identifier for the job queue.
+
+        Returns:
+            str: A unique job identifier.
+        """
+        # Generate a unique job ID.
+        job_id = str(uuid.uuid4())
+        # Create a new job instance and store it.
+        job = Job(job_id, data, peer_id, queue_id)
+        self.active_jobs[job_id] = job
+
+        # Create a future to await the job's result.
+        future = asyncio.get_event_loop().create_future()
+        self.job_futures[job_id] = future
+
+        # Notify that the job has been initialized.
+        await self.send_message_callback(
+            message_type=MessageTypes.SYSTEM_MESSAGE.value,
+            job_id=job_id,
+            peer_id=peer_id,
+            queue_id=queue_id,
+            text=f"Job initialized with id {job_id}",
+        )
+
+        # Start the job's asynchronous task.
+        job.task = asyncio.create_task(job_callable(data))
+        # Notify that the job has started.
+        await self.send_message_callback(
+            message_type=MessageTypes.SYSTEM_MESSAGE.value,
+            job_id=job_id,
+            peer_id=peer_id,
+            queue_id=queue_id,
+            text=f"Job started with id {job_id}",
+        )
+        return job_id
+
+    async def start_monitoring(self, interval: float = 1) -> None:
+        """
+        Continuously monitor active jobs, sending heartbeat and final messages.
+
+        Args:
+            interval (float): Time (in seconds) between monitoring checks.
+        """
+        while True:
+            finished_jobs: List[str] = []
+            # Iterate through a copy of active jobs to safely modify the dictionary.
+            for job_id, job in list(self.active_jobs.items()):
+                peer_id = job.peer_id
+                queue_id = job.queue_id
+                if job.task.done():
+                    try:
+                        # Get the result of the job.
+                        result = job.task.result()
+                    except Exception as e:
+                        logger.error(f"[JobManager] Error: {e}")
+                        result = f"Job error: {e}"
+
+                    # Set the result for the awaiting future.
+                    future = self.job_futures.get(job_id)
+                    if future and not future.done():
+                        future.set_result(result)
+
+                    # Notify that the job has finished.
+                    await self.send_message_callback(
+                        message_type=MessageTypes.FINAL_MESSAGE.value,
+                        job_id=job_id,
+                        peer_id=peer_id,
+                        queue_id=queue_id,
+                        text=result,
+                    )
+                    finished_jobs.append(job_id)
+                else:
+                    # Send a heartbeat message to indicate the job is still running.
+                    await self.send_message_callback(
+                        message_type=MessageTypes.HEARTBEAT.value,
+                        job_id=job_id,
+                        queue_id=queue_id,
+                        peer_id=peer_id,
+                        text="working",
+                    )
+
+            # Clean up finished jobs.
+            for job_id in finished_jobs:
+                del self.active_jobs[job_id]
+                if job_id in self.job_futures:
+                    del self.job_futures[job_id]
+
+            await asyncio.sleep(interval)
+
+    async def await_job_result(self, job_id: str, timeout: Optional[float] = None) -> Any:
+        """
+        Await the result of a specific job.
+
+        Args:
+            job_id (str): The unique job identifier.
+            timeout (Optional[float]): Optional timeout (in seconds).
+
+        Returns:
+            Any: The result of the job.
+
+        Raises:
+            ValueError: If the job does not exist.
+        """
+        future = self.job_futures.get(job_id)
+        if not future:
+            raise ValueError(f"Job with {job_id=} doesn't exist")
+        return await asyncio.wait_for(future, timeout=timeout)
+
+    def start(self, interval: float = 1) -> None:
+        """
+        Start monitoring active jobs.
+
+        Args:
+            interval (float): Time (in seconds) between monitoring checks.
+        """
+        self.monitoring_task = asyncio.create_task(self.start_monitoring(interval))
+
+    async def stop(self) -> None:
+        """
+        Stop monitoring active jobs by cancelling the monitoring task.
+        """
+        if self.monitoring_task:
+            self.monitoring_task.cancel()
+            try:
+                await self.monitoring_task
+            except asyncio.CancelledError:
+                pass
+
+    async def process_message(self, msg: str, from_id: str, queue_id: str) -> None:
+        """
+        Process an incoming message by creating a job using the first registered job callable.
+
+        Args:
+            msg (str): The message data.
+            from_id (str): Identifier of the peer that sent the message.
+            queue_id (str): Identifier for the job queue.
+        """
+        if self.jobs:
+            job_callable = self.jobs[0]
+            job_id = await self.create_job(msg, job_callable, from_id, queue_id)
+            await self.send_message_callback(
+                message_type=MessageTypes.SYSTEM_MESSAGE.value,
+                job_id=job_id,
+                peer_id=from_id,
+                queue_id=queue_id,
+                text=f"Ack: created job with id {job_id} for peer {from_id}",
+            )
+
+    def register_job(self, job_call_back: Callable[[Any], Awaitable[Any]]) -> None:
+        """
+        Register a new job callable.
+
+        Args:
+            job_call_back (Callable[[Any], Awaitable[Any]]):
+                Asynchronous callable that processes job data.
+        """
+        self.jobs.append(job_call_back)
+
+    def job_list(self) -> List[Callable[[Any], Awaitable[Any]]]:
+        """
+        Get the list of registered job callables.
+
+        Returns:
+            List[Callable[[Any], Awaitable[Any]]]: A list of job callables.
+        """
+        return list(self.jobs)
+
+    def job_active_list(self) -> Dict[str, Job]:
+        """
+        Get the dictionary of currently active jobs.
+
+        Returns:
+            Dict[str, Job]: A mapping of job IDs to Job instances.
+        """
+        return self.active_jobs

AIConnector/core/message_factory.py

@@ -0,0 +1,44 @@
+from typing import Any, Dict, Optional
+
+
+class MessageFactory:
+    """
+    Factory for generating message dictionaries with a specified type and sender.
+    """
+
+    def __init__(self, message_type: str, from_id: str, *args: Any, **kwargs: Any) -> None:
+        """
+        Initialize the MessageFactory.
+
+        Args:
+            message_type (str): The type of the message.
+            from_id (str): The identifier of the sender.
+            *args (Any): Additional positional arguments (unused).
+            **kwargs (Any): Additional keyword arguments to include in the message.
+        """
+        self.message_type = message_type
+        self.from_id = from_id
+        self.args = args
+        self.kwargs = kwargs
+
+    def generate_message(self, text: Optional[str] = None) -> Dict[str, Any]:
+        """
+        Generate a message dictionary.
+
+        Args:
+            text (Optional[str]): Optional text to include in the message.
+
+        Returns:
+            Dict[str, Any]: A dictionary containing the message data.
+        """
+        # Base message with type and sender.
+        message: Dict[str, Any] = {
+            "type": self.message_type,
+            "from_id": self.from_id,
+        }
+        # Optionally include text if provided.
+        if text is not None:
+            message["text"] = text
+        # Merge in any additional keyword arguments.
+        message.update(self.kwargs)
+        return message

AIConnector/discovery/azure_discovery_service.py

@@ -0,0 +1,206 @@
+import asyncio
+import json
+import logging
+import time
+import websockets
+from typing import Callable, Optional, Dict, Any, Tuple
+
+from azure.messaging.webpubsubservice import WebPubSubServiceClient
+from AIConnector.common.exceptions import ParameterException
+from AIConnector.discovery.base_discovery_service import BaseDiscoveryService
+
+logger = logging.getLogger(__name__)
+
+
+class AzureDiscoveryService(BaseDiscoveryService):
+    """
+    Service for peer discovery via Azure Web PubSub.
+
+    This service functions similarly to the UDP-based DiscoveryService,
+    but uses Azure for sending and receiving DISCOVERY messages.
+    """
+
+    def __init__(
+            self,
+            on_peer_discovered_callback: Optional[Callable[[Dict[str, Any], Tuple[str, int]], None]] = None,
+            on_peer_lost: Optional[Callable[[Dict[str, Any], Tuple[str, int]], None]] = None,
+            client_id: Optional[str] = None,
+            client_name: Optional[str] = None,
+            azure_endpoint_url: str = None,
+            azure_access_key: str = None,
+            azure_api_version: str = "1.0",
+            discovery_hub: str = "discovery_demo",
+            discovery_interval: float = 10,
+            heartbeat_interval: float = 10,
+    ) -> None:
+        """
+        Initialize the discovery service.
+
+        Args:
+            on_peer_discovered_callback: Callback invoked when a peer is discovered.
+            on_peer_lost: Callback invoked when a peer is lost.
+            client_id: Unique identifier of the client.
+            client_name: Display name of the client.
+            azure_endpoint_url: Azure Web PubSub endpoint URL (must start with "http://" or "https://").
+            azure_access_key: Azure access key.
+            azure_api_version: Azure API version.
+            discovery_hub: Name of the discovery hub.
+            discovery_interval: Interval (in seconds) for sending DISCOVERY messages.
+            heartbeat_interval: Interval (in seconds) for cleaning up stale peer entries.
+        """
+        if not azure_endpoint_url or not azure_endpoint_url.startswith(("http://", "https://")):
+            raise ParameterException(
+                f"Invalid azure_endpoint_url: {azure_endpoint_url}. It must start with 'http://' or 'https://'."
+            )
+        self.azure_endpoint_url = azure_endpoint_url
+        self.azure_access_key = azure_access_key
+        self.azure_api_version = azure_api_version
+        self.discovery_hub = discovery_hub
+        self.discovery_interval = discovery_interval
+        self.heartbeat_interval = heartbeat_interval
+
+        self.on_peer_discovered_callback = on_peer_discovered_callback
+        self.on_peer_lost_callback = on_peer_lost
+        self.client_id = client_id
+        self.client_name = client_name
+
+        self.running = False
+        self.allow_discovery = False
+        self.discovered_peers: Dict[str, Dict[str, Any]] = {}
+
+        # Initialize the service client for the discovery hub.
+        self.service_client = WebPubSubServiceClient.from_connection_string(
+            f"Endpoint={self.azure_endpoint_url};AccessKey={self.azure_access_key};ApiVersion={self.azure_api_version}",
+            hub=self.discovery_hub,
+        )
+        self.websocket = None
+        self._receive_task: Optional[asyncio.Task] = None
+        self._announce_task: Optional[asyncio.Task] = None
+        self._cleanup_task: Optional[asyncio.Task] = None
+
+    async def start(self, allow_discovery: bool = True) -> None:
+        """
+        Start the discovery service by establishing a WebSocket connection with the discovery hub
+        and starting background tasks for receiving DISCOVERY messages and periodic announcements.
+
+        Args:
+            allow_discovery: If True, enables peer discovery.
+        """
+        self.running = True
+        self.allow_discovery = allow_discovery
+        try:
+            token = self.service_client.get_client_access_token()
+            self.websocket = await websockets.connect(token["url"])
+            logger.info(f"[AzureDiscoveryService] Connected to Azure Web PubSub discovery hub: {self.discovery_hub}")
+
+            if allow_discovery:
+                self._announce_task = asyncio.create_task(self._periodic_announce())
+                logger.info("[AzureDiscoveryService] Azure Discovery Service started with peer discovery enabled")
+            else:
+                logger.info("[AzureDiscoveryService] Azure Discovery Service started with peer discovery disabled")
+
+            self._cleanup_task = asyncio.create_task(self._cleanup_stale_peers())
+            self._receive_task = asyncio.create_task(self._receive_messages())
+        except Exception as e:
+            logger.error(f"[AzureDiscoveryService] Discovery Service connection failed: {e}")
+
+    async def stop(self) -> None:
+        """
+        Stop the discovery service and clean up resources.
+        """
+        self.running = False
+        if self._receive_task:
+            self._receive_task.cancel()
+        if self._announce_task:
+            self._announce_task.cancel()
+        if self.websocket:
+            await self.websocket.close()
+        self.discovered_peers.clear()
+        logger.info("[AzureDiscoveryService] Azure Discovery Service stopped")
+
+    async def _receive_messages(self) -> None:
+        """
+        Receive incoming DISCOVERY messages and update the list of discovered peers.
+        """
+        while self.running and self.websocket:
+            try:
+                message = await self.websocket.recv()
+                msg = json.loads(message)
+                if isinstance(msg, str):
+                    msg = json.loads(msg)
+                peer_id = msg.get("client_id")
+                # Process DISCOVERY messages from other peers.
+                if msg.get("type") == "DISCOVERY" and peer_id != self.client_id:
+                    if peer_id:
+                        result_msg = {
+                            "timestamp": msg.get("timestamp", 0),
+                            "display_name": msg.get("client_name", "Unknown"),
+                            "client_id": peer_id,
+                            "metadata": msg.get("metadata", {})
+                        }
+                        self.discovered_peers[peer_id] = result_msg
+                        if self.on_peer_discovered_callback:
+                            # Use a dummy address since Azure does not provide a peer's IP.
+                            dummy_addr = ("0.0.0.0", 0)
+                            try:
+                                result = self.on_peer_discovered_callback(result_msg, dummy_addr)
+                                if asyncio.iscoroutine(result):
+                                    await result
+                            except Exception as e:
+                                logger.error(f"[DiscoveryService] Error in on_peer_discovered_callback: {e}")
+            except websockets.exceptions.ConnectionClosed:
+                logger.error("[AzureDiscoveryService] Discovery WebSocket connection closed")
+                break
+            except Exception as e:
+                logger.error(f"[AzureDiscoveryService] Error processing discovery message: {e}")
+
+    async def _periodic_announce(self) -> None:
+        """
+        Periodically send DISCOVERY messages through the discovery hub.
+        """
+        while self.running:
+            try:
+                discovery_msg = {
+                    "type": "DISCOVERY",
+                    "client_id": self.client_id,
+                    "client_name": self.client_name or "Unknown",
+                    "timestamp": time.time(),
+                }
+                self.service_client.send_to_all(json.dumps(discovery_msg))
+            except Exception as e:
+                logger.error(f"[AzureDiscoveryService] Error sending discovery message: {e}")
+            else:
+                logger.info("[AzureDiscoveryService] Discovery message sent")
+            await asyncio.sleep(self.discovery_interval)
+
+    async def _cleanup_stale_peers(self) -> None:
+        """
+        Periodically clean up stale peer entries.
+        A peer is considered stale if its timestamp is older than 3 times the heartbeat_interval.
+        """
+        while self.running:
+            try:
+                current_time = asyncio.get_event_loop().time()
+                stale_threshold = current_time - (self.heartbeat_interval * 3)
+
+                stale_peers = [
+                    peer_id for peer_id, info in self.discovered_peers.items()
+                    if info["timestamp"] < stale_threshold
+                ]
+
+                for peer_id in stale_peers:
+                    del self.discovered_peers[peer_id]
+                    if self.on_peer_lost_callback:
+                        try:
+                            result = self.on_peer_lost_callback(peer_id)
+                            if asyncio.iscoroutine(result):
+                                await result
+                        except Exception as e:
+                            logger.error(f"[DiscoveryService] Error in on_peer_disappear callback: {e}")
+
+                    logger.debug(f"[AzureDiscoveryService] Removed stale peer: {peer_id}")
+
+            except Exception as e:
+                logger.error(f"[AzureDiscoveryService] Error in cleanup task: {e}")
+
+            await asyncio.sleep(self.heartbeat_interval)

AIConnector/discovery/base_discovery_service.py

@@ -0,0 +1,27 @@
+from abc import ABC, abstractmethod
+
+
+class BaseDiscoveryService(ABC):
+    """
+    Abstract base class for discovery services.
+
+    Subclasses must implement methods to start and stop the discovery service.
+    """
+
+    @abstractmethod
+    def start(self, allow_discovery: bool = False) -> None:
+        """
+        Start the discovery service.
+
+        Args:
+            allow_discovery (bool, optional): Flag indicating whether peer discovery is enabled.
+                Defaults to False.
+        """
+        pass
+
+    @abstractmethod
+    def stop(self) -> None:
+        """
+        Stop the discovery service.
+        """
+        pass