Skip to content
112 linesCodeBlameRaw
1/**
2 * One model turn with tools: ask, run the tools it calls through the
3 * reply's `ToolBox`, give it the results, until it answers in text. Pure
4 * apart from `send`, so the loop's rails are tested on their own.
5 *
6 * Rails: the tool box allows at most `MAX_TOOL_CALLS` for the whole reply
7 * (consults included); once they are spent, or the turn has read
8 * `INPUT_BUDGET` input tokens, the model is asked to answer without tools.
9 * At most `MAX_ROUNDS` requests, whatever happens.
10 */
11import type { TokenPrice } from "../../runner/src/model-env.ts";
12import { type Tokens, costMicros } from "./budget.ts";
13import type { ToolBox } from "./tools.ts";
14
15/** The input tokens one turn may read before it must answer. */
16export const INPUT_BUDGET = 150_000;
17const MAX_ROUNDS = 12;
18/** A chat answer is short; this bounds the cost of one that is not. */
19export const MAX_OUTPUT_TOKENS = 2048;
20
21type Block = { type: string; text?: string; id?: string; name?: string; input?: unknown; [key: string]: unknown };
22export type ModelMessage = { role: "user" | "assistant"; content: string | Block[] };
23
24export type ModelAnswer = {
25 content?: Block[];
26 stop_reason?: string;
27 usage?: { input_tokens?: number; output_tokens?: number; cache_read_input_tokens?: number; cache_creation_input_tokens?: number };
28};
29
30/** Sends one Messages API request; throws when the model did not answer. */
31export type Send = (body: Record<string, unknown>) => Promise<ModelAnswer>;
32
33export type TurnResult = { text: string; tokens: Tokens; cost: number; rounds: number; stopped?: boolean };
34
35export function addTokens(a: Tokens, b: Tokens): Tokens {
36 return { input: a.input + b.input, output: a.output + b.output, cacheRead: a.cacheRead + b.cacheRead, cacheWrite: a.cacheWrite + b.cacheWrite };
37}
38
39export const NO_TOKENS: Tokens = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 };
40
41/** A session step's rails: more rounds, more input, longer answers than a chat reply. */
42export const SESSION_LIMITS = { rounds: 30, input: 400_000, output: 8192 } as const;
43
44export type TurnOptions = {
45 /** At most this many requests (default `MAX_ROUNDS`). */
46 maxRounds?: number;
47 /** Input tokens read before it must answer (default `INPUT_BUDGET`). */
48 inputBudget?: number;
49 /** The longest answer (default `MAX_OUTPUT_TOKENS`). */
50 maxOutput?: number;
51 /** Told what the model said in each round that also called tools: a session's transcript. */
52 onText?: (text: string) => void;
53 /** Asked before each round; true ends the turn with what it has (a session was stopped). */
54 stopped?: () => Promise<boolean>;
55 /**
56 * How hard the model reasons (`output_config.effort`), from the agent's
57 * effort setting; null or absent leaves it to the model.
58 */
59 effort?: string | null;
60};
61
62export async function runTurn(
63 send: Send,
64 input: { model: string; system: string; messages: ModelMessage[]; tools: ToolBox | null; price: TokenPrice | null } & TurnOptions,
65): Promise<TurnResult> {
66 const messages = [...input.messages];
67 let tokens = NO_TOKENS;
68 const maxRounds = input.maxRounds ?? MAX_ROUNDS;
69 const inputBudget = input.inputBudget ?? INPUT_BUDGET;
70 for (let round = 1; ; round++) {
71 if (round > 1 && input.stopped && (await input.stopped())) {
72 return { text: "", tokens, cost: costMicros(tokens, input.price), rounds: round - 1, stopped: true };
73 }
74 const definitions = input.tools?.definitions() ?? [];
75 const canUse = !!input.tools && definitions.length > 0 && !input.tools.spent && tokens.input + tokens.cacheRead < inputBudget && round < maxRounds;
76 const answer = await send({
77 model: input.model,
78 system: input.system,
79 messages,
80 max_tokens: input.maxOutput ?? MAX_OUTPUT_TOKENS,
81 ...(input.effort ? { output_config: { effort: input.effort } } : {}),
82 // Tools stay listed once the conversation has used them, so their
83 // results still read; past the rails, the model must answer in text.
84 ...(definitions.length ? { tools: definitions, tool_choice: { type: canUse ? "auto" : "none" } } : {}),
85 });
86 const usage = answer.usage ?? {};
87 tokens = addTokens(tokens, {
88 input: usage.input_tokens ?? 0,
89 output: usage.output_tokens ?? 0,
90 cacheRead: usage.cache_read_input_tokens ?? 0,
91 cacheWrite: usage.cache_creation_input_tokens ?? 0,
92 });
93 const content = answer.content ?? [];
94 const calls = content.filter((block) => block.type === "tool_use");
95 const text = content
96 .filter((block) => block.type === "text" && block.text)
97 .map((block) => block.text!.trim())
98 .join("\n\n")
99 .trim();
100 if (!calls.length || !input.tools || round >= maxRounds) {
101 return { text, tokens, cost: costMicros(tokens, input.price), rounds: round };
102 }
103 if (text) input.onText?.(text);
104 messages.push({ role: "assistant", content });
105 const results: Block[] = [];
106 for (const call of calls) {
107 const ran = await input.tools.run(String(call.name ?? ""), (call.input as Record<string, unknown>) ?? {});
108 results.push({ type: "tool_result", tool_use_id: call.id, content: ran.text });
109 }
110 messages.push({ role: "user", content: results });
111 }
112}