ruby_llm 0.1.0.pre35 → 0.1.0.pre36

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 586cc8473d06aebb4eb69ee646c8ad90250a522eac80d3471a22ef05e601b9de
-  data.tar.gz: 4f060d755a37373230d1364dda3fc8b74fabdf4128206eb4faf2b11fd02090ac
+  metadata.gz: ed17bc0b342342484bd1e92b57f3b88a74d77cd43daeb1a569b132188025bafd
+  data.tar.gz: '0970e337a393e85cff88a449e723aa7c3e3e189b1923e16cf619f95d1bf65b03'
 SHA512:
-  metadata.gz: 8e10930fd06b2dbdbaca81aa93f45cde41e8d1199173597e29d94af7ea073db470a91be9869f89b622aebc28d915510baf070ac6655ddbbe1591a94fbef7b2f8
-  data.tar.gz: de99a1e6173f25b3c3bc7222910f6d1c8771b4c4f1062e8970f9472c7f551543c345cdffb1b3733ebd510e2351c27b5b274d483ba7047a32d90884b62577cb18
+  metadata.gz: 1e8b35980c57cd61e10c50b3eb97dbedbe138e2720b11940646c74807ada4c8b260aff3cbee8ed84abfa4b998d3d2747ff5f2421731fcf73ea038306774a2dc9
+  data.tar.gz: f0211b49713b10e00f1070fa7abb00a151ee7084e3f5e14b047be2c567afe4fcd5ade1ce662697a776470ab4e6c21e98ab1c453526e36b4ec29d6f7f6b9f6c0c

data/lib/ruby_llm/models.json

@@ -11,8 +11,8 @@
     "supports_vision": false,
     "supports_functions": false,
     "supports_json_mode": false,
-    "input_price_per_million": 0.075,
-    "output_price_per_million": 0.3,
+    "input_price_per_million": 0.0,
+    "output_price_per_million": 0.0,
     "metadata": {
       "object": "model",
       "owned_by": "google"
@@ -23,8 +23,8 @@
     "created_at": "2023-08-21T18:16:55+02:00",
     "display_name": "Babbage 002",
     "provider": "openai",
-    "context_window": 4096,
-    "max_tokens": 4096,
+    "context_window": 16384,
+    "max_tokens": 16384,
     "type": "chat",
     "family": "babbage",
     "supports_vision": false,
@@ -80,7 +80,7 @@
     "created_at": "2023-07-11T00:00:00Z",
     "display_name": "Claude 2.0",
     "provider": "anthropic",
-    "context_window": 100000,
+    "context_window": 200000,
     "max_tokens": 4096,
     "type": "chat",
     "family": "claude2",
@@ -96,7 +96,7 @@
     "created_at": "2023-11-21T00:00:00Z",
     "display_name": "Claude 2.1",
     "provider": "anthropic",
-    "context_window": 100000,
+    "context_window": 200000,
     "max_tokens": 4096,
     "type": "chat",
     "family": "claude2",
@@ -116,7 +116,7 @@
     "max_tokens": 8192,
     "type": "chat",
     "family": "claude35_haiku",
-    "supports_vision": false,
+    "supports_vision": true,
     "supports_functions": true,
     "supports_json_mode": true,
     "input_price_per_million": 0.8,
@@ -161,9 +161,9 @@
     "display_name": "Claude 3.7 Sonnet",
     "provider": "anthropic",
     "context_window": 200000,
-    "max_tokens": 4096,
+    "max_tokens": 8192,
     "type": "chat",
-    "family": "claude2",
+    "family": "claude37_sonnet",
     "supports_vision": true,
     "supports_functions": true,
     "supports_json_mode": true,
@@ -262,8 +262,8 @@
     "created_at": "2023-08-21T18:11:41+02:00",
     "display_name": "Davinci 002",
     "provider": "openai",
-    "context_window": 4096,
-    "max_tokens": 4096,
+    "context_window": 16384,
+    "max_tokens": 16384,
     "type": "chat",
     "family": "davinci",
     "supports_vision": false,
@@ -857,7 +857,7 @@
     "family": "gemini20_flash_lite",
     "supports_vision": true,
     "supports_functions": false,
-    "supports_json_mode": true,
+    "supports_json_mode": false,
     "input_price_per_million": 0.075,
     "output_price_per_million": 0.3,
     "metadata": {
@@ -876,7 +876,7 @@
     "family": "gemini20_flash_lite",
     "supports_vision": true,
     "supports_functions": false,
-    "supports_json_mode": true,
+    "supports_json_mode": false,
     "input_price_per_million": 0.075,
     "output_price_per_million": 0.3,
     "metadata": {
@@ -895,7 +895,7 @@
     "family": "gemini20_flash_lite",
     "supports_vision": true,
     "supports_functions": false,
-    "supports_json_mode": true,
+    "supports_json_mode": false,
     "input_price_per_million": 0.075,
     "output_price_per_million": 0.3,
     "metadata": {
@@ -914,7 +914,7 @@
     "family": "gemini20_flash_lite",
     "supports_vision": true,
     "supports_functions": false,
-    "supports_json_mode": true,
+    "supports_json_mode": false,
     "input_price_per_million": 0.075,
     "output_price_per_million": 0.3,
     "metadata": {
@@ -1650,7 +1650,7 @@
     "display_name": "GPT-4o-Mini Realtime Preview",
     "provider": "openai",
     "context_window": 128000,
-    "max_tokens": 16384,
+    "max_tokens": 4096,
     "type": "chat",
     "family": "gpt4o_mini_realtime",
     "supports_vision": true,
@@ -1669,7 +1669,7 @@
     "display_name": "GPT-4o-Mini Realtime Preview 20241217",
     "provider": "openai",
     "context_window": 128000,
-    "max_tokens": 16384,
+    "max_tokens": 4096,
     "type": "chat",
     "family": "gpt4o_mini_realtime",
     "supports_vision": true,
@@ -1685,10 +1685,10 @@
   {
     "id": "gpt-4o-realtime-preview",
     "created_at": "2024-09-30T03:33:18+02:00",
-    "display_name": "GPT-4o Realtime Preview",
+    "display_name": "GPT-4o-Realtime Preview",
     "provider": "openai",
     "context_window": 128000,
-    "max_tokens": 16384,
+    "max_tokens": 4096,
     "type": "chat",
     "family": "gpt4o_realtime",
     "supports_vision": true,
@@ -1704,10 +1704,10 @@
   {
     "id": "gpt-4o-realtime-preview-2024-10-01",
     "created_at": "2024-09-24T00:49:26+02:00",
-    "display_name": "GPT-4o Realtime Preview 20241001",
+    "display_name": "GPT-4o-Realtime Preview 20241001",
     "provider": "openai",
     "context_window": 128000,
-    "max_tokens": 16384,
+    "max_tokens": 4096,
     "type": "chat",
     "family": "gpt4o_realtime",
     "supports_vision": true,
@@ -1723,10 +1723,10 @@
   {
     "id": "gpt-4o-realtime-preview-2024-12-17",
     "created_at": "2024-12-11T20:30:30+01:00",
-    "display_name": "GPT-4o Realtime Preview 20241217",
+    "display_name": "GPT-4o-Realtime Preview 20241217",
     "provider": "openai",
     "context_window": 128000,
-    "max_tokens": 16384,
+    "max_tokens": 4096,
     "type": "chat",
     "family": "gpt4o_realtime",
     "supports_vision": true,
@@ -1820,7 +1820,7 @@
     "created_at": "2024-09-06T20:56:48+02:00",
     "display_name": "O1-Mini",
     "provider": "openai",
-    "context_window": 200000,
+    "context_window": 128000,
     "max_tokens": 4096,
     "type": "chat",
     "family": "o1_mini",
@@ -1839,7 +1839,7 @@
     "created_at": "2024-09-06T20:56:19+02:00",
     "display_name": "O1-Mini 20240912",
     "provider": "openai",
-    "context_window": 200000,
+    "context_window": 128000,
     "max_tokens": 65536,
     "type": "chat",
     "family": "o1_mini",
@@ -1894,7 +1894,7 @@
   {
     "id": "omni-moderation-2024-09-26",
     "created_at": "2024-11-27T20:07:46+01:00",
-    "display_name": "Omni Moderation 20240926",
+    "display_name": "Omni-Moderation 20240926",
     "provider": "openai",
     "context_window": 4096,
     "max_tokens": 4096,
@@ -1913,7 +1913,7 @@
   {
     "id": "omni-moderation-latest",
     "created_at": "2024-11-15T17:47:45+01:00",
-    "display_name": "Omni Moderation Latest",
+    "display_name": "Omni-Moderation Latest",
     "provider": "openai",
     "context_window": 4096,
     "max_tokens": 4096,

data/lib/ruby_llm/providers/anthropic/capabilities.rb

@@ -7,45 +7,74 @@ module RubyLLM
       module Capabilities
         module_function
 
-        def determine_context_window(model_id)
-          case model_id
-          when /claude-3/ then 200_000
-          else 100_000
-          end
+        # Determines the context window size for a given model
+        # @param model_id [String] the model identifier
+        # @return [Integer] the context window size in tokens
+        def determine_context_window(_model_id)
+          # All Claude 3 and 3.5 and 3.7 models have 200K token context windows
+          200_000
         end
 
+        # Determines the maximum output tokens for a given model
+        # @param model_id [String] the model identifier
+        # @return [Integer] the maximum output tokens
         def determine_max_tokens(model_id)
           case model_id
+          when /claude-3-7-sonnet/ then 8_192 # Can be increased to 64K with extended thinking
           when /claude-3-5/ then 8_192
-          else 4_096
+          else 4_096 # Claude 3 Opus and Haiku
           end
         end
 
+        # Gets the input price per million tokens for a given model
+        # @param model_id [String] the model identifier
+        # @return [Float] the price per million tokens for input
         def get_input_price(model_id)
           PRICES.dig(model_family(model_id), :input) || default_input_price
         end
 
+        # Gets the output price per million tokens for a given model
+        # @param model_id [String] the model identifier
+        # @return [Float] the price per million tokens for output
         def get_output_price(model_id)
           PRICES.dig(model_family(model_id), :output) || default_output_price
         end
 
+        # Determines if a model supports vision capabilities
+        # @param model_id [String] the model identifier
+        # @return [Boolean] true if the model supports vision
         def supports_vision?(model_id)
-          return false if model_id.match?(/claude-3-5-haiku/)
-          return false if model_id.match?(/claude-[12]/)
-
-          true
+          # All Claude 3, 3.5, and 3.7 models support vision
+          !model_id.match?(/claude-[12]/)
         end
 
+        # Determines if a 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.include?('claude-3')
+          model_id.match?(/claude-3/)
         end
 
+        # Determines if a model supports JSON mode
+        # @param model_id [String] the model identifier
+        # @return [Boolean] true if the model supports JSON mode
         def supports_json_mode?(model_id)
-          model_id.include?('claude-3')
+          model_id.match?(/claude-3/)
+        end
+
+        # Determines if a model supports extended thinking
+        # @param model_id [String] the model identifier
+        # @return [Boolean] true if the model supports extended thinking
+        def supports_extended_thinking?(model_id)
+          model_id.match?(/claude-3-7-sonnet/)
         end
 
+        # Determines the model family for a given model ID
+        # @param model_id [String] the model identifier
+        # @return [Symbol] the model family identifier
         def model_family(model_id)
           case model_id
+          when /claude-3-7-sonnet/  then :claude37_sonnet
           when /claude-3-5-sonnet/  then :claude35_sonnet
           when /claude-3-5-haiku/   then :claude35_haiku
           when /claude-3-opus/      then :claude3_opus
@@ -55,23 +84,32 @@ module RubyLLM
           end
         end
 
+        # Returns the model type
+        # @param model_id [String] the model identifier (unused but kept for API consistency)
+        # @return [String] the model type, always 'chat' for Anthropic models
         def model_type(_)
           'chat'
         end
 
+        # Pricing information for Anthropic models (per million tokens)
         PRICES = {
-          claude35_sonnet: { input: 3.0, output: 15.0 }, # $3.00/$15.00 per million tokens
-          claude35_haiku: { input: 0.80, output: 4.0 }, # $0.80/$4.00 per million tokens
-          claude3_opus: { input: 15.0, output: 75.0 }, # $15.00/$75.00 per million tokens
-          claude3_sonnet: { input: 3.0, output: 15.0 }, # $3.00/$15.00 per million tokens
-          claude3_haiku: { input: 0.25, output: 1.25 }, # $0.25/$1.25 per million tokens
-          claude2: { input: 3.0,  output: 15.0 } # Default pricing for Claude 2.x models
+          claude37_sonnet: { input: 3.0, output: 15.0 },   # $3.00/$15.00 per million tokens
+          claude35_sonnet: { input: 3.0, output: 15.0 },   # $3.00/$15.00 per million tokens
+          claude35_haiku: { input: 0.80, output: 4.0 },    # $0.80/$4.00 per million tokens
+          claude3_opus: { input: 15.0, output: 75.0 },     # $15.00/$75.00 per million tokens
+          claude3_sonnet: { input: 3.0, output: 15.0 },    # $3.00/$15.00 per million tokens
+          claude3_haiku: { input: 0.25, output: 1.25 },    # $0.25/$1.25 per million tokens
+          claude2: { input: 3.0, output: 15.0 }            # Default pricing for Claude 2.x models
         }.freeze
 
+        # Default input price if model not found in PRICES
+        # @return [Float] default price per million tokens for input
         def default_input_price
           3.0
         end
 
+        # Default output price if model not found in PRICES
+        # @return [Float] default price per million tokens for output
         def default_output_price
           15.0
         end

data/lib/ruby_llm/providers/deepseek/capabilities.rb

@@ -7,6 +7,9 @@ module RubyLLM
       module Capabilities
         module_function
 
+        # Returns the context window size for the given model
+        # @param model_id [String] the model identifier
+        # @return [Integer] the context window size in tokens
         def context_window_for(model_id)
           case model_id
           when /deepseek-(?:chat|reasoner)/ then 64_000
@@ -14,6 +17,9 @@ module RubyLLM
           end
         end
 
+        # Returns the maximum number of tokens that can be generated
+        # @param model_id [String] the model identifier
+        # @return [Integer] the maximum number of tokens
         def max_tokens_for(model_id)
           case model_id
           when /deepseek-(?:chat|reasoner)/ then 8_192
@@ -21,30 +27,51 @@ module RubyLLM
           end
         end
 
+        # Returns the price per million tokens for input (cache miss)
+        # @param model_id [String] the model identifier
+        # @return [Float] the price per million tokens in USD
         def input_price_for(model_id)
           PRICES.dig(model_family(model_id), :input_miss) || default_input_price
         end
 
+        # Returns the price per million tokens for output
+        # @param model_id [String] the model identifier
+        # @return [Float] the price per million tokens in USD
         def output_price_for(model_id)
           PRICES.dig(model_family(model_id), :output) || default_output_price
         end
 
+        # Returns the price per million tokens for input with cache hit
+        # @param model_id [String] the model identifier
+        # @return [Float] the price per million tokens in USD
         def cache_hit_price_for(model_id)
           PRICES.dig(model_family(model_id), :input_hit) || default_cache_hit_price
         end
 
+        # Determines if the model supports vision capabilities
+        # @param model_id [String] the model identifier
+        # @return [Boolean] true if the model supports vision
         def supports_vision?(_model_id)
           false # DeepSeek models don't currently support vision
         end
 
+        # Determines if the model supports function calling
+        # @param model_id [String] the model identifier
+        # @return [Boolean] true if the model supports function calling
         def supports_functions?(model_id)
           model_id.match?(/deepseek-chat/) # Only deepseek-chat supports function calling
         end
 
+        # Determines if the model supports JSON mode
+        # @param model_id [String] the model identifier
+        # @return [Boolean] true if the model supports JSON mode
         def supports_json_mode?(model_id)
           model_id.match?(/deepseek-chat/) # Only deepseek-chat supports JSON mode
         end
 
+        # Returns a formatted display name for the model
+        # @param model_id [String] the model identifier
+        # @return [String] the formatted display name
         def format_display_name(model_id)
           case model_id
           when 'deepseek-chat' then 'DeepSeek V3'
@@ -56,10 +83,16 @@ module RubyLLM
           end
         end
 
+        # Returns the model type
+        # @param model_id [String] the model identifier
+        # @return [String] the model type (e.g., 'chat')
         def model_type(_model_id)
           'chat' # All DeepSeek models are chat models
         end
 
+        # Returns the model family
+        # @param model_id [String] the model identifier
+        # @return [Symbol] the model family
         def model_family(model_id)
           case model_id
           when /deepseek-chat/ then :chat
@@ -84,14 +117,20 @@ module RubyLLM
 
         private
 
+        # Default input price when model family can't be determined
+        # @return [Float] the default input price
         def default_input_price
           0.27 # Default to chat cache miss price
         end
 
+        # Default output price when model family can't be determined
+        # @return [Float] the default output price
         def default_output_price
           1.10 # Default to chat output price
         end
 
+        # Default cache hit price when model family can't be determined
+        # @return [Float] the default cache hit price
         def default_cache_hit_price
           0.07 # Default to chat cache hit price
         end

data/lib/ruby_llm/providers/gemini/capabilities.rb

@@ -7,25 +7,34 @@ module RubyLLM
       module Capabilities # rubocop:disable Metrics/ModuleLength
         module_function
 
+        # Returns the context window size (input token limit) for the given model
+        # @param model_id [String] the model identifier
+        # @return [Integer] the context window size in tokens
         def context_window_for(model_id)
           case model_id
           when /gemini-2\.0-flash/, /gemini-1\.5-flash/ then 1_048_576
           when /gemini-1\.5-pro/ then 2_097_152
-          when /text-embedding/, /embedding-001/ then 2_048
+          when /text-embedding-004/, /embedding-001/ then 2_048
           when /aqa/ then 7_168
           else 32_768 # Sensible default for unknown models
           end
         end
 
+        # Returns the maximum output tokens for the given model
+        # @param model_id [String] the model identifier
+        # @return [Integer] the maximum output tokens
         def max_tokens_for(model_id)
           case model_id
           when /gemini-2\.0-flash/, /gemini-1\.5/ then 8_192
-          when /text-embedding/, /embedding-001/ then 768 # Output dimension size for embeddings
+          when /text-embedding-004/, /embedding-001/ then 768 # Output dimension size for embeddings
           when /aqa/ then 1_024
           else 4_096 # Sensible default
           end
         end
 
+        # Returns the input price per million tokens for the given model
+        # @param model_id [String] the model identifier
+        # @return [Float] the price per million tokens in USD
         def input_price_for(model_id)
           base_price = PRICES.dig(pricing_family(model_id), :input) || default_input_price
           return base_price unless long_context_model?(model_id)
@@ -34,6 +43,9 @@ module RubyLLM
           context_length(model_id) > 128_000 ? base_price * 2 : base_price
         end
 
+        # Returns the output price per million tokens for the given model
+        # @param model_id [String] the model identifier
+        # @return [Float] the price per million tokens in USD
         def output_price_for(model_id)
           base_price = PRICES.dig(pricing_family(model_id), :output) || default_output_price
           return base_price unless long_context_model?(model_id)
@@ -42,6 +54,9 @@ module RubyLLM
           context_length(model_id) > 128_000 ? base_price * 2 : base_price
         end
 
+        # Determines if the model supports vision (image/video) inputs
+        # @param model_id [String] the model identifier
+        # @return [Boolean] true if the model supports vision inputs
         def supports_vision?(model_id)
           return false if model_id.match?(/text-embedding|embedding-001|aqa/)
           return false if model_id.match?(/gemini-1\.0/)
@@ -49,6 +64,9 @@ module RubyLLM
           model_id.match?(/gemini-[12]\.[05]/)
         end
 
+        # Determines if the model supports function calling
+        # @param model_id [String] the model identifier
+        # @return [Boolean] true if the model supports function calling
         def supports_functions?(model_id)
           return false if model_id.match?(/text-embedding|embedding-001|aqa/)
           return false if model_id.match?(/flash-lite/)
@@ -57,13 +75,20 @@ module RubyLLM
           model_id.match?(/gemini-[12]\.[05]-(?:pro|flash)(?!-lite)/)
         end
 
+        # Determines if the model supports JSON mode
+        # @param model_id [String] the model identifier
+        # @return [Boolean] true if the model supports JSON mode
         def supports_json_mode?(model_id)
           return false if model_id.match?(/text-embedding|embedding-001|aqa/)
           return false if model_id.match?(/gemini-1\.0/)
+          return false if model_id.match?(/gemini-2\.0-flash-lite/)
 
           model_id.match?(/gemini-\d/)
         end
 
+        # Formats the model ID into a human-readable display name
+        # @param model_id [String] the model identifier
+        # @return [String] the formatted display name
         def format_display_name(model_id)
           model_id
             .delete_prefix('models/')
@@ -76,20 +101,32 @@ module RubyLLM
             .strip
         end
 
+        # Determines if the model supports context caching
+        # @param model_id [String] the model identifier
+        # @return [Boolean] true if the model supports caching
         def supports_caching?(model_id)
           return false if model_id.match?(/flash-lite|gemini-1\.0/)
 
           model_id.match?(/gemini-[12]\.[05]/)
         end
 
+        # Determines if the model supports tuning
+        # @param model_id [String] the model identifier
+        # @return [Boolean] true if the model supports tuning
         def supports_tuning?(model_id)
           model_id.match?(/gemini-1\.5-flash/)
         end
 
+        # Determines if the model supports audio inputs
+        # @param model_id [String] the model identifier
+        # @return [Boolean] true if the model supports audio inputs
         def supports_audio?(model_id)
           model_id.match?(/gemini-[12]\.[05]/)
         end
 
+        # Returns the type of model (chat, embedding, image)
+        # @param model_id [String] the model identifier
+        # @return [String] the model type
         def model_type(model_id)
           case model_id
           when /text-embedding|embedding/ then 'embedding'
@@ -98,6 +135,9 @@ module RubyLLM
           end
         end
 
+        # Returns the model family identifier
+        # @param model_id [String] the model identifier
+        # @return [String] the model family identifier
         def model_family(model_id) # rubocop:disable Metrics/CyclomaticComplexity,Metrics/MethodLength
           case model_id
           when /gemini-2\.0-flash-lite/ then 'gemini20_flash_lite'
@@ -113,7 +153,10 @@ module RubyLLM
           end
         end
 
-        def pricing_family(model_id) # rubocop:disable Metrics/CyclomaticComplexity
+        # Returns the pricing family identifier for the model
+        # @param model_id [String] the model identifier
+        # @return [Symbol] the pricing family identifier
+        def pricing_family(model_id) # rubocop:disable Metrics/CyclomaticComplexity,Metrics/MethodLength
           case model_id
           when /gemini-2\.0-flash-lite/ then :flash_lite_2 # rubocop:disable Naming/VariableNumber
           when /gemini-2\.0-flash/ then :flash_2 # rubocop:disable Naming/VariableNumber
@@ -122,18 +165,26 @@ module RubyLLM
           when /gemini-1\.5-pro/ then :pro
           when /gemini-1\.0-pro/ then :pro_1_0 # rubocop:disable Naming/VariableNumber
           when /text-embedding|embedding/ then :embedding
+          when /aqa/ then :aqa
           else :base
           end
         end
 
+        # Determines if the model supports long context
+        # @param model_id [String] the model identifier
+        # @return [Boolean] true if the model supports long context
         def long_context_model?(model_id)
           model_id.match?(/gemini-1\.5-(?:pro|flash)/)
         end
 
+        # Returns the context length for the model
+        # @param model_id [String] the model identifier
+        # @return [Integer] the context length in tokens
         def context_length(model_id)
           context_window_for(model_id)
         end
 
+        # Pricing information for Gemini models (per 1M tokens in USD)
         PRICES = {
           flash_2: { # Gemini 2.0 Flash # rubocop:disable Naming/VariableNumber
             input: 0.10,
@@ -152,19 +203,22 @@ module RubyLLM
             input: 0.075,
             output: 0.30,
             cache: 0.01875,
-            cache_storage: 1.00
+            cache_storage: 1.00,
+            grounding_search: 35.00 # per 1K requests
           },
           flash_8b: { # Gemini 1.5 Flash 8B
             input: 0.0375,
             output: 0.15,
             cache: 0.01,
-            cache_storage: 0.25
+            cache_storage: 0.25,
+            grounding_search: 35.00 # per 1K requests
           },
           pro: { # Gemini 1.5 Pro
             input: 1.25,
             output: 5.0,
             cache: 0.3125,
-            cache_storage: 4.50
+            cache_storage: 4.50,
+            grounding_search: 35.00 # per 1K requests
           },
           pro_1_0: { # Gemini 1.0 Pro # rubocop:disable Naming/VariableNumber
             input: 0.50,
@@ -173,15 +227,23 @@ module RubyLLM
           embedding: { # Text Embedding models
             input: 0.00,
             output: 0.00
+          },
+          aqa: { # AQA model
+            input: 0.00,
+            output: 0.00
           }
         }.freeze
 
+        # Default input price for unknown models
+        # @return [Float] the default input price per million tokens
         def default_input_price
           0.075 # Default to Flash pricing
         end
 
+        # Default output price for unknown models
+        # @return [Float] the default output price per million tokens
         def default_output_price
-          0.30  # Default to Flash pricing
+          0.30 # Default to Flash pricing
         end
       end
     end

data/lib/ruby_llm/providers/openai/capabilities.rb

@@ -7,76 +7,113 @@ module RubyLLM
       module Capabilities # rubocop:disable Metrics/ModuleLength
         module_function
 
+        # Returns the context window size for the given model ID
+        # @param model_id [String] the model identifier
+        # @return [Integer] the context window size in tokens
         def context_window_for(model_id)
           case model_id
-          when /o[13]-mini/, /o3-mini-2025/         then 200_000
-          when /o1-2024/                            then 200_000
-          when /gpt-4o/, /gpt-4-turbo/             then 128_000
-          when /gpt-4-0[0-9]{3}/                   then 8_192
-          when /gpt-3.5-turbo-instruct/            then 4_096
-          when /gpt-3.5/                           then 16_385
+          when /o1-2024/, /o3-mini/, /o3-mini-2025/ then 200_000
+          when /gpt-4o/, /gpt-4o-mini/, /gpt-4-turbo/, /o1-mini/ then 128_000
+          when /gpt-4-0[0-9]{3}/ then 8_192
+          when /gpt-3.5-turbo-instruct/ then 4_096
+          when /gpt-3.5/ then 16_385
+          when /babbage-002/, /davinci-002/ then 16_384
           else 4_096
           end
         end
 
-        def max_tokens_for(model_id) # rubocop:disable Metrics/CyclomaticComplexity
+        # Returns the maximum output tokens for the given model ID
+        # @param model_id [String] the model identifier
+        # @return [Integer] the maximum output tokens
+        def max_tokens_for(model_id) # rubocop:disable Metrics/CyclomaticComplexity,Metrics/MethodLength
           case model_id
-          when /o1-2024/, /o3-mini/                then 100_000
-          when /o1-mini-2024/                      then 65_536
-          when /gpt-4o-2024-05-13/                then 4_096
-          when /gpt-4o/, /gpt-4o-mini/            then 16_384
-          when /gpt-4o-realtime/                  then 4_096
-          when /gpt-4-0[0-9]{3}/                  then 8_192
-          when /gpt-3.5-turbo/                    then 4_096
+          when /o1-2024/, /o3-mini/, /o3-mini-2025/ then 100_000
+          when /o1-mini-2024/ then 65_536
+          when /gpt-4o-2024-05-13/ then 4_096
+          when /gpt-4o-realtime/, /gpt-4o-mini-realtime/ then 4_096
+          when /gpt-4o/, /gpt-4o-mini/, /gpt-4o-audio/, /gpt-4o-mini-audio/ then 16_384
+          when /gpt-4-0[0-9]{3}/ then 8_192
+          when /gpt-4-turbo/, /gpt-3.5-turbo/ then 4_096
+          when /babbage-002/, /davinci-002/ then 16_384
           else 4_096
           end
         end
 
+        # Returns the input price per million tokens for the given model ID
+        # @param model_id [String] the model identifier
+        # @return [Float] the price per million tokens for input
         def input_price_for(model_id)
           PRICES.dig(model_family(model_id), :input) || default_input_price
         end
 
+        # Returns the output price per million tokens for the given model ID
+        # @param model_id [String] the model identifier
+        # @return [Float] the price per million tokens for output
         def output_price_for(model_id)
           PRICES.dig(model_family(model_id), :output) || default_output_price
         end
 
+        # Determines if the model supports vision capabilities
+        # @param model_id [String] the model identifier
+        # @return [Boolean] true if the model supports vision
         def supports_vision?(model_id)
           model_id.match?(/gpt-4o|o1/) || model_id.match?(/gpt-4-(?!0314|0613)/)
         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.include?('instruct')
         end
 
+        # Determines if the model supports audio input/output
+        # @param model_id [String] the model identifier
+        # @return [Boolean] true if the model supports audio
         def supports_audio?(model_id)
           model_id.match?(/audio-preview|realtime-preview|whisper|tts/)
         end
 
+        # Determines if the model supports JSON mode
+        # @param model_id [String] the model identifier
+        # @return [Boolean] true if the model supports JSON mode
         def supports_json_mode?(model_id)
           model_id.match?(/gpt-4-\d{4}-preview/) ||
             model_id.include?('turbo') ||
             model_id.match?(/gpt-3.5-turbo-(?!0301|0613)/)
         end
 
+        # Formats the model ID into a human-readable display name
+        # @param model_id [String] the model identifier
+        # @return [String] the formatted display name
         def format_display_name(model_id)
           model_id.then { |id| humanize(id) }
                   .then { |name| apply_special_formatting(name) }
         end
 
+        # Determines the type of model
+        # @param model_id [String] the model identifier
+        # @return [String] the model type (chat, embedding, image, audio, moderation)
         def model_type(model_id)
           case model_id
           when /text-embedding|embedding/ then 'embedding'
           when /dall-e/ then 'image'
           when /tts|whisper/ then 'audio'
-          when /omni-moderation/ then 'moderation'
+          when /omni-moderation|text-moderation/ then 'moderation'
           else 'chat'
           end
         end
 
+        # Determines if the model supports structured output
+        # @param model_id [String] the model identifier
+        # @return [Boolean] true if the model supports structured output
         def supports_structured_output?(model_id)
-          model_id.match?(/gpt-4o|o[13]-mini|o1/)
+          model_id.match?(/gpt-4o|o[13]-mini|o1|o3-mini/)
         end
 
+        # Determines the model family for pricing and capability lookup
+        # @param model_id [String] the model identifier
+        # @return [Symbol] the model family identifier
         def model_family(model_id) # rubocop:disable Metrics/AbcSize,Metrics/CyclomaticComplexity,Metrics/MethodLength
           case model_id
           when /o3-mini/ then 'o3_mini'
@@ -100,13 +137,14 @@ module RubyLLM
           when /tts-1-hd/ then 'tts1_hd'
           when /tts-1/ then 'tts1'
           when /whisper/ then 'whisper1'
-          when /omni-moderation/ then 'moderation'
+          when /omni-moderation|text-moderation/ then 'moderation'
           when /babbage/ then 'babbage'
           when /davinci/ then 'davinci'
           else 'other'
           end
         end
 
+        # Pricing information for OpenAI models (per million tokens unless otherwise specified)
         PRICES = {
           o1: { input: 15.0, cached_input: 7.5, output: 60.0 },
           o1_mini: { input: 1.10, cached_input: 0.55, output: 4.40 },
@@ -150,19 +188,27 @@ module RubyLLM
           embedding2: { price: 0.10 },
           davinci: { input: 2.0, output: 2.0 },
           babbage: { input: 0.40, output: 0.40 },
-          tts1: { price: 15.0 },
-          tts1_hd: { price: 30.0 },
-          whisper1: { price: 0.006 }
+          tts1: { price: 15.0 }, # per million characters
+          tts1_hd: { price: 30.0 }, # per million characters
+          whisper1: { price: 0.006 }, # per minute
+          moderation: { price: 0.0 } # free
         }.freeze
 
+        # Default input price when model-specific pricing is not available
+        # @return [Float] the default price per million tokens
         def default_input_price
           0.50
         end
 
+        # Default output price when model-specific pricing is not available
+        # @return [Float] the default price per million tokens
         def default_output_price
           1.50
         end
 
+        # Converts a model ID to a human-readable format
+        # @param id [String] the model identifier
+        # @return [String] the humanized model name
         def humanize(id)
           id.tr('-', ' ')
             .split(' ')
@@ -170,18 +216,25 @@ module RubyLLM
             .join(' ')
         end
 
+        # Applies special formatting rules to model names
+        # @param name [String] the humanized model name
+        # @return [String] the specially formatted model name
         def apply_special_formatting(name) # rubocop:disable Metrics/MethodLength
           name
             .gsub(/(\d{4}) (\d{2}) (\d{2})/, '\1\2\3')
             .gsub(/^Gpt /, 'GPT-')
             .gsub(/^O([13]) /, 'O\1-')
+            .gsub(/^O3 Mini/, 'O3-Mini')
+            .gsub(/^O1 Mini/, 'O1-Mini')
             .gsub(/^Chatgpt /, 'ChatGPT-')
             .gsub(/^Tts /, 'TTS-')
             .gsub(/^Dall E /, 'DALL-E-')
             .gsub(/3\.5 /, '3.5-')
             .gsub(/4 /, '4-')
-            .gsub(/4o (?=Mini|Preview|Turbo|Audio)/, '4o-')
+            .gsub(/4o (?=Mini|Preview|Turbo|Audio|Realtime)/, '4o-')
             .gsub(/\bHd\b/, 'HD')
+            .gsub(/Omni Moderation/, 'Omni-Moderation')
+            .gsub(/Text Moderation/, 'Text-Moderation')
         end
       end
     end

data/lib/ruby_llm/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module RubyLLM
-  VERSION = '0.1.0.pre35'
+  VERSION = '0.1.0.pre36'
 end

data/lib/tasks/models.rake

@@ -16,6 +16,9 @@ PROVIDER_DOCS = {
   },
   deepseek: {
     models: 'https://api-docs.deepseek.com/quick_start/pricing/'
+  },
+  anthropic: {
+    models: 'https://docs.anthropic.com/en/docs/about-claude/models/all-models'
   }
 }.freeze
 
@@ -85,8 +88,10 @@ namespace :models do # rubocop:disable Metrics/BlockLength
     end
   end
 
-  desc 'Update model capabilities modules by scraping provider documentation'
+  desc 'Update model capabilities modules by scraping provider documentation (use PROVIDER=name to update only one)'
   task :update_capabilities do # rubocop:disable Metrics/BlockLength
+    # Check if a specific provider was requested
+    target_provider = ENV['PROVIDER']&.to_sym
     require 'ruby_llm'
     require 'fileutils'
 
@@ -97,8 +102,15 @@ namespace :models do # rubocop:disable Metrics/BlockLength
       config.gemini_api_key = ENV.fetch('GEMINI_API_KEY')
     end
 
+    # Filter providers if a specific one was requested
+    providers_to_process = if target_provider && PROVIDER_DOCS.key?(target_provider)
+                             { target_provider => PROVIDER_DOCS[target_provider] }
+                           else
+                             PROVIDER_DOCS
+                           end
+
     # Process each provider
-    PROVIDER_DOCS.each do |provider, urls| # rubocop:disable Metrics/BlockLength
+    providers_to_process.each do |provider, urls| # rubocop:disable Metrics/BlockLength
       puts "Processing #{provider}..."
 
       # Initialize our AI assistants
@@ -175,12 +187,22 @@ namespace :models do # rubocop:disable Metrics/BlockLength
 
         response = claude.ask(code_prompt)
 
+        # Extract Ruby code from Claude's response
+        puts "  Extracting Ruby code from Claude's response..."
+        ruby_code = nil
+
+        # Look for Ruby code block
+        ruby_code = Regexp.last_match(1).strip if response.content =~ /```ruby\s*(.*?)```/m
+
+        # Verify we found Ruby code
+        raise "No Ruby code block found in Claude's response" if ruby_code.nil? || ruby_code.empty?
+
         # Save the file
         file_path = "lib/ruby_llm/providers/#{provider}/capabilities.rb"
         puts "  Writing #{file_path}..."
 
         FileUtils.mkdir_p(File.dirname(file_path))
-        File.write(file_path, response.content)
+        File.write(file_path, ruby_code)
       rescue StandardError => e
         raise "Failed to process #{provider}: #{e.message}"
       end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: ruby_llm
 version: !ruby/object:Gem::Version
-  version: 0.1.0.pre35
+  version: 0.1.0.pre36
 platform: ruby
 authors:
 - Carmine Paolino