chatbot-template/lib/ai/providers.ts

42 lines
1.2 KiB
TypeScript
Raw Normal View History

2025-03-04 17:25:46 -08:00
import {
customProvider,
extractReasoningMiddleware,
wrapLanguageModel,
2025-03-17 16:54:40 -07:00
} from 'ai';
import { openai } from '@ai-sdk/openai';
import { fireworks } from '@ai-sdk/fireworks';
import { xai } from '@ai-sdk/xai';
import { isTestEnvironment } from '../constants';
2025-03-04 17:25:46 -08:00
import {
artifactModel,
chatModel,
reasoningModel,
titleModel,
2025-03-17 16:54:40 -07:00
} from './models.test';
2025-03-04 17:25:46 -08:00
export const myProvider = isTestEnvironment
? customProvider({
languageModels: {
2025-03-17 16:54:40 -07:00
'chat-model-small': chatModel,
'chat-model-large': chatModel,
'chat-model-reasoning': reasoningModel,
'title-model': titleModel,
'artifact-model': artifactModel,
2025-03-04 17:25:46 -08:00
},
})
: customProvider({
languageModels: {
2025-03-17 16:54:40 -07:00
'chat-model-small': xai('grok-2-1212'),
'chat-model-reasoning': wrapLanguageModel({
model: fireworks('accounts/fireworks/models/deepseek-r1'),
middleware: extractReasoningMiddleware({ tagName: 'think' }),
2025-03-04 17:25:46 -08:00
}),
2025-03-17 16:54:40 -07:00
'title-model': xai('grok-2-1212'),
'artifact-model': xai('grok-2-1212'),
2025-03-04 17:25:46 -08:00
},
imageModels: {
2025-03-17 16:54:40 -07:00
'small-model': openai.image('dall-e-2'),
'large-model': openai.image('dall-e-3'),
2025-03-04 17:25:46 -08:00
},
});