95 lines
2.6 KiB
TypeScript
95 lines
2.6 KiB
TypeScript
import {
|
|
createAssistantMessageEventStream,
|
|
type AssistantMessage,
|
|
type AssistantMessageEvent,
|
|
type Context,
|
|
type Model,
|
|
} from "@mariozechner/pi-ai";
|
|
import { describe, expect, it } from "vitest";
|
|
import { createSafePiStream } from "../../src/main/worker/runtime/pi-safe-stream";
|
|
|
|
function createModel(): Model<"openai-responses"> {
|
|
return {
|
|
id: "mock",
|
|
name: "mock",
|
|
api: "openai-responses",
|
|
provider: "openai",
|
|
baseUrl: "https://example.invalid",
|
|
reasoning: false,
|
|
input: ["text"],
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
|
contextWindow: 8192,
|
|
maxTokens: 2048,
|
|
};
|
|
}
|
|
|
|
function createAssistantMessage(): AssistantMessage {
|
|
return {
|
|
role: "assistant",
|
|
content: [{ type: "text", text: "done" }],
|
|
api: "openai-responses",
|
|
provider: "openai",
|
|
model: "mock",
|
|
usage: {
|
|
input: 1,
|
|
output: 1,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
totalTokens: 2,
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
|
},
|
|
stopReason: "stop",
|
|
timestamp: Date.now(),
|
|
};
|
|
}
|
|
|
|
const emptyContext: Context = { messages: [] };
|
|
|
|
describe("pi safe stream", () => {
|
|
it("forwards successful stream events", async () => {
|
|
const safeStream = createSafePiStream(() => {
|
|
const stream = createAssistantMessageEventStream();
|
|
queueMicrotask(() => {
|
|
const message = createAssistantMessage();
|
|
stream.push({ type: "done", reason: "stop", message });
|
|
stream.end(message);
|
|
});
|
|
return stream;
|
|
});
|
|
|
|
const stream = safeStream(createModel(), emptyContext);
|
|
const events: AssistantMessageEvent[] = [];
|
|
for await (const event of stream) {
|
|
events.push(event);
|
|
}
|
|
|
|
expect(events.map((event) => event.type)).toEqual(["done"]);
|
|
await expect(stream.result()).resolves.toMatchObject({ stopReason: "stop" });
|
|
});
|
|
|
|
it("turns rejected stream setup into an error event", async () => {
|
|
const safeStream = createSafePiStream(async () => {
|
|
throw new Error("stream setup failed");
|
|
});
|
|
|
|
const stream = safeStream(createModel(), emptyContext);
|
|
const events: AssistantMessageEvent[] = [];
|
|
for await (const event of stream) {
|
|
events.push(event);
|
|
}
|
|
const result = await stream.result();
|
|
|
|
expect(events).toHaveLength(1);
|
|
expect(events[0]).toMatchObject({
|
|
type: "error",
|
|
reason: "error",
|
|
error: {
|
|
stopReason: "error",
|
|
errorMessage: "stream setup failed",
|
|
},
|
|
});
|
|
expect(result.stopReason).toBe("error");
|
|
expect(result.errorMessage).toBe("stream setup failed");
|
|
});
|
|
});
|