Merge pull request #446 from kubet/feat/improve-model-selector

backend/utils/constants.py

@@ -2,8 +2,7 @@ MODEL_ACCESS_TIERS = {
     "free": [
         "openrouter/deepseek/deepseek-chat",
         "openrouter/qwen/qwen3-235b-a22b",
-        "openrouter/google/gemini-2.5-flash-preview",
+        "openrouter/google/gemini-2.5-flash-preview-05-20",
-        # "openrouter/google/gemini-2.5-flash-preview:thinking",
     ],
     "tier_2_20": [
         "openrouter/deepseek/deepseek-chat",
@@ -114,10 +113,11 @@ MODEL_NAME_ALIASES = {
     # "deepseek-r1": "openrouter/deepseek/deepseek-r1",
     # "grok-3-mini": "xai/grok-3-mini-fast-beta",  # Commented out in constants.py
     "qwen3": "openrouter/qwen/qwen3-235b-a22b",  # Commented out in constants.py
+    "gemini-flash-2.5": "openrouter/google/gemini-2.5-flash-preview-05-20",
+    "gemini-2.5-flash:thinking":"openrouter/google/gemini-2.5-flash-preview-05-20:thinking",
     
-    
+    # "google/gemini-2.5-flash-preview":"openrouter/google/gemini-2.5-flash-preview",
-    "google/gemini-2.5-flash-preview":"openrouter/google/gemini-2.5-flash-preview",
+    # "google/gemini-2.5-flash-preview:thinking":"openrouter/google/gemini-2.5-flash-preview:thinking",
-    "google/gemini-2.5-flash-preview:thinking":"openrouter/google/gemini-2.5-flash-preview:thinking",
     "google/gemini-2.5-pro-preview":"openrouter/google/gemini-2.5-pro-preview",
     "deepseek/deepseek-chat-v3-0324":"openrouter/deepseek/deepseek-chat-v3-0324",
 

frontend/src/components/thread/chat-input/_use-model-selection.ts

@@ -86,7 +86,7 @@ export const MODELS = {
     lowQuality: false,
     description: 'Gemini 2.5 - Google\'s powerful versatile model'
   },
-  'gemini-2.5-flash-preview': { 
+  'gemini-flash-2.5:thinking': { 
     tier: 'premium', 
     priority: 90,
     recommended: true,
@@ -137,7 +137,7 @@ export const MODELS = {
     lowQuality: true,
     description: 'DeepSeek - Free tier model with good general capabilities'
   },
-  'google/gemini-2.5-flash-preview': { 
+  'gemini-flash-2.5': { 
     tier: 'free', 
     priority: 50,
     recommended: false,
@@ -247,11 +247,17 @@ export const useModelSelection = () => {
     ? 'active' 
     : 'no_subscription';
 
+  // Function to refresh custom models from localStorage
+  const refreshCustomModels = () => {
+    if (isLocalMode() && typeof window !== 'undefined') {
+      const freshCustomModels = getCustomModels();
+      setCustomModels(freshCustomModels);
+    }
+  };
+
   // Load custom models from localStorage
   useEffect(() => {
-    if (isLocalMode() && typeof window !== 'undefined') {
+    refreshCustomModels();
-      setCustomModels(getCustomModels());
-    }
   }, []);
 
   // Generate model options list with consistent structure
@@ -417,21 +423,37 @@ export const useModelSelection = () => {
 
   // Handle model selection change
   const handleModelChange = (modelId: string) => {
-    const modelOption = MODEL_OPTIONS.find(option => option.id === modelId);
+    console.log('handleModelChange', modelId);
+    
+    // Refresh custom models from localStorage to ensure we have the latest
+    if (isLocalMode()) {
+      refreshCustomModels();
+    }
+    
+    // First check if it's a custom model in local mode
     const isCustomModel = isLocalMode() && customModels.some(model => model.id === modelId);
     
-    // Check if model exists
+    // Then check if it's in standard MODEL_OPTIONS
+    const modelOption = MODEL_OPTIONS.find(option => option.id === modelId);
+    
+    // Check if model exists in either custom models or standard options
     if (!modelOption && !isCustomModel) {
-      console.warn('Model not found in options:', modelId);
+      console.warn('Model not found in options:', modelId, MODEL_OPTIONS, isCustomModel, customModels);
+      
+      // Reset to default model when the selected model is not found
+      const defaultModel = isLocalMode() ? DEFAULT_PREMIUM_MODEL_ID : DEFAULT_FREE_MODEL_ID;
+      setSelectedModel(defaultModel);
+      saveModelPreference(defaultModel);
       return;
     }
 
     // Check access permissions (except for custom models in local mode)
     if (!isCustomModel && !isLocalMode() && 
         !canAccessModel(subscriptionStatus, modelOption?.requiresSubscription ?? false)) {
+      console.warn('Model not accessible:', modelId);
       return;
     }
-    
+    console.log('setting selected model', modelId);
     setSelectedModel(modelId);
     saveModelPreference(modelId);
   };
@@ -444,12 +466,15 @@ export const useModelSelection = () => {
 
   return {
     selectedModel,
-    setSelectedModel: handleModelChange,
+    setSelectedModel: (modelId: string) => {
+      handleModelChange(modelId);
+    },
     subscriptionStatus,
     availableModels,
     allModels: MODEL_OPTIONS,  // Already pre-sorted
     customModels,
     getActualModelId,
+    refreshCustomModels,
     canAccessModel: (modelId: string) => {
       if (isLocalMode()) return true;
       const model = MODEL_OPTIONS.find(m => m.id === modelId);

frontend/src/components/thread/chat-input/chat-input.tsx

@@ -82,6 +82,7 @@ export const ChatInput = forwardRef<ChatInputHandles, ChatInputProps>(
       allModels: modelOptions,
       canAccessModel,
       getActualModelId,
+      refreshCustomModels,
     } = useModelSelection();
 
     const textareaRef = useRef<HTMLTextAreaElement | null>(null);
@@ -233,6 +234,7 @@ export const ChatInput = forwardRef<ChatInputHandles, ChatInputProps>(
                 modelOptions={modelOptions}
                 subscriptionStatus={subscriptionStatus}
                 canAccessModel={canAccessModel}
+                refreshCustomModels={refreshCustomModels}
               />
             </CardContent>
           </div>

frontend/src/components/thread/chat-input/message-input.tsx

@@ -36,7 +36,8 @@ interface MessageInputProps {
   onModelChange: (model: string) => void;
   modelOptions: any[];
   subscriptionStatus: SubscriptionStatus;
-  canAccessModel: (model: string) => boolean;
+  canAccessModel: (modelId: string) => boolean;
+  refreshCustomModels?: () => void;
 }
 
 export const MessageInput = forwardRef<HTMLTextAreaElement, MessageInputProps>(
@@ -66,6 +67,7 @@ export const MessageInput = forwardRef<HTMLTextAreaElement, MessageInputProps>(
       modelOptions,
       subscriptionStatus,
       canAccessModel,
+      refreshCustomModels,
     },
     ref,
   ) => {
@@ -148,6 +150,7 @@ export const MessageInput = forwardRef<HTMLTextAreaElement, MessageInputProps>(
               modelOptions={modelOptions}
               subscriptionStatus={subscriptionStatus}
               canAccessModel={canAccessModel}
+              refreshCustomModels={refreshCustomModels}
             />
             <Button
               type="submit"

frontend/src/components/thread/chat-input/model-selector.tsx

@@ -44,6 +44,7 @@ interface ModelSelectorProps {
   modelOptions: ModelOption[];
   canAccessModel: (modelId: string) => boolean;
   subscriptionStatus: SubscriptionStatus;
+  refreshCustomModels?: () => void;
 }
 
 export const ModelSelector: React.FC<ModelSelectorProps> = ({
@@ -52,6 +53,7 @@ export const ModelSelector: React.FC<ModelSelectorProps> = ({
   modelOptions,
   canAccessModel,
   subscriptionStatus,
+  refreshCustomModels,
 }) => {
   const [paywallOpen, setPaywallOpen] = useState(false);
   const [lockedModel, setLockedModel] = useState<string | null>(null);
@@ -96,17 +98,30 @@ export const ModelSelector: React.FC<ModelSelectorProps> = ({
     });
   });
 
-  // Then add custom models, overriding any with the same ID
+  // Then add custom models from the current customModels state (not from props)
-  currentCustomModels.forEach(model => {
+  // This ensures we're using the most up-to-date list of custom models
-    if (!modelMap.has(model.id)) {
+  if (isLocalMode()) {
-      modelMap.set(model.id, {
+    // Get current custom models from state (not from storage)
-        ...model,
+    customModels.forEach(model => {
-        requiresSubscription: false,
+      // Only add if it doesn't exist or mark it as a custom model if it does
-        top: false,
+      if (!modelMap.has(model.id)) {
-        isCustom: true
+        modelMap.set(model.id, {
-      });
+          id: model.id,
-    }
+          label: model.label || formatModelName(model.id),
-  });
+          requiresSubscription: false,
+          top: false,
+          isCustom: true
+        });
+      } else {
+        // If it already exists (rare case), mark it as a custom model
+        const existingModel = modelMap.get(model.id);
+        modelMap.set(model.id, {
+          ...existingModel,
+          isCustom: true
+        });
+      }
+    });
+  }
 
   // Convert map back to array
   const enhancedModelOptions = Array.from(modelMap.values());
@@ -255,10 +270,13 @@ export const ModelSelector: React.FC<ModelSelectorProps> = ({
     // First close the dialog to prevent UI issues
     closeCustomModelDialog();
 
+    // Create the new model object
+    const newModel = { id: modelId, label: modelLabel };
+
     // Update models array (add new or update existing)
     const updatedModels = dialogMode === 'add'
-      ? [...customModels, { id: modelId, label: modelLabel }]
+      ? [...customModels, newModel]
-      : customModels.map(model => model.id === editingModelId ? { id: modelId, label: modelLabel } : model);
+      : customModels.map(model => model.id === editingModelId ? newModel : model);
 
     // Save to localStorage first
     try {
@@ -270,6 +288,11 @@ export const ModelSelector: React.FC<ModelSelectorProps> = ({
     // Update state with new models
     setCustomModels(updatedModels);
 
+    // Refresh custom models in the parent hook if the function is available
+    if (refreshCustomModels) {
+      refreshCustomModels();
+    }
+
     // Handle model selection changes
     if (dialogMode === 'add') {
       // Always select newly added models
@@ -334,6 +357,11 @@ export const ModelSelector: React.FC<ModelSelectorProps> = ({
     // Update state with the new list
     setCustomModels(updatedCustomModels);
 
+    // Refresh custom models in the parent hook if the function is available
+    if (refreshCustomModels) {
+      refreshCustomModels();
+    }
+
     // Check if we need to change the selected model
     if (selectedModel === modelId) {
       const defaultModel = isLocalMode() ? DEFAULT_PREMIUM_MODEL_ID : DEFAULT_FREE_MODEL_ID;
@@ -345,18 +373,29 @@ export const ModelSelector: React.FC<ModelSelectorProps> = ({
       }
     }
 
-    // Force dropdown to close to ensure proper refresh on next open
+    // Force dropdown to close
     setIsOpen(false);
 
-    // Force a UI refresh by scheduling a state update after React completes this render cycle
+    // Update the modelMap and recreate enhancedModelOptions on next render
+    // This will force a complete refresh of the model list
     setTimeout(() => {
-      setIsOpen(false);
+      // Force React to fully re-evaluate the component with fresh data
-    }, 0);
+      setHighlightedIndex(-1);
+
+      // Reopen dropdown with fresh data if it was open
+      if (isOpen) {
+        setIsOpen(false);
+        setTimeout(() => setIsOpen(true), 50);
+      }
+    }, 10);
   };
 
   const renderModelOption = (opt: any, index: number) => {
-    // Custom models are always accessible in local mode
+    // More accurate check for custom models - use the actual customModels array
-    const isCustom = opt.isCustom || customModels.some(model => model.id === opt.id);
+    // from both the opt.isCustom flag and by checking if it exists in customModels
+    const isCustom = Boolean(opt.isCustom) ||
+      (isLocalMode() && customModels.some(model => model.id === opt.id));
+
     const accessible = isCustom ? true : canAccessModel(opt.id);
 
     // Fix the highlighting logic to use the index parameter instead of searching in filteredOptions
@@ -460,7 +499,7 @@ export const ModelSelector: React.FC<ModelSelectorProps> = ({
       setIsOpen(false);
       setTimeout(() => setIsOpen(true), 10);
     }
-  }, [customModels]);
+  }, [customModels, modelOptions]); // Also depend on modelOptions to refresh when parent changes
 
   return (
     <div className="relative">