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

AIConnector/discovery/discovery_service.py

@@ -0,0 +1,226 @@
+import asyncio
+import json
+import socket
+import struct
+import time
+import logging
+from typing import Callable, Tuple, Dict, Any, Optional
+
+from AIConnector.common.network import NetworkConfig
+
+logger = logging.getLogger(__name__)
+
+
+class DiscoveryProtocol(asyncio.DatagramProtocol):
+    """
+    Datagram protocol for handling UDP discovery messages.
+    """
+
+    def __init__(self, on_message_callback: Callable[[Dict[str, Any], Tuple[str, int]], None]) -> None:
+        """
+        Initialize the protocol with a callback for processing received messages.
+
+        Args:
+            on_message_callback: Callback function invoked with the parsed message and sender's address.
+        """
+        self.on_message_callback = on_message_callback
+
+    def datagram_received(self, data: bytes, addr: Tuple[str, int]) -> None:
+        """
+        Process an incoming UDP datagram.
+
+        Args:
+            data: The received data as bytes.
+            addr: A tuple containing the sender's IP address and port.
+        """
+        try:
+            msg = json.loads(data.decode("utf-8"))
+            # Process only DISCOVERY messages with a valid client_id.
+            if msg.get("type") == "DISCOVERY" and msg.get("client_id"):
+                asyncio.create_task(self.on_message_callback(msg, addr))
+        except Exception as e:
+            logger.debug(f"[DiscoveryProtocol] UDP parsing error: {e}")
+
+
+class DiscoveryService:
+    """
+    Service for broadcasting and listening for peer discovery messages using multicast UDP.
+
+    This service supports events for when peers appear (are discovered) and when peers are lost.
+    """
+
+    _PEER_DISAPPEAR_TIMEOUT = 30  # Seconds after which a peer is considered lost.
+
+    def __init__(
+            self,
+            on_peer_discovered_callback: Callable[[Dict[str, Any], Tuple[str, int]], None],
+            client_name: str,
+            port: int,
+            network_config: NetworkConfig,
+            on_peer_lost: Optional[Callable[[str], None]] = None,
+    ) -> None:
+        """
+        Initialize the DiscoveryService.
+
+        Args:
+            on_peer_discovered_callback: Callback invoked when a peer is discovered.
+            client_name: Display name of the client.
+            port: Port number used by the service.
+            network_config: Network configuration settings.
+            on_peer_lost: Optional callback invoked when a peer is lost.
+        """
+        self.on_peer_discovered_callback = on_peer_discovered_callback
+        self.on_peer_lost_callback = on_peer_lost
+
+        self.transport: Optional[asyncio.DatagramTransport] = None
+        self.running: bool = False
+        self.client_name: str = client_name
+        self.network_config: NetworkConfig = network_config
+        self.port: int = port
+
+        # Store peer information: key is client_id, value is a dict with last_seen timestamp and message data.
+        self.peers: Dict[str, Dict[str, Any]] = {}
+
+        self._cleanup_task: Optional[asyncio.Task] = None
+
+    @staticmethod
+    def get_socket(advertisement_port: int, advertisement_ip: str) -> socket.socket:
+        """
+        Create and configure a multicast UDP socket.
+
+        Args:
+            advertisement_port: The port to bind for receiving multicast messages.
+            advertisement_ip: The multicast IP address.
+
+        Returns:
+            A configured, non-blocking UDP socket.
+        """
+        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
+        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        try:
+            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
+        except AttributeError:
+            pass
+        sock.bind(("", advertisement_port))
+        mreq = struct.pack("4sl", socket.inet_aton(advertisement_ip), socket.INADDR_ANY)
+        sock.setsockopt(socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP, mreq)
+        sock.setblocking(False)
+        return sock
+
+    async def start(self, allow_discovery: bool = True) -> None:
+        """
+        Start the discovery service and begin periodic announcements.
+
+        Args:
+            allow_discovery: If True, the service actively announces and listens for discovery messages.
+        """
+        self.running = True
+
+        loop = asyncio.get_event_loop()
+        sock = self.get_socket(
+            self.network_config.advertisement_port,
+            self.network_config.advertisement_ip,
+        )
+
+        self.transport, _ = await loop.create_datagram_endpoint(
+            lambda: DiscoveryProtocol(self._process_discovery_message),
+            sock=sock
+        )
+        self._cleanup_task = asyncio.create_task(self._cleanup_stale_peers())
+
+        if allow_discovery:
+            asyncio.create_task(self._periodic_announce())
+
+    async def stop(self) -> None:
+        """
+        Stop the discovery service and close the underlying transport.
+        """
+        self.running = False
+        if self.transport:
+            self.transport.close()
+        if self._cleanup_task:
+            self._cleanup_task.cancel()
+            try:
+                await self._cleanup_task
+            except asyncio.CancelledError:
+                pass
+
+    async def _periodic_announce(self) -> None:
+        """
+        Periodically send discovery messages via multicast.
+        """
+        while self.running:
+            try:
+                msg = {
+                    "type": "DISCOVERY",
+                    "client_id": self.network_config.client_id,
+                    "display_name": self.client_name,
+                    "port": self.port,
+                }
+                encoded = json.dumps(msg).encode("utf-8")
+                self.transport.sendto(
+                    encoded,
+                    (
+                        self.network_config.advertisement_ip,
+                        self.network_config.advertisement_port,
+                    )
+                )
+            except Exception as e:
+                logger.debug(f"[DiscoveryService] Sending error: {e}")
+            await asyncio.sleep(self.network_config.advertisement_interval)
+
+    async def _process_discovery_message(self, msg: Dict[str, Any], addr: Tuple[str, int]) -> None:
+        """
+        Process an incoming discovery message.
+
+        Updates the last_seen timestamp for peers and invokes the on_peer_discovered_callback.
+
+        Args:
+            msg: The parsed discovery message.
+            addr: The sender's address.
+        """
+        peer_id = msg.get("client_id")
+        if not peer_id:
+            return
+
+        # Update or add peer information.
+        self.peers[peer_id] = {
+            "last_seen": time.time(),
+            "msg": msg,
+            "addr": addr,
+        }
+
+        if self.on_peer_discovered_callback:
+            try:
+                result = self.on_peer_discovered_callback(msg, addr)
+                if asyncio.iscoroutine(result):
+                    await result
+            except Exception as e:
+                logger.error(f"[DiscoveryService] Error in on_peer_discovered_callback: {e}")
+
+    async def _cleanup_stale_peers(self) -> None:
+        """
+        Periodically check the list of discovered peers.
+
+        If a peer's last update is older than _PEER_DISAPPEAR_TIMEOUT seconds,
+        it is considered lost and removed, triggering the on_peer_lost_callback.
+        """
+        while self.running:
+            try:
+                now = time.time()
+                stale_peers = [
+                    peer_id for peer_id, info in self.peers.items()
+                    if now - info.get("last_seen", now) > self._PEER_DISAPPEAR_TIMEOUT
+                ]
+                for peer_id in stale_peers:
+                    del self.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_lost callback: {e}")
+            except Exception as e:
+                logger.error(f"[DiscoveryService] Error in _cleanup_stale_peers: {e}")
+            await asyncio.sleep(1)

AIConnector/session.py

@@ -0,0 +1,274 @@
+import asyncio
+import uuid
+import json
+from typing import Any, Dict, Optional, Callable, Awaitable, List, Tuple
+from AIConnector.common.logger import Logger
+from AIConnector.common.exceptions import ParameterException
+from AIConnector.common.network import NetworkConfig, get_available_port, Config
+from AIConnector.core.chat_client import ChatClient
+
+
+class Session:
+    """
+    Represents a session for a client. This class manages configuration, unique client ID generation,
+    and initializes a ChatClient to handle peer-to-peer messaging.
+    """
+
+    FILENAME = "client_id.json"
+    ALLOWED_LOG_LEVELS = ("debug", "info", "warning", "error", "critical")
+
+    def __init__(
+        self,
+        client_name: str,
+        config: Config,
+        client_id: Optional[str] = None,
+        log_level: Optional[str] = "critical"
+    ) -> None:
+        """
+        Initialize a new Session.
+
+        Args:
+            client_name (str): The display name of the client.
+            config (Config): Network configuration settings.
+            client_id (Optional[str]): Optional pre-defined client identifier.
+            log_level (Optional[str]): Logging level. Defaults to "info".
+        """
+        self.client_name = client_name
+        self.__config = config
+        self.__client_id = client_id
+
+        self.chat_client: Optional[ChatClient] = None
+
+        # Callback lists for various peer events.
+        self._peer_connected_callbacks: List[Callable[[str], Awaitable[None]]] = []
+        self._peer_disconnected_callbacks: List[Callable[[str], Awaitable[None]]] = []
+        self._peer_discovered_callbacks: List[Callable[[str], Awaitable[None]]] = []
+        self._peer_lost_callbacks: List[Callable[[str], Awaitable[None]]] = []
+
+        self._send_lock = asyncio.Lock()
+
+        if log_level.lower() not in self.ALLOWED_LOG_LEVELS:
+            raise ParameterException(
+                f"Invalid log level '{log_level}'. Allowed values: {self.ALLOWED_LOG_LEVELS}"
+            )
+
+        Logger(log_level=log_level)
+
+    def __get_client_id(self) -> str:
+        """
+        Retrieve a unique client identifier from file, or generate and store a new one if necessary.
+
+        Returns:
+            str: The client identifier.
+        """
+        if self.__client_id:
+            return self.__client_id
+
+        try:
+            with open(self.FILENAME, "r", encoding="utf-8") as f:
+                data = json.load(f)
+        except FileNotFoundError:
+            data = {}
+
+        if stored_client_id := data.get(self.client_name):
+            self.__client_id = stored_client_id
+            return stored_client_id
+
+        # Generate a new client ID if not found.
+        new_id = str(uuid.uuid4())
+        data[self.client_name] = new_id
+
+        with open(self.FILENAME, "w", encoding="utf-8") as f:
+            json.dump(data, f, indent=4)
+
+        self.__client_id = new_id
+        return new_id
+
+    @property
+    def client_id(self) -> str:
+        """
+        Get the unique client identifier, generating it if needed.
+
+        Returns:
+            str: The client identifier.
+        """
+        if self.__client_id is None:
+            self.__client_id = self.__get_client_id()
+        return self.__client_id
+
+    @property
+    def config(self) -> NetworkConfig:
+        """
+        Retrieve the network configuration, ensuring the client_id is updated.
+
+        Returns:
+            NetworkConfig: The network configuration.
+        """
+        network_config = NetworkConfig(**self.__config.__dict__)
+        network_config.client_id = self.client_id
+        return network_config
+
+    async def start(self, message_handler: Optional[Callable[[str], Awaitable[str]]] = None) -> None:
+        """
+        Start the session by initializing and starting the ChatClient.
+        Optionally bind a message handler for incoming messages.
+
+        Args:
+            message_handler (Optional[Callable[[str], Awaitable[str]]]): Async function to process incoming messages.
+        """
+        self.chat_client = ChatClient(
+            client_name=self.client_name,
+            client_id=self.client_id,
+            network_config=self.config
+        )
+        self._register_chat_client_on_peer_connected_callback()
+        self._register_chat_client_on_peer_disconnected_callback()
+        self._register_chat_client_on_peer_discovered_callback()
+        self._register_chat_client_on_peer_lost_callback()
+
+        if message_handler:
+            await self.bind(message_handler)
+
+        # Attempt to start the chat client; if the port is unavailable, assign a new one.
+        while True:
+            try:
+                await self.chat_client.start()
+                break
+            except OSError:
+                self.config.webserver_port = get_available_port(
+                    self.config.webserver_port_min,
+                    self.config.webserver_port_max
+                )
+
+    async def stop(self) -> None:
+        """
+        Stop the session by stopping the ChatClient.
+        """
+        await self.chat_client.stop()
+
+    async def bind(self, handler: Callable[[str], Awaitable[str]]) -> None:
+        """
+        Bind a message handler to the ChatClient for processing incoming messages.
+
+        Args:
+            handler (Callable[[str], Awaitable[str]]): Async function that processes an incoming message.
+        """
+        if self.chat_client is None:
+            raise Exception("ChatClient is not started yet. Call start() before binding a handler.")
+        await self.chat_client.register_job(handler)
+
+    async def send(self, target_client_name: str, text: str, timeout: Optional[float] = None) -> Tuple[bool, str]:
+        """
+        Sends a message to a peer by display name and waits for a response.
+
+        Args:
+            target_client_name (str): The display name of the target peer.
+            text (str): The message text to send.
+            timeout (Optional[float], optional): The timeout in seconds for waiting for a response. Defaults to None.
+
+        Returns:
+            Tuple[bool, str]: A tuple where the first value indicates success (True if the message was sent successfully, False otherwise),
+            and the second value contains the response text from the remote peer.
+        """
+        if self.chat_client is None:
+            raise Exception("ChatClient is not started yet. Call start() before sending messages.")
+
+        async with self._send_lock:
+            # Wait for a peer with the specified display name.
+            peer = await self.chat_client.wait_for_peers(
+                target_client_name=target_client_name,
+                timeout=self.config.peer_discovery_timeout,
+                interval=self.config.peer_ping_interval,
+            )
+            if not peer:
+                raise Exception(f"Peer '{target_client_name}' not found.")
+
+            peer_id = peer["peer_id"]
+            queue_id = await self.chat_client.send_message(peer_id, text)
+            is_success, result = await self.chat_client.wait_for_result(queue_id, timeout=timeout)
+            await self.chat_client.close_connection(peer_id)
+            return is_success, result
+
+    async def get_client_list(self, timeout: int = 30) -> List[Dict[str, Any]]:
+        """
+        Retrieve a list of currently connected peers.
+
+        Args:
+            timeout (int): Maximum time to wait (in seconds) for peers.
+
+        Returns:
+            List[Dict[str, Any]]: A list of peer information dictionaries.
+        """
+        return await self.chat_client.wait_all_peers(timeout=timeout)
+
+    def _register_chat_client_on_peer_connected_callback(self) -> None:
+        """
+        Register all stored peer-connected callbacks with the ChatClient.
+        """
+        for registered_callback in self._peer_connected_callbacks:
+            self.chat_client.register_on_peer_connected_callback(callback=registered_callback)
+
+    def _register_chat_client_on_peer_disconnected_callback(self) -> None:
+        """
+        Register all stored peer-disconnected callbacks with the ChatClient.
+        """
+        for registered_callback in self._peer_disconnected_callbacks:
+            self.chat_client.register_on_peer_disconnected_callback(callback=registered_callback)
+
+    def register_on_peer_connected_callback(self, callback: Callable[[str], Awaitable[None]]) -> None:
+        """
+        Register an async callback to be called when a peer connects.
+
+        Args:
+            callback (Callable[[str], Awaitable[None]]): Async function accepting a peer ID.
+        """
+        self._peer_connected_callbacks.append(callback)
+        if self.chat_client is not None:
+            self._register_chat_client_on_peer_connected_callback()
+
+    def register_on_peer_disconnected_callback(self, callback: Callable[[str], Awaitable[None]]) -> None:
+        """
+        Register an async callback to be called when a peer disconnects.
+
+        Args:
+            callback (Callable[[str], Awaitable[None]]): Async function accepting a peer ID.
+        """
+        self._peer_disconnected_callbacks.append(callback)
+        if self.chat_client is not None:
+            self._register_chat_client_on_peer_disconnected_callback()
+
+    def _register_chat_client_on_peer_discovered_callback(self) -> None:
+        """
+        Register all stored peer-discovered callbacks with the ChatClient.
+        """
+        for registered_callback in self._peer_discovered_callbacks:
+            self.chat_client.register_on_peer_discovered_callback(callback=registered_callback)
+
+    def _register_chat_client_on_peer_lost_callback(self) -> None:
+        """
+        Register all stored peer-lost callbacks with the ChatClient.
+        """
+        for registered_callback in self._peer_lost_callbacks:
+            self.chat_client.register_on_peer_lost_callback(callback=registered_callback)
+
+    def register_on_peer_discovered_callback(self, callback: Callable[[str], Awaitable[None]]) -> None:
+        """
+        Register an async callback to be called when a peer is discovered.
+
+        Args:
+            callback (Callable[[str], Awaitable[None]]): Async function accepting a peer ID.
+        """
+        self._peer_discovered_callbacks.append(callback)
+        if self.chat_client is not None:
+            self._register_chat_client_on_peer_discovered_callback()
+
+    def register_on_peer_lost_callback(self, callback: Callable[[str], Awaitable[None]]) -> None:
+        """
+        Register an async callback to be called when a peer is lost.
+
+        Args:
+            callback (Callable[[str], Awaitable[None]]): Async function accepting a peer ID.
+        """
+        self._peer_lost_callbacks.append(callback)
+        if self.chat_client is not None:
+            self._register_chat_client_on_peer_lost_callback()

genai_protocol_lite-1.0.0.dist-info/METADATA

@@ -0,0 +1,186 @@
+Metadata-Version: 2.4
+Name: genai-protocol-lite
+Version: 1.0.0
+Summary: GenAI Python project for direct agents connection either locally via websockets or remote via Azure Web PubSub
+Home-page: https://github.com/genai-works-org/genai-protocol-lite
+Author: Yaroslav Oliinyk, Yaroslav Motalov, Ivan Kuzlo
+Author-email: yaroslav.oliinyk@chisw.com, yaroslav.motalov@chisw.com, ivan.kuzlo@chisw.com
+License: Apache License 2.0
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: wcwidth==0.2.13
+Requires-Dist: websockets~=15.0
+Requires-Dist: colorlog==6.9.0
+Requires-Dist: azure-messaging-webpubsubservice==1.2.1
+Provides-Extra: dev
+Requires-Dist: twine>=6.1.0; extra == "dev"
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: license-file
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+########################################
+
+```
+# Multi-Agent Communication Library
+
+This library provides a unified, session-based API for building distributed multi-agent communication systems. It
+supports two connection modes:
+
+- **Local Mode:** Uses UDP/WebSocket for local communication.
+- **Remote Mode:** Uses Azure Pub/Sub for cloud-based messaging.
+
+The `Session` class is the main entry point, allowing agents to start a session, send messages to peers, and process
+incoming messages asynchronously.
+
+## Features
+
+- **Dual Connection Modes:**
+  - **Local Mode:** Communication using UDP/WebSocket.
+  - **Remote Mode:** Communication via Azure Pub/Sub.
+- **Session Management:**
+  Persistent client identification with an optional display name.
+- **Asynchronous Messaging:**
+  Send requests to other clients and wait for their responses.
+- **Custom Workload Processing:**
+  Easily bind your own asynchronous message handler to process incoming requests.
+- **Flexible Configuration:**
+  Advanced settings available via the `config` parameter.
+
+## Installation
+
+```bash
+  python pip install multi-agent-communication-library
+```
+
+## Usage Examples
+
+### Agent Example (Providing Services)
+
+This example demonstrates an agent that provides a service. The agent binds a custom asynchronous message handler to
+process incoming requests with an abstract workload.
+
+```python
+import asyncio
+from library.session import Session
+
+async def process(message: str) -> str:
+    # Simulate an abstract workload (e.g., data processing or computation)
+    print(f"Processing workload for message: {message}")
+    await asyncio.sleep(2)  # Simulate time-consuming work
+    return f"Processed result for: {message}"
+
+async def main():
+    # Create a session in local mode (use "remote" for Azure Pub/Sub mode)
+    session = Session(
+        connection_type="local",
+        client_name="worker_agent"
+    )
+    # Start the session with the message handler bound to process incoming requests
+    await session.start(message_handler=process)
+    # Keep the agent running indefinitely to process incoming tasks
+    await asyncio.Future()
+
+if __name__ == '__main__':
+    asyncio.run(main())
+```
+
+### Client Example (Using Services)
+
+This example shows a client that sends a task request to a service provider identified by its `client_name`.
+
+```python
+import asyncio
+
+from library.session import Session
+
+async def main():
+    session = Session(
+        connection_type="local",
+        client_name="client"
+    )
+    await session.start()
+    request = "Ipsum tempor ut adipisicing do magna ut excepteur deserunt non irure veniam dolore."
+    result = await session.send("scrapper", request)
+    print("Result:", result)
+if __name__ == '__main__':
+    asyncio.run(main())
+
+```
+
+## API Reference
+
+### `Session` Class
+
+#### Constructor
+
+- **connection_type:**
+  - `'local'` – Uses UDP/WebSocket for local communication.
+  - `'remote'` – Uses Azure Pub/Sub for remote communication.
+- **client_name:** Optional display name for the client.
+- **client_id:** Optional unique identifier; if not provided, one will be generated.
+- **config:** Optional configuration dictionary for advanced settings.
+- **log_level:** Logging level (default is `'info'`).
+
+#### Methods
+
+- **`async start(message_handler: Optional[Callable[[str], Awaitable[str]]] = None) -> None`**
+  Starts the session. If a `message_handler` is provided (for local mode), it will be bound to process incoming
+  messages.
+- **`async send(target_client_name: str, text: str, timeout: Optional[float] = None) -> str`**
+  Sends a message to a peer identified by `target_client_name` and waits for a response
+
+## Configuration Parameters
+
+
+| Parameter                | Type            | Default Value | Description                                                                                        | Required In       |
+| ------------------------ | --------------- | ------------- | -------------------------------------------------------------------------------------------------- | ----------------- |
+| `webserver_port`         | `Optional[int]` | `None`        | Port number for the webserver. If not set, an available port from the autodiscovery range is used. | `local`           |
+| `azure_endpoint_url`     | `Optional[str]` | `None`        | URL endpoint for Azure connectivity.                                                               | `remote`          |
+| `azure_access_key`       | `Optional[str]` | `None`        | Access key for Azure services.                                                                     | `remote`          |
+| `advertisement_port`     | `int`           | `9999`        | Port used for network advertisement.                                                               | `local`           |
+| `advertisement_ip`       | `str`           | `"224.1.1.1"` | IP address for network advertisement.                                                              | `local`           |
+| `advertisement_interval` | `float`         | `15.0`        | Interval (in seconds) for sending advertisements.                                                  | `local`, `remote` |
+| `heartbeat_delay`        | `float`         | `5.0`         | Delay (in seconds) between heartbeat messages.                                                     | `local`, `remote` |
+| `webserver_port_min`     | `int`           | `5000`        | Minimum port number for autodiscovery.                                                             | `local`           |
+| `webserver_port_max`     | `int`           | `65000`       | Maximum port number for autodiscovery.                                                             | `local`           |
+| `peer_discovery_timeout` | `float`         | `30.0`        | Timeout (in seconds) for peer discovery.                                                           | `local`, `remote` |
+| `peer_ping_interval`     | `float`         | `5.0`         | Interval (in seconds) for pinging peers.                                                           | `local`, `remote` |
+| `azure_api_version`      | `str`           | `"1.0"`       | API version for Azure communications.                                                              | `remote`          |
+| `connection_type`        | `str`           | `"local"`     | Type of connection, either`"local"` or `"remote"`.                                                 | `local`, `remote` |
+
+### Notes:
+
+- For **remote** connections, `azure_endpoint_url` and `azure_access_key` are required.
+- If `webserver_port` is not set, it will be selected from the range [`webserver_port_min`, `webserver_port_max`].
+- `connection_type` must be either `"local"` or `"remote"`.
+
+## Customization
+
+- **Message Handler:**
+  Provide your own asynchronous function to process incoming messages. This function should accept a string message and
+  return an awaitable string result.
+- **Configuration Options:**
+  Use the `config` parameter in the `Session` constructor for advanced settings.
+
+## License
+
+[Specify your license here]
+
+---
+
+```
+
+```