Merge pull request #2 from Unsupervisedcom/azure-token-provider

Azure token provider

Author: CarterBland

README.md

@@ -218,6 +218,21 @@ To use the [Azure OpenAI Service](https://learn.microsoft.com/en-us/azure/cognit
 
 where `AZURE_OPENAI_URI` is e.g. `https://custom-domain.openai.azure.com/openai/deployments/gpt-35-turbo`
 
+##### Azure with Azure AD tokens
+
+To use Azure AD tokens you can configure the gem with a proc like this:
+
+```ruby
+    OpenAI.configure do |config|
+        config.azure_token_provider = ->() { your_code_caches_or_refreshes_token }
+        config.uri_base = ENV.fetch("AZURE_OPENAI_URI")
+        config.api_type = :azure
+        config.api_version = "2023-03-15-preview"
+    end
+```
+
+The azure_token_provider will be called on every request. This allows tokens to be cached and periodically refreshed by your custom code.
+
 #### Ollama
 
 Ollama allows you to run open-source LLMs, such as Llama 3, locally. It [offers chat compatibility](https://github.com/ollama/ollama/blob/main/docs/openai.md) with the OpenAI API.

lib/openai.rb

@@ -47,7 +47,8 @@ class Configuration
                   :organization_id,
                   :uri_base,
                   :request_timeout,
-                  :extra_headers
+                  :extra_headers,
+                  :azure_token_provider
 
     DEFAULT_API_VERSION = "v1".freeze
     DEFAULT_URI_BASE = "https://api.openai.com/".freeze
@@ -63,6 +64,7 @@ def initialize
       @uri_base = DEFAULT_URI_BASE
       @request_timeout = DEFAULT_REQUEST_TIMEOUT
       @extra_headers = {}
+      @azure_token_provider = nil
     end
   end
 

lib/openai/client.rb

@@ -12,6 +12,7 @@ class Client
       uri_base
       request_timeout
       extra_headers
+      azure_token_provider
     ].freeze
     attr_reader *CONFIG_KEYS, :faraday_middleware
 
@@ -25,6 +26,8 @@ def initialize(config = {}, &faraday_middleware)
         )
       end
       @faraday_middleware = faraday_middleware
+      validate_credential_config!
+      validate_azure_credential_provider!
     end
 
     def chat(parameters: {})
@@ -109,6 +112,31 @@ def beta(apis)
       end
     end
 
+    private
+
+    def validate_credential_config!
+      if @access_token && @azure_token_provider
+        raise ConfigurationError,
+              "Only one of OpenAI access token or Azure token provider can be set! See https://github.com/alexrudall/ruby-openai#usage"
+      end
+
+      return if @access_token || @azure_token_provider
+
+      raise ConfigurationError,
+            "OpenAI access token or Azure token provider missing! See https://github.com/alexrudall/ruby-openai#usage"
+    end
+
+    def validate_azure_credential_provider!
+      return if @azure_token_provider.nil?
+
+      unless @azure_token_provider.respond_to?(:to_proc)
+        raise ConfigurationError,
+              "OpenAI Azure AD token provider must be a Proc, Lambda, or respond to to_proc."
+      end
+
+      @azure_token_provider = @azure_token_provider&.to_proc
+    end
+
     def inspect
       vars = instance_variables.map do |var|
         value = instance_variable_get(var)

lib/openai/http_headers.rb

@@ -24,9 +24,18 @@ def openai_headers
 
     def azure_headers
       {
-        "Content-Type" => "application/json",
-        "api-key" => @access_token
-      }
+        "Content-Type" => "application/json"
+      }.merge(azure_auth_headers)
+    end
+
+    def azure_auth_headers
+      if @access_token
+        { "api-key" => @access_token }
+      elsif @azure_token_provider
+        { "Authorization" => "Bearer #{@azure_token_provider.call}" }
+      else
+        raise ConfigurationError, "access_token or azure_token_provider must be set."
+      end
     end
 
     def extra_headers

spec/fixtures/cassettes/http_json_post_with_azure_token_provider.yml

@@ -30,7 +30,6 @@
     let!(:c2) do
       OpenAI::Client.new(
         access_token: "access_token2",
-        organization_id: nil,
         request_timeout: 1,
         uri_base: "https://example.com/"
       )

spec/openai/client/client_spec.rb

@@ -120,6 +120,56 @@
     end
   end
 
+  describe ".json_post" do
+    context "with azure_token_provider" do
+      let(:token_provider) do
+        counter = 0
+        lambda do
+          counter += 1
+          "some dynamic token #{counter}"
+        end
+      end
+
+      let(:client) do
+        OpenAI::Client.new(
+          access_token: nil,
+          azure_token_provider: token_provider,
+          api_type: :azure,
+          uri_base: "https://custom-domain.openai.azure.com/openai/deployments/gpt-35-turbo",
+          api_version: "2024-02-01"
+        )
+      end
+
+      let(:cassette) { "http json post with azure token provider" }
+
+      it "calls the token provider on every request" do
+        expect(token_provider).to receive(:call).twice.and_call_original
+        VCR.use_cassette(cassette, record: :none) do
+          client.chat(
+            parameters: {
+              messages: [
+                {
+                  "role" => "user",
+                  "content" => "Hello world!"
+                }
+              ]
+            }
+          )
+          client.chat(
+            parameters: {
+              messages: [
+                {
+                  "role" => "user",
+                  "content" => "Who were the founders of Microsoft?"
+                }
+              ]
+            }
+          )
+        end
+      end
+    end
+  end
+
   describe ".to_json_stream" do
     context "with a proc" do
       let(:user_proc) { proc { |x| x } }
@@ -281,6 +331,28 @@
         expect(headers).to eq({ "Content-Type" => "application/json",
                                 "api-key" => OpenAI.configuration.access_token })
       }
+
+      context "with azure_token_provider" do
+        let(:token) { "some dynamic token" }
+        let(:token_provider) { -> { token } }
+
+        around do |ex|
+          old_access_token = OpenAI.configuration.access_token
+          OpenAI.configuration.access_token = nil
+          OpenAI.configuration.azure_token_provider = token_provider
+
+          ex.call
+        ensure
+          OpenAI.configuration.azure_token_provider = nil
+          OpenAI.configuration.access_token = old_access_token
+        end
+
+        it {
+          expect(token_provider).to receive(:call).once.and_call_original
+          expect(headers).to eq({ "Content-Type" => "application/json",
+                                  "Authorization" => "Bearer #{token}" })
+        }
+      end
     end
   end