Python SDK Rework

client_sdks/python/examples/example_producer.py

@@ -1,5 +1,6 @@
 import asyncio
 
+from broker.config import BrokerConfig
 from manager.config import ManagerConfig
 from publisher.client import PublisherClient
 
@@ -9,21 +10,35 @@
 # Setup the manager location configuration
 manager_config = ManagerConfig(url="http://localhost:3000")
 
+# Setup the broker configuration
+broker_config = BrokerConfig(url="amqp://user:password@localhost:5672")
+
 # Both the publisher and the worker need to know about the task kinds and
 # should have unified names for them.
+WORKER_KIND_NAME = "worker_kind"
 TASK_1_NAME = "task_1"
 TASK_2_NAME = "task_2"
 
 # APPLICATION CONFIGURATION ___________________________________________________
 
 # 1. Create a producer application
-worker_application = PublisherClient(manager_config)
+worker_application = PublisherClient(
+    manager_config=manager_config, broker_config=broker_config
+)
 
 
 # 2. Start the application
 async def main():
-    task1 = await worker_application.publish_task(TASK_1_NAME, {"data": "task_1_data"})
-    task2 = await worker_application.publish_task(TASK_2_NAME, {"data": "task_2_data"})
+    task1 = await worker_application.publish_task(
+        TASK_1_NAME,
+        WORKER_KIND_NAME,
+        {"data": "task_1_data"},
+    )
+    task2 = await worker_application.publish_task(
+        TASK_2_NAME,
+        WORKER_KIND_NAME,
+        {"data": "task_2_data"},
+    )
 
     print(f"Task 1: {task1}")
     print(f"Task 2: {task2}")

client_sdks/python/examples/example_worker.py

@@ -14,34 +14,37 @@
 
 # Both the publisher and the worker need to know about the task kinds and
 # should have unified names for them.
+WORKER_KIND_NAME = "worker_kind"
 TASK_1_NAME = "task_1"
 TASK_2_NAME = "task_2"
 
 # APPLICATION CONFIGURATION ___________________________________________________
 
-# 2. Configure the worker
-worker_config = WorkerApplicationConfig(
-    name="test_worker",
-    manager_config=manager_config,
-    broker_config=broker_config,
+# 1. Create a worker application
+worker_application = WorkerApplication(
+    config=WorkerApplicationConfig(
+        name="test_worker",
+        manager_config=manager_config,
+        broker_config=broker_config,
+        kind=WORKER_KIND_NAME,
+    ),
 )
 
-# 3. Create a worker application
-worker_application = WorkerApplication(worker_config)
 
-
-# 4. Create tasks and register them with the worker application
+# 2. Create tasks and register them with the worker application
 @worker_application.task(TASK_1_NAME)
 async def task_1(input_data: dict[Any, Any]) -> dict[Any, Any]:
     await asyncio.sleep(1)
     return input_data
 
 
 @worker_application.task(TASK_2_NAME)
-async def task_2(input_data: dict[Any, Any]) -> dict[Any, Any]:
+async def task_2(_: dict[Any, Any]) -> dict[Any, Any]:
     raise Exception("This is a test exception")
 
 
+# 3. Run the worker application
+
 if __name__ == "__main__":
     # Application can be run either as a standalone script or via the CLI.
     asyncio.run(worker_application.entrypoint())

client_sdks/python/pyproject.toml

@@ -16,9 +16,12 @@ classifiers = [
 requires-python = ">=3.12"
 dependencies = [
     "aio-pika>=9.5.3",
+    "aiohttp-retry>=2.9.1",
     "aiohttp>=3.11.8",
     "aioredis>=2.0.1",
+    "aioresponses>=0.7.8",
     "click>=8.1.7",
+    "pydantic>=2.10.5",
     "uuid>=1.30",
     "watchfiles>=1.0.3",
 ]
@@ -28,6 +31,9 @@ asyncio_mode = "auto"
 pythonpath = ["."]
 markers = [
     "bench: benchmark a function",
+    "target: current target test. Used locally.",
+    "unit: unit tests",
+    "service: tests that interact with external services",
 ]
 
 [build-system]

client_sdks/python/src/broker/__init__.py

@@ -1,28 +1,9 @@
+from broker.client import PublisherBrokerClient, WorkerBrokerClient, BaseBrokerClient
 from broker.config import BrokerConfig
-from broker.core import BrokerClient
 
-# TODO: Import only rabbit if tacoq[amqp] is installed
-from broker.rabbitmq import RabbitMQBroker
-
-
-def create_broker_instance(
-    config: BrokerConfig, exchange_name: str, worker_id: str
-) -> BrokerClient:
-    """Create appropriate broker client based on configuration.
-
-    ### Parameters
-    - `config`: Configuration for the broker connection
-    - `exchange_name`: Name of the exchange to use
-    - `worker_id`: Unique identifier for this worker instance
-
-    ### Returns
-    - `BrokerClient`: Configured broker client instance
-
-    ### Raises
-    - `ValueError`: If broker URL scheme is not supported
-    """
-
-    if config.url.startswith("amqp"):
-        return RabbitMQBroker(config, exchange_name, worker_id)
-    else:
-        raise ValueError(f"Unsupported broker URL: {config.url}")
+__all__ = [
+    "PublisherBrokerClient",
+    "WorkerBrokerClient",
+    "BaseBrokerClient",
+    "BrokerConfig",
+]

client_sdks/python/src/broker/client.py

@@ -0,0 +1,224 @@
+import json
+from typing import AsyncGenerator, Optional
+from broker.config import BrokerConfig
+from aio_pika import Message, connect_robust
+from models.task import Task
+from pydantic import BaseModel
+
+from aio_pika.abc import (
+    AbstractChannel,
+    AbstractQueue,
+    AbstractRobustConnection,
+    AbstractExchange,
+)
+
+# =========================================
+# Constants - Exchange and queue names
+# =========================================
+
+TASK_ASSIGNMENT_EXCHANGE = "task_assignment_exchange"
+""" Exchange for task assignments. """
+
+TASK_RESULT_EXCHANGE = "task_result_exchange"
+""" Exchange for task results. """
+
+# =========================================
+# Errors
+# =========================================
+
+
+class NoChannelError(Exception):
+    """Raised when a RabbitMQ client is not connected to the broker while
+    trying to perform an operation that requires a channel."""
+
+    pass
+
+
+class NotConnectedError(Exception):
+    """Raised when a RabbitMQ client is not connected to the broker while
+    trying to perform an operation that requires a connection."""
+
+    pass
+
+
+class QueueNotDeclaredError(Exception):
+    """Raised when a RabbitMQ client tries to use a queue that has not been
+    declared."""
+
+    pass
+
+
+class ExchangeNotDeclaredError(Exception):
+    """Raised when a RabbitMQ client tries to use an exchange that has not been
+    declared."""
+
+    pass
+
+
+## =========================================
+## Base Client
+## =========================================
+
+
+class BaseBrokerClient(BaseModel):
+    """RabbitMQ implementation of the broker interface."""
+
+    config: BrokerConfig
+    """ Configuration for the broker. """
+
+    _connection: Optional[AbstractRobustConnection] = None
+    """ The connection to the RabbitMQ server. """
+
+    _channel: Optional[AbstractChannel] = None
+    """ The channel to the RabbitMQ server. """
+
+    async def connect(self) -> None:
+        """Establish connection to RabbitMQ server and setup channel.
+
+        ### Raises
+        - `ConnectionError`: If connection to RabbitMQ fails
+        """
+
+        self._connection = await connect_robust(self.config.url)
+        self._channel = await self._connection.channel()
+
+    async def disconnect(self) -> None:
+        """Close RabbitMQ connection.
+
+        ### Raises
+        - `RabbitMQNotConnectedError`: If connection is not established
+        """
+
+        if self._connection is None:
+            raise NotConnectedError(
+                "Tried to disconnect from RabbitMQ, but connection was not established."
+            )
+
+        # Remove the exchanges
+        await self._connection.close()
+
+
+## =========================================
+## Publisher Client
+## =========================================
+
+
+class PublisherBrokerClient(BaseBrokerClient):
+    """RabbitMQ client for publishing tasks to workers.
+    Uses a fanout exchange to send tasks to both the manager queue
+    and the appropriate worker kind queue."""
+
+    _task_exchange: Optional[AbstractExchange] = None
+    """ The exchange for task assignments. """
+
+    async def connect(self) -> None:
+        await super().connect()
+
+        if self._channel is None:
+            raise NoChannelError(
+                "Tried to connect to RabbitMQ, but channel was not established."
+            )
+
+        # Connect to the exchange so to which we will publish tasks.
+        self._task_exchange = await self._channel.declare_exchange(
+            TASK_ASSIGNMENT_EXCHANGE,
+            passive=True,  # We only want to connect to the exchange if it already exists.
+        )
+
+    async def publish_task(self, routing_key: str, task: Task) -> None:
+        """Publish a task to both manager and worker queues via exchange and routing mechanisms.
+
+        ### Arguments
+        - `routing_key`: The routing key for the task. This is based on the worker kind. The publisher
+        client will know the routing key based on the requests it has made to the manager, who creates
+        the queues and binds them to the exchange.
+
+        ### Raises
+        - `RuntimeError`: If the exchange was not declared.
+        """
+
+        if self._task_exchange is None:
+            raise RuntimeError("Tried to publish task, but exchange was not declared.")
+
+        message = Message(body=task.model_dump_json().encode())
+
+        await self._task_exchange.publish(message, routing_key)
+
+
+## =========================================
+## Worker Client
+## =========================================
+
+
+class WorkerBrokerClient(BaseBrokerClient):
+    """RabbitMQ client for workers to consume tasks and publish results.
+    Each worker kind has its own queue for task assignments, but all workers
+    share a single queue for publishing results."""
+
+    task_assignment_queue_name: str
+    """ The name of the task assignment queue. """
+
+    _task_assignment_queue: Optional[AbstractQueue] = None
+    """ Queue for task assignments. """
+
+    _result_exchange: Optional[AbstractExchange] = None
+    """ Exchange for publishing results (shared by all workers). """
+
+    async def connect(self) -> None:
+        await super().connect()
+
+        if self._channel is None:
+            raise NoChannelError(
+                "Tried to connect to RabbitMQ, but channel was not established."
+            )
+
+        # =========================================
+        # Setup task assignment queue for this worker kind
+        # =========================================
+
+        self._task_assignment_queue = await self._channel.declare_queue(
+            self.task_assignment_queue_name,
+            passive=True,  # We only want to connect to the queue if it already exists.
+        )
+
+        # =========================================
+        # Setup result publishing infrastructure
+        # =========================================
+
+        # Set prefetch to one to enable fair dispatching
+        await self._channel.set_qos(prefetch_count=1)
+
+        # Setup result publishing infrastructure
+        self._result_exchange = await self._channel.declare_exchange(
+            TASK_RESULT_EXCHANGE,
+            passive=True,  # We only want to connect to the queue if it already exists.
+        )
+
+    async def listen(self) -> AsyncGenerator[Task, None]:
+        """Listen for tasks assigned to this worker's kind."""
+        if self._task_assignment_queue is None:
+            raise QueueNotDeclaredError(
+                "Tried to listen for tasks, but queue was not declared."
+            )
+
+        async for message in self._task_assignment_queue.iterator():
+            async with message.process():
+                yield Task(**json.loads(message.body.decode()))
+
+    async def publish_task_result(self, task: Task) -> None:
+        """Publish a task result to the shared results queue."""
+
+        # Check if the task has a result attached
+        if task.result is None:
+            raise ValueError(
+                "Tried to publish task result, but task has no result attached. How did it get to this point?"
+            )
+
+        if self._result_exchange is None:
+            raise ExchangeNotDeclaredError(
+                "Tried to publish task result, but exchange was not declared."
+            )
+
+        message = Message(body=task.model_dump_json().encode())
+
+        await self._result_exchange.publish(message, routing_key="")

client_sdks/python/src/broker/config.py

@@ -1,8 +1,8 @@
-from dataclasses import dataclass
+from pydantic import BaseModel
 
 
-@dataclass
-class BrokerConfig:
-    """Configuration for a broker."""
+class BrokerConfig(BaseModel):
+    """Configuration for a RabbitMQ broker."""
 
     url: str
+    """ The URL of the broker. """