Uses NATS to aggregate responses from workers publishing LLM inference to the subject inference.requests

llm_request_aggregator.py

import asyncio
import nats
import uuid


async def aggregate_inferences(nats_url, request_subject, data, timeout=10):
    nc = await nats.connect(nats_url)

    responses = []

    async def response_handler(msg):
        responses.append(msg.data.decode())

    inbox = f"_INBOX.{uuid.uuid4()}"
    subscription = await nc.subscribe(inbox, cb=response_handler)
    await nc.publish(request_subject, data.encode(), reply=inbox)
    await asyncio.sleep(timeout)
    await subscription.drain()
    await nc.close()
    return responses


if __name__ == "__main__":
    prompt = "Let's discuss the best way to run machine learning on the edge"
    results = asyncio.run(
        aggregate_inferences("nats://localhost:4222", "inference.requests", prompt)
    )
    print("Aggregated Results:", results)