53 lines
2.0 KiB
TypeScript
53 lines
2.0 KiB
TypeScript
import type { AiProvider, ChatMessage } from "@stc/contracts";
|
|
import type { ToolSet } from "./tools";
|
|
|
|
export type AssistantEvent =
|
|
| { type: "text"; text: string }
|
|
| { type: "tool"; name: string }
|
|
| { type: "error"; message: string }
|
|
| { type: "done"; text: string };
|
|
|
|
/**
|
|
* Runs the model with tools until it answers without calling a tool, or until maxSteps model calls.
|
|
* Mutates `messages` by appending the tool round trips, so callers can inspect them.
|
|
*/
|
|
export async function* runAssistant(o: {
|
|
provider: AiProvider;
|
|
messages: ChatMessage[];
|
|
tools: ToolSet;
|
|
signal?: AbortSignal;
|
|
maxSteps?: number;
|
|
}): AsyncIterable<AssistantEvent> {
|
|
const maxSteps = o.maxSteps ?? 6;
|
|
let answer = "";
|
|
for (let step = 0; step < maxSteps; step++) {
|
|
let text = "";
|
|
const calls: { id: string; name: string; input: unknown }[] = [];
|
|
try {
|
|
for await (const ev of o.provider.chat(o.messages, { tools: o.tools.specs, signal: o.signal })) {
|
|
if (ev.text) { text += ev.text; yield { type: "text", text: ev.text }; }
|
|
if (ev.toolCall) calls.push(ev.toolCall);
|
|
}
|
|
} catch (e) {
|
|
yield { type: "error", message: (e as Error).message };
|
|
return;
|
|
}
|
|
answer += text;
|
|
if (!calls.length) {
|
|
yield { type: "done", text: answer };
|
|
return;
|
|
}
|
|
if (text) o.messages.push({ role: "assistant", content: text });
|
|
for (const c of calls) o.messages.push({ role: "assistant", content: JSON.stringify(c.input ?? {}), toolCallId: c.id, name: c.name });
|
|
for (const c of calls) {
|
|
yield { type: "tool", name: c.name };
|
|
let result: string;
|
|
try { result = await o.tools.run(c.name, c.input); } catch (e) { result = `Tool error: ${(e as Error).message}`; }
|
|
o.messages.push({ role: "tool", content: result, toolCallId: c.id, name: c.name });
|
|
}
|
|
}
|
|
const note = "\n\n(I stopped after several lookups. Ask again to continue.)";
|
|
yield { type: "text", text: note };
|
|
yield { type: "done", text: answer + note };
|
|
}
|