File size: 1,948 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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
import { generateUUID } from '@/lib/utils';
import { tool, type UIMessageStreamWriter } from 'ai';
import { z } from 'zod';
import type { Session } from 'next-auth';
import {
artifactKinds,
documentHandlersByArtifactKind,
} from '@/lib/artifacts/server';
import type { ChatMessage } from '@/lib/types';
interface CreateDocumentProps {
session: Session;
dataStream: UIMessageStreamWriter<ChatMessage>;
}
export const createDocument = ({ session, dataStream }: CreateDocumentProps) =>
tool({
description:
'Create a document for a writing or content creation activities. This tool will call other functions that will generate the contents of the document based on the title and kind.',
inputSchema: z.object({
title: z.string(),
kind: z.enum(artifactKinds),
}),
execute: async ({ title, kind }) => {
const id = generateUUID();
dataStream.write({
type: 'data-kind',
data: kind,
transient: true,
});
dataStream.write({
type: 'data-id',
data: id,
transient: true,
});
dataStream.write({
type: 'data-title',
data: title,
transient: true,
});
dataStream.write({
type: 'data-clear',
data: null,
transient: true,
});
const documentHandler = documentHandlersByArtifactKind.find(
(documentHandlerByArtifactKind) =>
documentHandlerByArtifactKind.kind === kind,
);
if (!documentHandler) {
throw new Error(`No document handler found for kind: ${kind}`);
}
await documentHandler.onCreateDocument({
id,
title,
dataStream,
session,
});
dataStream.write({ type: 'data-finish', data: null, transient: true });
return {
id,
title,
kind,
content: 'A document was created and is now visible to the user.',
};
},
});
|