ruby_llm 0.1.0.pre4 → 0.1.0.pre5

data/lib/ruby_llm/models.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Models
+    module_function
+
+    def provider_for(model)
+      Provider.for(model)
+    end
+
+    def all
+      @all ||= begin
+        data = JSON.parse(File.read(File.expand_path('models.json', __dir__)))
+        data['models'].map { |model| ModelInfo.new(model.transform_keys(&:to_sym)) }
+      end
+    rescue Errno::ENOENT
+      [] # Return empty array if file doesn't exist yet
+    end
+
+    def find(model_id)
+      all.find { |m| m.id == model_id } or raise Error, "Unknown model: #{model_id}"
+    end
+
+    def chat_models
+      all.select { |m| m.type == 'chat' }
+    end
+
+    def embedding_models
+      all.select { |m| m.type == 'embedding' }
+    end
+
+    def audio_models
+      all.select { |m| m.type == 'audio' }
+    end
+
+    def image_models
+      all.select { |m| m.type == 'image' }
+    end
+
+    def by_family(family)
+      all.select { |m| m.family == family }
+    end
+
+    def default_model
+      'gpt-4o-mini'
+    end
+
+    def refresh!
+      @all = nil
+    end
+  end
+end

data/lib/ruby_llm/provider.rb

@@ -0,0 +1,99 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Provider
+    def self.included(base)
+      base.include(InstanceMethods)
+    end
+
+    module InstanceMethods
+      def complete(messages, tools: [], model: nil, &block)
+        # TODO: refactor
+        payload = build_payload(messages, tools, model: model, stream: block_given?)
+
+        content = String.new
+        model_id = nil
+        input_tokens = 0
+        output_tokens = 0
+        response = connection.post(completion_url, payload) do |req|
+          req.headers.merge! headers
+          if block_given?
+            req.options.on_data = handle_stream do |chunk|
+              model_id ||= chunk.model_id
+              content << (chunk.content || '')
+              input_tokens += chunk.input_tokens if chunk.input_tokens
+              output_tokens += chunk.output_tokens if chunk.output_tokens
+              block.call(chunk)
+            end
+          end
+        end
+
+        if block_given?
+          Message.new(
+            role: :assistant,
+            content: content,
+            model_id: model_id,
+            input_tokens: input_tokens.positive? ? input_tokens : nil,
+            output_tokens: output_tokens.positive? ? output_tokens : nil
+          )
+        else
+          parse_completion_response(response)
+        end
+      end
+
+      def list_models
+        response = connection.get(models_url) do |req|
+          req.headers.merge!(headers)
+        end
+
+        parse_list_models_response(response)
+      end
+
+      private
+
+      def connection
+        @connection ||= Faraday.new(api_base) do |f|
+          f.options.timeout = RubyLLM.config.request_timeout
+          f.request :json
+          f.response :json
+          f.adapter Faraday.default_adapter
+          f.use Faraday::Response::RaiseError
+          f.response :logger, RubyLLM.logger, { headers: false, bodies: true, errors: true, log_level: :debug }
+        end
+      end
+
+      def to_json_stream(&block)
+        parser = EventStreamParser::Parser.new
+        proc do |chunk, _bytes, _|
+          parser.feed(chunk) do |_type, data|
+            unless data == '[DONE]'
+              parsed_data = JSON.parse(data)
+              RubyLLM.logger.debug "chunk: #{parsed_data}"
+              block.call(parsed_data)
+            end
+          end
+        end
+      end
+    end
+
+    class << self
+      def register(name, provider_class)
+        providers[name.to_sym] = provider_class
+      end
+
+      def for(model)
+        model_info = Models.find(model)
+        provider_class = providers[model_info.provider.to_sym] or
+          raise Error, "No provider registered for #{model_info.provider}"
+
+        provider_class.new
+      end
+
+      private
+
+      def providers
+        @providers ||= {}
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/anthropic.rb

@@ -1,51 +1,82 @@
 # frozen_string_literal: true
 
-require 'time'
-
 module RubyLLM
   module Providers
-    class Anthropic < Base
-      def chat(messages, model: nil, temperature: 0.7, stream: false, tools: nil, &block)
-        payload = {
-          model: model || 'claude-3-5-sonnet-20241022',
+    class Anthropic
+      include Provider
+
+      private
+
+      def api_base
+        'https://api.anthropic.com'
+      end
+
+      def headers
+        {
+          'x-api-key' => RubyLLM.config.anthropic_api_key,
+          'anthropic-version' => '2023-06-01'
+        }
+      end
+
+      def completion_url
+        '/v1/messages'
+      end
+
+      def models_url
+        '/v1/models'
+      end
+
+      def build_payload(messages, tools, model:, temperature: 0.7, stream: false)
+        {
+          model: model,
           messages: format_messages(messages),
           temperature: temperature,
           stream: stream,
-          max_tokens: 4096
-        }
+          max_tokens: RubyLLM.models.find(model).max_tokens
+        }.tap do |payload|
+          payload[:tools] = tools.map { |t| function_for(t) } if tools.any?
+        end
+      end
 
-        payload[:tools] = tools.map { |tool| tool_to_anthropic(tool) } if tools&.any?
+      def parse_completion_response(response)
+        data = response.body
+        content_blocks = data['content'] || []
 
-        puts 'Sending payload to Anthropic:' if ENV['RUBY_LLM_DEBUG']
-        puts JSON.pretty_generate(payload) if ENV['RUBY_LLM_DEBUG']
+        text_content = content_blocks.find { |c| c['type'] == 'text' }&.fetch('text', '')
+        tool_use = content_blocks.find { |c| c['type'] == 'tool_use' }
 
-        if stream && block_given?
-          stream_chat_completion(payload, tools, &block)
+        if tool_use
+          Message.new(
+            role: :assistant,
+            content: text_content,
+            tool_calls: [
+              {
+                name: tool_use['name'],
+                arguments: JSON.generate(tool_use['input'] || {})
+              }
+            ]
+          )
         else
-          create_chat_completion(payload, tools)
+          Message.new(
+            role: :assistant,
+            content: text_content,
+            input_tokens: data['usage']['input_tokens'],
+            output_tokens: data['usage']['output_tokens'],
+            model_id: data['model']
+          )
         end
       end
 
-      def list_models
-        response = @connection.get('/v1/models') do |req|
-          req.headers['x-api-key'] = RubyLLM.configuration.anthropic_api_key
-          req.headers['anthropic-version'] = '2023-06-01'
-        end
-
-        raise RubyLLM::Error, "API error: #{parse_error_message(response)}" if response.status >= 400
-
-        capabilities = RubyLLM::ModelCapabilities::Anthropic.new
-        models_data = response.body['data'] || []
+      def parse_models_response(response)
+        capabilities = ModelCapabilities::Anthropic.new
 
-        models_data.map do |model|
+        (response.body['data'] || []).map do |model|
           ModelInfo.new(
             id: model['id'],
             created_at: Time.parse(model['created_at']),
             display_name: model['display_name'],
             provider: 'anthropic',
-            metadata: {
-              type: model['type']
-            },
+            metadata: { type: model['type'] },
             context_window: capabilities.determine_context_window(model['id']),
             max_tokens: capabilities.determine_max_tokens(model['id']),
             supports_vision: capabilities.supports_vision?(model['id']),
@@ -55,254 +86,72 @@ module RubyLLM
             output_price_per_million: capabilities.get_output_price(model['id'])
           )
         end
-      rescue Faraday::Error => e
-        handle_error(e)
       end
 
-      private
-
-      def tool_to_anthropic(tool)
-        # Get required fields and clean properties
-        required_fields = []
-        cleaned_properties = {}
-
-        tool.parameters.each do |name, props|
-          required_fields << name.to_s if props[:required]
-          cleaned_props = props.dup
-          cleaned_props.delete(:required)
-          cleaned_properties[name] = cleaned_props
+      def handle_stream(&block)
+        to_json_stream do |data|
+          block.call(
+            Chunk.new(
+              role: :assistant,
+              model_id: data.dig('message', 'model'),
+              content: data.dig('delta', 'text'),
+              input_tokens: data.dig('message', 'usage', 'input_tokens'),
+              output_tokens: data.dig('message', 'usage', 'output_tokens') || data.dig('usage', 'output_tokens')
+            )
+          )
         end
+      end
 
+      def function_for(tool)
         {
           name: tool.name,
           description: tool.description,
           input_schema: {
             type: 'object',
-            properties: cleaned_properties,
-            required: required_fields
+            properties: clean_parameters(tool.parameters),
+            required: required_parameters(tool.parameters)
           }
         }
       end
 
       def format_messages(messages)
         messages.map do |msg|
-          message = { role: msg.role == :user ? 'user' : 'assistant' }
-
-          message[:content] = if msg.tool_results
-                                [
-                                  {
-                                    type: 'tool_result',
-                                    tool_use_id: msg.tool_results[:tool_use_id],
-                                    content: msg.tool_results[:content],
-                                    is_error: msg.tool_results[:is_error]
-                                  }.compact
-                                ]
-                              else
-                                msg.content
-                              end
-
-          message
-        end
-      end
-
-      def create_chat_completion(payload, tools = nil, &block)
-        response = @connection.post('/v1/messages') do |req|
-          req.headers['x-api-key'] = RubyLLM.configuration.anthropic_api_key
-          req.headers['anthropic-version'] = '2023-06-01'
-          req.headers['Content-Type'] = 'application/json'
-          req.body = payload
-        end
-
-        puts 'Response from Anthropic:' if ENV['RUBY_LLM_DEBUG']
-        puts JSON.pretty_generate(response.body) if ENV['RUBY_LLM_DEBUG']
-
-        # Check for API errors first
-        check_for_api_error(response)
-
-        handle_response(response, tools, payload, &block)
-      rescue Faraday::Error => e
-        handle_error(e)
-      end
-
-      def stream_chat_completion(payload, tools = nil)
-        response = @connection.post('/v1/messages') do |req|
-          req.headers['x-api-key'] = RubyLLM.configuration.anthropic_api_key
-          req.headers['anthropic-version'] = '2023-06-01'
-          req.body = payload
-        end
-
-        messages = []
-        response.body.each_line do |line|
-          next if line.strip.empty?
-          next if line == 'data: [DONE]'
-
-          begin
-            data = JSON.parse(line.sub(/^data: /, ''))
-
-            message = case data['type']
-                      when 'content_block_delta'
-                        Message.new(role: :assistant, content: data['delta']['text']) if data['delta']['text']
-                      when 'tool_call'
-                        handle_tool_calls(data['tool_calls'], tools) do |result|
-                          Message.new(role: :assistant, content: result)
-                        end
-                      end
-
-            if message
-              messages << message
-              yield message if block_given?
-            end
-          rescue JSON::ParserError
-            next
-          end
-        end
-
-        messages
-      rescue Faraday::Error => e
-        handle_error(e)
-      end
-
-      def handle_response(response, tools, payload, &block)
-        data = response.body
-
-        content_parts = data['content'] || []
-        text_content = content_parts.find { |c| c['type'] == 'text' }&.fetch('text', '')
-        tool_use = content_parts.find { |c| c['type'] == 'tool_use' }
-
-        if tool_use && tools
-          # Tool call handling code...
-          tool_message = Message.new(
-            role: :assistant,
-            content: text_content,
-            tool_calls: [{
-              name: tool_use['name'],
-              arguments: JSON.generate(tool_use['input'] || {})
-            }]
-          )
-          yield tool_message if block_given?
-
-          tool = tools.find { |t| t.name == tool_use['name'] }
-          result = if tool
-                     begin
-                       tool_result = tool.call(tool_use['input'] || {})
-                       {
-                         tool_use_id: tool_use['id'],
-                         content: tool_result.to_s
-                       }
-                     rescue StandardError => e
-                       {
-                         tool_use_id: tool_use['id'],
-                         content: "Error executing tool #{tool.name}: #{e.message}",
-                         is_error: true
-                       }
-                     end
-                   end
-
-          result_message = Message.new(
-            role: :tool,
-            content: result[:content],
-            tool_results: result
-          )
-          yield result_message if block_given?
-
-          new_messages = payload[:messages] + [
-            { role: 'assistant', content: data['content'] },
+          if msg.tool_results
             {
-              role: 'user',
+              role: convert_role(msg.role),
               content: [
                 {
                   type: 'tool_result',
-                  tool_use_id: result[:tool_use_id],
-                  content: result[:content],
-                  is_error: result[:is_error]
+                  tool_use_id: msg.tool_results[:tool_use_id],
+                  content: msg.tool_results[:content],
+                  is_error: msg.tool_results[:is_error]
                 }.compact
               ]
             }
-          ]
-
-          final_response = create_chat_completion(
-            payload.merge(messages: new_messages),
-            tools,
-            &block
-          )
-
-          [tool_message, result_message] + final_response
-        else
-          token_usage = if data['usage']
-                          {
-                            input_tokens: data['usage']['input_tokens'],
-                            output_tokens: data['usage']['output_tokens'],
-                            total_tokens: data['usage']['input_tokens'] + data['usage']['output_tokens']
-                          }
-                        end
-
-          [Message.new(
-            role: :assistant,
-            content: text_content,
-            token_usage: token_usage,
-            model_id: data['model']
-          )]
-        end
-      end
-
-      def handle_tool_calls(tool_calls, tools)
-        return [] unless tool_calls && tools
-
-        tool_calls.map do |tool_call|
-          tool = tools.find { |t| t.name == tool_call['name'] }
-          next unless tool
-
-          begin
-            args = JSON.parse(tool_call['arguments'])
-            result = tool.call(args)
-            puts "Tool result: #{result}" if ENV['RUBY_LLM_DEBUG']
-            {
-              tool_use_id: tool_call['id'],
-              content: result.to_s
-            }
-          rescue JSON::ParserError, ArgumentError => e
-            puts "Error executing tool: #{e.message}" if ENV['RUBY_LLM_DEBUG']
+          else
             {
-              tool_use_id: tool_call['id'],
-              content: "Error executing tool #{tool.name}: #{e.message}",
-              is_error: true
+              role: convert_role(msg.role),
+              content: msg.content
             }
           end
-        end.compact
+        end
       end
 
-      def handle_api_error(error)
-        response_body = error.response[:body]
-        if response_body.is_a?(String)
-          begin
-            error_data = JSON.parse(response_body)
-            message = error_data.dig('error', 'message')
-            raise RubyLLM::Error, "API error: #{message}" if message
-          rescue JSON::ParserError
-            raise RubyLLM::Error, "API error: #{error.response[:status]}"
-          end
-        elsif response_body['error']
-          raise RubyLLM::Error, "API error: #{response_body['error']['message']}"
-        else
-          raise RubyLLM::Error, "API error: #{error.response[:status]}"
+      def convert_role(role)
+        case role
+        when :user then 'user'
+        else 'assistant'
         end
       end
 
-      def handle_error(error)
-        case error
-        when Faraday::TimeoutError
-          raise RubyLLM::Error, 'Request timed out'
-        when Faraday::ConnectionFailed
-          raise RubyLLM::Error, 'Connection failed'
-        when Faraday::ClientError
-          handle_api_error(error)
-        else
-          raise error
+      def clean_parameters(parameters)
+        parameters.transform_values do |props|
+          props.except(:required)
         end
       end
 
-      def api_base
-        'https://api.anthropic.com'
+      def required_parameters(parameters)
+        parameters.select { |_, props| props[:required] }.keys
       end
     end
   end