From 8b3c0d0c52beb9cd8c1a952ed1378297f710fd9c Mon Sep 17 00:00:00 2001
From: Tor Hagemann <hagemt@users.noreply.github.com>
Date: Thu, 3 Jun 2021 14:24:15 -0700
Subject: [PATCH] api.go: add constructor variant for organizations

---
 api.go | 33 ++++++++++++++++++++++++++-------
 1 file changed, 26 insertions(+), 7 deletions(-)

diff --git a/api.go b/api.go
index fec3af0..1911ce4 100644
--- a/api.go
+++ b/api.go
@@ -9,28 +9,47 @@ import (
 
 const apiURLv1 = "https://api.openai.com/v1"
 
+func newTransport() *http.Client {
+	return &http.Client{
+		Timeout: time.Minute,
+	}
+}
+
 // Client is OpenAI GPT-3 API client
 type Client struct {
 	BaseURL    string
-	authToken  string
 	HTTPClient *http.Client
+	authToken  string
+	idOrg      string
 }
 
 // NewClient creates new OpenAI API client
 func NewClient(authToken string) *Client {
 	return &Client{
-		BaseURL:   apiURLv1,
-		authToken: authToken,
-		HTTPClient: &http.Client{
-			Timeout: time.Minute,
-		},
+		BaseURL:    apiURLv1,
+		HTTPClient: newTransport(),
+		authToken:  authToken,
+		idOrg:      "",
+	}
+}
+
+// NewOrgClient creates new OpenAI API client for specified Organization ID
+func NewOrgClient(authToken, org string) *Client {
+	return &Client{
+		BaseURL:    apiURLv1,
+		HTTPClient: newTransport(),
+		authToken:  authToken,
+		idOrg:      org,
 	}
 }
 
 func (c *Client) sendRequest(req *http.Request, v interface{}) error {
-	req.Header.Set("Content-Type", "application/json; charset=utf-8")
 	req.Header.Set("Accept", "application/json; charset=utf-8")
 	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.authToken))
+	req.Header.Set("Content-Type", "application/json; charset=utf-8")
+	if len(c.idOrg) > 0 {
+		req.Header.Set("OpenAI-Organization", c.idOrg)
+	}
 
 	res, err := c.HTTPClient.Do(req)
 	if err != nil {