83 lines
2.0 KiB
TypeScript
83 lines
2.0 KiB
TypeScript
import { streamObject } from "ai";
|
|
import { z } from "zod";
|
|
import { codePrompt, updateDocumentPrompt } from "@/lib/ai/prompts";
|
|
import { myProvider } from "@/lib/ai/providers";
|
|
import { getSystemModel } from "@/lib/ai/dynamic-provider";
|
|
import { createDocumentHandler } from "@/lib/artifacts/server";
|
|
|
|
export const codeDocumentHandler = createDocumentHandler<"code">({
|
|
kind: "code",
|
|
onCreateDocument: async ({ title, dataStream }) => {
|
|
let draftContent = "";
|
|
|
|
// 优先使用数据库配置的模型
|
|
const model = (await getSystemModel("artifact")) ?? myProvider.languageModel("artifact-model");
|
|
|
|
const { fullStream } = streamObject({
|
|
model,
|
|
system: codePrompt,
|
|
prompt: title,
|
|
schema: z.object({
|
|
code: z.string(),
|
|
}),
|
|
});
|
|
|
|
for await (const delta of fullStream) {
|
|
const { type } = delta;
|
|
|
|
if (type === "object") {
|
|
const { object } = delta;
|
|
const { code } = object;
|
|
|
|
if (code) {
|
|
dataStream.write({
|
|
type: "data-codeDelta",
|
|
data: code ?? "",
|
|
transient: true,
|
|
});
|
|
|
|
draftContent = code;
|
|
}
|
|
}
|
|
}
|
|
|
|
return draftContent;
|
|
},
|
|
onUpdateDocument: async ({ document, description, dataStream }) => {
|
|
let draftContent = "";
|
|
|
|
// 优先使用数据库配置的模型
|
|
const model = (await getSystemModel("artifact")) ?? myProvider.languageModel("artifact-model");
|
|
|
|
const { fullStream } = streamObject({
|
|
model,
|
|
system: updateDocumentPrompt(document.content, "code"),
|
|
prompt: description,
|
|
schema: z.object({
|
|
code: z.string(),
|
|
}),
|
|
});
|
|
|
|
for await (const delta of fullStream) {
|
|
const { type } = delta;
|
|
|
|
if (type === "object") {
|
|
const { object } = delta;
|
|
const { code } = object;
|
|
|
|
if (code) {
|
|
dataStream.write({
|
|
type: "data-codeDelta",
|
|
data: code ?? "",
|
|
transient: true,
|
|
});
|
|
|
|
draftContent = code;
|
|
}
|
|
}
|
|
}
|
|
|
|
return draftContent;
|
|
},
|
|
});
|