Refactor getModels function to use options parameter for improved type safety

- Updated the getModels function to accept a single options parameter instead of multiple individual parameters.
- Introduced a discriminated union type for GetModelsOptions to enforce required properties based on the router type.
This commit is contained in:
slytechnical 2025-05-19 13:32:37 -05:00
parent cbed162462
commit bd1f93a755

View file

@ -30,22 +30,29 @@ async function readModels(router: RouterName): Promise<ModelRecord | undefined>
return exists ? JSON.parse(await fs.readFile(filePath, "utf8")) : undefined
}
/**
* Options for fetching models from different routers.
* This is a discriminated union type where the router property determines
* which other properties are required.
*/
export type GetModelsOptions =
| { router: "openrouter" }
| { router: "glama" }
| { router: "requesty"; apiKey?: string }
| { router: "unbound"; apiKey?: string }
| { router: "litellm"; apiKey: string; baseUrl: string }
/**
* Get models from the cache or fetch them from the provider and cache them.
* There are two caches:
* 1. Memory cache - This is a simple in-memory cache that is used to store models for a short period of time.
* 2. File cache - This is a file-based cache that is used to store models for a longer period of time.
*
* @param router - The router to fetch models from.
* @param apiKey - Optional API key for the provider.
* @param baseUrl - Optional base URL for the provider (currently used only for LiteLLM).
* @param options - Options for fetching models, including the router and any required parameters.
* @returns The models from the cache or the fetched models.
*/
export const getModels = async (
router: RouterName,
apiKey: string | undefined = undefined,
baseUrl: string | undefined = undefined,
): Promise<ModelRecord> => {
export const getModels = async (options: GetModelsOptions): Promise<ModelRecord> => {
const { router } = options
let models = memoryCache.get<ModelRecord>(router)
if (models) {
return models
@ -58,22 +65,18 @@ export const getModels = async (
break
case "requesty":
// Requesty models endpoint requires an API key for per-user custom policies
models = await getRequestyModels(apiKey)
models = await getRequestyModels(options.apiKey)
break
case "glama":
models = await getGlamaModels()
break
case "unbound":
// Unbound models endpoint requires an API key to fetch application specific models
models = await getUnboundModels(apiKey)
models = await getUnboundModels(options.apiKey)
break
case "litellm":
if (!baseUrl || !apiKey) {
// This case should ideally be handled by the caller if baseUrl is strictly required.
// However, for robustness, if called without baseUrl for litellm, it would fail in getLiteLLMModels or here.
throw new Error("Base URL and api key are required for LiteLLM models.")
}
models = await getLiteLLMModels(apiKey, baseUrl)
// Type safety ensures apiKey and baseUrl are always provided for litellm
models = await getLiteLLMModels(options.apiKey, options.baseUrl)
break
default:
// Ensures router is exhaustively checked if RouterName is a strict union