Parallelisation
in Agentic
Systems

# Hierarchical map-reduce for large document corpora
import asyncio
from typing import List
from anthropic import AsyncAnthropic

client = AsyncAnthropic()
CHUNK_SIZE = 8000   # tokens per chunk
REDUCE_GROUP = 10  # max summaries per reduce call

async def map_chunk(chunk: str, task: str) -> str:
    """Process a single chunk — the map step."""
    resp = await client.messages.create(
        model="claude-haiku-4-5-20251001",  # fast + cheap for map
        max_tokens=1024,
        messages=[{
            "role": "user",
            "content": f"Task: {task}\n\nDocument chunk:\n{chunk}\n\nProvide a structured intermediate result."
        }]
    )
    return resp.content[0].text

async def reduce_batch(intermediates: List[str], task: str, final: bool = False) -> str:
    """Reduce a batch of intermediate results."""
    model = "claude-sonnet-4-6" if final else "claude-haiku-4-5-20251001"
    combined = "\n\n---\n\n".join(intermediates)
    resp = await client.messages.create(
        model=model,
        max_tokens=2048,
        messages=[{
            "role": "user",
            "content": f"Task: {task}\n\nIntermediate results to synthesise:\n{combined}"
        }]
    )
    return resp.content[0].text

async def hierarchical_map_reduce(documents: List[str], task: str) -> str:
    # Phase 1: MAP — all chunks processed in parallel
    semaphore = asyncio.Semaphore(20)  # respect rate limits

    async def bounded_map(chunk):
        async with semaphore:
            return await map_chunk(chunk, task)

    intermediates = await asyncio.gather(
        *[bounded_map(doc) for doc in documents]
    )

    # Phase 2: HIERARCHICAL REDUCE — tree reduction
    current_level = list(intermediates)
    while len(current_level) > 1:
        groups = [
            current_level[i:i + REDUCE_GROUP]
            for i in range(0, len(current_level), REDUCE_GROUP)
        ]
        is_final = len(groups) == 1
        current_level = await asyncio.gather(
            *[reduce_batch(g, task, final=is_final) for g in groups]
        )

    return current_level[0]

import asyncio
from anthropic import AsyncAnthropic

client = AsyncAnthropic()

async def stream_and_extract(prompt: str) -> dict:
    """Pipeline: stream LLM response while concurrently extracting structure."""
    buffer = []
    extracted = {}

    async with client.messages.stream(
        model="claude-sonnet-4-6",
        max_tokens=2048,
        messages=[{"role": "user", "content": prompt}]
    ) as stream:
        async for chunk in stream.text_stream:
            buffer.append(chunk)
            current_text = "".join(buffer)

            # Pipeline stage: extract JSON as it streams, don't wait for completion
            if "action" in current_text and not extracted.get("action"):
                extracted["action"] = parse_partial_json(current_text, "action")
                # Fire tool call immediately — don't wait for stream to end
                asyncio.create_task(dispatch_tool(extracted["action"]))

    return extracted

async def pipeline_rag_agent(query: str) -> str:
    # Stage 1+2 OVERLAP: embed query while pre-warming retrieval connection
    embed_task   = asyncio.create_task(embed_query(query))
    warmup_task  = asyncio.create_task(warmup_retriever())
    embedding, _ = await asyncio.gather(embed_task, warmup_task)

    # Stage 2+3 OVERLAP: retrieve while assembling prompt template
    retrieve_task = asyncio.create_task(retrieve_contexts(embedding))
    template_task = asyncio.create_task(build_prompt_template(query))
    contexts, template = await asyncio.gather(retrieve_task, template_task)

    # Stage 3: LLM streams — extraction is pipelined mid-stream
    prompt = template.format(contexts=contexts)
    result = await stream_and_extract(prompt)
    return result

import asyncio
from dataclasses import dataclass

async def speculative_route(query: str) -> str:
    """Launch fast and slow paths simultaneously; return first valid result."""

    async def fast_path() -> str:
        # Try cache first — returns instantly or raises CacheMiss
        result = await semantic_cache_lookup(query)
        if not result:
            raise ValueError("cache miss")
        return result

    async def medium_path() -> str:
        # Try RAG + small model — ~800ms
        context = await retrieve_and_summarise(query)
        return await haiku_generate(query, context)

    async def full_path() -> str:
        # Full agentic pipeline with tool calls — 3-6s
        return await full_agent_pipeline(query)

    # Stagger starts: fast fires immediately, medium after 50ms, full after 200ms
    # This avoids burning tokens on slow path if fast path will win
    tasks = {
        asyncio.create_task(fast_path()): "fast",
        asyncio.create_task(_delayed(medium_path, 0.05)): "medium",
        asyncio.create_task(_delayed(full_path, 0.20)): "full",
    }

    pending = set(tasks.keys())
    while pending:
        done, pending = await asyncio.wait(
            pending, return_when=asyncio.FIRST_COMPLETED
        )
        for task in done:
            if not task.exception():
                # Cancel all remaining tasks — we have our answer
                for p in pending: p.cancel()
                return task.result()

    raise RuntimeError("All speculative paths failed")

async def _delayed(coro_fn, delay: float):
    await asyncio.sleep(delay)
    return await coro_fn()

# THE FOUR CORE ASYNC PARALLELISATION PRIMITIVES

# 1. asyncio.gather — N independent coroutines, all results needed
results = await asyncio.gather(
    llm_call(prompt_a),
    retrieve_context(query),
    call_tool("database", params),
    return_exceptions=True  # don't abort on single failure
)

# 2. asyncio.wait FIRST_COMPLETED — speculative, race to first valid result
done, pending = await asyncio.wait(
    [task_a, task_b, task_c],
    return_when=asyncio.FIRST_COMPLETED
)
winner = done.pop().result()
for p in pending: p.cancel()

# 3. asyncio.Semaphore — bounded concurrency (respect rate limits!)
sem = asyncio.Semaphore(20)  # max 20 concurrent LLM calls
async def bounded_call(prompt):
    async with sem:
        return await llm_call(prompt)
results = await asyncio.gather(*[bounded_call(p) for p in prompts])

# 4. asyncio.Queue — producer/consumer pipeline between stages
queue: asyncio.Queue = asyncio.Queue(maxsize=100)

async def producer():
    async for item in document_stream():
        await queue.put(item)  # blocks if queue is full (backpressure)
    await queue.put(None)   # sentinel

async def consumer():
    while True:
        item = await queue.get()
        if item is None: break
        await process(item)

# Run producer and N consumers concurrently
await asyncio.gather(
    producer(),
    *[consumer() for _ in range(10)]
)

from langgraph.graph import StateGraph, END
from typing import TypedDict, Annotated, List
import operator, asyncio

class AgentState(TypedDict):
    query: str
    db_result: str
    rag_result: str
    api_result: str
    # Annotated with operator.add means results are MERGED not overwritten
    tool_outputs: Annotated[List[str], operator.add]
    final_answer: str

# Define parallel nodes — each is an independent async function
async def query_database(state: AgentState) -> dict:
    result = await db_lookup(state["query"])
    return {"db_result": result, "tool_outputs": [result]}

async def retrieve_context(state: AgentState) -> dict:
    result = await vector_search(state["query"])
    return {"rag_result": result, "tool_outputs": [result]}

async def call_external_api(state: AgentState) -> dict:
    result = await api_fetch(state["query"])
    return {"api_result": result, "tool_outputs": [result]}

async def synthesise(state: AgentState) -> dict:
    # This node only runs after all parallel branches complete (fan-in)
    combined = "\n".join(state["tool_outputs"])
    answer = await llm_synthesise(state["query"], combined)
    return {"final_answer": answer}

# BUILD THE GRAPH — topology encodes parallelism
graph = StateGraph(AgentState)
graph.add_node("db",  query_database)
graph.add_node("rag", retrieve_context)
graph.add_node("api", call_external_api)
graph.add_node("synth", synthesise)

# Fan-out from START to three parallel branches
graph.set_entry_point("db")
graph.set_entry_point("rag")   # LangGraph executes all entry points in parallel
graph.set_entry_point("api")

# Fan-in: all branches must complete before synth runs
graph.add_edge("db",  "synth")
graph.add_edge("rag", "synth")
graph.add_edge("api", "synth")
graph.add_edge("synth", END)

app = graph.compile()
result = await app.ainvoke({"query": "What is the customer's current balance?"})

import redis.asyncio as redis
import time, asyncio

class TokenBucketLimiter:
    """Distributed token bucket — shared across all agent pods via Redis."""
    def __init__(self, r: redis.Redis, key: str,
                 capacity: int, refill_rate: float):
        self.r = r
        self.key = key          # e.g. "rate:claude:tokens"
        self.capacity = capacity # max tokens (e.g. 100_000 TPM)
        self.refill_rate = refill_rate  # tokens per second

    async def acquire(self, tokens: int) -> bool:
        """Atomically acquire tokens. Returns True if granted, False if throttled."""
        now = time.time()
        # Lua script runs atomically in Redis — no race conditions
        script = """
        local key, capacity, rate, tokens, now = KEYS[1], ARGV[1], ARGV[2], ARGV[3], ARGV[4]
        local bucket = redis.call('HMGET', key, 'tokens', 'last_refill')
        local current = tonumber(bucket[1]) or capacity
        local last_refill = tonumber(bucket[2]) or now
        local elapsed = now - last_refill
        local refilled = math.min(capacity, current + elapsed * rate)
        if refilled >= tokens then
            redis.call('HMSET', key, 'tokens', refilled - tokens, 'last_refill', now)
            redis.call('EXPIRE', key, 3600)
            return 1
        end
        return 0
        """
        result = await self.r.eval(
            script, 1, self.key,
            self.capacity, self.refill_rate, tokens, now
        )
        return bool(result)

    async def acquire_or_wait(self, tokens: int, max_wait: float = 30.0):
        deadline = time.time() + max_wait
        while time.time() < deadline:
            if await self.acquire(tokens):
                return
            await asyncio.sleep(0.1)
        raise TimeoutError(f"Rate limit: could not acquire {tokens} tokens in {max_wait}s")