plano/demos/use_cases/vercel-ai-sdk/artifacts/sheet/server.ts

81 lines
1.9 KiB
TypeScript

import { streamObject } from "ai";
import { z } from "zod";
import { sheetPrompt, updateDocumentPrompt } from "@/lib/ai/prompts";
import { getArtifactModel } from "@/lib/ai/providers";
import { createDocumentHandler } from "@/lib/artifacts/server";
export const sheetDocumentHandler = createDocumentHandler<"sheet">({
kind: "sheet",
onCreateDocument: async ({ title, dataStream }) => {
let draftContent = "";
const { fullStream } = streamObject({
model: getArtifactModel(),
system: sheetPrompt,
prompt: title,
schema: z.object({
csv: z.string().describe("CSV data"),
}),
});
for await (const delta of fullStream) {
const { type } = delta;
if (type === "object") {
const { object } = delta;
const { csv } = object;
if (csv) {
dataStream.write({
type: "data-sheetDelta",
data: csv,
transient: true,
});
draftContent = csv;
}
}
}
dataStream.write({
type: "data-sheetDelta",
data: draftContent,
transient: true,
});
return draftContent;
},
onUpdateDocument: async ({ document, description, dataStream }) => {
let draftContent = "";
const { fullStream } = streamObject({
model: getArtifactModel(),
system: updateDocumentPrompt(document.content, "sheet"),
prompt: description,
schema: z.object({
csv: z.string(),
}),
});
for await (const delta of fullStream) {
const { type } = delta;
if (type === "object") {
const { object } = delta;
const { csv } = object;
if (csv) {
dataStream.write({
type: "data-sheetDelta",
data: csv,
transient: true,
});
draftContent = csv;
}
}
}
return draftContent;
},
});