Speculative Tool Execution: Cut Agent Latency 48% in Production

async with client.messages.stream(
    model="claude-opus-4-7",
    max_tokens=4096,
    messages=history,
    tools=tool_schemas,
    extra_headers={"anthropic-beta": "fine-grained-tool-streaming-2025-05-14"},
) as stream:
    async for event in stream:
        if event.type == "content_block_start" and event.content_block.type == "tool_use":
            tool_name = event.content_block.name
            tool_use_id = event.content_block.id
            # SPECULATION TRIGGER. Name is final. Args still streaming.
            asyncio.create_task(speculate(tool_name, tool_use_id, history))

"""
paste_harness.py
A client-side speculative tool execution harness over the Anthropic SDK.

Architecture:
  - BigramPredictor: P(next_tool | prev_tool, prev_args_shape)
  - EligibilityPolicy: per-tool full | warmup | forbid
  - PASTEHarness: streams a turn, fires speculation on content_block_start,
    promotes on commit, cancels orphans on mismatch.
"""
import asyncio
import json
import time
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Any, Awaitable, Callable, Literal

from anthropic import AsyncAnthropic
from anthropic.types import ToolUseBlock

Eligibility = Literal["full", "warmup", "forbid"]


def canonical(d: dict) -> str:
    """Stable canonical form of a tool input dict for cache keying."""
    return json.dumps(d, sort_keys=True, separators=(",", ":"))


def _shape(d: dict) -> str:
    """Argument-shape key. Strips values, keeps keys + value-types."""
    if not isinstance(d, dict):
        return type(d).__name__
    return "{" + ",".join(f"{k}:{type(v).__name__}" for k, v in sorted(d.items())) + "}"


# ---------------------------------------------------------------------------
# Predictor
# ---------------------------------------------------------------------------

class BigramPredictor:
    """Frequency-counted predictor: P(next_tool | prev_tool, prev_args_shape).

    Symbolic arg derivation: for known templates we copy fields from the
    previous call (e.g. arxiv_search.id -> fetch_url.url). Heuristics live
    in `_derive_args`. Replace with a small LM if you want better recall.
    """

    def __init__(self) -> None:
        self.counts: dict[tuple[str, str], dict[tuple[str, str], int]] = (
            defaultdict(lambda: defaultdict(int))
        )

    def observe(
        self,
        prev_tool: str,
        prev_args: dict,
        next_tool: str,
        next_args: dict,
    ) -> None:
        key = (prev_tool, _shape(prev_args))
        self.counts[key][(next_tool, _shape(next_args))] += 1

    def predict(
        self,
        prev_tool: str,
        prev_args: dict,
        history: list[dict],
        top_k: int = 3,
    ) -> list[tuple[str, dict, float]]:
        key = (prev_tool, _shape(prev_args))
        bucket = self.counts.get(key)
        if not bucket:
            return []
        total = sum(bucket.values())
        ranked = sorted(bucket.items(), key=lambda kv: -kv[1])[:top_k]
        out: list[tuple[str, dict, float]] = []
        for (next_tool, next_shape), count in ranked:
            args = self._derive_args(next_tool, prev_tool, prev_args, history)
            if args is None:
                continue
            out.append((next_tool, args, count / total))
        return out

    @staticmethod
    def _derive_args(
        next_tool: str,
        prev_tool: str,
        prev_args: dict,
        history: list[dict],
    ) -> dict | None:
        # Symbolic templates. Add yours here.
        if prev_tool == "arxiv_search" and next_tool == "fetch_url":
            q = prev_args.get("query", "")
            return {"url": f"https://arxiv.org/abs/{q}"} if q else None
        if prev_tool == "fetch_url" and next_tool == "extract_text":
            return {"url": prev_args.get("url")} if prev_args.get("url") else None
        if prev_tool == "extract_text" and next_tool == "summarize":
            return {"source_url": prev_args.get("url")} if prev_args.get("url") else None
        if prev_tool == "web_search" and next_tool == "fetch_url":
            return None  # too many candidates, do not speculate blindly
        return None


# ---------------------------------------------------------------------------
# Stats
# ---------------------------------------------------------------------------

@dataclass
class Stats:
    fired: int = 0
    committed: int = 0
    wasted: int = 0
    dollars_wasted: float = 0.0
    p_hit: float = field(default=0.0)

    def update(self) -> None:
        self.p_hit = self.committed / self.fired if self.fired else 0.0


# ---------------------------------------------------------------------------
# Harness
# ---------------------------------------------------------------------------

class PASTEHarness:
    def __init__(
        self,
        client: AsyncAnthropic,
        tools: dict[str, Callable[..., Awaitable[Any]]],
        tool_schemas: list[dict],
        policy: dict[str, Eligibility],
        predictor: BigramPredictor,
        model: str = "claude-opus-4-7",
        cost_per_call: dict[str, float] | None = None,
    ) -> None:
        self.client = client
        self.tools = tools
        self.tool_schemas = tool_schemas
        self.policy = policy
        self.predictor = predictor
        self.model = model
        self.cost_per_call = cost_per_call or {}
        self.stats = Stats()
        # Speculative result cache: (tool_name, canonical_args) -> Future[result]
        self._cache: dict[tuple[str, str], asyncio.Future] = {}
        # tool_use_id -> Future[result]; bound only after commit
        self._committed: dict[str, asyncio.Future] = {}
        self._last_call: tuple[str, dict] | None = None

    # ---- public ----------------------------------------------------------

    async def run(self, user_msg: str, max_turns: int = 16) -> str:
        history: list[dict] = [{"role": "user", "content": user_msg}]
        for _ in range(max_turns):
            assistant_blocks, stop_reason = await self._stream_turn(history)
            history.append({"role": "assistant", "content": assistant_blocks})
            tool_uses = [b for b in assistant_blocks if isinstance(b, ToolUseBlock)
                         or (isinstance(b, dict) and b.get("type") == "tool_use")]
            if not tool_uses or stop_reason == "end_turn":
                self._cancel_orphans()
                return _final_text(assistant_blocks)
            results = await asyncio.gather(*(self._resolve(tu) for tu in tool_uses))
            history.append({
                "role": "user",
                "content": [
                    {"type": "tool_result", "tool_use_id": tu_id, "content": json.dumps(res)}
                    for (tu_id, res) in results
                ],
            })
            self._cancel_orphans()
        return "max_turns_exceeded"

    # ---- streaming -------------------------------------------------------

    async def _stream_turn(self, history: list[dict]) -> tuple[list, str]:
        async with self.client.messages.stream(
            model=self.model,
            max_tokens=4096,
            messages=history,
            tools=self.tool_schemas,
            extra_headers={"anthropic-beta": "fine-grained-tool-streaming-2025-05-14"},
        ) as stream:
            async for event in stream:
                if (
                    event.type == "content_block_start"
                    and getattr(event.content_block, "type", None) == "tool_use"
                ):
                    name = event.content_block.name
                    if self.policy.get(name, "forbid") != "forbid":
                        # Fire before args finish streaming.
                        asyncio.create_task(self._speculate(name, history))
            final = await stream.get_final_message()
        return list(final.content), final.stop_reason

    # ---- speculation -----------------------------------------------------

    async def _speculate(self, observed_name: str, history: list[dict]) -> None:
        if self._last_call is None:
            return
        prev_name, prev_args = self._last_call
        candidates = self.predictor.predict(prev_name, prev_args, history)
        for cand_name, cand_args, conf in candidates:
            if cand_name != observed_name:
                continue
            if self.policy.get(cand_name, "forbid") == "forbid":
                continue
            key = (cand_name, canonical(cand_args))
            if key in self._cache:
                return
            fut: asyncio.Future = asyncio.get_event_loop().create_future()
            self._cache[key] = fut
            self.stats.fired += 1

            async def _exec() -> None:
                try:
                    if self.policy.get(cand_name) == "warmup":
                        # DNS prewarm, connection prewarm, no side effect.
                        result = await _warmup(cand_name, cand_args, self.tools)
                    else:
                        result = await self.tools[cand_name](**cand_args)
                    if not fut.done():
                        fut.set_result(result)
                except Exception as exc:
                    if not fut.done():
                        fut.set_exception(exc)

            asyncio.create_task(_exec())
            return  # one speculation per trigger

    # ---- commit / discard ------------------------------------------------

    async def _resolve(self, tu) -> tuple[str, Any]:
        name = tu.name if hasattr(tu, "name") else tu["name"]
        args = tu.input if hasattr(tu, "input") else tu["input"]
        tu_id = tu.id if hasattr(tu, "id") else tu["id"]
        key = (name, canonical(args))
        fut = self._cache.pop(key, None)
        if fut is not None:
            self.stats.committed += 1
            self._last_call = (name, args)
            try:
                return tu_id, await fut
            except Exception as exc:
                return tu_id, {"error": repr(exc)}
        # No cache hit. Run synchronously.
        result = await self.tools[name](**args)
        # Train the predictor on the actual transition.
        if self._last_call is not None:
            self.predictor.observe(self._last_call[0], self._last_call[1], name, args)
        self._last_call = (name, args)
        return tu_id, result

    def _cancel_orphans(self) -> None:
        for (name, _), fut in list(self._cache.items()):
            if not fut.done():
                fut.cancel()
            self.stats.wasted += 1
            self.stats.dollars_wasted += self.cost_per_call.get(name, 0.0)
        self._cache.clear()
        self.stats.update()


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

async def _warmup(name: str, args: dict, tools: dict) -> Any:
    # Hook for DNS/connection prewarm. Default: no-op.
    return {"warmup": True, "tool": name}


def _final_text(blocks: list) -> str:
    parts = []
    for b in blocks:
        text = getattr(b, "text", None) if not isinstance(b, dict) else b.get("text")
        if text:
            parts.append(text)
    return "\n".join(parts)


# ---------------------------------------------------------------------------
# Example: a small read-only research stack.
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import os

    async def web_search(query: str) -> dict:
        await asyncio.sleep(0.4)
        return {"hits": [f"https://example.com/{query}"]}

    async def arxiv_search(query: str) -> dict:
        await asyncio.sleep(0.5)
        return {"id": "2603.18897", "title": "PASTE"}

    async def fetch_url(url: str) -> dict:
        await asyncio.sleep(0.6)
        return {"url": url, "html": "<html>...</html>"}

    async def extract_text(url: str) -> dict:
        await asyncio.sleep(0.2)
        return {"url": url, "text": "Speculation reduces task time 48.5% ..."}

    async def summarize(source_url: str) -> dict:
        await asyncio.sleep(0.3)
        return {"summary": f"summary of {source_url}"}

    async def synthesize(notes: list) -> dict:
        await asyncio.sleep(0.4)
        return {"answer": "synthesized"}

    tools = {
        "web_search": web_search,
        "arxiv_search": arxiv_search,
        "fetch_url": fetch_url,
        "extract_text": extract_text,
        "summarize": summarize,
        "synthesize": synthesize,
    }
    tool_schemas = [
        {"name": "web_search", "description": "Search the web.",
         "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]}},
        {"name": "arxiv_search", "description": "Search arxiv.",
         "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]}},
        {"name": "fetch_url", "description": "Fetch a URL.",
         "input_schema": {"type": "object", "properties": {"url": {"type": "string"}}, "required": ["url"]}},
        {"name": "extract_text", "description": "Extract text from a URL.",
         "input_schema": {"type": "object", "properties": {"url": {"type": "string"}}, "required": ["url"]}},
        {"name": "summarize", "description": "Summarize a source.",
         "input_schema": {"type": "object", "properties": {"source_url": {"type": "string"}}, "required": ["source_url"]}},
        {"name": "synthesize", "description": "Synthesize notes.",
         "input_schema": {"type": "object", "properties": {"notes": {"type": "array"}}, "required": ["notes"]}},
    ]
    policy: dict[str, Eligibility] = {
        "web_search": "full",
        "arxiv_search": "full",
        "fetch_url": "full",
        "extract_text": "full",
        "summarize": "full",
        "synthesize": "forbid",  # final answer, no point speculating
    }
    cost_per_call = {n: 0.001 for n in tools}

    async def main() -> None:
        client = AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
        harness = PASTEHarness(client, tools, tool_schemas, policy,
                               BigramPredictor(), cost_per_call=cost_per_call)
        out = await harness.run("Summarize the PASTE paper and the UMD speculative tool-calls paper.")
        print(out)
        print(harness.stats)

    asyncio.run(main())