Enable use of Azure Endpoints for OpenAI

overmindtech · Dec 3, 2024 · 76322f0 · 76322f0
1 parent 5600877
commit 76322f0
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/llm/openai.go b/llm/openai.go
@@ -17,9 +17,15 @@ import (
 
 // Creates a new OpenAI Provider. The user must pass in the API key, the model
 // to use, and the name. The name is used in the subsequent names of all the
-// assistants that are created
-func NewOpenAIProvider(apiKey string, model string, name string, jsonMode bool) *openAIProvider {
-	cfg := openai.DefaultConfig(apiKey)
+// assistants that are created. Set `azureBaseURL` to enable azure mode
+func NewOpenAIProvider(apiKey string, azureBaseURL string, model string, name string, jsonMode bool) *openAIProvider {
+	var cfg openai.ClientConfig
+	if azureBaseURL != "" {
+		cfg = openai.DefaultAzureConfig(apiKey, azureBaseURL)
+	} else {
+		cfg = openai.DefaultConfig(apiKey)
+	}
+
 	cfg.HTTPClient = otelhttp.DefaultClient
 
 	return &openAIProvider{

diff --git a/llm/openai_test.go b/llm/openai_test.go
@@ -16,7 +16,7 @@ func TestNewOpenAIProvider(t *testing.T) {
 		t.Skip("OPENAI_API_KEY not set")
 	}
 
-	openaiProvider := NewOpenAIProvider(key, openai.GPT4oMini, t.Name(), false)
+	openaiProvider := NewOpenAIProvider(key, "", openai.GPT4oMini, t.Name(), false)
 
 	// Assert that the result matches the provider interface
 	var _ Provider = openaiProvider