58 lines
2.2 KiB
TypeScript
58 lines
2.2 KiB
TypeScript
import type { UIMessageChunk } from "ai";
|
|
|
|
/** Canonical parts are projected independently of an attached browser response. */
|
|
export function createAnswerStreamProjection(options: {
|
|
messageId: string;
|
|
sessionId: string;
|
|
write: (chunk: UIMessageChunk) => void;
|
|
onChange: () => void;
|
|
}) {
|
|
const parts: Record<string, unknown>[] = [];
|
|
let segment = 0;
|
|
let textId = "answer";
|
|
let currentText: Record<string, unknown> | undefined;
|
|
let finished = false;
|
|
const endText = () => {
|
|
if (!currentText) return;
|
|
currentText.state = "done";
|
|
options.write({ type: "text-end", id: textId });
|
|
currentText = undefined;
|
|
textId = `answer-${++segment}`;
|
|
};
|
|
const appendText = (delta: string) => {
|
|
if (!delta || finished) return;
|
|
if (!currentText) {
|
|
currentText = { type: "text", text: "", state: "streaming" };
|
|
parts.push(currentText);
|
|
options.write({ type: "text-start", id: textId });
|
|
}
|
|
currentText.text = `${String(currentText.text)}${delta}`;
|
|
options.write({ type: "text-delta", id: textId, delta });
|
|
options.write({ type: "data-stream_token", transient: true, data: {
|
|
session_id: options.sessionId, message_id: options.messageId, content: delta,
|
|
} } as UIMessageChunk);
|
|
options.onChange();
|
|
};
|
|
const appendData = (event: string, data: Record<string, unknown>) => {
|
|
if (finished) return;
|
|
if (event === "answer_ui") endText();
|
|
const type = `data-${event}`;
|
|
const id = [data.id, data.request_id, data.envelope_id, data.block_id].find((value): value is string => typeof value === "string");
|
|
const existing = id === undefined ? -1 : parts.findIndex((part) => part.type === type && part.id === id);
|
|
const part = id ? { type, id, data } : { type, data };
|
|
if (existing >= 0) parts[existing] = part;
|
|
else parts.push(part);
|
|
options.write(part as UIMessageChunk);
|
|
options.onChange();
|
|
};
|
|
const finish = (finishReason: "stop" | "error" = "stop") => {
|
|
if (finished) return;
|
|
if (!parts.length) appendText("Agent 已完成处理。");
|
|
endText();
|
|
finished = true;
|
|
options.onChange();
|
|
options.write({ type: "finish", finishReason });
|
|
};
|
|
return { parts, appendText, appendData, finish, get finished() { return finished; } };
|
|
}
|