feat: Add with_headers method to Chat for custom HTTP headers

Enables use of provider-specific features like beta headers while
maintaining proper header precedence for security.

Closes #85, #105

Author: crmne

lib/ruby_llm/active_record/acts_as.rb

@@ -138,6 +138,11 @@ def with_params(...)
         self
       end
 
+      def with_headers(...)
+        to_llm.with_headers(...)
+        self
+      end
+
       def with_schema(...)
         to_llm.with_schema(...)
         self

lib/ruby_llm/chat.rb

@@ -11,7 +11,7 @@ module RubyLLM
   class Chat
     include Enumerable
 
-    attr_reader :model, :messages, :tools, :params, :schema
+    attr_reader :model, :messages, :tools, :params, :headers, :schema
 
     def initialize(model: nil, provider: nil, assume_model_exists: false, context: nil)
       if assume_model_exists && !provider
@@ -26,6 +26,7 @@ def initialize(model: nil, provider: nil, assume_model_exists: false, context: n
       @messages = []
       @tools = {}
       @params = {}
+      @headers = {}
       @schema = nil
       @on = {
         new_message: nil,
@@ -87,6 +88,11 @@ def with_params(**params)
       self
     end
 
+    def with_headers(**headers)
+      @headers = headers
+      self
+    end
+
     def with_schema(schema, force: false)
       unless force || @model.structured_output?
         raise UnsupportedStructuredOutputError, "Model #{@model.id} doesn't support structured output"
@@ -135,6 +141,7 @@ def complete(&) # rubocop:disable Metrics/PerceivedComplexity
         temperature: @temperature,
         model: @model.id,
         params: @params,
+        headers: @headers,
         schema: @schema,
         &wrap_streaming_block(&)
       )

lib/ruby_llm/provider.rb

@@ -40,7 +40,7 @@ def configuration_requirements
       self.class.configuration_requirements
     end
 
-    def complete(messages, tools:, temperature:, model:, params: {}, schema: nil, &) # rubocop:disable Metrics/ParameterLists
+    def complete(messages, tools:, temperature:, model:, params: {}, headers: {}, schema: nil, &) # rubocop:disable Metrics/ParameterLists
       normalized_temperature = maybe_normalize_temperature(temperature, model)
 
       payload = Utils.deep_merge(
@@ -56,9 +56,9 @@ def complete(messages, tools:, temperature:, model:, params: {}, schema: nil, &)
       )
 
       if block_given?
-        stream_response @connection, payload, &
+        stream_response @connection, payload, headers, &
       else
-        sync_response @connection, payload
+        sync_response @connection, payload, headers
       end
     end
 
@@ -208,8 +208,11 @@ def maybe_normalize_temperature(temperature, _model_id)
       temperature
     end
 
-    def sync_response(connection, payload)
-      response = connection.post completion_url, payload
+    def sync_response(connection, payload, additional_headers = {})
+      response = connection.post completion_url, payload do |req|
+        # Merge additional headers, with existing headers taking precedence
+        req.headers = additional_headers.merge(req.headers) unless additional_headers.empty?
+      end
       parse_completion_response response
     end
   end

lib/ruby_llm/providers/bedrock/chat.rb

@@ -7,10 +7,12 @@ class Bedrock
       module Chat
         module_function
 
-        def sync_response(connection, payload)
+        def sync_response(connection, payload, additional_headers = {})
           signature = sign_request("#{connection.connection.url_prefix}#{completion_url}", payload:)
           response = connection.post completion_url, payload do |req|
             req.headers.merge! build_headers(signature.headers, streaming: block_given?)
+            # Merge additional headers, with existing headers taking precedence
+            req.headers = additional_headers.merge(req.headers) unless additional_headers.empty?
           end
           Anthropic::Chat.parse_completion_response response
         end

lib/ruby_llm/providers/bedrock/streaming/base.rb

@@ -29,12 +29,14 @@ def stream_url
             "model/#{@model_id}/invoke-with-response-stream"
           end
 
-          def stream_response(connection, payload, &block)
+          def stream_response(connection, payload, additional_headers = {}, &block)
             signature = sign_request("#{connection.connection.url_prefix}#{stream_url}", payload:)
             accumulator = StreamAccumulator.new
 
             response = connection.post stream_url, payload do |req|
               req.headers.merge! build_headers(signature.headers, streaming: block_given?)
+              # Merge additional headers, with existing headers taking precedence
+              req.headers = additional_headers.merge(req.headers) unless additional_headers.empty?
               req.options.on_data = handle_stream do |chunk|
                 accumulator.add chunk
                 block.call chunk

lib/ruby_llm/streaming.rb

@@ -8,10 +8,12 @@ module RubyLLM
   module Streaming
     module_function
 
-    def stream_response(connection, payload, &block)
+    def stream_response(connection, payload, additional_headers = {}, &block)
       accumulator = StreamAccumulator.new
 
       response = connection.post stream_url, payload do |req|
+        # Merge additional headers, with existing headers taking precedence
+        req.headers = additional_headers.merge(req.headers) unless additional_headers.empty?
         if req.options.respond_to?(:on_data)
           # Handle Faraday 2.x streaming with on_data method
           req.options.on_data = handle_stream do |chunk|

spec/fixtures/vcr_cassettes/RubyLLM_Chat/_with_headers/with_Anthropic_beta_headers/works_with_anthropic-beta_header_for_fine-grained_tool_streaming.yml

@@ -131,6 +131,33 @@ def execute(expression:)
     end
   end
 
+  describe 'custom headers' do
+    it 'supports with_headers for custom HTTP headers' do
+      chat = Chat.create!(model_id: model)
+
+      result = chat.with_headers('X-Custom-Header' => 'test-value')
+      expect(result).to eq(chat) # Should return self for chaining
+
+      # Verify the headers are passed through to the underlying chat
+      llm_chat = chat.instance_variable_get(:@chat)
+      expect(llm_chat.headers).to eq('X-Custom-Header' => 'test-value')
+    end
+
+    it 'allows chaining with_headers with other methods' do
+      chat = Chat.create!(model_id: model)
+
+      result = chat
+               .with_temperature(0.5)
+               .with_headers('X-Test' => 'value')
+               .with_tool(Calculator)
+
+      expect(result).to eq(chat)
+
+      llm_chat = chat.instance_variable_get(:@chat)
+      expect(llm_chat.headers).to eq('X-Test' => 'value')
+    end
+  end
+
   describe 'error handling' do
     it 'destroys empty assistant messages on API failure' do
       chat = Chat.create!(model_id: model)

spec/ruby_llm/active_record/acts_as_spec.rb

@@ -0,0 +1,108 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe RubyLLM::Chat do
+  include_context 'with configured RubyLLM'
+  describe '#with_headers' do
+    it 'stores headers' do
+      chat = RubyLLM.chat.with_headers('X-Custom-Header' => 'value')
+      expect(chat.headers).to eq('X-Custom-Header' => 'value')
+    end
+
+    it 'returns self for chaining' do
+      chat = RubyLLM.chat
+      expect(chat.with_headers('X-Test' => 'test')).to eq(chat)
+    end
+
+    it 'passes headers to provider complete method' do
+      chat = RubyLLM.chat
+      provider = chat.instance_variable_get(:@provider)
+
+      allow(provider).to receive(:complete).and_return(
+        RubyLLM::Message.new(role: :assistant, content: 'Test response')
+      )
+
+      chat.with_headers('X-Custom' => 'header').ask('Test')
+
+      expect(provider).to have_received(:complete).with(
+        anything,
+        hash_including(headers: { 'X-Custom' => 'header' })
+      )
+    end
+
+    it 'allows chaining with other methods' do
+      chat = RubyLLM.chat
+                    .with_temperature(0.5)
+                    .with_headers('X-Test' => 'value')
+                    .with_params(max_tokens: 100)
+
+      expect(chat.headers).to eq('X-Test' => 'value')
+      expect(chat.params).to eq(max_tokens: 100)
+      expect(chat.instance_variable_get(:@temperature)).to eq(0.5)
+    end
+
+    context 'with Anthropic beta headers' do
+      it 'works with anthropic-beta header for fine-grained tool streaming', :vcr do
+        skip unless ENV['ANTHROPIC_API_KEY']
+
+        chat = RubyLLM.chat(model: 'claude-3-5-haiku-20241022', provider: 'anthropic')
+                      .with_headers('anthropic-beta' => 'fine-grained-tool-streaming-2025-05-14')
+
+        response = chat.ask('Say "beta headers work"')
+        expect(response.content).to include('beta headers work')
+      end
+    end
+
+    context 'with header precedence' do
+      it 'user headers do not override provider headers' do
+        chat = RubyLLM.chat
+        connection = chat.instance_variable_get(:@connection)
+        provider = chat.instance_variable_get(:@provider)
+
+        # Mock provider headers
+        allow(provider).to receive_messages(
+          headers: {
+            'X-Api-Key' => 'provider-key',
+            'Content-Type' => 'application/json'
+          },
+          parse_completion_response: RubyLLM::Message.new(role: :assistant, content: 'Test')
+        )
+
+        # Set user headers that try to override provider headers
+        chat.with_headers(
+          'X-Api-Key' => 'user-key',
+          'X-Custom' => 'user-value'
+        )
+
+        # Mock the connection.post to verify header merging
+        allow(connection).to receive(:post) do |_url, _payload, &block|
+          req = instance_double(Faraday::Request)
+          initial_headers = {
+            'X-Api-Key' => 'provider-key',
+            'Content-Type' => 'application/json'
+          }
+
+          allow(req).to receive(:headers).and_return(initial_headers)
+          allow(req).to receive(:headers=) do |merged_headers|
+            # Provider headers should take precedence
+            expect(merged_headers['X-Api-Key']).to eq('provider-key')
+            expect(merged_headers['Content-Type']).to eq('application/json')
+            # User headers should be added
+            expect(merged_headers['X-Custom']).to eq('user-value')
+          end
+
+          block&.call(req)
+
+          instance_double(Faraday::Response,
+                          body: {
+                            'content' => [{ 'type' => 'text', 'text' => 'Test' }],
+                            'model' => 'test-model'
+                          })
+        end
+
+        chat.ask('Test')
+      end
+    end
+  end
+end