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);
}