diff --git a/api_internal_test.go b/api_internal_test.go index 83dcafc..9651ad4 100644 --- a/api_internal_test.go +++ b/api_internal_test.go @@ -42,6 +42,7 @@ func TestRequestAuthHeader(t *testing.T) { APIType APIType HeaderKey string Token string + OrgID string Expect string }{ { @@ -49,6 +50,15 @@ func TestRequestAuthHeader(t *testing.T) { "", "Authorization", "dummy-token-openai", + "", + "Bearer dummy-token-openai", + }, + { + "OpenAIOrg", + APITypeOpenAI, + "Authorization", + "dummy-token-openai", + "dummy-org-openai", "Bearer dummy-token-openai", }, { @@ -56,6 +66,7 @@ func TestRequestAuthHeader(t *testing.T) { APITypeOpenAI, "Authorization", "dummy-token-openai", + "", "Bearer dummy-token-openai", }, { @@ -63,6 +74,7 @@ func TestRequestAuthHeader(t *testing.T) { APITypeAzureAD, "Authorization", "dummy-token-azure", + "", "Bearer dummy-token-azure", }, { @@ -70,6 +82,7 @@ func TestRequestAuthHeader(t *testing.T) { APITypeAzure, AzureAPIKeyHeader, "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) { az := DefaultConfig(c.Token) az.APIType = c.APIType + az.OrgID = c.OrgID cli := NewClientWithConfig(az) req, err := cli.newStreamRequest(context.Background(), "POST", "/chat/completions", nil) diff --git a/client.go b/client.go index e17ded2..368947b 100644 --- a/client.go +++ b/client.go @@ -134,6 +134,9 @@ func (c *Client) newStreamRequest( // OpenAI or Azure AD authentication 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 }