ruby_llm 0.1.0.pre3 → 0.1.0.pre5

data/lib/ruby_llm/model_capabilities/openai.rb

@@ -2,95 +2,47 @@
 
 module RubyLLM
   module ModelCapabilities
-    class OpenAI < Base
-      def determine_context_window(model_id)
-        case model_id
-        when /gpt-4o/, /o1/, /gpt-4-turbo/
-          128_000
-        when /gpt-4-0[0-9]{3}/
-          8_192
-        when /gpt-3.5-turbo-instruct/
-          4_096
-        when /gpt-3.5/
-          16_385
-        else
-          4_096
-        end
-      end
+    module OpenAI
+      extend self
 
-      def determine_max_tokens(model_id)
+      def context_window_for(model_id)
         case model_id
-        when /o1-2024-12-17/
-          100_000
-        when /o1-mini-2024-09-12/
-          65_536
-        when /o1-preview-2024-09-12/
-          32_768
-        when /gpt-4o/, /gpt-4-turbo/
-          16_384
-        when /gpt-4-0[0-9]{3}/
-          8_192
-        when /gpt-3.5-turbo/
-          4_096
-        else
-          4_096
+        when /gpt-4o/, /o1/, /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
+        else                                   4_096
         end
       end
 
-      def get_input_price(model_id)
+      def max_tokens_for(model_id)
         case model_id
-        when /o1-2024/
-          15.0    # $15.00 per million tokens
-        when /o1-mini/
-          3.0     # $3.00 per million tokens
-        when /gpt-4o-realtime-preview/
-          5.0     # $5.00 per million tokens
-        when /gpt-4o-mini-realtime-preview/
-          0.60    # $0.60 per million tokens
-        when /gpt-4o-mini/
-          0.15    # $0.15 per million tokens
-        when /gpt-4o/
-          2.50    # $2.50 per million tokens
-        when /gpt-4-turbo/
-          10.0    # $10.00 per million tokens
-        when /gpt-3.5/
-          0.50    # $0.50 per million tokens
-        else
-          0.50    # Default to GPT-3.5 pricing
+        when /o1-2024-12-17/              then 100_000
+        when /o1-mini-2024-09-12/         then 65_536
+        when /o1-preview-2024-09-12/      then 32_768
+        when /gpt-4o/, /gpt-4-turbo/      then 16_384
+        when /gpt-4-0[0-9]{3}/           then 8_192
+        when /gpt-3.5-turbo/             then 4_096
+        else                                  4_096
         end
       end
 
-      def get_output_price(model_id)
-        case model_id
-        when /o1-2024/
-          60.0    # $60.00 per million tokens
-        when /o1-mini/
-          12.0    # $12.00 per million tokens
-        when /gpt-4o-realtime-preview/
-          20.0    # $20.00 per million tokens
-        when /gpt-4o-mini-realtime-preview/
-          2.40    # $2.40 per million tokens
-        when /gpt-4o-mini/
-          0.60    # $0.60 per million tokens
-        when /gpt-4o/
-          10.0    # $10.00 per million tokens
-        when /gpt-4-turbo/
-          30.0    # $30.00 per million tokens
-        when /gpt-3.5/
-          1.50    # $1.50 per million tokens
-        else
-          1.50    # Default to GPT-3.5 pricing
-        end
+      def input_price_for(model_id)
+        PRICES.dig(model_family(model_id), :input) || default_input_price
       end
 
-      def supports_functions?(model_id)
-        !model_id.include?('instruct')
+      def output_price_for(model_id)
+        PRICES.dig(model_family(model_id), :output) || default_output_price
       end
 
       def supports_vision?(model_id)
         model_id.include?('vision') || model_id.match?(/gpt-4-(?!0314|0613)/)
       end
 
+      def supports_functions?(model_id)
+        !model_id.include?('instruct')
+      end
+
       def supports_json_mode?(model_id)
         model_id.match?(/gpt-4-\d{4}-preview/) ||
           model_id.include?('turbo') ||
@@ -98,23 +50,64 @@ module RubyLLM
       end
 
       def format_display_name(model_id)
-        # First replace hyphens with spaces
-        name = model_id.tr('-', ' ')
-
-        # Capitalize each word
-        name = name.split(' ').map { |word| word.capitalize }.join(' ')
-
-        # Apply specific formatting rules
-        name.gsub(/(\d{4}) (\d{2}) (\d{2})/, '\1\2\3') # Convert dates to YYYYMMDD
-            .gsub(/^Gpt /, 'GPT-')
-            .gsub(/^O1 /, 'O1-')
-            .gsub(/^Chatgpt /, 'ChatGPT-')
-            .gsub(/^Tts /, 'TTS-')
-            .gsub(/^Dall E /, 'DALL-E-')
-            .gsub(/3\.5 /, '3.5-')
-            .gsub(/4 /, '4-')
-            .gsub(/4o (?=Mini|Preview|Turbo)/, '4o-')
-            .gsub(/\bHd\b/, 'HD')
+        model_id.then { |id| humanize(id) }
+                .then { |name| apply_special_formatting(name) }
+      end
+
+      private
+
+      def model_family(model_id)
+        case model_id
+        when /o1-2024/                   then :o1
+        when /o1-mini/                   then :o1_mini
+        when /gpt-4o-realtime-preview/   then :gpt4o_realtime
+        when /gpt-4o-mini-realtime/      then :gpt4o_mini_realtime
+        when /gpt-4o-mini/               then :gpt4o_mini
+        when /gpt-4o/                    then :gpt4o
+        when /gpt-4-turbo/               then :gpt4_turbo
+        when /gpt-3.5/                   then :gpt35
+        else :gpt35
+        end
+      end
+
+      PRICES = {
+        o1: { input: 15.0, output: 60.0 },
+        o1_mini: { input: 3.0, output: 12.0 },
+        gpt4o_realtime: { input: 5.0, output: 20.0 },
+        gpt4o_mini_realtime: { input: 0.60, output: 2.40 },
+        gpt4o_mini: { input: 0.15, output: 0.60 },
+        gpt4o: { input: 2.50, output: 10.0 },
+        gpt4_turbo: { input: 10.0, output: 30.0 },
+        gpt35: { input: 0.50, output: 1.50 }
+      }.freeze
+
+      def default_input_price
+        0.50
+      end
+
+      def default_output_price
+        1.50
+      end
+
+      def humanize(id)
+        id.tr('-', ' ')
+          .split(' ')
+          .map(&:capitalize)
+          .join(' ')
+      end
+
+      def apply_special_formatting(name)
+        name
+          .gsub(/(\d{4}) (\d{2}) (\d{2})/, '\1\2\3')
+          .gsub(/^Gpt /, 'GPT-')
+          .gsub(/^O1 /, 'O1-')
+          .gsub(/^Chatgpt /, 'ChatGPT-')
+          .gsub(/^Tts /, 'TTS-')
+          .gsub(/^Dall E /, 'DALL-E-')
+          .gsub(/3\.5 /, '3.5-')
+          .gsub(/4 /, '4-')
+          .gsub(/4o (?=Mini|Preview|Turbo)/, '4o-')
+          .gsub(/\bHd\b/, 'HD')
       end
     end
   end

data/lib/ruby_llm/model_info.rb

@@ -1,42 +1,51 @@
 # frozen_string_literal: true
 
+require 'time'
+
 module RubyLLM
   class ModelInfo
     attr_reader :id, :created_at, :display_name, :provider, :metadata,
                 :context_window, :max_tokens, :supports_vision, :supports_functions,
                 :supports_json_mode, :input_price_per_million, :output_price_per_million
 
-    def initialize(id:, created_at:, display_name:, provider:, context_window:, max_tokens:, supports_vision:,
-                   supports_functions:, supports_json_mode:, input_price_per_million:, output_price_per_million:, metadata: {})
-      @id = id
-      @created_at = created_at
-      @display_name = display_name
-      @provider = provider
-      @metadata = metadata
-      @context_window = context_window
-      @max_tokens = max_tokens
-      @supports_vision = supports_vision
-      @supports_functions = supports_functions
-      @supports_json_mode = supports_json_mode
-      @input_price_per_million = input_price_per_million
-      @output_price_per_million = output_price_per_million
+    def initialize(data)
+      @id = data[:id]
+      @created_at = data[:created_at].is_a?(String) ? Time.parse(data[:created_at]) : data[:created_at]
+      @display_name = data[:display_name]
+      @provider = data[:provider]
+      @context_window = data[:context_window]
+      @max_tokens = data[:max_tokens]
+      @supports_vision = data[:supports_vision]
+      @supports_functions = data[:supports_functions]
+      @supports_json_mode = data[:supports_json_mode]
+      @input_price_per_million = data[:input_price_per_million]
+      @output_price_per_million = data[:output_price_per_million]
+      @metadata = data[:metadata] || {}
     end
 
     def to_h
       {
         id: id,
-        created_at: created_at,
+        created_at: created_at.iso8601,
         display_name: display_name,
         provider: provider,
-        metadata: metadata,
         context_window: context_window,
         max_tokens: max_tokens,
         supports_vision: supports_vision,
         supports_functions: supports_functions,
         supports_json_mode: supports_json_mode,
         input_price_per_million: input_price_per_million,
-        output_price_per_million: output_price_per_million
+        output_price_per_million: output_price_per_million,
+        metadata: metadata
       }
     end
+
+    def type
+      metadata['type']
+    end
+
+    def family
+      metadata['family']
+    end
   end
 end