feat: add reasoning model (#750)

Co-authored-by: Matt Apperson <me@mattapperson.com>
This commit is contained in:
Jeremy 2025-02-03 20:33:15 +05:30 committed by GitHub
parent 76804269c4
commit c61d4f91d4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 342 additions and 209 deletions

View file

@ -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({