From 47f634994b8be5a7b8194c61de47091aab6515ce Mon Sep 17 00:00:00 2001
From: Austin Spires <austin@fastly.com>
Date: Sat, 15 Mar 2025 22:20:10 -0600
Subject: [PATCH] Add OpenAI URL override and custom req headers

This commit introduces the configuration primitives necessary for
semantic caching proxy services for OpenAI requests.Specifically,
the ability to override the OpenAI root API URL path and the
ability to add custom headers to the request payload in the form
of a hash.

This commit does not introduce tests, and should not be merged
until a subsequent commit lands with proper testing.
---
 bin/console                      | 2 ++
 lib/ruby_llm/configuration.rb    | 4 +++-
 lib/ruby_llm/provider.rb         | 2 ++
 lib/ruby_llm/providers/openai.rb | 3 ++-
 4 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/bin/console b/bin/console
index e43cf3d..1095622 100755
--- a/bin/console
+++ b/bin/console
@@ -12,6 +12,8 @@ RubyLLM.configure do |config|
   config.anthropic_api_key = ENV.fetch('ANTHROPIC_API_KEY', nil)
   config.gemini_api_key = ENV.fetch('GEMINI_API_KEY', nil)
   config.deepseek_api_key = ENV.fetch('DEEPSEEK_API_KEY', nil)
+
+  config.openai_base_url_override = ENV.fetch('OPENAI_BASE_URL', nil)
 end
 
 IRB.start(__FILE__)
diff --git a/lib/ruby_llm/configuration.rb b/lib/ruby_llm/configuration.rb
index 72a878a..c9075de 100644
--- a/lib/ruby_llm/configuration.rb
+++ b/lib/ruby_llm/configuration.rb
@@ -18,7 +18,9 @@ class Configuration
                   :default_embedding_model,
                   :default_image_model,
                   :request_timeout,
-                  :max_retries
+                  :max_retries,
+                  :openai_base_url_override,
+                  :additional_headers
 
     def initialize
       @request_timeout = 120
diff --git a/lib/ruby_llm/provider.rb b/lib/ruby_llm/provider.rb
index 6e2ad48..c1463e6 100644
--- a/lib/ruby_llm/provider.rb
+++ b/lib/ruby_llm/provider.rb
@@ -60,7 +60,9 @@ def stream_response(payload, &block)
       end
 
       def post(url, payload)
+        additional_headers = RubyLLM.config.additional_headers
         connection.post url, payload do |req|
+          req.headers.merge! additional_headers unless additional_headers.nil?
           req.headers.merge! headers
           yield req if block_given?
         end
diff --git a/lib/ruby_llm/providers/openai.rb b/lib/ruby_llm/providers/openai.rb
index f78062e..83d8883 100644
--- a/lib/ruby_llm/providers/openai.rb
+++ b/lib/ruby_llm/providers/openai.rb
@@ -29,7 +29,8 @@ def self.extended(base)
       module_function
 
       def api_base
-        'https://api.openai.com/v1'
+        return 'https://api.openai.com/v1' unless RubyLLM.config.openai_base_url_override
+        return RubyLLM.config.openai_base_url_override
       end
 
       def headers