diff --git a/app/(chat)/api/chat/route.ts b/app/(chat)/api/chat/route.ts index ec404ec8b..2ca0a0ebc 100644 --- a/app/(chat)/api/chat/route.ts +++ b/app/(chat)/api/chat/route.ts @@ -120,7 +120,7 @@ export async function POST(request: Request) { sendReasoning: true, }); }, - onError: (error) => { + onError: () => { return 'Oops, an error occured!'; }, }); diff --git a/blocks/code.tsx b/blocks/code/client.tsx similarity index 100% rename from blocks/code.tsx rename to blocks/code/client.tsx diff --git a/blocks/code/server.ts b/blocks/code/server.ts new file mode 100644 index 000000000..b5e163b83 --- /dev/null +++ b/blocks/code/server.ts @@ -0,0 +1,73 @@ +import { z } from 'zod'; +import { streamObject } from 'ai'; +import { myProvider } from '@/lib/ai/models'; +import { codePrompt, updateDocumentPrompt } from '@/lib/ai/prompts'; +import { createDocumentHandler } from '@/lib/blocks/server'; + +export const codeDocumentHandler = createDocumentHandler<'code'>({ + kind: 'code', + onCreateDocument: async ({ title, dataStream }) => { + let draftContent = ''; + + const { fullStream } = streamObject({ + model: myProvider.languageModel('block-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.writeData({ + type: 'code-delta', + content: code ?? '', + }); + + draftContent = code; + } + } + } + + return draftContent; + }, + onUpdateDocument: async ({ document, description, dataStream }) => { + let draftContent = ''; + + const { fullStream } = streamObject({ + model: myProvider.languageModel('block-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.writeData({ + type: 'code-delta', + content: code ?? '', + }); + + draftContent = code; + } + } + } + + return draftContent; + }, +}); diff --git a/blocks/image.tsx b/blocks/image/client.tsx similarity index 100% rename from blocks/image.tsx rename to blocks/image/client.tsx diff --git a/blocks/image/server.ts b/blocks/image/server.ts new file mode 100644 index 000000000..6c6adce6c --- /dev/null +++ b/blocks/image/server.ts @@ -0,0 +1,54 @@ +import { myProvider } from '@/lib/ai/models'; +import { createDocumentHandler } from '@/lib/blocks/server'; +import { saveDocument } from '@/lib/db/queries'; +import { experimental_generateImage } from 'ai'; + +export const imageDocumentHandler = createDocumentHandler<'image'>({ + kind: 'image', + onCreateDocument: async ({ id, title, dataStream, session }) => { + let draftContent = ''; + + const { image } = await experimental_generateImage({ + model: myProvider.imageModel('small-model'), + prompt: title, + n: 1, + }); + + draftContent = image.base64; + + dataStream.writeData({ + type: 'image-delta', + content: image.base64, + }); + + if (session?.user?.id) { + await saveDocument({ + id, + title, + kind: 'image', + content: draftContent, + userId: session.user.id, + }); + } + + return draftContent; + }, + onUpdateDocument: async ({ description, dataStream }) => { + let draftContent = ''; + + const { image } = await experimental_generateImage({ + model: myProvider.imageModel('image-model'), + prompt: description, + n: 1, + }); + + draftContent = image.base64; + + dataStream.writeData({ + type: 'image-delta', + content: image.base64, + }); + + return draftContent; + }, +}); diff --git a/blocks/sheet.tsx b/blocks/sheet/client.tsx similarity index 100% rename from blocks/sheet.tsx rename to blocks/sheet/client.tsx diff --git a/blocks/sheet/server.ts b/blocks/sheet/server.ts new file mode 100644 index 000000000..d2de86e81 --- /dev/null +++ b/blocks/sheet/server.ts @@ -0,0 +1,78 @@ +import { myProvider } from '@/lib/ai/models'; +import { sheetPrompt, updateDocumentPrompt } from '@/lib/ai/prompts'; +import { createDocumentHandler } from '@/lib/blocks/server'; +import { streamObject } from 'ai'; +import { z } from 'zod'; + +export const sheetDocumentHandler = createDocumentHandler<'sheet'>({ + kind: 'sheet', + onCreateDocument: async ({ title, dataStream }) => { + let draftContent = ''; + + const { fullStream } = streamObject({ + model: myProvider.languageModel('block-model'), + 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.writeData({ + type: 'sheet-delta', + content: csv, + }); + + draftContent = csv; + } + } + } + + dataStream.writeData({ + type: 'sheet-delta', + content: draftContent, + }); + + return draftContent; + }, + onUpdateDocument: async ({ document, description, dataStream }) => { + let draftContent = ''; + + const { fullStream } = streamObject({ + model: myProvider.languageModel('block-model'), + 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.writeData({ + type: 'sheet-delta', + content: csv, + }); + + draftContent = csv; + } + } + } + + return draftContent; + }, +}); diff --git a/blocks/text.tsx b/blocks/text/client.tsx similarity index 99% rename from blocks/text.tsx rename to blocks/text/client.tsx index 7678d7a2b..360bb5160 100644 --- a/blocks/text.tsx +++ b/blocks/text/client.tsx @@ -12,7 +12,7 @@ import { } from '@/components/icons'; import { Suggestion } from '@/lib/db/schema'; import { toast } from 'sonner'; -import { getSuggestions } from './actions'; +import { getSuggestions } from '../actions'; interface TextBlockMetadata { suggestions: Array; diff --git a/blocks/text/server.ts b/blocks/text/server.ts new file mode 100644 index 000000000..5618598db --- /dev/null +++ b/blocks/text/server.ts @@ -0,0 +1,70 @@ +import { smoothStream, streamText } from 'ai'; +import { myProvider } from '@/lib/ai/models'; +import { createDocumentHandler } from '@/lib/blocks/server'; +import { updateDocumentPrompt } from '@/lib/ai/prompts'; + +export const textDocumentHandler = createDocumentHandler<'text'>({ + kind: 'text', + onCreateDocument: async ({ title, dataStream }) => { + let draftContent = ''; + + const { fullStream } = streamText({ + model: myProvider.languageModel('block-model'), + system: + 'Write about the given topic. Markdown is supported. Use headings wherever appropriate.', + experimental_transform: smoothStream({ chunking: 'word' }), + prompt: title, + }); + + for await (const delta of fullStream) { + const { type } = delta; + + if (type === 'text-delta') { + const { textDelta } = delta; + + draftContent += textDelta; + + dataStream.writeData({ + type: 'text-delta', + content: textDelta, + }); + } + } + + return draftContent; + }, + onUpdateDocument: async ({ document, description, dataStream }) => { + let draftContent = ''; + + const { fullStream } = streamText({ + model: myProvider.languageModel('block-model'), + system: updateDocumentPrompt(document.content, 'text'), + experimental_transform: smoothStream({ chunking: 'word' }), + prompt: description, + experimental_providerMetadata: { + openai: { + prediction: { + type: 'content', + content: document.content, + }, + }, + }, + }); + + for await (const delta of fullStream) { + const { type } = delta; + + if (type === 'text-delta') { + const { textDelta } = delta; + + draftContent += textDelta; + dataStream.writeData({ + type: 'text-delta', + content: textDelta, + }); + } + } + + return draftContent; + }, +}); diff --git a/components/block-actions.tsx b/components/block-actions.tsx index b660d704f..47bf1c377 100644 --- a/components/block-actions.tsx +++ b/components/block-actions.tsx @@ -91,6 +91,7 @@ export const BlockActions = memo(PureBlockActions, (prevProps, nextProps) => { if (prevProps.currentVersionIndex !== nextProps.currentVersionIndex) return false; if (prevProps.isCurrentVersion !== nextProps.isCurrentVersion) return false; + if (prevProps.block.content !== nextProps.block.content) return false; return true; }); diff --git a/components/block.tsx b/components/block.tsx index 890ad8f68..cf6298fbd 100644 --- a/components/block.tsx +++ b/components/block.tsx @@ -26,11 +26,11 @@ import { BlockCloseButton } from './block-close-button'; import { BlockMessages } from './block-messages'; import { useSidebar } from './ui/sidebar'; import { useBlock } from '@/hooks/use-block'; -import { textBlock } from '@/blocks/text'; -import { imageBlock } from '@/blocks/image'; -import { codeBlock } from '@/blocks/code'; +import { imageBlock } from '@/blocks/image/client'; +import { codeBlock } from '@/blocks/code/client'; +import { sheetBlock } from '@/blocks/sheet/client'; +import { textBlock } from '@/blocks/text/client'; import equal from 'fast-deep-equal'; -import { sheetBlock } from '@/blocks/sheet'; export const blockDefinitions = [textBlock, codeBlock, imageBlock, sheetBlock]; export type BlockKind = (typeof blockDefinitions)[number]['kind']; diff --git a/lib/ai/tools/create-document.ts b/lib/ai/tools/create-document.ts index 9f755d640..46c5434c1 100644 --- a/lib/ai/tools/create-document.ts +++ b/lib/ai/tools/create-document.ts @@ -1,17 +1,8 @@ import { generateUUID } from '@/lib/utils'; -import { - DataStreamWriter, - experimental_generateImage, - smoothStream, - streamObject, - streamText, - tool, -} from 'ai'; +import { DataStreamWriter, tool } from 'ai'; import { z } from 'zod'; -import { codePrompt, sheetPrompt } from '../prompts'; -import { saveDocument } from '@/lib/db/queries'; import { Session } from 'next-auth'; -import { myProvider } from '../models'; +import { blockKinds, documentHandlersByBlockKind } from '@/lib/blocks/server'; interface CreateDocumentProps { session: Session; @@ -24,11 +15,15 @@ export const createDocument = ({ session, dataStream }: CreateDocumentProps) => '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.', parameters: z.object({ title: z.string(), - kind: z.enum(['text', 'code', 'image', 'sheet']), + kind: z.enum(blockKinds), }), execute: async ({ title, kind }) => { const id = generateUUID(); - let draftText = ''; + + dataStream.writeData({ + type: 'kind', + content: kind, + }); dataStream.writeData({ type: 'id', @@ -40,129 +35,28 @@ export const createDocument = ({ session, dataStream }: CreateDocumentProps) => content: title, }); - dataStream.writeData({ - type: 'kind', - content: kind, - }); - dataStream.writeData({ type: 'clear', content: '', }); - if (kind === 'text') { - const { fullStream } = streamText({ - model: myProvider.languageModel('block-model'), - system: - 'Write about the given topic. Markdown is supported. Use headings wherever appropriate.', - experimental_transform: smoothStream({ chunking: 'word' }), - prompt: title, - }); - - for await (const delta of fullStream) { - const { type } = delta; - - if (type === 'text-delta') { - const { textDelta } = delta; - - draftText += textDelta; - dataStream.writeData({ - type: 'text-delta', - content: textDelta, - }); - } - } - - dataStream.writeData({ type: 'finish', content: '' }); - } else if (kind === 'code') { - const { fullStream } = streamObject({ - model: myProvider.languageModel('block-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; + const documentHandler = documentHandlersByBlockKind.find( + (documentHandlerByBlockKind) => + documentHandlerByBlockKind.kind === kind, + ); - if (code) { - dataStream.writeData({ - type: 'code-delta', - content: code ?? '', - }); - - draftText = code; - } - } - } - - dataStream.writeData({ type: 'finish', content: '' }); - } else if (kind === 'image') { - const { image } = await experimental_generateImage({ - model: myProvider.imageModel('small-model'), - prompt: title, - n: 1, - }); - - draftText = image.base64; - - dataStream.writeData({ - type: 'image-delta', - content: image.base64, - }); - - dataStream.writeData({ type: 'finish', content: '' }); - } else if (kind === 'sheet') { - const { fullStream } = streamObject({ - model: myProvider.languageModel('block-model'), - 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.writeData({ - type: 'sheet-delta', - content: csv, - }); - - draftText = csv; - } - } - } - - dataStream.writeData({ - type: 'sheet-delta', - content: draftText, - }); - - dataStream.writeData({ type: 'finish', content: '' }); + if (!documentHandler) { + throw new Error(`No document handler found for kind: ${kind}`); } - if (session.user?.id) { - await saveDocument({ - id, - title, - kind, - content: draftText, - userId: session.user.id, - }); - } + await documentHandler.onCreateDocument({ + id, + title, + dataStream, + session, + }); + + dataStream.writeData({ type: 'finish', content: '' }); return { id, diff --git a/lib/ai/tools/update-document.ts b/lib/ai/tools/update-document.ts index 8d937d866..2ba4abd45 100644 --- a/lib/ai/tools/update-document.ts +++ b/lib/ai/tools/update-document.ts @@ -1,16 +1,8 @@ -import { - DataStreamWriter, - experimental_generateImage, - smoothStream, - streamObject, - streamText, - tool, -} from 'ai'; +import { DataStreamWriter, tool } from 'ai'; import { Session } from 'next-auth'; import { z } from 'zod'; import { getDocumentById, saveDocument } from '@/lib/db/queries'; -import { updateDocumentPrompt } from '../prompts'; -import { myProvider } from '../models'; +import { documentHandlersByBlockKind } from '@/lib/blocks/server'; interface UpdateDocumentProps { session: Session; @@ -35,129 +27,28 @@ export const updateDocument = ({ session, dataStream }: UpdateDocumentProps) => }; } - const { content: currentContent } = document; - let draftText = ''; - dataStream.writeData({ type: 'clear', content: document.title, }); - if (document.kind === 'text') { - const { fullStream } = streamText({ - model: myProvider.languageModel('block-model'), - system: updateDocumentPrompt(currentContent, 'text'), - experimental_transform: smoothStream({ chunking: 'word' }), - prompt: description, - experimental_providerMetadata: { - openai: { - prediction: { - type: 'content', - content: currentContent, - }, - }, - }, - }); - - for await (const delta of fullStream) { - const { type } = delta; - - if (type === 'text-delta') { - const { textDelta } = delta; - - draftText += textDelta; - dataStream.writeData({ - type: 'text-delta', - content: textDelta, - }); - } - } - - dataStream.writeData({ type: 'finish', content: '' }); - } else if (document.kind === 'code') { - const { fullStream } = streamObject({ - model: myProvider.languageModel('block-model'), - system: updateDocumentPrompt(currentContent, '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.writeData({ - type: 'code-delta', - content: code ?? '', - }); - - draftText = code; - } - } - } - - dataStream.writeData({ type: 'finish', content: '' }); - } else if (document.kind === 'image') { - const { image } = await experimental_generateImage({ - model: myProvider.imageModel('image-model'), - prompt: description, - n: 1, - }); + const documentHandler = documentHandlersByBlockKind.find( + (documentHandlerByBlockKind) => + documentHandlerByBlockKind.kind === document.kind, + ); - draftText = image.base64; - - dataStream.writeData({ - type: 'image-delta', - content: image.base64, - }); - - dataStream.writeData({ type: 'finish', content: '' }); - } else if (document.kind === 'sheet') { - const { fullStream } = streamObject({ - model: myProvider.languageModel('block-model'), - system: updateDocumentPrompt(currentContent, '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.writeData({ - type: 'sheet-delta', - content: csv, - }); - - draftText = csv; - } - } - } - - dataStream.writeData({ type: 'finish', content: '' }); + if (!documentHandler) { + throw new Error(`No document handler found for kind: ${document.kind}`); } - if (session.user?.id) { - await saveDocument({ - id, - title: document.title, - content: draftText, - kind: document.kind, - userId: session.user.id, - }); - } + await documentHandler.onUpdateDocument({ + document, + description, + dataStream, + session, + }); + + dataStream.writeData({ type: 'finish', content: '' }); return { id, diff --git a/lib/blocks/server.ts b/lib/blocks/server.ts new file mode 100644 index 000000000..98b037b4a --- /dev/null +++ b/lib/blocks/server.ts @@ -0,0 +1,99 @@ +import { codeDocumentHandler } from '@/blocks/code/server'; +import { imageDocumentHandler } from '@/blocks/image/server'; +import { sheetDocumentHandler } from '@/blocks/sheet/server'; +import { textDocumentHandler } from '@/blocks/text/server'; +import { BlockKind } from '@/components/block'; +import { DataStreamWriter } from 'ai'; +import { Document } from '../db/schema'; +import { saveDocument } from '../db/queries'; +import { Session } from 'next-auth'; + +export interface SaveDocumentProps { + id: string; + title: string; + kind: BlockKind; + content: string; + userId: string; +} + +export interface CreateDocumentCallbackProps { + id: string; + title: string; + dataStream: DataStreamWriter; + session: Session; +} + +export interface UpdateDocumentCallbackProps { + document: Document; + description: string; + dataStream: DataStreamWriter; + session: Session; +} + +export interface DocumentHandler { + kind: T; + onCreateDocument: (args: CreateDocumentCallbackProps) => Promise; + onUpdateDocument: (args: UpdateDocumentCallbackProps) => Promise; +} + +export function createDocumentHandler(config: { + kind: T; + onCreateDocument: (params: CreateDocumentCallbackProps) => Promise; + onUpdateDocument: (params: UpdateDocumentCallbackProps) => Promise; +}): DocumentHandler { + return { + kind: config.kind, + onCreateDocument: async (args: CreateDocumentCallbackProps) => { + const draftContent = await config.onCreateDocument({ + id: args.id, + title: args.title, + dataStream: args.dataStream, + session: args.session, + }); + + if (args.session?.user?.id) { + await saveDocument({ + id: args.id, + title: args.title, + content: draftContent, + kind: config.kind, + userId: args.session.user.id, + }); + } + + return; + }, + onUpdateDocument: async (args: UpdateDocumentCallbackProps) => { + const draftContent = await config.onUpdateDocument({ + document: args.document, + description: args.description, + dataStream: args.dataStream, + session: args.session, + }); + + if (args.session?.user?.id) { + await saveDocument({ + id: args.document.id, + title: args.document.title, + content: draftContent, + kind: config.kind, + userId: args.session.user.id, + }); + } + + return; + }, + }; +} + +/* + * Use this array to define the document handlers for each block kind. + */ +export const documentHandlersByBlockKind: Array = [ + textDocumentHandler, + codeDocumentHandler, + imageDocumentHandler, + sheetDocumentHandler, +]; + +export const blockKinds = ['text', 'code', 'image', 'sheet'] as const;