Skip to content

Commit

Permalink
Merge branch 'main' into jrmy/live-sync-2025-01-23
Browse files Browse the repository at this point in the history
  • Loading branch information
@jeremyphilemon
jeremyphilemon committed Jan 23, 2025
2 parents 5e51a26 + 5e8cddc commit 5171cbf
Show file tree
Hide file tree
Showing 17 changed files with 647 additions and 411 deletions.
355 changes: 16 additions & 339 deletions app/(chat)/api/chat/route.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,36 +2,31 @@ import {
type Message,
convertToCoreMessages,
createDataStreamResponse,
streamObject,
smoothStream,
streamText,
} from 'ai';
import { z } from 'zod';

import { auth } from '@/app/(auth)/auth';
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,
sanitizeResponseMessages,
} 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;

Expand All @@ -48,7 +43,6 @@ const blocksTools: AllowedTools[] = [
];

const weatherTools: AllowedTools[] = ['getWeather'];

const allTools: AllowedTools[] = [...blocksTools, ...weatherTools];

export async function POST(request: Request) {
Expand Down Expand Up @@ -121,336 +115,19 @@ export async function POST(request: Request) {
messages: coreMessages,
maxSteps: 5,
experimental_activeTools: allTools,
experimental_transform: smoothStream({ chunking: 'word' }),
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}&current=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. 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']),
}),
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: '' });
}

if (session && session.user && session.user.id) {
await saveDocument({
id,
title,
kind,
content: draftText,
userId: session.user.id,
});
} else {
await saveDocument({
id,
title,
kind,
content: draftText,
userId: process.env.ANON_USERID!,
});
}

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: '' });
}

if (session && session.user && session.user.id) {
await saveDocument({
id,
title: document.title,
content: draftText,
kind: document.kind,
userId: session.user.id,
});
} else {
await saveDocument({
id,
title: document.title,
kind: document.kind,
content: draftText,
userId: process.env.ANON_USERID!,
});
}

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 && session.user && session.user.id) {
const userId = session.user.id;

await saveSuggestions({
suggestions: suggestions.map((suggestion) => ({
...suggestion,
userId,
createdAt: new Date(),
documentCreatedAt: document.createdAt,
})),
});
} else {
await saveSuggestions({
suggestions: suggestions.map((suggestion) => ({
...suggestion,
userId: process.env.ANON_USERID!,
createdAt: new Date(),
documentCreatedAt: document.createdAt,
})),
});
}

return {
id: documentId,
title: document.title,
kind: document.kind,
message: 'Suggestions have been added to the document',
};
},
},
getWeather,
createDocument: createDocument({ session, dataStream, model }),
updateDocument: updateDocument({ session, dataStream, model }),
requestSuggestions: requestSuggestions({
session,
dataStream,
model,
}),
},
onFinish: async ({ response }) => {
if (session && session.user && session.user.id) {
if (session?.user?.id) {
try {
const responseMessagesWithoutIncompleteToolCalls =
sanitizeResponseMessages(response.messages);
Expand Down
Loading

0 comments on commit 5171cbf

Please sign in to comment.