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

AIConnector/common/exceptions.py

@@ -0,0 +1,13 @@
+class ParameterException(ValueError):
+    """Exception raised for errors in the input parameters."""
+    pass
+
+
+class PeerDisconnectedException(Exception):
+    """Exception indicating that the peer has disconnected."""
+    pass
+
+
+class AgentConnectionError(Exception):
+    """Exception indicating that we couldn't connect to agent."""
+    pass

AIConnector/common/logger.py

@@ -0,0 +1,57 @@
+import logging
+from colorlog import ColoredFormatter
+
+
+class Logger:
+    LOG_LEVEL = {
+        "critical": (logging.CRITICAL, "%(asctime)s %(levelname)-8s [%(name)s] %(message)s"),
+        "error": (logging.ERROR, "%(asctime)s %(levelname)-8s [%(name)s] %(message)s"),
+        "warning": (logging.WARNING, "%(asctime)s %(levelname)-8s [%(name)s] %(message)s"),
+        "info": (logging.INFO, "%(asctime)s %(levelname)-8s [%(name)s] %(message)s"),
+        "debug": (logging.DEBUG, f"%(pathname)s:%(lineno)d:\n%(asctime)s %(levelname)-8s [%(name)s] %(message)s"),
+        "notset": (logging.NOTSET, "%(asctime)s %(levelname)-8s [%(name)s] %(message)s")
+    }
+
+    def __init__(self, log_level):
+        self._date_fmt = "%d-%m-%Y:%H:%M:%S"
+        self.log_level = self._get_log_level(log_level)
+        self.log_format = self._get_log_format(log_level)
+        self._config()
+
+    def _get_log_level(self, log_level):
+        return self.LOG_LEVEL[log_level][0]
+
+    def _get_log_format(self, log_level):
+        return self.LOG_LEVEL[log_level][1]
+
+    def _config(self):
+        """
+        Configures logging using a colored formatter.
+        """
+        # Get the root logger
+        root_logger = logging.getLogger()
+
+        # Remove existing handlers, if any
+        for handler in root_logger.handlers[:]:
+            root_logger.removeHandler(handler)
+
+        # Create a console handler
+        stream_handler = logging.StreamHandler()
+
+        # Define the color scheme for different logging levels
+        formatter = ColoredFormatter(
+            "%(log_color)s" + self.log_format,
+            datefmt=self._date_fmt,
+            log_colors={
+                'DEBUG': 'blue',
+                'INFO': 'green',
+                'WARNING': 'yellow',
+                'ERROR': 'red',
+                'CRITICAL': 'red,bg_white',
+            }
+        )
+        stream_handler.setFormatter(formatter)
+
+        # Configure the root logger
+        root_logger.setLevel(self.log_level)
+        root_logger.addHandler(stream_handler)

AIConnector/common/message.py

@@ -0,0 +1,14 @@
+from enum import Enum
+
+
+class MessageTypes(Enum):
+    """
+    Enum representing different types of messages exchanged between peers.
+    """
+    HELLO = "hello"
+    TEXT = "text"
+    SYSTEM_MESSAGE = "system_message"
+    HEARTBEAT = "heartbeat"
+    FINAL_MESSAGE = "final_message"
+    ERROR = "error"
+    JOB_LIST = "job_list"

AIConnector/common/network.py

@@ -0,0 +1,146 @@
+import socket
+
+from dataclasses import dataclass
+from typing import Optional
+
+from AIConnector.common.exceptions import ParameterException
+
+
+@dataclass
+class Config:
+    """
+    Data class representing the Class network configuration settings.
+
+    Attributes:
+        webserver_port (Optional[int]): Port number for the webserver.
+            If not provided, an available port within the autodiscovery range will be used.
+            Required in: 'local' and 'remote'.
+        azure_endpoint_url (Optional[str]): URL endpoint for Azure connectivity.
+                                            Required in: 'remote'.
+        azure_access_key (Optional[str]): Access key for Azure services.
+                                          Required in: 'remote'.
+        advertisement_port (int): Port used for network advertisement. Defaults to 9999.
+                                  Required in: 'local' and 'remote'.
+        advertisement_ip (str): IP address for network advertisement. Defaults to "224.1.1.1".
+                                Required in: 'local' and 'remote'.
+        advertisement_interval (float): Interval (in seconds) for sending advertisements. Defaults to 15.0.
+                                        Required in: 'local' and 'remote'.
+        heartbeat_delay (float): Delay (in seconds) between heartbeat messages. Defaults to 5.0.
+                                 Required in: 'local' and 'remote'.
+        webserver_port_min (int): Minimum port number for autodiscovery. Defaults to 5000.
+                                                Required in: 'local'.
+        webserver_port_max (int): Maximum port number for autodiscovery. Defaults to 7000.
+                                                Required in: 'local'.
+        peer_discovery_timeout (float): Timeout (in seconds) for peer discovery. Defaults to 30.0.
+                                        Required in: 'local' and 'remote'.
+        peer_ping_interval (float): Interval (in seconds) for pinging peers. Defaults to 5.0.
+                                    Required in: 'local' and 'remote'.
+        azure_api_version (str): API version for Azure communications. Defaults to "1.0".
+                                 Required in: 'remote'.
+        connection_type (str): Type of connection, either "local" or "remote". Defaults to "local".
+                               Required in: 'local' and 'remote'.
+    """
+    webserver_port: Optional[int] = None
+
+    azure_endpoint_url: Optional[str] = None
+    azure_access_key: Optional[str] = None
+
+    advertisement_port: int = 9999
+    advertisement_ip: str = "224.1.1.1"
+    advertisement_interval: float = 15.0
+    heartbeat_delay: float = 15.0
+    webserver_port_min: int = 5000
+    webserver_port_max: int = 65000
+    peer_discovery_timeout: float = 30.0
+    peer_ping_interval: float = 5.0
+    azure_api_version: str = "1.0"
+    connection_type: str = "local"
+
+    def __post_init__(self):
+        """
+        Validate the network configuration parameters immediately after initialization.
+
+        Raises:
+            ParameterException: If any parameter is invalid.
+        """
+        # Validate connection type
+        if self.connection_type not in ("local", "remote"):
+            raise ParameterException("Invalid connection type. Only 'local' and 'remote' are allowed.")
+
+        # For remote connections, ensure Azure settings are provided
+        if self.connection_type == "remote":
+            if not (self.azure_endpoint_url and self.azure_access_key):
+                raise ParameterException(
+                    "Azure endpoint URL and access key must be provided for 'remote' sessions."
+                )
+
+        # Define a list of fields to validate (field name, value, expected type)
+        fields_to_validate = (
+            ("advertisement_port", self.advertisement_port, int),
+            ("webserver_port_min", self.webserver_port_min, int),
+            ("webserver_port_max", self.webserver_port_max, int),
+            ("advertisement_interval", self.advertisement_interval, float),
+            ("heartbeat_delay", self.heartbeat_delay, float),
+            ("peer_discovery_timeout", self.peer_discovery_timeout, float),
+            ("peer_ping_interval", self.peer_ping_interval, float),
+        )
+
+        # Check each field for correct type and positive value
+        for field_name, value, expected_type in fields_to_validate:
+            if not isinstance(value, expected_type):
+                raise ParameterException(f"Parameter '{field_name}' must be of type {expected_type.__name__}.")
+            if value <= 0:
+                raise ParameterException(f"Parameter '{field_name}' must be greater than 0.")
+
+        # Ensure autodiscovery port range is valid
+        if self.webserver_port_min > self.webserver_port_max:
+            raise ParameterException(
+                "webserver_port_min cannot be greater than webserver_port_max."
+            )
+
+        # If webserver_port is not set, determine an available port from the autodiscovery range
+        if not self.webserver_port:
+            self.webserver_port = get_available_port(
+                self.webserver_port_min,
+                self.webserver_port_max
+            )
+
+
+@dataclass
+class NetworkConfig(Config):
+    """
+    Data class representing the App network configuration settings.
+
+    Attributes:
+        client_id (Optional[str]): Identifier for the client.
+                                   Required in: 'local' and 'remote'.
+    """
+    client_id: Optional[str] = None
+
+
+def get_available_port(start_port: int = 5000, end_port: int = 7000) -> Optional[int]:
+    """
+    Find and return the first available network port in the given range (inclusive).
+
+    The function iterates through the port range and attempts to bind a temporary
+    socket to each port. If the bind is successful, the port is considered available.
+
+    Args:
+        start_port (int): The beginning of the port range to check. Defaults to 5000.
+        end_port (int): The end of the port range to check. Defaults to 7000.
+
+    Returns:
+        Optional[int]: The first available port number, or None if no port is available.
+    """
+    for port in range(start_port, end_port + 1):
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            try:
+                # Attempt to bind the socket to the current port on all interfaces.
+                sock.bind(("0.0.0.0", port))
+                # Successful bind indicates the port is free.
+                return port
+            except OSError:
+                # Port is already in use; try the next one.
+                continue
+    # Return None if no available port is found in the specified range.
+    return None

AIConnector/connector/azure_connector.py

@@ -0,0 +1,205 @@
+import asyncio
+import json
+import logging
+import time
+from typing import Callable, Optional, Dict, Any
+
+import websockets
+
+from azure.messaging.webpubsubservice import WebPubSubServiceClient
+from AIConnector.connector.base_connector import BaseConnector
+
+logger = logging.getLogger(__name__)
+
+
+class AzureConnector(BaseConnector):
+    """
+    AzureConnector implements message exchange via Azure Web PubSub,
+    similar to how WsConnector operates for local connections.
+    """
+
+    def __init__(
+        self,
+        client_id: str,
+        on_message_callback: Callable[[Dict[str, Any]], Any],
+        azure_endpoint_url: str,
+        azure_access_key: str,
+        azure_api_version: str,
+        messaging_hub: str = "messaging_demo",
+        reconnect_attempts: int = 2,
+        client_name: Optional[str] = None,
+        on_peer_disconnected: Optional[Callable[[str], Any]] = None,
+        on_peer_connected: Optional[Callable[[str], Any]] = None,
+    ) -> None:
+        """
+        Initialize the AzureConnector.
+
+        Args:
+            client_id (str): Identifier for the client.
+            on_message_callback (Callable[[Dict[str, Any]], Any]): Callback function to handle incoming messages.
+            azure_endpoint_url (str): Azure Web PubSub endpoint URL.
+            azure_access_key (str): Access key for Azure Web PubSub.
+            azure_api_version (str): API version for Azure services.
+            messaging_hub (str, optional): Name of the messaging hub. Defaults to "messaging_demo".
+            reconnect_attempts (int, optional): Number of reconnection attempts on failure. Defaults to 2.
+            client_name (Optional[str], optional): Name of the client. Defaults to None.
+            on_peer_disconnected (Optional[Callable[[str], Any]], optional): Callback when a peer disconnects. Defaults to None.
+            on_peer_connected (Optional[Callable[[str], Any]], optional): Callback when a peer connects. Defaults to None.
+        """
+        self.client_id = client_id
+        self.on_message_callback = on_message_callback
+        self.on_peer_disconnected = on_peer_disconnected
+        self.on_peer_connected = on_peer_connected
+        self.client_name = client_name
+        self.reconnect_attempts = reconnect_attempts
+        self.messaging_hub = messaging_hub
+
+        self.is_connected = False
+        self.running = False
+        self.active_connections: Dict[str, websockets.WebSocketClientProtocol] = {}
+
+        # Initialize the service client for the messaging hub.
+        connection_string = f"Endpoint={azure_endpoint_url};AccessKey={azure_access_key};ApiVersion={azure_api_version}"
+        self.service_client = WebPubSubServiceClient.from_connection_string(
+            connection_string,
+            hub=messaging_hub,
+        )
+        self.websocket = None
+
+    async def start_server(self, host: str = None, port: int = None, allow_discovery: bool = False) -> None:
+        """
+        Establish a WebSocket connection with the messaging hub.
+
+        The host and port parameters are provided for interface compatibility,
+        but they are not used in AzureConnector.
+
+        Args:
+            host (str, optional): Not used in AzureConnector.
+            port (int, optional): Not used in AzureConnector.
+            allow_discovery (bool, optional): Not used in AzureConnector.
+        """
+        self.running = True
+        try:
+            logger.info(f"Client ID: {self.client_id}")
+            token = self.service_client.get_client_access_token(user_id=self.client_id)
+            self.websocket = await websockets.connect(token["url"])
+            self.is_connected = True
+            logger.info(f"[AzureConnector] Connected to Azure Web PubSub messaging hub: {self.messaging_hub}")
+
+            # Start receiving messages asynchronously.
+            asyncio.create_task(self._receive_messages())
+
+        except Exception as e:
+            logger.error(f"[AzureConnector] Connection failed: {e}")
+            await self._attempt_reconnect()
+
+    async def stop_server(self) -> None:
+        """
+        Stop the WebSocket connection with the messaging hub.
+        """
+        self.running = False
+        self.is_connected = False
+        if self.websocket:
+            await self.websocket.close()
+            logger.info("[AzureConnector] Messaging WebSocket connection closed")
+
+    async def send_message(self, peer_id: str, msg: Dict[str, Any]) -> None:
+        """
+        Send a message to a specific peer via the messaging hub.
+
+        Args:
+            peer_id (str): The identifier of the target peer.
+            msg (Dict[str, Any]): The message payload to send.
+        """
+        if not self.is_connected:
+            logger.error("[AzureConnector] Cannot send message: not connected")
+            return
+
+        try:
+            message = {
+                "type": "MESSAGE",
+                "from_id": self.client_id,
+                "from_name": self.client_name or "Unknown",
+                "to_id": peer_id,
+                "timestamp": time.time(),
+                "payload": msg,
+            }
+            # Send the message as a JSON string.
+            self.service_client.send_to_user(peer_id, json.dumps(message))
+        except Exception as e:
+            logger.error(f"[AzureConnector] Error sending message to {peer_id}: {e}")
+        else:
+            logger.info(f"[AzureConnector] Message sent to {peer_id}")
+
+    async def connect_to_peer(self, ip: str, port: int, peer_id: str) -> None:
+        """
+        For Azure, a direct connection to peers is not required.
+        Instead, a HELLO message is sent to simulate a handshake.
+
+        Args:
+            ip (str): Not used in AzureConnector.
+            port (int): Not used in AzureConnector.
+            peer_id (str): The identifier of the peer to connect to.
+        """
+        hello_msg = {
+            "type": "HELLO",
+            "from_id": self.client_id,
+            "from_name": self.client_name or "Unknown",
+            "timestamp": time.time(),
+        }
+        await self.send_message(peer_id, hello_msg)
+
+    async def _receive_messages(self) -> None:
+        """
+        Receive incoming messages via the messaging hub.
+        """
+        while self.running and self.websocket:
+            try:
+                message = await self.websocket.recv()
+                msg = json.loads(message)
+                # If the received message is a string, try to parse it again.
+                if isinstance(msg, str):
+                    msg = json.loads(msg)
+                payload = msg["payload"]
+                # If a HELLO message is received, update active connections and trigger callback.
+                if payload.get("type").lower() == "hello":
+                    peer_id = payload.get("from_id")
+                    if peer_id:
+                        self.active_connections[peer_id] = self.websocket
+                        if self.on_peer_connected:
+                            await self.on_peer_connected(peer_id)
+
+                # Process the incoming message with the provided callback.
+                await self.on_message_callback(payload)
+            except websockets.exceptions.ConnectionClosed:
+                logger.error("[AzureConnector] Messaging WebSocket connection closed")
+                await self._attempt_reconnect()
+                break
+            except Exception as e:
+                logger.error(f"[AzureConnector] Error processing message: {e}")
+
+    async def _attempt_reconnect(self) -> None:
+        """
+        Attempt to reconnect to Azure Web PubSub after a connection loss.
+        """
+        for attempt in range(self.reconnect_attempts):
+            try:
+                logger.info(f"[AzureConnector] Reconnection attempt {attempt + 1}/{self.reconnect_attempts}")
+                await self.start_server()
+                if self.is_connected:
+                    return
+            except Exception as e:
+                logger.error(f"[AzureConnector] Reconnection attempt failed: {e}")
+            # Exponential backoff before the next attempt.
+            await asyncio.sleep(2 ** attempt)
+        logger.error("[AzureConnector] Failed to reconnect after multiple attempts")
+
+
+    async def disconnect_from_peer(self, peer_id: str) -> None:
+        """
+        Close connection
+
+        Args:
+           peer_id (str): Unique identifier of the lost peer.
+       """
+        pass

AIConnector/connector/base_connector.py

@@ -0,0 +1,51 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+
+class BaseConnector(ABC):
+    """
+    Abstract base class for network connectors.
+
+    Classes inheriting from BaseConnector must implement the methods
+    to start and stop the server, as well as to send messages.
+    """
+
+    @abstractmethod
+    async def start_server(self, host: str, port: int, allow_discovery: bool) -> None:
+        """
+        Start the server and establish necessary connections.
+
+        Args:
+            host (str): The host address to bind the server.
+            port (int): The port number to bind the server.
+            allow_discovery (bool): Flag indicating if discovery is allowed.
+        """
+        pass
+
+    @abstractmethod
+    async def stop_server(self) -> None:
+        """
+        Stop the server and close all active connections.
+        """
+        pass
+
+    @abstractmethod
+    async def send_message(self, peer_id: str, msg: Dict[str, Any]) -> None:
+        """
+        Send a message to a specified peer.
+
+        Args:
+            peer_id (str): Identifier of the target peer.
+            msg (Dict[str, Any]): The message payload to be sent.
+        """
+        pass
+
+    @abstractmethod
+    async def disconnect_from_peer(self, peer_id: str) -> None:
+        """
+        Disconnect from peer by peer id
+
+        Args:
+            peer_id (str): Identifier of the target peer.
+        """
+        pass