Try to make model capabilities more robust. Add to model update task

Author: OriPekelman

lib/ruby_llm/providers/mistral/capabilities.rb

@@ -53,19 +53,34 @@ def output_price_for(model_id)
         def supports_vision?(model_id)
           # Explicitly match the known vision-capable models
           vision_models = [
-            'pixtral-12b-latest',
-            'pixtral-large-latest',
-            'mistral-medium-latest',
-            'mistral-small-latest'
+            'pixtral-12b',
+            'pixtral-large',
+            'mistral-medium',
+            'mistral-small'
           ]
-          vision_models.any? { |id| model_id.include?(id) }
+          model_id = model_id.to_s.split('/').last
+          result = vision_models.any? { |id| model_id.match?(/^#{Regexp.escape(id)}/) }
+          result
         end
 
         # Determines if the model supports function calling
         # @param model_id [String] the model identifier
         # @return [Boolean] true if the model supports functions
         def supports_functions?(model_id)
-          !model_id.match?(/embed|moderation/)
+          function_calling_models = [
+            'mistral-large',
+            'mistral-medium',
+            'mistral-small',
+            'codestral',
+            'ministral-8b',
+            'ministral-3b',
+            'pixtral-12b',
+            'pixtral-large',
+            'mistral-nemo'
+          ]
+          model_id = model_id.to_s.split('/').last
+          result = function_calling_models.any? { |id| model_id.match?(/^#{Regexp.escape(id)}/) }
+          result
         end
 
         # Determines if the model supports audio input/output
@@ -168,6 +183,14 @@ def apply_special_formatting(name)
             .gsub("Mathstral ", "Mathstral-")
             .gsub("Embed ", "Embed-")
         end
+
+        def capabilities_for(model_id)
+          capabilities = ['streaming']
+          capabilities << 'function_calling' if supports_functions?(model_id)
+          capabilities << 'structured_output' if supports_json_mode?(model_id)
+          # Add more as needed (e.g., 'batch', 'caching', etc.)
+          capabilities
+        end
       end
     end
   end

lib/ruby_llm/providers/mistral/models.rb

@@ -33,11 +33,7 @@ def parse_model(model, slug, capabilities)
             },
             context_window: capabilities.context_window_for(id),
             max_tokens: capabilities.max_tokens_for(id),
-            supports_vision: capabilities.supports_vision?(id),
-            supports_functions: capabilities.supports_functions?(id),
-            supports_json_mode: capabilities.supports_json_mode?(id),
-            input_price_per_million: capabilities.input_price_for(id),
-            output_price_per_million: capabilities.output_price_for(id),
+            capabilities: capabilities.capabilities_for(id),
           )
         end
       end

lib/tasks/models_update.rake

@@ -23,6 +23,7 @@ def configure_from_env
     config.gemini_api_key = ENV.fetch('GEMINI_API_KEY', nil)
     config.deepseek_api_key = ENV.fetch('DEEPSEEK_API_KEY', nil)
     config.openrouter_api_key = ENV.fetch('OPENROUTER_API_KEY', nil)
+    config.mistral_api_key = ENV.fetch('MISTRAL_API_KEY', nil)
     configure_bedrock(config)
     config.request_timeout = 30
   end