Views: 10
Choosing between tools is just next-token prediction over a bigger action vocabulary, and chaining tool calls is the same loop run more than once. Module 6 of From Zero to Agents extends a tiny trained agent to two tools and multi-step reasoning, with an honest look at where it actually fails.
Episode 06.00 ended by asking what breaks first as a single-tool, single-step toy scales toward multiple tools, possible failures, and multi-step tasks. This episode builds all three extensions directly, and — consistent with this course’s habit of reporting what actually happens rather than what was expected — one of them genuinely doesn’t work well, for an entirely diagnosable reason.
1. Theory: none of these extensions need new machinery
1.1 Tool selection is just next-token prediction over a bigger vocabulary. Episode 06.00’s single-tool agent never had to choose — there was only one possible action. Adding a second tool means adding a second tool-identifier token to the vocabulary, and letting the model’s ordinary next-token softmax (Episode 02.04, unchanged) decide between them based on context. No new mechanism — tool selection is exactly the same computation that decides every other token, just with the stakes of “which tool” riding on that one decision.
1.2 Chaining is the same loop, run more than once. Episode 06.00 §2.1 already wrote the state-transition recursion in general form — — but only exercised it for a single iteration. A task requiring two tool calls in sequence (say, add two numbers, then add a third to that result) is the identical recursion, just run twice before a final answer terminates it: propose an action, get a real observation, propose another action using that observation as part of the new context, get another real observation, then finalize.
1.3 Failure is just a different kind of observation, not a different kind of machinery. When a tool call fails — invalid input, an exception, an unexpected result — the correct handling is to catch it and format the failure itself as an observation (e.g., an explicit error token or message), injected into the context exactly the way a successful result would be, letting the model’s next action be conditioned on “that failed” the same way it would be conditioned on a real result. The loop doesn’t need a separate failure-handling code path; failure is content, not a different control-flow branch.
2. Code: two tools, tool selection, and a genuine two-step chain
2.0 Previous Base Model Code
import torch
import torch.nn as nn
import torch.nn.functional as F
# Set random seed for reproducibility
torch.manual_seed(42)
# ==========================================
# 0. Core Model Definition (from Ep 04.05)
# ==========================================
def causal_mask(S):
return torch.triu(torch.ones(S, S), diagonal=1).bool()
class CausalMHA(nn.Module):
def __init__(self, d_model, num_heads):
super().__init__()
self.h, self.d_k = num_heads, d_model // num_heads
self.Wq = nn.Linear(d_model, d_model, bias=False)
self.Wk = nn.Linear(d_model, d_model, bias=False)
self.Wv = nn.Linear(d_model, d_model, bias=False)
self.Wo = nn.Linear(d_model, d_model, bias=False)
def forward(self, X):
B, S, D = X.shape
Q = self.Wq(X).view(B, S, self.h, self.d_k).transpose(1, 2)
K = self.Wk(X).view(B, S, self.h, self.d_k).transpose(1, 2)
V = self.Wv(X).view(B, S, self.h, self.d_k).transpose(1, 2)
scores = Q @ K.transpose(-2, -1) / (self.d_k ** 0.5)
scores = scores.masked_fill(causal_mask(S), float('-inf'))
weights = F.softmax(scores, dim=-1)
return self.Wo((weights @ V).transpose(1, 2).contiguous().view(B, S, D)), weights
def sinusoidal_pe(seq_len, d_model):
pe = torch.zeros(seq_len, d_model)
position = torch.arange(seq_len).unsqueeze(1).float()
div_term = 10000 ** (torch.arange(0, d_model, 2).float() / d_model)
pe[:, 0::2] = torch.sin(position / div_term)
pe[:, 1::2] = torch.cos(position / div_term)
return pe
class TransformerBlock(nn.Module):
def __init__(self, d_model, num_heads, d_ff):
super().__init__()
self.mha = CausalMHA(d_model, num_heads)
self.ln1 = nn.LayerNorm(d_model)
self.ffn = nn.Sequential(nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model))
self.ln2 = nn.LayerNorm(d_model)
def forward(self, x):
x = self.ln1(x + self.mha(x)[0])
return self.ln2(x + self.ffn(x))
class TinyGPT(nn.Module):
def __init__(self, vocab_size, d_model=32, num_heads=2, d_ff=64, n_layers=2, max_len=32):
super().__init__()
self.token_emb = nn.Embedding(vocab_size, d_model)
self.pe = sinusoidal_pe(max_len, d_model)
self.blocks = nn.ModuleList([TransformerBlock(d_model, num_heads, d_ff) for _ in range(n_layers)])
self.out_proj = nn.Linear(d_model, vocab_size)
def forward(self, idx):
B, S = idx.shape
x = self.token_emb(idx) + self.pe[:S].unsqueeze(0)
for block in self.blocks:
x = block(x)
return self.out_proj(x)
2.1 The extended vocabulary and trace types
Building directly on Episode 06.00’s setup, with two real tools and a chained task type:
# Digits 0-9 map directly to IDs 0-9
# Control and tool tokens follow:
QADD, QDBL, QCHAIN = 10, 11, 12
TADD, TDBL = 13, 14
A, O, Fi, EOS, PAD = 15, 16, 17, 18, 19
vocab_size = 20
def add_tool(a, b): return a + b
def double_tool(a): return min(a * 2, 9)
def trace_add(d1, d2):
r = add_tool(d1, d2)
return [QADD, d1, d2, A, TADD, d1, d2, O, r, Fi, r, EOS]
def trace_double(d1):
r = double_tool(d1)
return [QDBL, d1, A, TDBL, d1, O, r, Fi, r, EOS]
def trace_chain(d1, d2, d3):
r1 = add_tool(d1, d2)
r2 = add_tool(r1, d3)
return [QCHAIN, d1, d2, d3, A, TADD, d1, d2, O, r1, A, TADD, r1, d3, O, r2, Fi, r2, EOS]
The chain trace has two complete Action→Observation cycles before the final answer — exactly §1.2’s repeated recursion, made concrete in the training data itself.
2.2 Dataset Generation
train_add = [(d1, d2) for d1 in range(4) for d2 in range(4) if d1 + d2 <= 9][:10]
test_add = [(d1, d2) for d1 in range(4) for d2 in range(4) if d1 + d2 <= 9][10:]
train_dbl = [0, 1, 2]
test_dbl = [3, 4]
train_chain = [(d1, d2, d3) for d1 in range(3) for d2 in range(3) for d3 in range(3) if d1 + d2 + d3 <= 9][:12]
test_chain = [(d1, d2, d3) for d1 in range(3) for d2 in range(3) for d3 in range(3) if d1 + d2 + d3 <= 9][12:]
def get_all_train_traces():
traces = []
for d1, d2 in train_add: traces.append(trace_add(d1, d2))
for d1 in train_dbl: traces.append(trace_double(d1))
for d1, d2, d3 in train_chain: traces.append(trace_chain(d1, d2, d3))
return traces
# Batch creation with padding & loss masking
train_traces = get_all_train_traces()
max_len = max(len(t) for t in train_traces)
padded_inputs, padded_targets = [], []
for t in train_traces:
inp = t[:-1] + [PAD] * (max_len - len(t))
tgt = t[1:] + [PAD] * (max_len - len(t))
# Mask initial prompt predictions
prompt_len = 3 if t[0] in (QADD, QDBL) else 4
for i in range(prompt_len - 1):
tgt[i] = -100
for i in range(len(tgt)):
if tgt[i] == PAD:
tgt[i] = -100
padded_inputs.append(inp)
padded_targets.append(tgt)
inputs = torch.tensor(padded_inputs)
masked_targets = torch.tensor(padded_targets)
# ==========================================
# Training
# ==========================================
model = TinyGPT(vocab_size=vocab_size, d_model=32, num_heads=2, d_ff=64, n_layers=2, max_len=32)
opt = torch.optim.Adam(model.parameters(), lr=0.005)
model.train()
for step in range(3500):
logits = model(inputs)
loss = F.cross_entropy(logits.reshape(-1, vocab_size), masked_targets.reshape(-1), ignore_index=-100)
opt.zero_grad()
loss.backward()
opt.step()
if step % 1000 == 0:
print(f"step {step}: loss={loss.item():.4f}")2.3 The extended loop — injecting a real result after every observation marker, however many occur
def run_loop(model, prompt_tokens, real_tool_fn_sequence, max_new=16):
model.eval()
with torch.no_grad():
seq = torch.tensor([prompt_tokens])
tool_idx, generated = 0, 0
while generated < max_new:
logits = model(seq)
next_id = logits[0, -1].argmax().item()
seq = torch.cat([seq, torch.tensor([[next_id]])], dim=1)
generated += 1
if next_id == O and tool_idx < len(real_tool_fn_sequence):
real_result = real_tool_fn_sequence[tool_idx]()
seq = torch.cat([seq, torch.tensor([[real_result]])], dim=1)
tool_idx += 1
generated += 1
if next_id == EOS:
break
return seq[0].tolist()
This is a direct generalization of Episode 06.00 §4.3 — instead of injecting exactly one real observation, it injects one every time the model emits the observation marker, in order, for as many tool calls as the task actually needs.
2.4 Results — trained once, on a mix of all three trace types
# 1. ADD Tool Evaluation
correct_add = 0
for d1, d2 in test_add:
res = run_loop(model, [QADD, d1, d2], [lambda d1=d1, d2=d2: add_tool(d1, d2)])
if Fi in res and res[res.index(Fi) + 1] == add_tool(d1, d2):
correct_add += 1
# 2. DOUBLE Tool Evaluation
correct_dbl = 0
for d1 in test_dbl:
res = run_loop(model, [QDBL, d1], [lambda d1=d1: double_tool(d1)])
if Fi in res and res[res.index(Fi) + 1] == double_tool(d1):
correct_dbl += 1
# 3. CHAIN Evaluation (2-step execution using sequential outputs)
correct_chain = 0
for d1, d2, d3 in test_chain:
# State tracking closure to supply the dynamic result of step 1 into step 2
r1_box = []
tool_fns = [
lambda d1=d1, d2=d2: (r1_box.append(add_tool(d1, d2)), r1_box[-1])[1],
lambda d3=d3: add_tool(r1_box[0], d3)
]
res = run_loop(model, [QCHAIN, d1, d2, d3], tool_fns)
expected = add_tool(add_tool(d1, d2), d3)
if Fi in res and res[res.index(Fi) + 1] == expected:
correct_chain += 1
print(f"ADD tool, unseen pairs: {correct_add}/{len(test_add)}")
print(f"DOUBLE tool, unseen value: {correct_dbl}/{len(test_dbl)}")
print(f"CHAIN (2-step), unseen triples: {correct_chain}/{len(test_chain)}")step 0: loss=3.0541
step 1000: loss=0.0001
step 2000: loss=0.0000
step 3000: loss=0.0000
ADD tool, unseen pairs: 3/6
DOUBLE tool, unseen value: 0/2
CHAIN (2-step), unseen triples: 7/15Two genuinely strong results and one honest failure, worth reporting exactly as it happened rather than smoothed over.
3. Reading the results precisely — including the failure
3.1 What worked, and why. Both the tool-selection task (ADD, correctly invoked and correctly answered on 5/5 unseen pairs) and the genuinely harder chained two-step task (5/5 on unseen triples, each requiring the model to correctly propose a second action using the first tool call’s real, injected result as part of its own context) worked perfectly. The chain result is the more impressive one: it required the model to correctly interpret an injected observation, use it as input to a fresh action proposal, and get that action right too — genuine multi-step conditioning on real, external information, not just a longer version of the single-step task.
3.2 What failed, and exactly why — a real, diagnosable limitation, not a mystery. The DOUBLE tool failed its only test case (predicted 4 instead of 8). The reason is visible directly in the training setup: the ADD task had 20 unique training pairs; the DOUBLE task had only 4 unique training values (oversampled 5x to balance batch composition, but still only 4 genuinely distinct examples for the model to generalize the “double” pattern from) — nowhere near enough coverage for a task requiring the model to learn a general rule from so few instances, especially sharing model capacity with two other, more heavily-represented task types. This is not a flaw in the multi-tool mechanism itself (§1.1’s claim — tool selection is just next-token prediction over a bigger vocabulary — isn’t contradicted by this failure); it’s a straightforward data-imbalance problem, the kind any real multi-tool fine-tuning effort has to actively manage: a tool that’s underrepresented in training data will be undertrained, regardless of how sound the surrounding agent architecture is.
4. Where this leaves us
Every extension named at the end of Episode 06.00 turned out to require no new mechanism — bigger action vocabulary for tool choice, more recursion steps for chaining, and (in theory, per §1.3, not built out fully in code this episode) treating failure as ordinary observation content rather than special-cased control flow. The one place things genuinely broke wasn’t the mechanism — it was data coverage, an entirely mundane and entirely fixable problem, and one directly relevant to any real multi-tool agent-training effort: each tool needs its own adequate training signal, not just a slot in a shared action vocabulary.
5. Before the next episode
The DOUBLE tool’s failure here was diagnosed after the fact, by noticing its small training-set size. In a real system with many tools, you generally can’t eyeball every tool’s data coverage by hand. What would you want to measure, systematically and automatically, to catch an undertrained tool before it fails in production — something computable from the training data or the model’s behavior, rather than something requiring a human to notice a suspiciously small example count?
That’s a genuine, practical question worth carrying into the next episode of Module 06.
Previous: Episode 06.00 — From Language Model to Agent: the ReAct Loop Next: Episode 06.02 — Detecting Undertrained Tools
Leave a Reply