From 9ee7dff064a9412c069a793ff62096b8db4aa43d Mon Sep 17 00:00:00 2001
From: Krzysztof Cwalina <kcwalina@microsoft.com>
Date: Fri, 21 Jun 2024 08:47:02 -0700
Subject: [PATCH] Added support for OrganizationId and ProjectId (#80)

* added support for OrganizationId and ProjectId

* changed header names to consts
---
 src/Custom/OpenAIClient.cs        | 38 +++++++++++++++++++------------
 src/Custom/OpenAIClientOptions.cs | 10 ++++++++
 2 files changed, 33 insertions(+), 15 deletions(-)

diff --git a/src/Custom/OpenAIClient.cs b/src/Custom/OpenAIClient.cs
index e7ffd76e..a9718090 100644
--- a/src/Custom/OpenAIClient.cs
+++ b/src/Custom/OpenAIClient.cs
@@ -6,17 +6,13 @@
 using OpenAI.Files;
 using OpenAI.FineTuning;
 using OpenAI.Images;
-using OpenAI.LegacyCompletions;
 using OpenAI.Models;
 using OpenAI.Moderations;
 using OpenAI.VectorStores;
 using System;
 using System.ClientModel;
 using System.ClientModel.Primitives;
-using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
-using System.Reflection;
-using System.Runtime.Versioning;
 
 namespace OpenAI;
 
@@ -220,12 +216,13 @@ internal static ClientPipeline CreatePipeline(ApiKeyCredential credential, OpenA
     {
         return ClientPipeline.Create(
             options ?? new(),
-            perCallPolicies: [],
+            perCallPolicies: [
+                CreateAddBetaFeatureHeaderPolicy(),
+                CreateAddCustomHeadersPolicy(options),
+            ],
             perTryPolicies:
             [
-                ApiKeyAuthenticationPolicy.CreateHeaderApiKeyPolicy(credential, AuthorizationHeader, AuthorizationApiKeyPrefix),
-                CreateAddBetaFeatureHeaderPolicy(),
-                CreateAddUserAgentHeaderPolicy(options),
+                ApiKeyAuthenticationPolicy.CreateHeaderApiKeyPolicy(credential, AuthorizationHeader, AuthorizationApiKeyPrefix)
             ],
             beforeTransportPolicies: []);
     }
@@ -263,29 +260,40 @@ private static PipelinePolicy CreateAddBetaFeatureHeaderPolicy()
     {
         return new GenericActionPipelinePolicy((message) =>
         {
-            if (message?.Request?.Headers?.TryGetValue(OpenAIBetaFeatureHeaderKey, out string _) == false)
+            if (message?.Request?.Headers?.TryGetValue(OpenAIBetaFeatureHeaderName, out string _) == false)
             {
-                message.Request.Headers.Set(OpenAIBetaFeatureHeaderKey, OpenAIBetaAssistantsV1HeaderValue);
+                message.Request.Headers.Set(OpenAIBetaFeatureHeaderName, OpenAIBetaAssistantsV1HeaderValue);
             }
         });
     }
 
-    private static PipelinePolicy CreateAddUserAgentHeaderPolicy(OpenAIClientOptions options = null)
+    private static PipelinePolicy CreateAddCustomHeadersPolicy(OpenAIClientOptions options = null)
     {
         TelemetryDetails telemetryDetails = new(typeof(OpenAIClientOptions).Assembly, options?.ApplicationId);
         return new GenericActionPipelinePolicy((message) =>
         {
-            if (message?.Request?.Headers?.TryGetValue(UserAgentHeaderKey, out string _) == false)
+            if (message?.Request?.Headers?.TryGetValue(UserAgentHeaderName, out string _) == false)
             {
-                message.Request.Headers.Set(UserAgentHeaderKey, telemetryDetails.ToString());
+                message.Request.Headers.Set(UserAgentHeaderName, telemetryDetails.ToString());
+            }
+
+            if (!string.IsNullOrEmpty(options.OrganizationId))
+            {
+                message.Request.Headers.Set(OpenAIOrganizationHeaderName, options.OrganizationId);
+            }
+            if (!string.IsNullOrEmpty(options.ProjectId))
+            {
+                message.Request.Headers.Set(OpenAIProjectHeaderName, options.ProjectId);
             }
         });
     }
 
-    private const string OpenAIBetaFeatureHeaderKey = "OpenAI-Beta";
+    private const string OpenAIBetaFeatureHeaderName = "OpenAI-Beta";
+    private const string OpenAIOrganizationHeaderName = "OpenAI-Organization";
+    private const string OpenAIProjectHeaderName = "OpenAI-Project";
     private const string OpenAIBetaAssistantsV1HeaderValue = "assistants=v2";
     private const string OpenAIEndpointEnvironmentVariable = "OPENAI_ENDPOINT";
     private const string OpenAIApiKeyEnvironmentVariable = "OPENAI_API_KEY";
     private const string OpenAIV1Endpoint = "https://api.openai.com/v1";
-    private const string UserAgentHeaderKey = "User-Agent";
+    private const string UserAgentHeaderName = "User-Agent";
 }
diff --git a/src/Custom/OpenAIClientOptions.cs b/src/Custom/OpenAIClientOptions.cs
index b5da9a71..5835d0fd 100644
--- a/src/Custom/OpenAIClientOptions.cs
+++ b/src/Custom/OpenAIClientOptions.cs
@@ -18,4 +18,14 @@ public partial class OpenAIClientOptions : ClientPipelineOptions
     /// An optional application ID to use as part of the request User-Agent header.
     /// </summary>
     public string ApplicationId { get; init; }
+
+    /// <summary>
+    /// An optional ID added to OpenAI-Organization header
+    /// </summary>
+    public string OrganizationId { get; init; }
+
+    /// <summary>
+    /// An optional ID added to OpenAI-Project header
+    /// </summary>
+    public string ProjectId { get; init; }
 }