feat: add reasoning model (#750)

Co-authored-by: Matt Apperson <me@mattapperson.com>

lib/ai/tools/create-document.ts

@@ -8,23 +8,17 @@ import {
   tool,
 } from 'ai';
 import { z } from 'zod';
-import { customModel, imageGenerationModel } from '..';
 import { codePrompt, sheetPrompt } from '../prompts';
 import { saveDocument } from '@/lib/db/queries';
 import { Session } from 'next-auth';
-import { Model } from '../models';
+import { myProvider } from '../models';
 
 interface CreateDocumentProps {
-  model: Model;
   session: Session;
   dataStream: DataStreamWriter;
 }
 
-export const createDocument = ({
-  model,
-  session,
-  dataStream,
-}: CreateDocumentProps) =>
+export const createDocument = ({ session, dataStream }: CreateDocumentProps) =>
   tool({
     description:
       '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.',
@@ -58,7 +52,7 @@ export const createDocument = ({
 
       if (kind === 'text') {
         const { fullStream } = streamText({
-          model: customModel(model.apiIdentifier),
+          model: myProvider.languageModel('block-model'),
           system:
             'Write about the given topic. Markdown is supported. Use headings wherever appropriate.',
           experimental_transform: smoothStream({ chunking: 'word' }),
@@ -82,7 +76,7 @@ export const createDocument = ({
         dataStream.writeData({ type: 'finish', content: '' });
       } else if (kind === 'code') {
         const { fullStream } = streamObject({
-          model: customModel(model.apiIdentifier),
+          model: myProvider.languageModel('block-model'),
           system: codePrompt,
           prompt: title,
           schema: z.object({
@@ -111,7 +105,7 @@ export const createDocument = ({
         dataStream.writeData({ type: 'finish', content: '' });
       } else if (kind === 'image') {
         const { image } = await experimental_generateImage({
-          model: imageGenerationModel,
+          model: myProvider.imageModel('small-model'),
           prompt: title,
           n: 1,
         });
@@ -126,7 +120,7 @@ export const createDocument = ({
         dataStream.writeData({ type: 'finish', content: '' });
       } else if (kind === 'sheet') {
         const { fullStream } = streamObject({
-          model: customModel(model.apiIdentifier),
+          model: myProvider.languageModel('block-model'),
           system: sheetPrompt,
           prompt: title,
           schema: z.object({

lib/ai/tools/request-suggestions.ts

@@ -1,20 +1,17 @@
 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';
+import { myProvider } from '../models';
 
 interface RequestSuggestionsProps {
-  model: Model;
   session: Session;
   dataStream: DataStreamWriter;
 }
 
 export const requestSuggestions = ({
-  model,
   session,
   dataStream,
 }: RequestSuggestionsProps) =>
@@ -39,7 +36,7 @@ export const requestSuggestions = ({
       > = [];
 
       const { elementStream } = streamObject({
-        model: customModel(model.apiIdentifier),
+        model: myProvider.languageModel('block-model'),
         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,

lib/ai/tools/update-document.ts

@@ -6,24 +6,18 @@ import {
   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';
+import { myProvider } from '../models';
 
 interface UpdateDocumentProps {
-  model: Model;
   session: Session;
   dataStream: DataStreamWriter;
 }
 
-export const updateDocument = ({
-  model,
-  session,
-  dataStream,
-}: UpdateDocumentProps) =>
+export const updateDocument = ({ session, dataStream }: UpdateDocumentProps) =>
   tool({
     description: 'Update a document with the given description.',
     parameters: z.object({
@@ -51,7 +45,7 @@ export const updateDocument = ({
 
       if (document.kind === 'text') {
         const { fullStream } = streamText({
-          model: customModel(model.apiIdentifier),
+          model: myProvider.languageModel('block-model'),
           system: updateDocumentPrompt(currentContent, 'text'),
           experimental_transform: smoothStream({ chunking: 'word' }),
           prompt: description,
@@ -82,7 +76,7 @@ export const updateDocument = ({
         dataStream.writeData({ type: 'finish', content: '' });
       } else if (document.kind === 'code') {
         const { fullStream } = streamObject({
-          model: customModel(model.apiIdentifier),
+          model: myProvider.languageModel('block-model'),
           system: updateDocumentPrompt(currentContent, 'code'),
           prompt: description,
           schema: z.object({
@@ -111,7 +105,7 @@ export const updateDocument = ({
         dataStream.writeData({ type: 'finish', content: '' });
       } else if (document.kind === 'image') {
         const { image } = await experimental_generateImage({
-          model: imageGenerationModel,
+          model: myProvider.imageModel('image-model'),
           prompt: description,
           n: 1,
         });
@@ -126,7 +120,7 @@ export const updateDocument = ({
         dataStream.writeData({ type: 'finish', content: '' });
       } else if (document.kind === 'sheet') {
         const { fullStream } = streamObject({
-          model: customModel(model.apiIdentifier),
+          model: myProvider.languageModel('block-model'),
           system: updateDocumentPrompt(currentContent, 'sheet'),
           prompt: description,
           schema: z.object({