feat: add image block type (#709)

This commit is contained in:
Jeremy 2025-01-15 19:59:29 +05:30 committed by GitHub
parent 6c23a8a604
commit df449a408a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 229 additions and 92 deletions

View file

@ -2,13 +2,14 @@ import {
type Message,
convertToCoreMessages,
createDataStreamResponse,
experimental_generateImage,
streamObject,
streamText,
} from 'ai';
import { z } from 'zod';
import { auth } from '@/app/(auth)/auth';
import { customModel } from '@/lib/ai';
import { customModel, imageGenerationModel } from '@/lib/ai';
import { models } from '@/lib/ai/models';
import {
codePrompt,
@ -124,10 +125,10 @@ export async function POST(request: Request) {
},
createDocument: {
description:
'Create a document for a writing activity. This tool will call other functions that will generate the contents of the document based on the title and kind.',
'Create a document for a writing or content creation activities like image generation. 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']),
kind: z.enum(['text', 'code', 'image']),
}),
execute: async ({ title, kind }) => {
const id = generateUUID();
@ -204,6 +205,21 @@ export async function POST(request: Request) {
}
}
dataStream.writeData({ type: 'finish', content: '' });
} else if (kind === 'image') {
const { image } = await experimental_generateImage({
model: imageGenerationModel,
prompt: title,
n: 1,
});
draftText = image.base64;
dataStream.writeData({
type: 'image-delta',
content: image.base64,
});
dataStream.writeData({ type: 'finish', content: '' });
}
@ -309,6 +325,21 @@ export async function POST(request: Request) {
}
}
dataStream.writeData({ type: 'finish', content: '' });
} else if (document.kind === 'image') {
const { image } = await experimental_generateImage({
model: imageGenerationModel,
prompt: description,
n: 1,
});
draftText = image.base64;
dataStream.writeData({
type: 'image-delta',
content: image.base64,
});
dataStream.writeData({ type: 'finish', content: '' });
}