From c2b58e77ed93ee12b97c419885fe4f5d13639f16 Mon Sep 17 00:00:00 2001
From: biubiu7 <panjiajia.coder@gmail.com>
Date: Thu, 20 Apr 2023 21:08:29 +0800
Subject: [PATCH] Fixing missing OrgId assignment issue in the
 ChatCompletionStream (#268)

* Update client.go

* Add test

---------

Co-authored-by: panjiajia <panjiajia@haoshuo.com>
---
 api_internal_test.go | 14 ++++++++++++++
 client.go            |  3 +++
 2 files changed, 17 insertions(+)

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
 }