37 lines
1 KiB
TypeScript
37 lines
1 KiB
TypeScript
|
import { createOpenAICompatible } from "@ai-sdk/openai-compatible";
|
||
|
|
import { customProvider } from "ai";
|
||
|
|
import { isTestEnvironment } from "../constants";
|
||
|
|
import { titleModel } from "./models";
|
||
|
|
|
||
|
|
const litellm = createOpenAICompatible({
|
||
|
|
name: "egbe-litellm",
|
||
|
|
baseURL: process.env.EGBE_AI_API_URL ?? "https://proxy.gcp.egbe.dev/v1",
|
||
|
|
apiKey: process.env.EGBE_AI_API_KEY ?? "",
|
||
|
|
});
|
||
|
|
|
||
|
|
export const myProvider = isTestEnvironment
|
||
|
|
? (() => {
|
||
|
|
const { chatModel, titleModel: mockTitle } = require("./models.mock");
|
||
|
|
return customProvider({
|
||
|
|
languageModels: {
|
||
|
|
"chat-model": chatModel,
|
||
|
|
"title-model": mockTitle,
|
||
|
|
},
|
||
|
|
});
|
||
|
|
})()
|
||
|
|
: null;
|
||
|
|
|
||
|
|
export function getLanguageModel(modelId: string) {
|
||
|
|
if (isTestEnvironment && myProvider) {
|
||
|
|
return myProvider.languageModel(modelId);
|
||
|
|
}
|
||
|
|
return litellm.chatModel(modelId);
|
||
|
|
}
|
||
|
|
|
||
|
|
export function getTitleModel() {
|
||
|
|
if (isTestEnvironment && myProvider) {
|
||
|
|
return myProvider.languageModel("title-model");
|
||
|
|
}
|
||
|
|
return litellm.chatModel(titleModel.id);
|
||
|
|
}
|