0

Multi-step task with an LLM tool-use loop

by · about 3 hours ago· ⑂ 0 remixes

Question: Multi-step task with an LLM tool-use loop

Companies: Anthropic

Difficulty: Medium

Problem statement

In a notebook with docs open, solve a multi-step problem by writing code that calls a language model as a building block: design the prompts and the loop that invokes tools and folds results back in. Starter API code is given; judgment about decomposition matters more than syntax.

class Tool:
    def __init__(self, name, description, input_schema, fn): ...


def run_agent(client, model, system, user_message, tools, max_turns=8, max_tokens=1024):
    """Messages-API tool loop. Returns (final_text, messages)."""


def map_prompts(client, model, system, inputs, render, parse, workers=8, retries=1):
    """One validated call per input, in parallel."""

Approach

What the round grades. You get a notebook, the API docs, a starter client, and a task too big for one prompt (for example: "given 40 support tickets and a small product database, find the five most common root causes and the affected customers"). The grade is about judgment: what you give the model, what you keep in plain code, how you check the model's output, and whether the loop terminates and fails loudly. Syntax help from the docs is expected.

Decide what is code and what is model. A rule that holds up well in this round:

  • Plain code for anything exact: parsing, joins, counting, sorting, arithmetic, deduplication. Models are slow and occasionally wrong at these, and code is free and testable.
  • The model for anything fuzzy: reading free text, classifying into categories it has to infer, summarizing, choosing which tool to call next.
  • Tools connect the two: when the model needs a fact or an exact computation, it asks for a tool and your code answers.

Two shapes, and when to use each.

  1. A fixed pipeline (prompt chain). You know the steps ahead of time: classify each ticket (one call each, in parallel), group the labels in code, then one call to name and summarize each group. Cheapest, easiest to test, easiest to debug. Prefer it whenever the steps are known.
  2. A tool-use loop (agent). The model decides the next step: it calls lookup_customer, sees the plan tier, then decides to call get_incidents. Use it when the path depends on what the model finds. It costs more tokens and needs guardrails.

Most good solutions are a pipeline with one small agent inside it, not one giant agent.

The loop, step by step (Messages API shape):

  1. Send the system prompt, the conversation so far, and the tool definitions (name, description, JSON schema for input).
  2. If the response's stop_reason is "tool_use", run every tool_use block in it (the model may ask for several at once). Append the assistant message as is, then one user message holding a tool_result block per call, each with the matching tool_use_id.
  3. A tool that fails returns its error as the result with is_error: true. The model usually corrects itself (wrong id format, missing field). Never crash the loop on a tool error.
  4. Stop when stop_reason is "end_turn"; return the text. Stop with an error after max_turns, so a confused model cannot loop forever and spend money. Treat "max_tokens" as a failure to surface, not a final answer.

Prompt design that matters here. Tool descriptions are prompts: say what the tool returns, its units, and when not to use it. Ask for the final answer in a fixed format (JSON with named fields) and validate it in code; on a validation error, send the error back once and ask for a fix. Put the instructions and the output schema in the system prompt and the data in the user turn, inside clear tags (<ticket id="17">...</ticket>) so the model cannot confuse data with instructions.

Numbers. 40 tickets × ~600 input tokens for classification = 24k input tokens, run with 8 concurrent calls: about 5 rounds of ~2 s, so ~10 s. The same work as one agent that reads tickets through a tool: 40+ turns, each resending the growing conversation, easily 500k input tokens and minutes. That gap is the argument for the pipeline.

Code

import json
from concurrent.futures import ThreadPoolExecutor


class Tool:
    def __init__(self, name, description, input_schema, fn):
        self.name, self.description, self.input_schema, self.fn = name, description, input_schema, fn

    def spec(self):
        return {"name": self.name, "description": self.description, "input_schema": self.input_schema}


def _get(block, key):
    return block[key] if isinstance(block, dict) else getattr(block, key)


def run_agent(client, model, system, user_message, tools, max_turns=8, max_tokens=1024):
    """Tool-use loop. Returns (final_text, messages). Raises on runaway loops."""
    by_name = {t.name: t for t in tools}
    messages = [{"role": "user", "content": user_message}]
    for _ in range(max_turns):
        response = client.messages.create(
            model=model, system=system, max_tokens=max_tokens,
            messages=messages, tools=[t.spec() for t in tools],
        )
        content = list(response.content)
        messages.append({"role": "assistant", "content": content})
        if response.stop_reason == "end_turn":
            return "".join(_get(b, "text") for b in content if _get(b, "type") == "text"), messages
        if response.stop_reason != "tool_use":
            raise RuntimeError(f"stopped early: {response.stop_reason}")
        results = []
        for block in content:
            if _get(block, "type") != "tool_use":
                continue
            results.append(_call_tool(by_name, block))
        messages.append({"role": "user", "content": results})
    raise RuntimeError(f"no final answer after {max_turns} turns")


def _call_tool(by_name, block):
    name, args, call_id = _get(block, "name"), _get(block, "input"), _get(block, "id")
    tool = by_name.get(name)
    try:
        if tool is None:
            raise KeyError(f"unknown tool {name!r}; available: {sorted(by_name)}")
        output = tool.fn(**args)
        text = output if isinstance(output, str) else json.dumps(output)
        return {"type": "tool_result", "tool_use_id": call_id, "content": text}
    except Exception as e:  # the model sees the error and can retry
        return {"type": "tool_result", "tool_use_id": call_id, "content": f"{type(e).__name__}: {e}", "is_error": True}


def map_prompts(client, model, system, inputs, render, parse, workers=8, retries=1):
    """Fixed-pipeline step: one call per input, in parallel, each output validated.

    render(item) -> user text; parse(text) -> value or raises ValueError.
    """
    def one(item):
        messages = [{"role": "user", "content": render(item)}]
        for attempt in range(retries + 1):
            response = client.messages.create(model=model, system=system, max_tokens=512, messages=messages)
            text = "".join(_get(b, "text") for b in response.content if _get(b, "type") == "text")
            try:
                return parse(text)
            except ValueError as e:
                if attempt == retries:
                    raise
                messages += [{"role": "assistant", "content": text},
                             {"role": "user", "content": f"That output was invalid ({e}). Reply again in the required format only."}]

    with ThreadPoolExecutor(max_workers=workers) as pool:
        return list(pool.map(one, inputs))

With the real SDK, client = anthropic.Anthropic() works unchanged: client.messages.create takes these arguments, and response blocks have .type, .text, .id, .name and .input. _get lets the same code run on plain dicts in tests.

A sketch of the full pipeline for the ticket task:

LABEL_SYSTEM = (
    "You label one support ticket with its root cause. Reply with JSON only: "
    '{"cause": "<2-4 word label>", "customer_id": "<id or null>"}'
)


def parse_label(text):
    data = json.loads(text)  # json.JSONDecodeError is a ValueError
    if not isinstance(data.get("cause"), str):
        raise ValueError("missing cause")
    return data


# labels = map_prompts(client, MODEL, LABEL_SYSTEM, tickets,
#                      render=lambda t: f"<ticket id=\"{t['id']}\">{t['body']}</ticket>",
#                      parse=parse_label)
# Group with collections.Counter in code; ask the model once to merge
# near-synonym labels ("login bug" vs "auth failure"), then count again.

Complexity

  • Calls: the pipeline makes N + (number of groups) calls; the agent makes one call per turn.
  • Tokens: an agent resends the whole conversation every turn, so input tokens grow roughly with turns² × average turn size. A pipeline's cost is linear in N. Prompt caching of the shared system prompt and tool definitions cuts the repeated part substantially.
  • Latency: the pipeline's map step runs in parallel, so wall clock is about (N / workers) × call latency. An agent loop is strictly sequential.

Test cases

from types import SimpleNamespace as NS

from solution import Tool, run_agent


class FakeClient:
    """Plays back scripted responses and records every request."""

    def __init__(self, responses):
        self.responses, self.requests = list(responses), []
        self.messages = self

    def create(self, **kwargs):
        self.requests.append({**kwargs, "messages": list(kwargs["messages"])})  # copy: the loop keeps appending
        return self.responses.pop(0)


def text(t):
    return NS(type="text", text=t)


def call(id_, name, **args):
    return NS(type="tool_use", id=id_, name=name, input=args)


def lookup_tool(calls):
    def lookup(customer_id):
        calls.append(customer_id)
        if customer_id == "bad":
            raise ValueError("no such customer")
        return {"id": customer_id, "plan": "pro"}
    return Tool("lookup", "Look up a customer by id.", {"type": "object"}, lookup)


def test_runs_parallel_tool_calls_and_returns_final_text():
    calls = []
    client = FakeClient([
        NS(stop_reason="tool_use", content=[text("checking"), call("t1", "lookup", customer_id="a"), call("t2", "lookup", customer_id="b")]),
        NS(stop_reason="end_turn", content=[text("both are pro")]),
    ])
    answer, messages = run_agent(client, "m", "sys", "who?", [lookup_tool(calls)])
    assert answer == "both are pro"
    assert calls == ["a", "b"]
    results = client.requests[1]["messages"][-1]["content"]
    assert [r["tool_use_id"] for r in results] == ["t1", "t2"]


def test_tool_errors_go_back_to_the_model():
    client = FakeClient([
        NS(stop_reason="tool_use", content=[call("t1", "lookup", customer_id="bad"), call("t2", "nope")]),
        NS(stop_reason="end_turn", content=[text("sorry")]),
    ])
    run_agent(client, "m", "sys", "q", [lookup_tool([])])
    r1, r2 = client.requests[1]["messages"][-1]["content"]
    assert r1["is_error"] and "no such customer" in r1["content"]
    assert r2["is_error"] and "unknown tool" in r2["content"]


def test_loop_is_bounded():
    looping = [NS(stop_reason="tool_use", content=[call(f"t{i}", "lookup", customer_id="a")]) for i in range(5)]
    try:
        run_agent(FakeClient(looping), "m", "sys", "q", [lookup_tool([])], max_turns=3)
    except RuntimeError as e:
        assert "3 turns" in str(e)
    else:
        assert False, "expected RuntimeError"


def test_max_tokens_is_a_failure():
    client = FakeClient([NS(stop_reason="max_tokens", content=[text("half an ans")])])
    try:
        run_agent(client, "m", "sys", "q", [])
    except RuntimeError as e:
        assert "max_tokens" in str(e)
    else:
        assert False
import json
from types import SimpleNamespace as NS

from solution import map_prompts


class Scripted:
    def __init__(self, replies):
        self.replies, self.calls = dict(replies), []
        self.messages = self

    def create(self, **kwargs):
        prompt = kwargs["messages"][0]["content"]
        self.calls.append(len(kwargs["messages"]))
        return NS(stop_reason="end_turn", content=[NS(type="text", text=self.replies[prompt].pop(0))])


def parse(text):
    data = json.loads(text)
    if "cause" not in data:
        raise ValueError("missing cause")
    return data["cause"]


def test_map_validates_and_retries_once():
    client = Scripted({"t1": ['{"cause": "login"}'], "t2": ["not json", '{"cause": "billing"}']})
    out = map_prompts(client, "m", "sys", ["t1", "t2"], render=lambda x: x, parse=parse, workers=2)
    assert out == ["login", "billing"]
    assert sorted(client.calls) == [1, 1, 3]  # the retry carries the bad answer and the error

Follow-ups

  1. Cost and latency. Cache the shared prefix (system prompt and tool definitions) with prompt caching, use a small fast model for the map step and a larger one only for the final synthesis, and run the map step concurrently within the rate limit (a semaphore plus retry with backoff on rate-limit errors).
  2. Context growth. Long agent runs fill the context window with tool output. Truncate large tool results (return the first rows and a count), summarize old turns, or move data out of the conversation and give the model a tool to page through it.
  3. Evaluating it. Build 20–50 labeled examples, score the pipeline end to end and per step (label accuracy, grouping quality), and rerun after every prompt change. Read failures by hand before tuning anything.
  4. Safety of tools. Tools with side effects (sending email, writing data) need a confirmation step or a dry-run mode, and an allow-list of arguments. Data from tickets or web pages can contain instructions; keep it inside tags and never let it pick which side-effecting tool runs without a check in code.
  5. Determinism. For tests, record real responses once and replay them, as the fake client does. In production, log every request and response with a run id so a bad answer can be traced to the step that caused it.

Interviewer expectations

Mid-level (L4)

  • Writes a working tool-use loop: runs tool calls, returns results with matching ids, stops on end_turn.
  • Keeps exact work (counting, joins) in code rather than asking the model to do arithmetic.
  • Gets a correct answer to the task on a few examples.

Senior (L5)

  • Decomposes the task into a pipeline and justifies where a model is needed; handles parallel tool calls, tool errors, and a turn limit.
  • Asks for structured output, validates it, and retries with the error.
  • Checks results on a small labeled set instead of eyeballing one run.

Staff (L6+)

  • Reasons about token cost and latency with numbers and picks the pipeline, model sizes and caching accordingly.
  • Designs evaluation and logging so prompt changes can be compared, and treats tool inputs from untrusted data as a security boundary.
  • Explains when an agent loop is worth its cost and how to bound it (turns, tokens, time, side effects).

Sign in to comment, vote or remix.