Fixing missing OrgId assignment issue in the ChatCompletionStream (#268)

* Update client.go

* Add test

---------

Co-authored-by: panjiajia <panjiajia@haoshuo.com>
This commit is contained in:
biubiu7
2023-04-20 21:08:29 +08:00
committed by GitHub
parent ecdea45b67
commit c2b58e77ed
2 changed files with 17 additions and 0 deletions

View File

@@ -42,6 +42,7 @@ func TestRequestAuthHeader(t *testing.T) {
APIType APIType APIType APIType
HeaderKey string HeaderKey string
Token string Token string
OrgID string
Expect string Expect string
}{ }{
{ {
@@ -49,6 +50,15 @@ func TestRequestAuthHeader(t *testing.T) {
"", "",
"Authorization", "Authorization",
"dummy-token-openai", "dummy-token-openai",
"",
"Bearer dummy-token-openai",
},
{
"OpenAIOrg",
APITypeOpenAI,
"Authorization",
"dummy-token-openai",
"dummy-org-openai",
"Bearer dummy-token-openai", "Bearer dummy-token-openai",
}, },
{ {
@@ -56,6 +66,7 @@ func TestRequestAuthHeader(t *testing.T) {
APITypeOpenAI, APITypeOpenAI,
"Authorization", "Authorization",
"dummy-token-openai", "dummy-token-openai",
"",
"Bearer dummy-token-openai", "Bearer dummy-token-openai",
}, },
{ {
@@ -63,6 +74,7 @@ func TestRequestAuthHeader(t *testing.T) {
APITypeAzureAD, APITypeAzureAD,
"Authorization", "Authorization",
"dummy-token-azure", "dummy-token-azure",
"",
"Bearer dummy-token-azure", "Bearer dummy-token-azure",
}, },
{ {
@@ -70,6 +82,7 @@ func TestRequestAuthHeader(t *testing.T) {
APITypeAzure, APITypeAzure,
AzureAPIKeyHeader, AzureAPIKeyHeader,
"dummy-api-key-here", "dummy-api-key-here",
"",
"dummy-api-key-here", "dummy-api-key-here",
}, },
} }
@@ -78,6 +91,7 @@ func TestRequestAuthHeader(t *testing.T) {
t.Run(c.Name, func(t *testing.T) { t.Run(c.Name, func(t *testing.T) {
az := DefaultConfig(c.Token) az := DefaultConfig(c.Token)
az.APIType = c.APIType az.APIType = c.APIType
az.OrgID = c.OrgID
cli := NewClientWithConfig(az) cli := NewClientWithConfig(az)
req, err := cli.newStreamRequest(context.Background(), "POST", "/chat/completions", nil) req, err := cli.newStreamRequest(context.Background(), "POST", "/chat/completions", nil)

View File

@@ -134,6 +134,9 @@ func (c *Client) newStreamRequest(
// OpenAI or Azure AD authentication // OpenAI or Azure AD authentication
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.config.authToken)) req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.config.authToken))
} }
if c.config.OrgID != "" {
req.Header.Set("OpenAI-Organization", c.config.OrgID)
}
return req, nil return req, nil
} }