mirror of
https://github.com/run-llama/ai-chatbot.git
synced 2026-07-01 21:14:02 -04:00
refactor: replace StreamData with createDataStreamResponse (#628)
Co-authored-by: Ali Mirlou <alimirlou@gmail.com>
This commit is contained in:
Jeremy
+326
-324
@@ -1,8 +1,7 @@
|
||||
import {
|
||||
type Message,
|
||||
StreamData,
|
||||
convertToCoreMessages,
|
||||
generateObject,
|
||||
createDataStreamResponse,
|
||||
streamObject,
|
||||
streamText,
|
||||
} from 'ai';
|
||||
@@ -94,358 +93,361 @@ export async function POST(request: Request) {
|
||||
],
|
||||
});
|
||||
|
||||
const streamingData = new StreamData();
|
||||
return createDataStreamResponse({
|
||||
execute: (dataStream) => {
|
||||
dataStream.writeData({
|
||||
type: 'user-message-id',
|
||||
content: userMessageId,
|
||||
});
|
||||
|
||||
streamingData.append({
|
||||
type: 'user-message-id',
|
||||
content: userMessageId,
|
||||
});
|
||||
const result = streamText({
|
||||
model: customModel(model.apiIdentifier),
|
||||
system: systemPrompt,
|
||||
messages: coreMessages,
|
||||
maxSteps: 5,
|
||||
experimental_activeTools: allTools,
|
||||
tools: {
|
||||
getWeather: {
|
||||
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 result = streamText({
|
||||
model: customModel(model.apiIdentifier),
|
||||
system: systemPrompt,
|
||||
messages: coreMessages,
|
||||
maxSteps: 5,
|
||||
experimental_activeTools: allTools,
|
||||
tools: {
|
||||
getWeather: {
|
||||
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;
|
||||
},
|
||||
},
|
||||
createDocument: {
|
||||
description: 'Create a document for a writing activity.',
|
||||
parameters: z.object({
|
||||
title: z.string(),
|
||||
kind: z.enum(['text', 'code']),
|
||||
}),
|
||||
execute: async ({ title, kind }) => {
|
||||
const id = generateUUID();
|
||||
let draftText = '';
|
||||
|
||||
const weatherData = await response.json();
|
||||
return weatherData;
|
||||
},
|
||||
},
|
||||
createDocument: {
|
||||
description: 'Create a document for a writing activity.',
|
||||
parameters: z.object({
|
||||
title: z.string(),
|
||||
kind: z.enum(['text', 'code']),
|
||||
}),
|
||||
execute: async ({ title, kind }) => {
|
||||
const id = generateUUID();
|
||||
let draftText = '';
|
||||
dataStream.writeData({
|
||||
type: 'id',
|
||||
content: id,
|
||||
});
|
||||
|
||||
streamingData.append({
|
||||
type: 'id',
|
||||
content: id,
|
||||
});
|
||||
dataStream.writeData({
|
||||
type: 'title',
|
||||
content: title,
|
||||
});
|
||||
|
||||
streamingData.append({
|
||||
type: 'title',
|
||||
content: title,
|
||||
});
|
||||
dataStream.writeData({
|
||||
type: 'kind',
|
||||
content: kind,
|
||||
});
|
||||
|
||||
streamingData.append({
|
||||
type: 'kind',
|
||||
content: kind,
|
||||
});
|
||||
dataStream.writeData({
|
||||
type: 'clear',
|
||||
content: '',
|
||||
});
|
||||
|
||||
streamingData.append({
|
||||
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,
|
||||
});
|
||||
|
||||
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;
|
||||
|
||||
for await (const delta of fullStream) {
|
||||
const { type } = delta;
|
||||
if (type === 'text-delta') {
|
||||
const { textDelta } = delta;
|
||||
|
||||
if (type === 'text-delta') {
|
||||
const { textDelta } = delta;
|
||||
draftText += textDelta;
|
||||
dataStream.writeData({
|
||||
type: 'text-delta',
|
||||
content: textDelta,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
draftText += textDelta;
|
||||
streamingData.append({
|
||||
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: '' });
|
||||
}
|
||||
|
||||
if (session.user?.id) {
|
||||
await saveDocument({
|
||||
id,
|
||||
title,
|
||||
kind,
|
||||
content: draftText,
|
||||
userId: session.user.id,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
streamingData.append({ 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) {
|
||||
streamingData.append({
|
||||
type: 'code-delta',
|
||||
content: code ?? '',
|
||||
});
|
||||
|
||||
draftText = code;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
streamingData.append({ 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 = '';
|
||||
|
||||
streamingData.append({
|
||||
type: 'clear',
|
||||
content: document.title,
|
||||
});
|
||||
|
||||
if (document.kind === 'text') {
|
||||
const { fullStream } = streamText({
|
||||
model: customModel(model.apiIdentifier),
|
||||
system: updateDocumentPrompt(currentContent),
|
||||
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;
|
||||
streamingData.append({
|
||||
type: 'text-delta',
|
||||
content: textDelta,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
streamingData.append({ type: 'finish', content: '' });
|
||||
} else if (document.kind === 'code') {
|
||||
const { fullStream } = streamObject({
|
||||
model: customModel(model.apiIdentifier),
|
||||
system: updateDocumentPrompt(currentContent),
|
||||
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) {
|
||||
streamingData.append({
|
||||
type: 'code-delta',
|
||||
content: code ?? '',
|
||||
});
|
||||
|
||||
draftText = code;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
streamingData.append({ 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'),
|
||||
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 the suggestion'),
|
||||
.describe('The description of changes that need to be made'),
|
||||
}),
|
||||
});
|
||||
execute: async ({ id, description }) => {
|
||||
const document = await getDocumentById({ id });
|
||||
|
||||
for await (const element of elementStream) {
|
||||
const suggestion = {
|
||||
originalText: element.originalSentence,
|
||||
suggestedText: element.suggestedSentence,
|
||||
description: element.description,
|
||||
id: generateUUID(),
|
||||
documentId: documentId,
|
||||
isResolved: false,
|
||||
};
|
||||
if (!document) {
|
||||
return {
|
||||
error: 'Document not found',
|
||||
};
|
||||
}
|
||||
|
||||
streamingData.append({
|
||||
type: 'suggestion',
|
||||
content: suggestion,
|
||||
});
|
||||
const { content: currentContent } = document;
|
||||
let draftText = '';
|
||||
|
||||
suggestions.push(suggestion);
|
||||
}
|
||||
dataStream.writeData({
|
||||
type: 'clear',
|
||||
content: document.title,
|
||||
});
|
||||
|
||||
if (session.user?.id) {
|
||||
const userId = session.user.id;
|
||||
if (document.kind === 'text') {
|
||||
const { fullStream } = streamText({
|
||||
model: customModel(model.apiIdentifier),
|
||||
system: updateDocumentPrompt(currentContent),
|
||||
prompt: description,
|
||||
experimental_providerMetadata: {
|
||||
openai: {
|
||||
prediction: {
|
||||
type: 'content',
|
||||
content: currentContent,
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
await saveSuggestions({
|
||||
suggestions: suggestions.map((suggestion) => ({
|
||||
...suggestion,
|
||||
userId,
|
||||
createdAt: new Date(),
|
||||
documentCreatedAt: document.createdAt,
|
||||
})),
|
||||
});
|
||||
}
|
||||
for await (const delta of fullStream) {
|
||||
const { type } = delta;
|
||||
|
||||
return {
|
||||
id: documentId,
|
||||
title: document.title,
|
||||
kind: document.kind,
|
||||
message: 'Suggestions have been added to the document',
|
||||
};
|
||||
},
|
||||
},
|
||||
},
|
||||
onFinish: async ({ response }) => {
|
||||
if (session.user?.id) {
|
||||
try {
|
||||
const responseMessagesWithoutIncompleteToolCalls =
|
||||
sanitizeResponseMessages(response.messages);
|
||||
if (type === 'text-delta') {
|
||||
const { textDelta } = delta;
|
||||
|
||||
await saveMessages({
|
||||
messages: responseMessagesWithoutIncompleteToolCalls.map(
|
||||
(message) => {
|
||||
const messageId = generateUUID();
|
||||
|
||||
if (message.role === 'assistant') {
|
||||
streamingData.appendMessageAnnotation({
|
||||
messageIdFromServer: messageId,
|
||||
});
|
||||
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),
|
||||
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: '' });
|
||||
}
|
||||
|
||||
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 {
|
||||
id: messageId,
|
||||
chatId: id,
|
||||
role: message.role,
|
||||
content: message.content,
|
||||
createdAt: new Date(),
|
||||
error: 'Document not found',
|
||||
};
|
||||
},
|
||||
),
|
||||
});
|
||||
} catch (error) {
|
||||
console.error('Failed to save chat');
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
streamingData.close();
|
||||
},
|
||||
experimental_telemetry: {
|
||||
isEnabled: true,
|
||||
functionId: 'stream-text',
|
||||
},
|
||||
});
|
||||
const suggestions: Array<
|
||||
Omit<Suggestion, 'userId' | 'createdAt' | 'documentCreatedAt'>
|
||||
> = [];
|
||||
|
||||
return result.toDataStreamResponse({
|
||||
data: streamingData,
|
||||
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) {
|
||||
try {
|
||||
const responseMessagesWithoutIncompleteToolCalls =
|
||||
sanitizeResponseMessages(response.messages);
|
||||
|
||||
await saveMessages({
|
||||
messages: responseMessagesWithoutIncompleteToolCalls.map(
|
||||
(message) => {
|
||||
const messageId = generateUUID();
|
||||
|
||||
if (message.role === 'assistant') {
|
||||
dataStream.writeMessageAnnotation({
|
||||
messageIdFromServer: messageId,
|
||||
});
|
||||
}
|
||||
|
||||
return {
|
||||
id: messageId,
|
||||
chatId: id,
|
||||
role: message.role,
|
||||
content: message.content,
|
||||
createdAt: new Date(),
|
||||
};
|
||||
},
|
||||
),
|
||||
});
|
||||
} catch (error) {
|
||||
console.error('Failed to save chat');
|
||||
}
|
||||
}
|
||||
},
|
||||
experimental_telemetry: {
|
||||
isEnabled: true,
|
||||
functionId: 'stream-text',
|
||||
},
|
||||
});
|
||||
|
||||
result.mergeIntoDataStream(dataStream);
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user