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",
+]