72 lines
2.1 KiB
TypeScript
72 lines
2.1 KiB
TypeScript
import type { AssistantMessage } from "@mariozechner/pi-ai";
|
|
import { describe, expect, it, vi } from "vitest";
|
|
import {
|
|
THINKING_TAG_CASES,
|
|
createStubSessionHarness,
|
|
} from "./pi-embedded-subscribe.e2e-harness.js";
|
|
import { subscribeEmbeddedPiSession } from "./pi-embedded-subscribe.js";
|
|
|
|
describe("subscribeEmbeddedPiSession", () => {
|
|
function createReasoningBlockReplyHarness() {
|
|
const { session, emit } = createStubSessionHarness();
|
|
const onBlockReply = vi.fn();
|
|
|
|
subscribeEmbeddedPiSession({
|
|
session,
|
|
runId: "run",
|
|
onBlockReply,
|
|
blockReplyBreak: "message_end",
|
|
reasoningMode: "on",
|
|
});
|
|
|
|
return { emit, onBlockReply };
|
|
}
|
|
|
|
function expectReasoningAndAnswerCalls(onBlockReply: ReturnType<typeof vi.fn>) {
|
|
expect(onBlockReply).toHaveBeenCalledTimes(2);
|
|
expect(onBlockReply.mock.calls[0][0].text).toBe("Reasoning:\n_Because it helps_");
|
|
expect(onBlockReply.mock.calls[1][0].text).toBe("Final answer");
|
|
}
|
|
|
|
it("emits reasoning as a separate message when enabled", () => {
|
|
const { emit, onBlockReply } = createReasoningBlockReplyHarness();
|
|
|
|
const assistantMessage = {
|
|
role: "assistant",
|
|
content: [
|
|
{ type: "thinking", thinking: "Because it helps" },
|
|
{ type: "text", text: "Final answer" },
|
|
],
|
|
} as AssistantMessage;
|
|
|
|
emit({ type: "message_end", message: assistantMessage });
|
|
|
|
expectReasoningAndAnswerCalls(onBlockReply);
|
|
});
|
|
it.each(THINKING_TAG_CASES)(
|
|
"promotes <%s> tags to thinking blocks at write-time",
|
|
({ open, close }) => {
|
|
const { emit, onBlockReply } = createReasoningBlockReplyHarness();
|
|
|
|
const assistantMessage = {
|
|
role: "assistant",
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `${open}\nBecause it helps\n${close}\n\nFinal answer`,
|
|
},
|
|
],
|
|
} as AssistantMessage;
|
|
|
|
emit({ type: "message_end", message: assistantMessage });
|
|
|
|
expectReasoningAndAnswerCalls(onBlockReply);
|
|
|
|
expect(assistantMessage.content).toEqual([
|
|
{ type: "thinking", thinking: "Because it helps" },
|
|
{ type: "text", text: "Final answer" },
|
|
]);
|
|
},
|
|
);
|
|
});
|