Merge pull request #340 from overmindtech/azuuuuuure

Accept azureModelName configuration
overmindtech · Dec 9, 2024 · da33b2f · da33b2f
2 parents 8b110f6 + 7336910
commit da33b2f
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/llm/openai.go b/llm/openai.go
@@ -18,10 +18,13 @@ 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. Set `azureBaseURL` to enable azure mode
-func NewOpenAIProvider(apiKey string, azureBaseURL string, model string, name string, jsonMode bool) *openAIProvider {
+func NewOpenAIProvider(apiKey string, azureBaseURL string, model, azureModelName string, name string, jsonMode bool) *openAIProvider {
 	var cfg openai.ClientConfig
 	if azureBaseURL != "" {
 		cfg = openai.DefaultAzureConfig(apiKey, azureBaseURL)
+		cfg.AzureModelMapperFunc = func(m string) string {
+			return azureModelName
+		}
 	} else {
 		cfg = openai.DefaultConfig(apiKey)
 	}

diff --git a/llm/openai_test.go b/llm/openai_test.go
@@ -18,7 +18,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