feat: add azure openai support (#214)

* feat: add azure openai support * chore: refine config * chore: make config options like the python one * chore: adjust config struct field order * test: fix tests * style: make the linter happy * fix: support Azure API Key authentication in sendRequest * chore: check error in CreateChatCompletionStream * chore: pass tests * chore: try pass tests again * chore: change ClientConfig back due to this lib does not like WithXxx config style * chore: revert fix to CreateChatCompletionStream() due to cause tests not pass * chore: at least add some comment about the required fields * chore: re order ClientConfig fields * chore: add DefaultAzure() * chore: set default api_version the same as py one "2023-03-15-preview" * style: fixup typo * test: add api_internal_test.go * style: make lint happy * chore: add constant AzureAPIKeyHeader * chore: use AzureAPIKeyHeader for api-key header, fix azure base url auto trim suffix / * test: add TestAzureFullURL, TestRequestAuthHeader and TestOpenAIFullURL * test: simplify TestRequestAuthHeader * test: refine TestOpenAIFullURL * chore: refine comments * feat: DefaultAzureConfig
2023-04-04 16:05:20 +08:00
parent bee0656174
commit 8677fb4bb4
3 changed files with 198 additions and 9 deletions
--- a/api_internal_test.go
+++ b/api_internal_test.go
@@ -0,0 +1,133 @@
+package openai
+
+import (
+	"context"
+	"testing"
+)
+
+func TestOpenAIFullURL(t *testing.T) {
+	cases := []struct {
+		Name   string
+		Suffix string
+		Expect string
+	}{
+		{
+			"ChatCompletionsURL",
+			"/chat/completions",
+			"https://api.openai.com/v1/chat/completions",
+		},
+		{
+			"CompletionsURL",
+			"/completions",
+			"https://api.openai.com/v1/completions",
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.Name, func(t *testing.T) {
+			az := DefaultConfig("dummy")
+			cli := NewClientWithConfig(az)
+			actual := cli.fullURL(c.Suffix)
+			if actual != c.Expect {
+				t.Errorf("Expected %s, got %s", c.Expect, actual)
+			}
+			t.Logf("Full URL: %s", actual)
+		})
+	}
+}
+
+func TestRequestAuthHeader(t *testing.T) {
+	cases := []struct {
+		Name      string
+		APIType   APIType
+		HeaderKey string
+		Token     string
+		Expect    string
+	}{
+		{
+			"OpenAIDefault",
+			"",
+			"Authorization",
+			"dummy-token-openai",
+			"Bearer dummy-token-openai",
+		},
+		{
+			"OpenAI",
+			APITypeOpenAI,
+			"Authorization",
+			"dummy-token-openai",
+			"Bearer dummy-token-openai",
+		},
+		{
+			"AzureAD",
+			APITypeAzureAD,
+			"Authorization",
+			"dummy-token-azure",
+			"Bearer dummy-token-azure",
+		},
+		{
+			"Azure",
+			APITypeAzure,
+			AzureAPIKeyHeader,
+			"dummy-api-key-here",
+			"dummy-api-key-here",
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.Name, func(t *testing.T) {
+			az := DefaultConfig(c.Token)
+			az.APIType = c.APIType
+
+			cli := NewClientWithConfig(az)
+			req, err := cli.newStreamRequest(context.Background(), "POST", "/chat/completions", nil)
+			if err != nil {
+				t.Errorf("Failed to create request: %v", err)
+			}
+			actual := req.Header.Get(c.HeaderKey)
+			if actual != c.Expect {
+				t.Errorf("Expected %s, got %s", c.Expect, actual)
+			}
+			t.Logf("%s: %s", c.HeaderKey, actual)
+		})
+	}
+}
+
+func TestAzureFullURL(t *testing.T) {
+	cases := []struct {
+		Name    string
+		BaseURL string
+		Engine  string
+		Expect  string
+	}{
+		{
+			"AzureBaseURLWithSlashAutoStrip",
+			"https://httpbin.org/",
+			"chatgpt-demo",
+			"https://httpbin.org/" +
+				"openai/deployments/chatgpt-demo" +
+				"/chat/completions?api-version=2023-03-15-preview",
+		},
+		{
+			"AzureBaseURLWithoutSlashOK",
+			"https://httpbin.org",
+			"chatgpt-demo",
+			"https://httpbin.org/" +
+				"openai/deployments/chatgpt-demo" +
+				"/chat/completions?api-version=2023-03-15-preview",
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.Name, func(t *testing.T) {
+			az := DefaultAzureConfig("dummy", c.BaseURL, c.Engine)
+			cli := NewClientWithConfig(az)
+			// /openai/deployments/{engine}/chat/completions?api-version={api_version}
+			actual := cli.fullURL("/chat/completions")
+			if actual != c.Expect {
+				t.Errorf("Expected %s, got %s", c.Expect, actual)
+			}
+			t.Logf("Full URL: %s", actual)
+		})
+	}
+}