refactor: use tool() and organize by file (#721)
Co-authored-by: Marcus Schiesser <mail@marcusschiesser.de>
This commit is contained in:
Jeremy
GitHub
parent
371242d683
commit
3f9d379a6a
5 changed files with 412 additions and 345 deletions
|
|
@ -2,30 +2,19 @@ import {
|
|||
type Message,
|
||||
convertToCoreMessages,
|
||||
createDataStreamResponse,
|
||||
experimental_generateImage,
|
||||
streamObject,
|
||||
streamText,
|
||||
} from 'ai';
|
||||
import { z } from 'zod';
|
||||
|
||||
import { auth } from '@/app/(auth)/auth';
|
||||
import { customModel, imageGenerationModel } from '@/lib/ai';
|
||||
import { customModel } from '@/lib/ai';
|
||||
import { models } from '@/lib/ai/models';
|
||||
import {
|
||||
codePrompt,
|
||||
systemPrompt,
|
||||
updateDocumentPrompt,
|
||||
} from '@/lib/ai/prompts';
|
||||
import { systemPrompt } from '@/lib/ai/prompts';
|
||||
import {
|
||||
deleteChatById,
|
||||
getChatById,
|
||||
getDocumentById,
|
||||
saveChat,
|
||||
saveDocument,
|
||||
saveMessages,
|
||||
saveSuggestions,
|
||||
} from '@/lib/db/queries';
|
||||
import type { Suggestion } from '@/lib/db/schema';
|
||||
import {
|
||||
generateUUID,
|
||||
getMostRecentUserMessage,
|
||||
|
|
@ -33,6 +22,10 @@ import {
|
|||
} from '@/lib/utils';
|
||||
|
||||
import { generateTitleFromUserMessage } from '../../actions';
|
||||
import { createDocument } from '@/lib/ai/tools/create-document';
|
||||
import { updateDocument } from '@/lib/ai/tools/update-document';
|
||||
import { requestSuggestions } from '@/lib/ai/tools/request-suggestions';
|
||||
import { getWeather } from '@/lib/ai/tools/get-weather';
|
||||
|
||||
export const maxDuration = 60;
|
||||
|
||||
|
|
@ -49,7 +42,6 @@ const blocksTools: AllowedTools[] = [
|
|||
];
|
||||
|
||||
const weatherTools: AllowedTools[] = ['getWeather'];
|
||||
|
||||
const allTools: AllowedTools[] = [...blocksTools, ...weatherTools];
|
||||
|
||||
export async function POST(request: Request) {
|
||||
|
|
@ -108,337 +100,14 @@ export async function POST(request: Request) {
|
|||
maxSteps: 5,
|
||||
experimental_activeTools: allTools,
|
||||
tools: {
|
||||
getWeather: {
|
||||
description: 'Get the current weather at a location',
|
||||
parameters: z.object({
|
||||
latitude: z.number(),
|
||||
longitude: z.number(),
|
||||
getWeather,
|
||||
createDocument: createDocument({ session, dataStream, model }),
|
||||
updateDocument: updateDocument({ session, dataStream, model }),
|
||||
requestSuggestions: requestSuggestions({
|
||||
session,
|
||||
dataStream,
|
||||
model,
|
||||
}),
|
||||
execute: async ({ latitude, longitude }) => {
|
||||
const response = await fetch(
|
||||
`https://api.open-meteo.com/v1/forecast?latitude=${latitude}&longitude=${longitude}¤t=temperature_2m&hourly=temperature_2m&daily=sunrise,sunset&timezone=auto`,
|
||||
);
|
||||
|
||||
const weatherData = await response.json();
|
||||
return weatherData;
|
||||
},
|
||||
},
|
||||
createDocument: {
|
||||
description:
|
||||
'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', 'image']),
|
||||
}),
|
||||
execute: async ({ title, kind }) => {
|
||||
const id = generateUUID();
|
||||
let draftText = '';
|
||||
|
||||
dataStream.writeData({
|
||||
type: 'id',
|
||||
content: id,
|
||||
});
|
||||
|
||||
dataStream.writeData({
|
||||
type: 'title',
|
||||
content: title,
|
||||
});
|
||||
|
||||
dataStream.writeData({
|
||||
type: 'kind',
|
||||
content: kind,
|
||||
});
|
||||
|
||||
dataStream.writeData({
|
||||
type: 'clear',
|
||||
content: '',
|
||||
});
|
||||
|
||||
if (kind === 'text') {
|
||||
const { fullStream } = streamText({
|
||||
model: customModel(model.apiIdentifier),
|
||||
system:
|
||||
'Write about the given topic. Markdown is supported. Use headings wherever appropriate.',
|
||||
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: customModel(model.apiIdentifier),
|
||||
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 ?? '',
|
||||
});
|
||||
|
||||
draftText = code;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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: '' });
|
||||
}
|
||||
|
||||
if (session.user?.id) {
|
||||
await saveDocument({
|
||||
id,
|
||||
title,
|
||||
kind,
|
||||
content: draftText,
|
||||
userId: session.user.id,
|
||||
});
|
||||
}
|
||||
|
||||
return {
|
||||
id,
|
||||
title,
|
||||
kind,
|
||||
content:
|
||||
'A document was created and is now visible to the user.',
|
||||
};
|
||||
},
|
||||
},
|
||||
updateDocument: {
|
||||
description: 'Update a document with the given description.',
|
||||
parameters: 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',
|
||||
};
|
||||
}
|
||||
|
||||
const { content: currentContent } = document;
|
||||
let draftText = '';
|
||||
|
||||
dataStream.writeData({
|
||||
type: 'clear',
|
||||
content: document.title,
|
||||
});
|
||||
|
||||
if (document.kind === 'text') {
|
||||
const { fullStream } = streamText({
|
||||
model: customModel(model.apiIdentifier),
|
||||
system: updateDocumentPrompt(currentContent, 'text'),
|
||||
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: customModel(model.apiIdentifier),
|
||||
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: imageGenerationModel,
|
||||
prompt: description,
|
||||
n: 1,
|
||||
});
|
||||
|
||||
draftText = image.base64;
|
||||
|
||||
dataStream.writeData({
|
||||
type: 'image-delta',
|
||||
content: image.base64,
|
||||
});
|
||||
|
||||
dataStream.writeData({ type: 'finish', content: '' });
|
||||
}
|
||||
|
||||
if (session.user?.id) {
|
||||
await saveDocument({
|
||||
id,
|
||||
title: document.title,
|
||||
content: draftText,
|
||||
kind: document.kind,
|
||||
userId: session.user.id,
|
||||
});
|
||||
}
|
||||
|
||||
return {
|
||||
id,
|
||||
title: document.title,
|
||||
kind: document.kind,
|
||||
content: 'The document has been updated successfully.',
|
||||
};
|
||||
},
|
||||
},
|
||||
requestSuggestions: {
|
||||
description: 'Request suggestions for a document',
|
||||
parameters: z.object({
|
||||
documentId: z
|
||||
.string()
|
||||
.describe('The ID of the document to request edits'),
|
||||
}),
|
||||
execute: async ({ documentId }) => {
|
||||
const document = await getDocumentById({ id: documentId });
|
||||
|
||||
if (!document || !document.content) {
|
||||
return {
|
||||
error: 'Document not found',
|
||||
};
|
||||
}
|
||||
|
||||
const suggestions: Array<
|
||||
Omit<Suggestion, 'userId' | 'createdAt' | 'documentCreatedAt'>
|
||||
> = [];
|
||||
|
||||
const { elementStream } = streamObject({
|
||||
model: customModel(model.apiIdentifier),
|
||||
system:
|
||||
'You are a help writing assistant. Given a piece of writing, please offer suggestions to improve the piece of writing and describe the change. It is very important for the edits to contain full sentences instead of just words. Max 5 suggestions.',
|
||||
prompt: document.content,
|
||||
output: 'array',
|
||||
schema: z.object({
|
||||
originalSentence: z
|
||||
.string()
|
||||
.describe('The original sentence'),
|
||||
suggestedSentence: z
|
||||
.string()
|
||||
.describe('The suggested sentence'),
|
||||
description: z
|
||||
.string()
|
||||
.describe('The description of the suggestion'),
|
||||
}),
|
||||
});
|
||||
|
||||
for await (const element of elementStream) {
|
||||
const suggestion = {
|
||||
originalText: element.originalSentence,
|
||||
suggestedText: element.suggestedSentence,
|
||||
description: element.description,
|
||||
id: generateUUID(),
|
||||
documentId: documentId,
|
||||
isResolved: false,
|
||||
};
|
||||
|
||||
dataStream.writeData({
|
||||
type: 'suggestion',
|
||||
content: suggestion,
|
||||
});
|
||||
|
||||
suggestions.push(suggestion);
|
||||
}
|
||||
|
||||
if (session.user?.id) {
|
||||
const userId = session.user.id;
|
||||
|
||||
await saveSuggestions({
|
||||
suggestions: suggestions.map((suggestion) => ({
|
||||
...suggestion,
|
||||
userId,
|
||||
createdAt: new Date(),
|
||||
documentCreatedAt: document.createdAt,
|
||||
})),
|
||||
});
|
||||
}
|
||||
|
||||
return {
|
||||
id: documentId,
|
||||
title: document.title,
|
||||
kind: document.kind,
|
||||
message: 'Suggestions have been added to the document',
|
||||
};
|
||||
},
|
||||
},
|
||||
},
|
||||
onFinish: async ({ response }) => {
|
||||
if (session.user?.id) {
|
||||
|
|
|
|||
144
lib/ai/tools/create-document.ts
Normal file
144
lib/ai/tools/create-document.ts
Normal file
|
|
@ -0,0 +1,144 @@
|
|||
import { generateUUID } from '@/lib/utils';
|
||||
import {
|
||||
DataStreamWriter,
|
||||
experimental_generateImage,
|
||||
streamObject,
|
||||
streamText,
|
||||
tool,
|
||||
} from 'ai';
|
||||
import { z } from 'zod';
|
||||
import { customModel, imageGenerationModel } from '..';
|
||||
import { codePrompt } from '../prompts';
|
||||
import { saveDocument } from '@/lib/db/queries';
|
||||
import { Session } from 'next-auth';
|
||||
import { Model } from '../models';
|
||||
|
||||
interface CreateDocumentProps {
|
||||
model: Model;
|
||||
session: Session;
|
||||
dataStream: DataStreamWriter;
|
||||
}
|
||||
|
||||
export const createDocument = ({
|
||||
model,
|
||||
session,
|
||||
dataStream,
|
||||
}: CreateDocumentProps) =>
|
||||
tool({
|
||||
description:
|
||||
'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', 'image']),
|
||||
}),
|
||||
execute: async ({ title, kind }) => {
|
||||
const id = generateUUID();
|
||||
let draftText = '';
|
||||
|
||||
dataStream.writeData({
|
||||
type: 'id',
|
||||
content: id,
|
||||
});
|
||||
|
||||
dataStream.writeData({
|
||||
type: 'title',
|
||||
content: title,
|
||||
});
|
||||
|
||||
dataStream.writeData({
|
||||
type: 'kind',
|
||||
content: kind,
|
||||
});
|
||||
|
||||
dataStream.writeData({
|
||||
type: 'clear',
|
||||
content: '',
|
||||
});
|
||||
|
||||
if (kind === 'text') {
|
||||
const { fullStream } = streamText({
|
||||
model: customModel(model.apiIdentifier),
|
||||
system:
|
||||
'Write about the given topic. Markdown is supported. Use headings wherever appropriate.',
|
||||
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: customModel(model.apiIdentifier),
|
||||
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 ?? '',
|
||||
});
|
||||
|
||||
draftText = code;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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: '' });
|
||||
}
|
||||
|
||||
if (session.user?.id) {
|
||||
await saveDocument({
|
||||
id,
|
||||
title,
|
||||
kind,
|
||||
content: draftText,
|
||||
userId: session.user.id,
|
||||
});
|
||||
}
|
||||
|
||||
return {
|
||||
id,
|
||||
title,
|
||||
kind,
|
||||
content: 'A document was created and is now visible to the user.',
|
||||
};
|
||||
},
|
||||
});
|
||||
18
lib/ai/tools/get-weather.ts
Normal file
18
lib/ai/tools/get-weather.ts
Normal file
|
|
@ -0,0 +1,18 @@
|
|||
import { tool } from 'ai';
|
||||
import { z } from 'zod';
|
||||
|
||||
export const getWeather = tool({
|
||||
description: 'Get the current weather at a location',
|
||||
parameters: z.object({
|
||||
latitude: z.number(),
|
||||
longitude: z.number(),
|
||||
}),
|
||||
execute: async ({ latitude, longitude }) => {
|
||||
const response = await fetch(
|
||||
`https://api.open-meteo.com/v1/forecast?latitude=${latitude}&longitude=${longitude}¤t=temperature_2m&hourly=temperature_2m&daily=sunrise,sunset&timezone=auto`,
|
||||
);
|
||||
|
||||
const weatherData = await response.json();
|
||||
return weatherData;
|
||||
},
|
||||
});
|
||||
92
lib/ai/tools/request-suggestions.ts
Normal file
92
lib/ai/tools/request-suggestions.ts
Normal file
|
|
@ -0,0 +1,92 @@
|
|||
import { z } from 'zod';
|
||||
import { Model } from '../models';
|
||||
import { Session } from 'next-auth';
|
||||
import { DataStreamWriter, streamObject, tool } from 'ai';
|
||||
import { getDocumentById, saveSuggestions } from '@/lib/db/queries';
|
||||
import { Suggestion } from '@/lib/db/schema';
|
||||
import { customModel } from '..';
|
||||
import { generateUUID } from '@/lib/utils';
|
||||
|
||||
interface RequestSuggestionsProps {
|
||||
model: Model;
|
||||
session: Session;
|
||||
dataStream: DataStreamWriter;
|
||||
}
|
||||
|
||||
export const requestSuggestions = ({
|
||||
model,
|
||||
session,
|
||||
dataStream,
|
||||
}: RequestSuggestionsProps) =>
|
||||
tool({
|
||||
description: 'Request suggestions for a document',
|
||||
parameters: z.object({
|
||||
documentId: z
|
||||
.string()
|
||||
.describe('The ID of the document to request edits'),
|
||||
}),
|
||||
execute: async ({ documentId }) => {
|
||||
const document = await getDocumentById({ id: documentId });
|
||||
|
||||
if (!document || !document.content) {
|
||||
return {
|
||||
error: 'Document not found',
|
||||
};
|
||||
}
|
||||
|
||||
const suggestions: Array<
|
||||
Omit<Suggestion, 'userId' | 'createdAt' | 'documentCreatedAt'>
|
||||
> = [];
|
||||
|
||||
const { elementStream } = streamObject({
|
||||
model: customModel(model.apiIdentifier),
|
||||
system:
|
||||
'You are a help writing assistant. Given a piece of writing, please offer suggestions to improve the piece of writing and describe the change. It is very important for the edits to contain full sentences instead of just words. Max 5 suggestions.',
|
||||
prompt: document.content,
|
||||
output: 'array',
|
||||
schema: z.object({
|
||||
originalSentence: z.string().describe('The original sentence'),
|
||||
suggestedSentence: z.string().describe('The suggested sentence'),
|
||||
description: z.string().describe('The description of the suggestion'),
|
||||
}),
|
||||
});
|
||||
|
||||
for await (const element of elementStream) {
|
||||
const suggestion = {
|
||||
originalText: element.originalSentence,
|
||||
suggestedText: element.suggestedSentence,
|
||||
description: element.description,
|
||||
id: generateUUID(),
|
||||
documentId: documentId,
|
||||
isResolved: false,
|
||||
};
|
||||
|
||||
dataStream.writeData({
|
||||
type: 'suggestion',
|
||||
content: suggestion,
|
||||
});
|
||||
|
||||
suggestions.push(suggestion);
|
||||
}
|
||||
|
||||
if (session.user?.id) {
|
||||
const userId = session.user.id;
|
||||
|
||||
await saveSuggestions({
|
||||
suggestions: suggestions.map((suggestion) => ({
|
||||
...suggestion,
|
||||
userId,
|
||||
createdAt: new Date(),
|
||||
documentCreatedAt: document.createdAt,
|
||||
})),
|
||||
});
|
||||
}
|
||||
|
||||
return {
|
||||
id: documentId,
|
||||
title: document.title,
|
||||
kind: document.kind,
|
||||
message: 'Suggestions have been added to the document',
|
||||
};
|
||||
},
|
||||
});
|
||||
144
lib/ai/tools/update-document.ts
Normal file
144
lib/ai/tools/update-document.ts
Normal file
|
|
@ -0,0 +1,144 @@
|
|||
import {
|
||||
DataStreamWriter,
|
||||
experimental_generateImage,
|
||||
streamObject,
|
||||
streamText,
|
||||
tool,
|
||||
} from 'ai';
|
||||
import { Model } from '../models';
|
||||
import { Session } from 'next-auth';
|
||||
import { z } from 'zod';
|
||||
import { getDocumentById, saveDocument } from '@/lib/db/queries';
|
||||
import { customModel, imageGenerationModel } from '..';
|
||||
import { updateDocumentPrompt } from '../prompts';
|
||||
|
||||
interface UpdateDocumentProps {
|
||||
model: Model;
|
||||
session: Session;
|
||||
dataStream: DataStreamWriter;
|
||||
}
|
||||
|
||||
export const updateDocument = ({
|
||||
model,
|
||||
session,
|
||||
dataStream,
|
||||
}: UpdateDocumentProps) =>
|
||||
tool({
|
||||
description: 'Update a document with the given description.',
|
||||
parameters: 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',
|
||||
};
|
||||
}
|
||||
|
||||
const { content: currentContent } = document;
|
||||
let draftText = '';
|
||||
|
||||
dataStream.writeData({
|
||||
type: 'clear',
|
||||
content: document.title,
|
||||
});
|
||||
|
||||
if (document.kind === 'text') {
|
||||
const { fullStream } = streamText({
|
||||
model: customModel(model.apiIdentifier),
|
||||
system: updateDocumentPrompt(currentContent, 'text'),
|
||||
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: customModel(model.apiIdentifier),
|
||||
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: imageGenerationModel,
|
||||
prompt: description,
|
||||
n: 1,
|
||||
});
|
||||
|
||||
draftText = image.base64;
|
||||
|
||||
dataStream.writeData({
|
||||
type: 'image-delta',
|
||||
content: image.base64,
|
||||
});
|
||||
|
||||
dataStream.writeData({ type: 'finish', content: '' });
|
||||
}
|
||||
|
||||
if (session.user?.id) {
|
||||
await saveDocument({
|
||||
id,
|
||||
title: document.title,
|
||||
content: draftText,
|
||||
kind: document.kind,
|
||||
userId: session.user.id,
|
||||
});
|
||||
}
|
||||
|
||||
return {
|
||||
id,
|
||||
title: document.title,
|
||||
kind: document.kind,
|
||||
content: 'The document has been updated successfully.',
|
||||
};
|
||||
},
|
||||
});
|
||||
Loading…
Add table
Add a link
Reference in a new issue