bolt.diy/app/utils/chatStreamController.ts

45 lines
1.6 KiB
TypeScript

// Define the ChatStreamController class for writing messages to a readable
// stream which will be decoded by the react/ai Chat API. There does not seem
// to be functionality exported from the associated packages to do this so
// for now we do it manually after reverse engineering the protocol.
import type { ChatAnthropicInfo } from "~/lib/.server/llm/chat-anthropic";
export interface ChatFileChange {
filePath: string;
contents: string;
}
export class ChatStreamController {
private controller: ReadableStreamDefaultController;
private encoder: TextEncoder;
constructor(controller: ReadableStreamDefaultController) {
this.controller = controller;
this.encoder = new TextEncoder();
}
writeText(text: string) {
const data = this.encoder.encode(`0:${JSON.stringify(text)}\n`);
this.controller.enqueue(data);
}
writeFileChanges(title: string, fileChanges: ChatFileChange[]) {
let text = `<boltArtifact title="${title}">`;
for (const fileChange of fileChanges) {
text += `<boltAction type="file" filePath="${fileChange.filePath}">${fileChange.contents}</boltAction>`;
}
text += "</boltArtifact>";
this.writeText(text);
}
writeAnnotation(type: string, value: any) {
const data = this.encoder.encode(`8:[{"type":"${type}","value":${JSON.stringify(value)}}]\n`);
this.controller.enqueue(data);
}
writeUsage({ chatInfo, completionTokens, promptTokens }: { chatInfo: ChatAnthropicInfo, completionTokens: number, promptTokens: number }) {
this.writeAnnotation("usage", { chatInfo, completionTokens, promptTokens, totalTokens: completionTokens + promptTokens });
}
}