@marimo-team/islands 0.15.1-dev20 → 0.15.1-dev23

package/src/core/ai/__tests__/model-registry.test.ts

@@ -78,21 +78,45 @@ describe("AiModelRegistry", () => {
       const displayedModels = ["openai/gpt-4", "anthropic/claude-3-sonnet"];
       const registry = AiModelRegistry.create({ displayedModels });
 
+      const ids = [...registry.getModelsMap().keys()];
+      expect(ids).toEqual(["openai/gpt-4", "anthropic/claude-3-sonnet"]);
       expect(registry.getCustomModels()).toEqual(new Set());
       expect(registry.getDisplayedModels()).toEqual(new Set(displayedModels));
     });
 
     it("should create registry with both custom and displayed models", () => {
       const customModels = ["openai/custom-gpt"];
-      const displayedModels = ["openai/gpt-4", "anthropic/claude-3-sonnet"];
+      const displayedModels = ["openai/custom-gpt"];
       const registry = AiModelRegistry.create({
         customModels,
         displayedModels,
       });
 
+      const ids = [...registry.getModelsMap().keys()];
+      expect(ids).toEqual(["openai/custom-gpt"]);
       expect(registry.getCustomModels()).toEqual(new Set(customModels));
       expect(registry.getDisplayedModels()).toEqual(new Set(displayedModels));
     });
+
+    it("should create registry with non-existent displayed_model", () => {
+      const customModels = ["openai/custom-gpt"];
+      const displayedModels = ["something-wrong/model-id"];
+      const registry = AiModelRegistry.create({
+        customModels,
+        displayedModels,
+      });
+
+      const ids = [...registry.getModelsMap().keys()];
+      // Include custom and all default ones.
+      expect(ids).toEqual([
+        "openai/custom-gpt",
+        "openai/gpt-4",
+        "anthropic/claude-3-sonnet",
+        "google/gemini-pro",
+        "openai/multi-model",
+        "anthropic/multi-model",
+      ]);
+    });
   });
 
   describe("getModelsByProvider", () => {

package/src/core/ai/model-registry.ts

@@ -7,6 +7,7 @@ import type {
 } from "@marimo-team/llm-info";
 import { models } from "@marimo-team/llm-info/models.json";
 import { providers } from "@marimo-team/llm-info/providers.json";
+import { Logger } from "@/utils/Logger";
 import { MultiMap } from "@/utils/multi-map";
 import { once } from "@/utils/once";
 import type { ProviderId } from "./ids/ids";
@@ -14,6 +15,7 @@ import { AiModelId, type QualifiedModelId, type ShortModelId } from "./ids/ids";
 
 export interface AiModel extends AiModelType {
   roles: Role[];
+  model: ShortModelId;
   providers: ProviderId[];
   /** Whether this is a custom model. */
   custom: boolean;
@@ -25,6 +27,7 @@ const getKnownModelMap = once((): ReadonlyMap<QualifiedModelId, AiModel> => {
     const modelId = model.model as ShortModelId;
     const modelInfo: AiModel = {
       ...model,
+      model: model.model as ShortModelId,
       roles: model.roles.map((role) => role as Role),
       providers: model.providers as ProviderId[],
       custom: false,
@@ -84,14 +87,42 @@ export class AiModelRegistry {
    * Builds the maps of models by provider and custom models.
    */
   private buildMaps() {
-    const displayedModels = this.displayedModels;
+    let result = AiModelRegistry.buildMapsFromConfig({
+      displayedModels: this.displayedModels,
+      customModels: this.customModels,
+    });
+
+    // If we got zero results, then build the maps with no displayedModels
+    // This can happen if displayedModels is configured to non existent models
+    if (result.modelsMap.size === 0) {
+      Logger.error(
+        "The configured displayed_models have filtered out all registered models. Reverting back to showing all models.",
+        [...this.displayedModels],
+      );
+
+      result = AiModelRegistry.buildMapsFromConfig({
+        displayedModels: new Set(),
+        customModels: this.customModels,
+      });
+    }
+
+    this.modelsByProviderMap = result.modelsByProviderMap;
+    this.modelsMap = result.modelsMap;
+  }
+
+  private static buildMapsFromConfig(opts: {
+    customModels: ReadonlySet<QualifiedModelId>;
+    displayedModels: ReadonlySet<QualifiedModelId>;
+  }) {
+    const { displayedModels, customModels } = opts;
     const hasDisplayedModels = displayedModels.size > 0;
     const knownModelMap = getKnownModelMap();
     const customModelsMap = new Map<QualifiedModelId, AiModel>();
 
     let modelsMap = new Map<QualifiedModelId, AiModel>();
+    const modelsByProviderMap = new MultiMap<ProviderId, AiModel>();
 
-    for (const model of this.customModels) {
+    for (const model of customModels) {
       if (hasDisplayedModels && !displayedModels.has(model)) {
         continue;
       }
@@ -121,15 +152,16 @@ export class AiModelRegistry {
     }
 
     // Set custom models first, then known models
+    // Known models will overwrite custom models (which is desired)
     modelsMap = new Map([...customModelsMap, ...modelsMap]);
 
     // Group by provider
     for (const [qualifiedModelId, model] of modelsMap.entries()) {
       const modelId = AiModelId.parse(qualifiedModelId);
-      this.modelsByProviderMap.add(modelId.providerId, model);
+      modelsByProviderMap.add(modelId.providerId, model);
     }
 
-    this.modelsMap = modelsMap;
+    return { modelsByProviderMap, modelsMap };
   }
 
   getDisplayedModels() {

package/src/hooks/useDebounce.ts

@@ -84,7 +84,8 @@ export function useDebounceControlledState<T>(opts: {
   };
 }
 
-export function useDebouncedCallback<T extends (...args: unknown[]) => unknown>(
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export function useDebouncedCallback<T extends (...args: any[]) => unknown>(
   callback: T,
   delay: number,
 ) {