Add Replicate provider

Author: josefarias

bin/console

@@ -23,6 +23,8 @@ RubyLLM.configure do |config|
   config.openai_api_key = ENV.fetch('OPENAI_API_KEY', nil)
   config.openrouter_api_key = ENV.fetch('OPENROUTER_API_KEY', nil)
   config.perplexity_api_key = ENV.fetch('PERPLEXITY_API_KEY', nil)
+  config.replicate_api_key = ENV.fetch('REPLICATE_API_KEY', nil)
+  config.replicate_webhook_url = ENV.fetch('REPLICATE_WEBHOOK_URL', nil)
   config.vertexai_location = ENV.fetch('GOOGLE_CLOUD_LOCATION', nil)
   config.vertexai_project_id = ENV.fetch('GOOGLE_CLOUD_PROJECT', nil)
 end

docs/_getting_started/configuration.md

@@ -59,6 +59,8 @@ RubyLLM.configure do |config|
   config.mistral_api_key = ENV['MISTRAL_API_KEY']
   config.perplexity_api_key = ENV['PERPLEXITY_API_KEY']
   config.openrouter_api_key = ENV['OPENROUTER_API_KEY']
+  config.replicate_api_key = ENV['REPLICATE_API_KEY']
+  config.replicate_webhook_url = ENV['REPLICATE_WEBHOOK_URL']
 
   # Local providers
   config.ollama_api_base = 'http://localhost:11434/v1'
@@ -363,4 +365,4 @@ Now that you've configured RubyLLM, you're ready to:
 
 - [Start chatting with AI models]({% link _core_features/chat.md %})
 - [Work with different providers and models]({% link _advanced/models.md %})
-- [Set up Rails integration]({% link _advanced/rails.md %})
+- [Set up Rails integration]({% link _advanced/rails.md %})

lib/ruby_llm.rb

@@ -93,6 +93,7 @@ def logger
 RubyLLM::Provider.register :openrouter, RubyLLM::Providers::OpenRouter
 RubyLLM::Provider.register :perplexity, RubyLLM::Providers::Perplexity
 RubyLLM::Provider.register :vertexai, RubyLLM::Providers::VertexAI
+RubyLLM::Provider.register :replicate, RubyLLM::Providers::Replicate
 
 if defined?(Rails::Railtie)
   require 'ruby_llm/railtie'

lib/ruby_llm/configuration.rb

@@ -23,6 +23,9 @@ class Configuration
                   :gpustack_api_base,
                   :gpustack_api_key,
                   :mistral_api_key,
+                  :replicate_api_key,
+                  :replicate_webhook_url,
+                  :replicate_webhook_events_filter,
                   # Default models
                   :default_model,
                   :default_embedding_model,

lib/ruby_llm/image.rb

@@ -36,14 +36,15 @@ def self.paint(prompt, # rubocop:disable Metrics/ParameterLists
                    provider: nil,
                    assume_model_exists: false,
                    size: '1024x1024',
-                   context: nil)
+                   context: nil,
+                   model_params: {})
       config = context&.config || RubyLLM.config
       model ||= config.default_image_model
       model, provider_instance = Models.resolve(model, provider: provider, assume_exists: assume_model_exists,
                                                        config: config)
       model_id = model.id
 
-      provider_instance.paint(prompt, model: model_id, size:)
+      provider_instance.paint(prompt, model: model_id, size:, **model_params)
     end
   end
 end

lib/ruby_llm/provider.rb

@@ -76,8 +76,8 @@ def moderate(input, model:)
       parse_moderation_response(response, model:)
     end
 
-    def paint(prompt, model:, size:)
-      payload = render_image_payload(prompt, model:, size:)
+    def paint(prompt, model:, size:, **params)
+      payload = render_image_payload(prompt, model:, size:, **params)
       response = @connection.post images_url, payload
       parse_image_response(response, model:)
     end

lib/ruby_llm/providers/replicate.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    # Replicate API integration
+    class Replicate < Provider
+      include Replicate::Capabilities
+      include Replicate::Images
+      include Replicate::Models
+
+      def api_base
+        'https://api.replicate.com'
+      end
+
+      def headers
+        {
+          'Authorization' => "Bearer #{@config.replicate_api_key}",
+          'Content-Type' => 'application/json'
+        }
+      end
+
+      class << self
+        def capabilities
+          Replicate::Capabilities
+        end
+
+        def configuration_requirements
+          %i[replicate_api_key replicate_webhook_url]
+        end
+      end
+    end
+  end
+end

lib/ruby_llm/providers/replicate/capabilities.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class Replicate
+      # Determines capabilities for Replicate models
+      module Capabilities
+        module_function
+
+        def name_from(model_data)
+          "replicate/#{model_data['owner']}/#{model_data['name']}"
+        end
+
+        def metadata_from(model_data)
+          input_schema = model_data.dig('latest_version', 'openapi_schema', 'components', 'schemas', 'Input')
+
+          {
+            url: model_data['url'],
+            description: model_data['description'],
+            license_url: model_data['license_url'],
+            is_official: model_data['is_official'],
+            supported_parameters: input_schema['properties'].keys - ['prompt'],
+            latest_version_created_at: model_data.dig('latest_version', 'created_at')
+          }
+        end
+      end
+    end
+  end
+end

lib/ruby_llm/providers/replicate/images.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class Replicate
+      # Image generation methods for the Replicate API implementation
+      module Images
+        attr_reader :model
+
+        def images_url
+          if official_model?
+            "/v1/models/#{canonical_model_name}/predictions"
+          else
+            '/v1/predictions'
+          end
+        end
+
+        def render_image_payload(prompt, model:, size:, **params)
+          RubyLLM.logger.debug "Use `model_params` for model-dependent sizing instead of size #{size}." if size
+
+          self.model_id = model
+
+          {}.tap do |payload|
+            payload[:webhook] = @config.replicate_webhook_url
+            payload[:version] = model.id unless official_model?
+            payload[:input] = { prompt: prompt }.merge(params)
+
+            if @config.replicate_webhook_events_filter
+              payload[:webhook_events_filter] = @config.replicate_webhook_events_filter
+            end
+          end
+        end
+
+        def parse_image_response(response, **)
+          response
+        end
+
+        private
+
+        def model_id=(id)
+          @model_id = id
+          @model = Models.find(@model_id, 'replicate')
+        end
+
+        def official_model?
+          model.metadata['is_official'] == true
+        end
+
+        def canonical_model_name
+          model.name.split('/')[1..].join('/')
+        end
+      end
+    end
+  end
+end

lib/ruby_llm/providers/replicate/models.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class Replicate
+      # Models methods of the Replicate API integration
+      module Models
+        def list_models(**)
+          response = @connection.get models_url
+          parse_list_models_response response
+        end
+
+        def models_url
+          '/v1/collections/text-to-image'
+        end
+
+        def parse_list_models_response(response)
+          Array(response.body['models']).map do |model_data|
+            Model::Info.new(
+              id: model_data.dig('latest_version', 'id'),
+              name: capabilities.name_from(model_data),
+              provider: 'replicate',
+              created_at: model_data['created_at'],
+              modalities: { input: ['text'], output: ['text'] },
+              capabilities: ['image_generation'],
+              metadata: capabilities.metadata_from(model_data)
+            )
+          end
+        end
+      end
+    end
+  end
+end