Remove beta client duplication

The current implementation is breaking any attempt of connection pool usage.

Author: madwork

lib/middleware/beta.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+module OpenAI
+  class BetaMiddleware < Faraday::Middleware
+    BETA_REGEX = %r{
+      \A/#{OpenAI.configuration.api_version}
+      /(assistants|batches|threads|vector_stores)
+    }ix.freeze
+
+    def on_request(env)
+      return unless env[:url].path.match?(BETA_REGEX)
+
+      env[:request_headers].merge!(
+        {
+          "OpenAI-Beta" => "assistants=v2"
+        }
+      )
+    end
+  end
+end

lib/openai.rb

@@ -90,3 +90,5 @@ def self.rough_token_count(content = "")
     [1, estimate].max
   end
 end
+
+require_relative "middleware/beta"

lib/openai/assistants.rb

@@ -1,9 +1,7 @@
 module OpenAI
   class Assistants
-    BETA_VERSION = "v2".freeze
-
     def initialize(client:)
-      @client = client.beta(assistants: OpenAI::Assistants::BETA_VERSION)
+      @client = client
     end
 
     def list

lib/openai/batches.rb

@@ -1,7 +1,7 @@
 module OpenAI
   class Batches
     def initialize(client:)
-      @client = client.beta(assistants: OpenAI::Assistants::BETA_VERSION)
+      @client = client
     end
 
     def list(parameters: {})

lib/openai/client.rb

@@ -107,12 +107,6 @@ def azure?
       @api_type&.to_sym == :azure
     end
 
-    def beta(apis)
-      dup.tap do |client|
-        client.add_headers("OpenAI-Beta": apis.map { |k, v| "#{k}=#{v}" }.join(";"))
-      end
-    end
-
     private
 
     attr_reader :connection, :multipart_connection
@@ -122,6 +116,7 @@ def build_connection(multipart: false)
         faraday.options[:timeout] = @request_timeout
         faraday.request(:multipart) if multipart
         faraday.use MiddlewareErrors if @log_errors
+        faraday.use BetaMiddleware
         faraday.response :raise_error
         faraday.response :json
       end

lib/openai/images.rb

@@ -1,6 +1,6 @@
 module OpenAI
   class Images
-    def initialize(client: nil)
+    def initialize(client:)
       @client = client
     end
 

lib/openai/messages.rb

@@ -1,7 +1,7 @@
 module OpenAI
   class Messages
     def initialize(client:)
-      @client = client.beta(assistants: OpenAI::Assistants::BETA_VERSION)
+      @client = client
     end
 
     def list(thread_id:, parameters: {})

lib/openai/run_steps.rb

@@ -1,7 +1,7 @@
 module OpenAI
   class RunSteps
     def initialize(client:)
-      @client = client.beta(assistants: OpenAI::Assistants::BETA_VERSION)
+      @client = client
     end
 
     def list(thread_id:, run_id:, parameters: {})

lib/openai/runs.rb

@@ -1,7 +1,7 @@
 module OpenAI
   class Runs
     def initialize(client:)
-      @client = client.beta(assistants: OpenAI::Assistants::BETA_VERSION)
+      @client = client
     end
 
     def list(thread_id:, parameters: {})

lib/openai/threads.rb

@@ -1,7 +1,7 @@
 module OpenAI
   class Threads
     def initialize(client:)
-      @client = client.beta(assistants: OpenAI::Assistants::BETA_VERSION)
+      @client = client
     end
 
     def retrieve(id:)

lib/openai/vector_store_file_batches.rb

@@ -1,7 +1,7 @@
 module OpenAI
   class VectorStoreFileBatches
     def initialize(client:)
-      @client = client.beta(assistants: OpenAI::Assistants::BETA_VERSION)
+      @client = client
     end
 
     def list(vector_store_id:, id:, parameters: {})

lib/openai/vector_store_files.rb

@@ -1,7 +1,7 @@
 module OpenAI
   class VectorStoreFiles
     def initialize(client:)
-      @client = client.beta(assistants: OpenAI::Assistants::BETA_VERSION)
+      @client = client
     end
 
     def list(vector_store_id:, parameters: {})

lib/openai/vector_stores.rb

@@ -1,7 +1,7 @@
 module OpenAI
   class VectorStores
     def initialize(client:)
-      @client = client.beta(assistants: OpenAI::Assistants::BETA_VERSION)
+      @client = client
     end
 
     def list(parameters: {})

spec/openai/client/client_spec.rb

@@ -44,7 +44,9 @@
       expect(c0.uri_base).to eq(OpenAI::Configuration::DEFAULT_URI_BASE)
       expect(c0.send(:headers).values).to include("Bearer #{c0.access_token}")
       expect(c0.send(:headers).values).to include(c0.organization_id)
-      expect(c0.send(:connection).options.timeout).to eq(OpenAI::Configuration::DEFAULT_REQUEST_TIMEOUT)
+      expect(c0.send(:connection).options.timeout).to eq(
+        OpenAI::Configuration::DEFAULT_REQUEST_TIMEOUT
+      )
       expect(c0.send(:uri, path: "")).to include(OpenAI::Configuration::DEFAULT_URI_BASE)
       expect(c0.send(:headers).values).to include("X-Default")
       expect(c0.send(:headers).values).not_to include("X-Test")
@@ -112,14 +114,6 @@
     end
   end
 
-  context "when using beta APIs" do
-    let(:client) { OpenAI::Client.new.beta(assistants: "v2") }
-
-    it "sends the appropriate header value" do
-      expect(client.send(:headers)["OpenAI-Beta"]).to eq "assistants=v2"
-    end
-  end
-
   context "with a block" do
     let(:client) do
       OpenAI::Client.new do |client|