Create new Azure OpenAI client for each call

solliancenet · Sep 11, 2024 · ebbb544 · ebbb544
1 parent 0ebc22d
commit ebbb544
Showing 1 changed file with 23 additions and 13 deletions.
diff --git a/src/dotnet/Gateway/Services/GatewayCore.cs b/src/dotnet/Gateway/Services/GatewayCore.cs
@@ -233,17 +233,9 @@ private async Task<Dictionary<string, object>> CreateOpenAIAgentCapability(strin
                     StringComparison.OrdinalIgnoreCase) == 0)
                 ?? throw new GatewayException($"The Gateway service is not configured to use the {endpoint} endpoint.");
 
-            var azureOpenAIClient = new AzureOpenAIClient(
-                new Uri(azureOpenAIAccount.Endpoint),
-                DefaultAuthentication.AzureCredential,
-                new AzureOpenAIClientOptions
-                {
-                    NetworkTimeout = TimeSpan.FromSeconds(1000)
-                });
-
             if (createAssistant)
             {
-                var assistantClient = azureOpenAIClient.GetAssistantClient();
+                var assistantClient = GetAzureOpenAIAssistantClient(azureOpenAIAccount.Endpoint);
 
                 var prompt = GetRequiredParameterValue<string>(parameters, OpenAIAgentCapabilityParameterNames.AssistantPrompt);
                 var modelDeploymentName = GetRequiredParameterValue<string>(parameters, OpenAIAgentCapabilityParameterNames.ModelDeploymentName);
@@ -271,8 +263,8 @@ private async Task<Dictionary<string, object>> CreateOpenAIAgentCapability(strin
 
             if (createAssistantThread)
             {
-                var assistantClient = azureOpenAIClient.GetAssistantClient();
-                var vectorStoreClient = azureOpenAIClient.GetVectorStoreClient();
+                var assistantClient = GetAzureOpenAIAssistantClient(azureOpenAIAccount.Endpoint);
+                var vectorStoreClient = GetAzureOpenAIVectorStoreClient(azureOpenAIAccount.Endpoint);
 
                 var vectorStoreResult = await vectorStoreClient.CreateVectorStoreAsync(new VectorStoreCreationOptions
                 {
@@ -304,7 +296,7 @@ private async Task<Dictionary<string, object>> CreateOpenAIAgentCapability(strin
 
             if (createAssistantFile)
             {
-                var fileClient = azureOpenAIClient.GetFileClient();
+                var fileClient = GetAzureOpenAIFileClient(azureOpenAIAccount.Endpoint);
 
                 var attachmentObjectId = GetRequiredParameterValue<string>(parameters, OpenAIAgentCapabilityParameterNames.AttachmentObjectId);
                 var attachmentFile = await _attachmentResourceProvider.GetResource<AttachmentFile>(attachmentObjectId, userIdentity, new ResourceProviderOptions { LoadContent = true });
@@ -320,7 +312,7 @@ private async Task<Dictionary<string, object>> CreateOpenAIAgentCapability(strin
 
             if (addAssistantFileToVectorStore)
             {
-                var vectorStoreClient = azureOpenAIClient.GetVectorStoreClient();
+                var vectorStoreClient = GetAzureOpenAIVectorStoreClient(azureOpenAIAccount.Endpoint);
                 var vectorStoreId = GetRequiredParameterValue<string>(parameters, OpenAIAgentCapabilityParameterNames.AssistantVectorStoreId);
 
                 var vectorizationResult = await vectorStoreClient.AddFileToVectorStoreAsync(vectorStoreId, fileId);
@@ -366,5 +358,23 @@ private T GetRequiredParameterValue<T>(Dictionary<string, object> parameters, st
                 ? ((JsonElement)parameterValueObject!).Deserialize<T>()
                     ?? throw new GatewayException($"Could not load required parameter {parameterName}.", StatusCodes.Status400BadRequest)
                 : throw new GatewayException($"The required parameter {parameterName} was not found.");
+
+        private AzureOpenAIClient GetAzureOpenAIClient(string endpoint) =>
+            new AzureOpenAIClient(
+                new Uri(endpoint),
+                DefaultAuthentication.AzureCredential,
+                new AzureOpenAIClientOptions
+                {
+                    NetworkTimeout = TimeSpan.FromSeconds(1000)
+                });
+
+        private AssistantClient GetAzureOpenAIAssistantClient(string endpoint) =>
+            GetAzureOpenAIClient(endpoint).GetAssistantClient();
+
+        private VectorStoreClient GetAzureOpenAIVectorStoreClient(string endpoint) =>
+            GetAzureOpenAIClient(endpoint).GetVectorStoreClient();
+
+        private FileClient GetAzureOpenAIFileClient(string endpoint) =>
+            GetAzureOpenAIClient(endpoint).GetFileClient();
     }
 }