Provider registry enables use of other providers

Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    raix (1.0.2)
+    raix (1.0.3)
       activesupport (>= 6.0)
       faraday-retry (~> 2.0)
       open_router (~> 0.2)

README.md

@@ -111,7 +111,7 @@ The second (optional) module that you can add to your Ruby classes after `ChatCo
 
 When the AI responds with tool function calls instead of a text message, Raix automatically:
 1. Executes the requested tool functions
-2. Adds the function results to the conversation transcript  
+2. Adds the function results to the conversation transcript
 3. Sends the updated transcript back to the AI for another completion
 4. Repeats this process until the AI responds with a regular text message
 
@@ -739,6 +739,13 @@ You can add an initializer to your application's `config/initializers` directory
 
 You will also need to configure the OpenRouter API access token as per the instructions here: https://github.com/OlympiaAI/open_router?tab=readme-ov-file#quickstart
 
+### Custom Providers
+
+You may register custom providers instead of either of the above by registering an object that responds
+to `request(params:, model:, messages:)` and returns a OpenAI compatible chat completion response.
+
+See the [OpenAIProvider implementation](/OlympiaAI/raix/blob/main/lib/raix/providers/openai_provider.rb) for the OpenAI compatible provider implementation.
+
 ### Global vs class level configuration
 
 You can either configure Raix globally or at the class level. The global configuration is set in the initializer as shown above. You can however also override all configuration options of the `Configuration` class on the class level with the

lib/raix/chat_completion.rb

@@ -146,7 +146,7 @@ def chat_completion(params: {}, loop: false, json: false, raw: false, openai: ni
         response = if openai
                      openai_request(params:, model: openai, messages:)
                    else
-                     openrouter_request(params:, model:, messages:)
+                     provider_request(params:, model:, messages:)
                    end
         retry_count = 0
         content = nil
@@ -174,7 +174,7 @@ def chat_completion(params: {}, loop: false, json: false, raw: false, openai: ni
             response = if openai
                          openai_request(params:, model: openai, messages:)
                        else
-                         openrouter_request(params:, model:, messages:)
+                         provider_request(params:, model:, messages:)
                        end
 
             # Process the final response
@@ -220,7 +220,7 @@ def chat_completion(params: {}, loop: false, json: false, raw: false, openai: ni
             response = if openai
                          openai_request(params:, model: openai, messages:)
                        else
-                         openrouter_request(params:, model:, messages:)
+                         provider_request(params:, model:, messages:)
                        end
 
             content = response.dig("choices", 0, "message", "content")
@@ -308,41 +308,23 @@ def filtered_tools(tool_names)
     end
 
     def openai_request(params:, model:, messages:)
-      if params[:prediction]
-        params.delete(:max_completion_tokens)
-      else
-        params[:max_completion_tokens] ||= params[:max_tokens]
-        params.delete(:max_tokens)
-      end
-
       params[:stream] ||= stream.presence
-      params[:stream_options] = { include_usage: true } if params[:stream]
-
-      params.delete(:temperature) if model.start_with?("o") || model.include?("gpt-5")
+      provider = configuration.provider(:openai)
+      raise "OpenAI provider not configured. Use configuration.openai_client = OpenAI::Client.new" unless provider
 
-      configuration.openai_client.chat(parameters: params.compact.merge(model:, messages:))
+      provider.request(params:, model:, messages:)
     end
 
-    def openrouter_request(params:, model:, messages:)
-      # max_completion_tokens is not supported by OpenRouter
-      params.delete(:max_completion_tokens)
+    def provider_request(params:, model:, messages:)
+      params[:stream] ||= stream.presence
 
-      retry_count = 0
+      # If openrouter is set, use it and pass provider as a parameter
+      # Otherwise, use provider to select the provider from the registry
+      provider = configuration.provider(:openrouter) || configuration.provider(params.delete(:provider))
 
-      params.delete(:temperature) if model.start_with?("openai/o") || model.include?("gpt-5")
+      raise "No provider configured." unless provider
 
-      begin
-        configuration.openrouter_client.complete(messages, model:, extras: params.compact, stream:)
-      rescue OpenRouter::ServerError => e
-        if e.message.include?("retry")
-          warn "Retrying OpenRouter request... (#{retry_count} attempts) #{e.message}"
-          retry_count += 1
-          sleep 1 * retry_count # backoff
-          retry if retry_count < 5
-        end
-
-        raise e
-      end
+      provider.request(params:, model:, messages:)
     end
   end
 end

lib/raix/configuration.rb

@@ -1,5 +1,8 @@
 # frozen_string_literal: true
 
+require_relative "providers/open_router_provider"
+require_relative "providers/openai_provider"
+
 module Raix
   # The Configuration class holds the configuration options for the Raix gem.
   class Configuration
@@ -30,11 +33,25 @@ def self.attr_accessor_with_fallback(method_name)
     # is normally set in each class that includes the ChatCompletion module.
     attr_accessor_with_fallback :model
 
+    attr_writer :openrouter_client, :openai_client
+
     # The openrouter_client option determines the default client to use for communication.
-    attr_accessor_with_fallback :openrouter_client
+    def openrouter_client
+      value = @openrouter_client
+      return value if value
+      return unless fallback
+
+      fallback.openrouter_client
+    end
 
     # The openai_client option determines the OpenAI client to use for communication.
-    attr_accessor_with_fallback :openai_client
+    def openai_client
+      value = @openai_client
+      return value if value
+      return unless fallback
+
+      fallback.openai_client
+    end
 
     # The max_tool_calls option determines the maximum number of tool calls
     # before forcing a text response to prevent excessive function invocations.
@@ -54,12 +71,41 @@ def initialize(fallback: nil)
       self.model = DEFAULT_MODEL
       self.max_tool_calls = DEFAULT_MAX_TOOL_CALLS
       self.fallback = fallback
+      @providers = {}
     end
 
     def client?
-      !!(openrouter_client || openai_client)
+      !!(openrouter_client || openai_client || @providers.any?)
+    end
+
+    def register_provider(name, client)
+      @providers[name] = client
     end
 
+    # Find the provider to use based on the name, if given.
+    # Fall back to the next registered provider if no name is provided.
+    # We must use the openai_client and openrouter_client methods so that the
+    # previous fallback behavior is preserved.
+    def provider(name = nil)
+      # Prioritize use of registered providers before using openai_client or openrouter_client.
+      return @providers[name] if name && @providers.key?(name)
+
+      # if openai is specified explicitly, use openai_client.
+      # if openrouter_client is set, use it for backwards compatibility.
+      # finally, use the named or first registered provider.
+      if name == :openai
+        openai_client ? Providers::OpenAIProvider.new(openai_client) : nil
+      elsif name == :openrouter || openrouter_client
+        openrouter_client ? Providers::OpenRouterProvider.new(openrouter_client) : nil
+      elsif @providers.any?
+        @providers.values.first
+      elsif fallback
+        fallback.provider(name)
+      end
+    end
+
+    attr_reader :providers
+
     private
 
     attr_accessor :fallback

lib/raix/providers/open_router_provider.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+module Raix
+  module Providers
+    # A wrapper around the OpenRouter client interface to make it compatible with the provider interface.
+    class OpenRouterProvider
+      attr_reader :client
+
+      def initialize(client)
+        @client = client
+      end
+
+      def request(params:, model:, messages:)
+        params = params.dup
+
+        # max_completion_tokens is not supported by OpenRouter
+        params.delete(:max_completion_tokens)
+
+        retry_count = 0
+
+        params.delete(:temperature) if model.start_with?("openai/o") || model.include?("gpt-5")
+
+        stream = params.delete(:stream)
+
+        begin
+          client.complete(messages, model:, extras: params.compact, stream:)
+        rescue ::OpenRouter::ServerError => e
+          if e.message.include?("retry")
+            warn "Retrying OpenRouter request... (#{retry_count} attempts) #{e.message}"
+            retry_count += 1
+            sleep 1 * retry_count # backoff
+            retry if retry_count < 5
+          end
+
+          raise e
+        end
+      end
+    end
+  end
+end

lib/raix/providers/openai_provider.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module Raix
+  module Providers
+    # A wrapper around the OpenAI client to make it compatible with the provider interface.
+    class OpenAIProvider
+      attr_reader :client
+
+      def initialize(client)
+        @client = client
+      end
+
+      def request(params:, model:, messages:)
+        params = params.dup
+
+        if params[:prediction]
+          params.delete(:max_completion_tokens)
+        else
+          params[:max_completion_tokens] ||= params[:max_tokens]
+          params.delete(:max_tokens)
+        end
+
+        params[:stream_options] = { include_usage: true } if params[:stream]
+
+        params.delete(:temperature) if model.start_with?("o") || model.include?("gpt-5")
+
+        client.chat(parameters: params.compact.merge(model:, messages:))
+      end
+    end
+  end
+end

spec/raix/chat_completion_spec.rb

@@ -22,6 +22,18 @@ def initialize
   end
 end
 
+class TestOveriddenConfiguration
+  include Raix::ChatCompletion
+
+  # Override the configuration accessor to make testing non-global
+  attr_accessor :configuration
+
+  def initialize
+    self.model = "test-model"
+    transcript << { user: "What is the meaning of life?" }
+  end
+end
+
 RSpec.describe MeaningOfLife, :vcr do
   subject { described_class.new }
 
@@ -67,3 +79,55 @@ def initialize
     subject.chat_completion
   end
 end
+
+RSpec.describe "Provider parameter behavior" do
+  context "when openrouter_client is set" do
+    it "passes provider as a parameter to openrouter" do
+      mock_openrouter = instance_double("OpenRouter::Client")
+      expect(mock_openrouter).to receive(:complete).with(
+        anything,
+        model: "test-model",
+        extras: hash_including(provider: "anthropic"),
+        stream: anything
+      ).and_return("choices" => [{ "message" => { "content" => "42" } }])
+
+      chat_client = TestOveriddenConfiguration.new
+      chat_client.configuration = Raix::Configuration.new
+      chat_client.configuration.openrouter_client = mock_openrouter
+      chat_client.provider = "anthropic"
+
+      expect(chat_client.chat_completion).to eq("42")
+    end
+  end
+
+  context "when openrouter_client is not set" do
+    it "uses provider parameter to select the registered provider" do
+      mock_provider = instance_double("CustomProvider")
+      expect(mock_provider).to receive(:request).with(
+        params: hash_not_including(:provider),
+        model: "test-model",
+        messages: anything
+      ).and_return("choices" => [{ "message" => { "content" => "42" } }])
+
+      chat_client = TestOveriddenConfiguration.new
+      chat_client.configuration = Raix::Configuration.new
+      chat_client.configuration.register_provider(:custom, mock_provider)
+      chat_client.provider = :custom
+
+      expect(chat_client.chat_completion).to eq("42")
+    end
+  end
+
+  context "when openrouter_client is not set and provider is not found" do
+    it "raises error" do
+      chat_client = TestOveriddenConfiguration.new
+      chat_client.configuration = Raix::Configuration.new
+      chat_client.provider = :nonexistent
+
+      expect { chat_client.chat_completion }.to raise_error(
+        RuntimeError,
+        "No provider configured."
+      )
+    end
+  end
+end