From b6317ecf3466e3036d3c61b4e57e1be9df0af751 Mon Sep 17 00:00:00 2001
From: Vincent Durand <vincent.durand@madwork.org>
Date: Thu, 18 Jul 2024 10:22:26 +0200
Subject: [PATCH 1/3] Add persistent client connections

---
 lib/openai/client.rb              | 23 +++++++++++++++++++++--
 lib/openai/http.rb                | 24 +++++-------------------
 spec/openai/client/client_spec.rb | 10 ++++------
 3 files changed, 30 insertions(+), 27 deletions(-)

diff --git a/lib/openai/client.rb b/lib/openai/client.rb
index ac05d33f..4e1ce1b5 100644
--- a/lib/openai/client.rb
+++ b/lib/openai/client.rb
@@ -12,7 +12,7 @@ class Client
       request_timeout
       extra_headers
     ].freeze
-    attr_reader *CONFIG_KEYS, :faraday_middleware
+    attr_reader *CONFIG_KEYS
 
     def initialize(config = {}, &faraday_middleware)
       CONFIG_KEYS.each do |key|
@@ -23,7 +23,12 @@ def initialize(config = {}, &faraday_middleware)
           config[key].nil? ? OpenAI.configuration.send(key) : config[key]
         )
       end
-      @faraday_middleware = faraday_middleware
+
+      @connection = build_connection
+      faraday_middleware&.call(@connection)
+
+      @multipart_connection = build_connection(multipart: true)
+      faraday_middleware&.call(@multipart_connection)
     end
 
     def chat(parameters: {})
@@ -107,5 +112,19 @@ def beta(apis)
         client.add_headers("OpenAI-Beta": apis.map { |k, v| "#{k}=#{v}" }.join(";"))
       end
     end
+
+    private
+
+    attr_reader :connection, :multipart_connection
+
+    def build_connection(multipart: false)
+      Faraday.new do |faraday|
+        faraday.options[:timeout] = @request_timeout
+        faraday.request(:multipart) if multipart
+        faraday.use MiddlewareErrors if @log_errors
+        faraday.response :raise_error
+        faraday.response :json
+      end
+    end
   end
 end
diff --git a/lib/openai/http.rb b/lib/openai/http.rb
index d04cca5f..18a6a022 100644
--- a/lib/openai/http.rb
+++ b/lib/openai/http.rb
@@ -7,32 +7,32 @@ module HTTP
     include HTTPHeaders
 
     def get(path:, parameters: nil)
-      parse_jsonl(conn.get(uri(path: path), parameters) do |req|
+      parse_jsonl(connection.get(uri(path: path), parameters) do |req|
         req.headers = headers
       end&.body)
     end
 
     def post(path:)
-      parse_jsonl(conn.post(uri(path: path)) do |req|
+      parse_jsonl(connection.post(uri(path: path)) do |req|
         req.headers = headers
       end&.body)
     end
 
     def json_post(path:, parameters:)
-      conn.post(uri(path: path)) do |req|
+      connection.post(uri(path: path)) do |req|
         configure_json_post_request(req, parameters)
       end&.body
     end
 
     def multipart_post(path:, parameters: nil)
-      conn(multipart: true).post(uri(path: path)) do |req|
+      multipart_connection.post(uri(path: path)) do |req|
         req.headers = headers.merge({ "Content-Type" => "multipart/form-data" })
         req.body = multipart_parameters(parameters)
       end&.body
     end
 
     def delete(path:)
-      conn.delete(uri(path: path)) do |req|
+      connection.delete(uri(path: path)) do |req|
         req.headers = headers
       end&.body
     end
@@ -70,20 +70,6 @@ def to_json_stream(user_proc:)
       end
     end
 
-    def conn(multipart: false)
-      connection = Faraday.new do |f|
-        f.options[:timeout] = @request_timeout
-        f.request(:multipart) if multipart
-        f.use MiddlewareErrors if @log_errors
-        f.response :raise_error
-        f.response :json
-      end
-
-      @faraday_middleware&.call(connection)
-
-      connection
-    end
-
     def uri(path:)
       if azure?
         base = File.join(@uri_base, path)
diff --git a/spec/openai/client/client_spec.rb b/spec/openai/client/client_spec.rb
index daa1652a..b336525f 100644
--- a/spec/openai/client/client_spec.rb
+++ b/spec/openai/client/client_spec.rb
@@ -44,7 +44,7 @@
       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(:conn).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")
@@ -55,7 +55,7 @@
       expect(c1.request_timeout).to eq(60)
       expect(c1.uri_base).to eq("https://oai.hconeai.com/")
       expect(c1.send(:headers).values).to include(c1.access_token)
-      expect(c1.send(:conn).options.timeout).to eq(60)
+      expect(c1.send(:connection).options.timeout).to eq(60)
       expect(c1.send(:uri, path: "")).to include("https://oai.hconeai.com/")
       expect(c1.send(:headers).values).not_to include("X-Default")
       expect(c1.send(:headers).values).to include("X-Test")
@@ -67,7 +67,7 @@
       expect(c2.uri_base).to eq("https://example.com/")
       expect(c2.send(:headers).values).to include("Bearer #{c2.access_token}")
       expect(c2.send(:headers).values).to include(c2.organization_id)
-      expect(c2.send(:conn).options.timeout).to eq(1)
+      expect(c2.send(:connection).options.timeout).to eq(1)
       expect(c2.send(:uri, path: "")).to include("https://example.com/")
       expect(c2.send(:headers).values).to include("X-Default")
       expect(c2.send(:headers).values).not_to include("X-Test")
@@ -128,9 +128,7 @@
     end
 
     it "sets the logger" do
-      connection = Faraday.new
-      client.faraday_middleware.call(connection)
-      expect(connection.builder.handlers).to include Faraday::Response::Logger
+      expect(client.send(:connection).builder.handlers).to include Faraday::Response::Logger
     end
   end
 end

From 6cfabd6929888788d8d86808a5c1e01ce57a57a7 Mon Sep 17 00:00:00 2001
From: Vincent Durand <vincent.durand@madwork.org>
Date: Wed, 17 Jul 2024 15:26:38 +0200
Subject: [PATCH 2/3] Remove beta client duplication

The current implementation is breaking any attempt of connection pool usage.
---
 lib/middleware/beta.rb                  | 20 ++++++++++++++++++++
 lib/openai.rb                           |  2 ++
 lib/openai/assistants.rb                |  4 +---
 lib/openai/batches.rb                   |  2 +-
 lib/openai/client.rb                    |  7 +------
 lib/openai/images.rb                    |  2 +-
 lib/openai/messages.rb                  |  2 +-
 lib/openai/run_steps.rb                 |  2 +-
 lib/openai/runs.rb                      |  2 +-
 lib/openai/threads.rb                   |  2 +-
 lib/openai/vector_store_file_batches.rb |  2 +-
 lib/openai/vector_store_files.rb        |  2 +-
 lib/openai/vector_stores.rb             |  2 +-
 spec/openai/client/client_spec.rb       | 12 +++---------
 14 files changed, 36 insertions(+), 27 deletions(-)
 create mode 100644 lib/middleware/beta.rb

diff --git a/lib/middleware/beta.rb b/lib/middleware/beta.rb
new file mode 100644
index 00000000..1c22fb34
--- /dev/null
+++ b/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
diff --git a/lib/openai.rb b/lib/openai.rb
index 0fad1b2e..25638ac8 100644
--- a/lib/openai.rb
+++ b/lib/openai.rb
@@ -90,3 +90,5 @@ def self.rough_token_count(content = "")
     [1, estimate].max
   end
 end
+
+require_relative "middleware/beta"
diff --git a/lib/openai/assistants.rb b/lib/openai/assistants.rb
index c7fbc2d0..f665cfd3 100644
--- a/lib/openai/assistants.rb
+++ b/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
diff --git a/lib/openai/batches.rb b/lib/openai/batches.rb
index 86624081..ccb47458 100644
--- a/lib/openai/batches.rb
+++ b/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: {})
diff --git a/lib/openai/client.rb b/lib/openai/client.rb
index 4e1ce1b5..2b84b562 100644
--- a/lib/openai/client.rb
+++ b/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
diff --git a/lib/openai/images.rb b/lib/openai/images.rb
index 391e3c52..237af84e 100644
--- a/lib/openai/images.rb
+++ b/lib/openai/images.rb
@@ -1,6 +1,6 @@
 module OpenAI
   class Images
-    def initialize(client: nil)
+    def initialize(client:)
       @client = client
     end
 
diff --git a/lib/openai/messages.rb b/lib/openai/messages.rb
index 1e335fd7..63cfe5f5 100644
--- a/lib/openai/messages.rb
+++ b/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: {})
diff --git a/lib/openai/run_steps.rb b/lib/openai/run_steps.rb
index 0a436dfe..d312fad0 100644
--- a/lib/openai/run_steps.rb
+++ b/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: {})
diff --git a/lib/openai/runs.rb b/lib/openai/runs.rb
index e1af4601..b2b0bf9f 100644
--- a/lib/openai/runs.rb
+++ b/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: {})
diff --git a/lib/openai/threads.rb b/lib/openai/threads.rb
index ed494aac..c2fd9a35 100644
--- a/lib/openai/threads.rb
+++ b/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:)
diff --git a/lib/openai/vector_store_file_batches.rb b/lib/openai/vector_store_file_batches.rb
index c5aa7e52..ceaf19e6 100644
--- a/lib/openai/vector_store_file_batches.rb
+++ b/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: {})
diff --git a/lib/openai/vector_store_files.rb b/lib/openai/vector_store_files.rb
index fc338612..bd791597 100644
--- a/lib/openai/vector_store_files.rb
+++ b/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: {})
diff --git a/lib/openai/vector_stores.rb b/lib/openai/vector_stores.rb
index e16410ee..253ebcd3 100644
--- a/lib/openai/vector_stores.rb
+++ b/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: {})
diff --git a/spec/openai/client/client_spec.rb b/spec/openai/client/client_spec.rb
index b336525f..86844225 100644
--- a/spec/openai/client/client_spec.rb
+++ b/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|

From 475b25e15c9fcf43d9733246e9c6f3cf122b7484 Mon Sep 17 00:00:00 2001
From: Vincent Durand <vincent.durand@madwork.org>
Date: Thu, 18 Jul 2024 21:41:16 +0200
Subject: [PATCH 3/3] Add MiddlewareErrors to own file

---
 lib/middleware/beta.rb   |  2 +-
 lib/middleware/errors.rb | 19 +++++++++++++++++++
 lib/openai.rb            | 19 ++-----------------
 lib/openai/client.rb     |  2 +-
 4 files changed, 23 insertions(+), 19 deletions(-)
 create mode 100644 lib/middleware/errors.rb

diff --git a/lib/middleware/beta.rb b/lib/middleware/beta.rb
index 1c22fb34..ba58bc66 100644
--- a/lib/middleware/beta.rb
+++ b/lib/middleware/beta.rb
@@ -1,7 +1,7 @@
 # frozen_string_literal: true
 
 module OpenAI
-  class BetaMiddleware < Faraday::Middleware
+  class MiddlewareBeta < Faraday::Middleware
     BETA_REGEX = %r{
       \A/#{OpenAI.configuration.api_version}
       /(assistants|batches|threads|vector_stores)
diff --git a/lib/middleware/errors.rb b/lib/middleware/errors.rb
new file mode 100644
index 00000000..d58e436c
--- /dev/null
+++ b/lib/middleware/errors.rb
@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module OpenAI
+  class MiddlewareErrors < Faraday::Middleware
+    def call(env)
+      @app.call(env)
+    rescue Faraday::Error => e
+      raise e unless e.response.is_a?(Hash)
+
+      logger = Logger.new($stdout)
+      logger.formatter = proc do |_severity, _datetime, _progname, msg|
+        "\033[31mOpenAI HTTP Error (spotted in ruby-openai #{VERSION}): #{msg}\n\033[0m"
+      end
+      logger.error(e.response[:body])
+
+      raise e
+    end
+  end
+end
diff --git a/lib/openai.rb b/lib/openai.rb
index 25638ac8..5a584072 100644
--- a/lib/openai.rb
+++ b/lib/openai.rb
@@ -23,21 +23,8 @@ module OpenAI
   class Error < StandardError; end
   class ConfigurationError < Error; end
 
-  class MiddlewareErrors < Faraday::Middleware
-    def call(env)
-      @app.call(env)
-    rescue Faraday::Error => e
-      raise e unless e.response.is_a?(Hash)
-
-      logger = Logger.new($stdout)
-      logger.formatter = proc do |_severity, _datetime, _progname, msg|
-        "\033[31mOpenAI HTTP Error (spotted in ruby-openai #{VERSION}): #{msg}\n\033[0m"
-      end
-      logger.error(e.response[:body])
-
-      raise e
-    end
-  end
+  autoload :MiddlewareErrors, "middleware/errors"
+  autoload :MiddlewareBeta, "middleware/beta"
 
   class Configuration
     attr_accessor :access_token,
@@ -90,5 +77,3 @@ def self.rough_token_count(content = "")
     [1, estimate].max
   end
 end
-
-require_relative "middleware/beta"
diff --git a/lib/openai/client.rb b/lib/openai/client.rb
index 2b84b562..fa5cd304 100644
--- a/lib/openai/client.rb
+++ b/lib/openai/client.rb
@@ -116,7 +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.use MiddlewareBeta
         faraday.response :raise_error
         faraday.response :json
       end