54 lines
2.4 KiB
TypeScript
54 lines
2.4 KiB
TypeScript
import type { AiProvider, ChatMessage, ToolSpec } from "@stc/contracts";
|
|
|
|
/** SSE body from a list of [event, data] pairs, split into awkward chunks to exercise buffering. */
|
|
export function sseBody(events: [string | null, unknown][], chunk = 7): ReadableStream<Uint8Array> {
|
|
const text = events.map(([e, d]) => `${e ? `event: ${e}\n` : ""}data: ${typeof d === "string" ? d : JSON.stringify(d)}\n\n`).join("");
|
|
const bytes = new TextEncoder().encode(text);
|
|
let i = 0;
|
|
return new ReadableStream({
|
|
pull(c) {
|
|
if (i >= bytes.length) return c.close();
|
|
c.enqueue(bytes.slice(i, i + chunk));
|
|
i += chunk;
|
|
},
|
|
});
|
|
}
|
|
|
|
export type Captured = { url: string; init: RequestInit; body: any };
|
|
/** A fake fetch that records requests and answers from a queue of responses. */
|
|
export function fakeFetch(responses: (() => Response)[]) {
|
|
const calls: Captured[] = [];
|
|
const f = (async (url: string | URL | Request, init?: RequestInit) => {
|
|
calls.push({ url: String(url), init: init ?? {}, body: init?.body ? JSON.parse(String(init.body)) : null });
|
|
const next = responses.shift();
|
|
if (!next) throw new Error("fakeFetch: no response queued");
|
|
return next();
|
|
}) as unknown as typeof fetch;
|
|
return { fetch: f, calls };
|
|
}
|
|
|
|
export async function collect(it: AsyncIterable<{ text?: string; toolCall?: { id: string; name: string; input: unknown } }>) {
|
|
let text = "";
|
|
const toolCalls: { id: string; name: string; input: unknown }[] = [];
|
|
for await (const ev of it) {
|
|
if (ev.text) text += ev.text;
|
|
if (ev.toolCall) toolCalls.push(ev.toolCall);
|
|
}
|
|
return { text, toolCalls };
|
|
}
|
|
|
|
/** Scripted fake provider: each call to chat() plays the next turn and records the messages it got. */
|
|
export function fakeProvider(turns: ((messages: ChatMessage[], tools?: ToolSpec[]) => { text?: string; toolCalls?: { id: string; name: string; input: unknown }[] })[]) {
|
|
const seen: { messages: ChatMessage[]; tools?: ToolSpec[] }[] = [];
|
|
const provider: AiProvider = {
|
|
id: "fake",
|
|
async *chat(messages, opts) {
|
|
seen.push({ messages: structuredClone(messages), tools: opts.tools });
|
|
const turn = turns[Math.min(seen.length - 1, turns.length - 1)]!(messages, opts.tools);
|
|
if (turn.text) for (const part of turn.text.match(/.{1,8}/gs) ?? []) yield { text: part };
|
|
for (const tc of turn.toolCalls ?? []) yield { toolCall: tc };
|
|
},
|
|
};
|
|
return { provider, seen };
|
|
}
|