Skip to content
Sign up

flagon-io/g1t

public
0
106 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
57export async function runTurn(
58 send: Send,
59 input: { model: string; system: string; messages: ModelMessage[]; tools: ToolBox | null; price: TokenPrice | null } & TurnOptions,
60): Promise<TurnResult> {
61 const messages = [...input.messages];
62 let tokens = NO_TOKENS;
63 const maxRounds = input.maxRounds ?? MAX_ROUNDS;
64 const inputBudget = input.inputBudget ?? INPUT_BUDGET;
65 for (let round = 1; ; round++) {
66 if (round > 1 && input.stopped && (await input.stopped())) {
67 return { text: "", tokens, cost: costMicros(tokens, input.price), rounds: round - 1, stopped: true };
68 }
69 const definitions = input.tools?.definitions() ?? [];
70 const canUse = !!input.tools && definitions.length > 0 && !input.tools.spent && tokens.input + tokens.cacheRead < inputBudget && round < maxRounds;
71 const answer = await send({
72 model: input.model,
73 system: input.system,
74 messages,
75 max_tokens: input.maxOutput ?? MAX_OUTPUT_TOKENS,
76 // Tools stay listed once the conversation has used them, so their
77 // results still read; past the rails, the model must answer in text.
78 ...(definitions.length ? { tools: definitions, tool_choice: { type: canUse ? "auto" : "none" } } : {}),
79 });
80 const usage = answer.usage ?? {};
81 tokens = addTokens(tokens, {
82 input: usage.input_tokens ?? 0,
83 output: usage.output_tokens ?? 0,
84 cacheRead: usage.cache_read_input_tokens ?? 0,
85 cacheWrite: usage.cache_creation_input_tokens ?? 0,
86 });
87 const content = answer.content ?? [];
88 const calls = content.filter((block) => block.type === "tool_use");
89 const text = content
90 .filter((block) => block.type === "text" && block.text)
91 .map((block) => block.text!.trim())
92 .join("\n\n")
93 .trim();
94 if (!calls.length || !input.tools || round >= maxRounds) {
95 return { text, tokens, cost: costMicros(tokens, input.price), rounds: round };
96 }
97 if (text) input.onText?.(text);
98 messages.push({ role: "assistant", content });
99 const results: Block[] = [];
100 for (const call of calls) {
101 const ran = await input.tools.run(String(call.name ?? ""), (call.input as Record<string, unknown>) ?? {});
102 results.push({ type: "tool_result", tool_use_id: call.id, content: ran.text });
103 }
104 messages.push({ role: "user", content: results });
105 }
106}