From bd1f93a75516e7650d50049382e47d979bf6a664 Mon Sep 17 00:00:00 2001
From: slytechnical <slytechnical@linuto.com>
Date: Mon, 19 May 2025 13:32:37 -0500
Subject: [PATCH] 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.
---
 src/api/providers/fetchers/modelCache.ts | 35 +++++++++++++-----------
 1 file changed, 19 insertions(+), 16 deletions(-)

diff --git a/src/api/providers/fetchers/modelCache.ts b/src/api/providers/fetchers/modelCache.ts
index 7ff625b6cb..175a6b2e15 100644
--- a/src/api/providers/fetchers/modelCache.ts
+++ b/src/api/providers/fetchers/modelCache.ts
@@ -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