From d5fd66f07cc8162250e7ce0cceb8d28a8cbdd157 Mon Sep 17 00:00:00 2001
From: Enrico Ros <enrico.ros@gmail.com>
Date: Wed, 2 Oct 2024 15:13:44 -0700
Subject: [PATCH] AutoConf w/ rerank of Vendors

---
 src/common/logic/autoconf.ts         | 15 ++++++++++++---
 src/common/stores/llms/store-llms.ts | 21 +++++++++++++++++++++
 2 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/src/common/logic/autoconf.ts b/src/common/logic/autoconf.ts
index 21006366d..ac00ee356 100644
--- a/src/common/logic/autoconf.ts
+++ b/src/common/logic/autoconf.ts
@@ -6,8 +6,8 @@ import { findAllModelVendors } from '~/modules/llms/vendors/vendors.registry';
 import { getBackendCapabilities } from '~/modules/backend/store-backend-capabilities';
 import { llmsUpdateModelsForServiceOrThrow } from '~/modules/llms/llm.client';
 
-import type { DModelsService } from '~/common/stores/llms/modelsservice.types';
-import { llmsStoreState } from '~/common/stores/llms/store-llms';
+import type { DModelsService, DModelsServiceId } from '~/common/stores/llms/modelsservice.types';
+import { llmsStoreActions, llmsStoreState } from '~/common/stores/llms/store-llms';
 
 
 interface AutoConfStore {
@@ -50,6 +50,9 @@ const autoConfVanillaStore = createVanillaStore<AutoConfStore>()(persist((_set,
     let configurableVendors = findAllModelVendors()
       .filter(vendor => vendor.hasBackendCapKey && backendCaps[vendor.hasBackendCapKey]);
 
+    // List to keep track of the service IDs in order
+    const configuredServiceIds: DModelsServiceId[] = [];
+
     // Sequentially auto-configure each vendor
     await configurableVendors.reduce(async (promiseChain, vendor) => {
       return promiseChain
@@ -63,11 +66,14 @@ const autoConfVanillaStore = createVanillaStore<AutoConfStore>()(persist((_set,
             // create and append the model service, assuming the backend configuration will be successful
             service = createModelsServiceForVendor(vendor.id, modelsServices);
             addService(service);
-            // re-find it now what's added
+            // re-find it now that it's added
             service = llmsStoreState().sources.find(_s => _s.id === service.id)!;
           } else
             service = firstServiceForVendor;
 
+          // keep track of the configured service IDs
+          configuredServiceIds.push(service.id);
+
           // auto-configure this service
           await llmsUpdateModelsForServiceOrThrow(service.id, true);
         })
@@ -81,6 +87,9 @@ const autoConfVanillaStore = createVanillaStore<AutoConfStore>()(persist((_set,
         });
     }, Promise.resolve());
 
+    // Re-rank the LLMs based on the order of configured services
+    llmsStoreActions().rerankLLMsByServices(configuredServiceIds);
+
     // end configuration
     _set({ isConfiguring: false, isConfigurationDone: true });
   },
diff --git a/src/common/stores/llms/store-llms.ts b/src/common/stores/llms/store-llms.ts
index d01d67387..b45d61480 100644
--- a/src/common/stores/llms/store-llms.ts
+++ b/src/common/stores/llms/store-llms.ts
@@ -30,6 +30,7 @@ interface LlmsActions {
 
   setLLMs: (llms: DLLM[], serviceId: DModelsServiceId, deleteExpiredVendorLlms: boolean, keepUserEdits: boolean) => void;
   removeLLM: (id: DLLMId) => void;
+  rerankLLMsByServices: (serviceIdOrder: DModelsServiceId[]) => void;
   updateLLM: (id: DLLMId, partial: Partial<DLLM>) => void;
   updateLLMOptions: <TLLMOptions>(id: DLLMId, partialOptions: Partial<TLLMOptions>) => void;
 
@@ -102,6 +103,26 @@ export const useModelsStore = create<LlmsState & LlmsActions>()(persist(
         };
       }),
 
+    rerankLLMsByServices: (serviceIdOrder: DModelsServiceId[]) =>
+      set(state => {
+        // Create a mapping of service IDs to their index in the provided order
+        const serviceIdToIndex = serviceIdOrder.reduce((acc, sId, idx) => {
+          acc[sId] = idx;
+          return acc;
+        }, {} as Record<DModelsServiceId, number>);
+
+        // Sort the LLMs based on the order of their service IDs
+        const orderedLlms = [...state.llms].sort((a, b) => {
+          const aIndex = serviceIdToIndex[a.sId] ?? Number.MAX_SAFE_INTEGER;
+          const bIndex = serviceIdToIndex[b.sId] ?? Number.MAX_SAFE_INTEGER;
+          return aIndex - bIndex;
+        });
+
+        return {
+          llms: orderedLlms,
+        };
+      }),
+
     updateLLM: (id: DLLMId, partial: Partial<DLLM>) =>
       set(state => ({
         llms: state.llms.map((llm: DLLM): DLLM =>