62 lines
1.6 KiB
TypeScript
62 lines
1.6 KiB
TypeScript
import { gateway } from "@ai-sdk/gateway";
|
|
import {
|
|
customProvider,
|
|
extractReasoningMiddleware,
|
|
wrapLanguageModel,
|
|
} from "ai";
|
|
import { isTestEnvironment } from "../constants";
|
|
|
|
const THINKING_SUFFIX_REGEX = /-thinking$/;
|
|
|
|
export const myProvider = isTestEnvironment
|
|
? (() => {
|
|
const {
|
|
artifactModel,
|
|
chatModel,
|
|
reasoningModel,
|
|
titleModel,
|
|
} = require("./models.mock");
|
|
return customProvider({
|
|
languageModels: {
|
|
"chat-model": chatModel,
|
|
"chat-model-reasoning": reasoningModel,
|
|
"title-model": titleModel,
|
|
"artifact-model": artifactModel,
|
|
},
|
|
});
|
|
})()
|
|
: null;
|
|
|
|
export function getLanguageModel(modelId: string) {
|
|
if (isTestEnvironment && myProvider) {
|
|
return myProvider.languageModel(modelId);
|
|
}
|
|
|
|
const isReasoningModel =
|
|
modelId.includes("reasoning") || modelId.endsWith("-thinking");
|
|
|
|
if (isReasoningModel) {
|
|
const gatewayModelId = modelId.replace(THINKING_SUFFIX_REGEX, "");
|
|
|
|
return wrapLanguageModel({
|
|
model: gateway.languageModel(gatewayModelId),
|
|
middleware: extractReasoningMiddleware({ tagName: "thinking" }),
|
|
});
|
|
}
|
|
|
|
return gateway.languageModel(modelId);
|
|
}
|
|
|
|
export function getTitleModel() {
|
|
if (isTestEnvironment && myProvider) {
|
|
return myProvider.languageModel("title-model");
|
|
}
|
|
return gateway.languageModel("google/gemini-2.5-flash-lite");
|
|
}
|
|
|
|
export function getArtifactModel() {
|
|
if (isTestEnvironment && myProvider) {
|
|
return myProvider.languageModel("artifact-model");
|
|
}
|
|
return gateway.languageModel("anthropic/claude-haiku-4.5");
|
|
}
|