File size: 1,767 Bytes
867b17d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
import { tool, type UIMessageStreamWriter } from 'ai';
import type { Session } from 'next-auth';
import { z } from 'zod';
import { getDocumentById } from '@/lib/db/queries';
import { documentHandlersByArtifactKind } from '@/lib/artifacts/server';
import type { ChatMessage } from '@/lib/types';
interface UpdateDocumentProps {
session: Session;
dataStream: UIMessageStreamWriter<ChatMessage>;
}
export const updateDocument = ({ session, dataStream }: UpdateDocumentProps) =>
tool({
description: 'Update a document with the given description.',
inputSchema: z.object({
id: z.string().describe('The ID of the document to update'),
description: z
.string()
.describe('The description of changes that need to be made'),
}),
execute: async ({ id, description }) => {
const document = await getDocumentById({ id });
if (!document) {
return {
error: 'Document not found',
};
}
dataStream.write({
type: 'data-clear',
data: null,
transient: true,
});
const documentHandler = documentHandlersByArtifactKind.find(
(documentHandlerByArtifactKind) =>
documentHandlerByArtifactKind.kind === document.kind,
);
if (!documentHandler) {
throw new Error(`No document handler found for kind: ${document.kind}`);
}
await documentHandler.onUpdateDocument({
document,
description,
dataStream,
session,
});
dataStream.write({ type: 'data-finish', data: null, transient: true });
return {
id,
title: document.title,
kind: document.kind,
content: 'The document has been updated successfully.',
};
},
});
|