Add support for OpenAI Organization ID header

openai/client.go

@@ -78,21 +78,29 @@ func (c *CreateEmbeddingResponse) String() string {
 type OpenAIClient struct {
 	BaseURL string
 	APIKey  string
+	OrgID   string
 	Client  *http.Client
 }
 
-func NewOpenAIClient(apiKey string) *OpenAIClient {
-	return &OpenAIClient{
+func NewOpenAIClient(apiKey string, opts ...Option) *OpenAIClient {
+	client := &OpenAIClient{
 		BaseURL: "https://api.openai.com/v1/",
 		Client:  &http.Client{},
 		APIKey:  apiKey,
 	}
+	applyClientOptions(client, opts...)
+
+	return client
 }
 
 func (c *OpenAIClient) SetAPIKey(apiKey string) {
 	c.APIKey = apiKey
 }
 
+func (c *OpenAIClient) SetOrgID(orgID string) {
+	c.OrgID = orgID
+}
+
 func (c *OpenAIClient) SetBaseURL(baseURL string) {
 	c.BaseURL = baseURL
 }
@@ -118,6 +126,11 @@ func (c *OpenAIClient) CreateEmbedding(ctx context.Context, req *CreateEmbedding
 	httpReq.Header.Set("Content-Type", "application/json")
 	httpReq.Header.Set("Authorization", "Bearer "+c.getAPIKey())
 
+	// OpenAI Organization ID (Optional)
+	if c.OrgID != "" {
+		httpReq.Header.Set("OpenAI-Organization", c.OrgID)
+	}
+
 	resp, err := c.Client.Do(httpReq)
 	if err != nil {
 		return nil, err
@@ -147,9 +160,9 @@ type OpenAIEmbeddingFunction struct {
 	apiClient *OpenAIClient
 }
 
-func NewOpenAIEmbeddingFunction(apiKey string) *OpenAIEmbeddingFunction {
+func NewOpenAIEmbeddingFunction(apiKey string, opts ...Option) *OpenAIEmbeddingFunction {
 	cli := &OpenAIEmbeddingFunction{
-		apiClient: NewOpenAIClient(apiKey),
+		apiClient: NewOpenAIClient(apiKey, opts...),
 	}
 
 	return cli

openai/client_test.go

@@ -20,9 +20,10 @@ func Test_openai_client(t *testing.T) {
 		}
 		apiKey = os.Getenv("OPENAI_API_KEY")
 	}
-	ef := NewOpenAIEmbeddingFunction(apiKey)
 
 	t.Run("Test DefaultApiService Add", func(t *testing.T) {
+		ef := NewOpenAIEmbeddingFunction(apiKey)
+
 		documents := []string{
 			"Document 1 content here",
 			"Document 2 content here",
@@ -33,5 +34,18 @@ func Test_openai_client(t *testing.T) {
 		require.NotNil(t, resp)
 		fmt.Printf("resp: %v\n", resp)
 		// assert.Equal(t, 201, httpRes.StatusCode)
+		require.Empty(t, ef.apiClient.OrgID)
+	})
+
+	t.Run("Test Adding Organization Id with NewOpenAIClient", func(t *testing.T) {
+		apiClient := NewOpenAIClient(apiKey, WithOpenAIOrganizationID("org-123"))
+
+		require.Equal(t, "org-123", apiClient.OrgID)
+	})
+
+	t.Run("Test Adding Organization Id with NewOpenAIEmbeddingFunction", func(t *testing.T) {
+		ef := NewOpenAIEmbeddingFunction(apiKey, WithOpenAIOrganizationID("org-123"))
+
+		require.Equal(t, "org-123", ef.apiClient.OrgID)
 	})
 }

openai/options.go

@@ -0,0 +1,17 @@
+package openai
+
+// Option is a function type that can be used to modify the client.
+type Option func(c *OpenAIClient)
+
+// WithOpenAIOrganizationID is an option for setting the OpenAI org id.
+func WithOpenAIOrganizationID(openAiAPIKey string) Option {
+	return func(c *OpenAIClient) {
+		c.SetOrgID(openAiAPIKey)
+	}
+}
+
+func applyClientOptions(c *OpenAIClient, opts ...Option) {
+	for _, opt := range opts {
+		opt(c)
+	}
+}