Merge pull request #1272 from solliancenet/gg-core-examples-aimodel

Update Core.Examples to create agents based on AIModel
solliancenet · Jul 25, 2024 · e753746 · e753746
2 parents 7240659 + 7cb9614
commit e753746
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 3 deletions.
diff --git a/tests/dotnet/Core.Examples/Interfaces/IManagementAPITestManager.cs b/tests/dotnet/Core.Examples/Interfaces/IManagementAPITestManager.cs
@@ -17,7 +17,8 @@ public interface IManagementAPITestManager
     /// </summary>
     /// <param name="agentName">The name of the agent and its dependencies to retrieve from the test catalog and create.</param>
     /// <returns>The created agent.</returns>
-    Task<AgentBase> CreateAgent(string agentName, string? indexingProfileName = null, string? textEmbeddingProfileName = null, string? textPartitioningProfileName = null);
+    Task<AgentBase> CreateAgent(string agentName, string? indexingProfileName = null, string? textEmbeddingProfileName = null, 
+        string? textPartitioningProfileName = null, string? apiEndpointName = null, string? apiEndpointUrl = null, string? aiModelName = null);
 
     /// <summary>
     /// Deletes an agent and its dependencies.

diff --git a/tests/dotnet/Core.Examples/Services/ManagementAPITestManager.cs b/tests/dotnet/Core.Examples/Services/ManagementAPITestManager.cs
@@ -11,7 +11,6 @@
 using FoundationaLLM.Core.Examples.Catalogs;
 using FoundationaLLM.Core.Examples.Exceptions;
 using FoundationaLLM.Core.Examples.Interfaces;
-using FoundationaLLM.Core.Examples.Setup;
 using Microsoft.Extensions.Options;
 using System.Text;
 using System.Text.Json;
@@ -115,6 +114,36 @@ public async Task CreateTextPartitioningProfile(string textPartitioningProfileNa
             throw new FoundationaLLMException($"Failed to create text partitioning profile: {textPartitioningProfileName}.");
         }
 
+        public async Task<string> CreateAPIEndpointConfiguration(string apiEndpointName, string apiEndpointUrl)
+        {
+            var apiEndpointProfile = APIEndpointConfigurationCatalog.GetAllAPIEndpointConfigurations().FirstOrDefault(a => a.Name == apiEndpointName);
+
+            if (apiEndpointProfile == null)
+            {
+                throw new InvalidOperationException($"The api endpoint profile {apiEndpointName} was not found.");
+            }
+
+            apiEndpointProfile.Url = apiEndpointUrl;
+            var response = await managementClient.Configuration.UpsertAPIEndpointConfiguration(apiEndpointProfile);
+
+            return GetObjectId(response); 
+        }
+
+        public async Task<string> CreateAIModel(string aiModelName, string endpointObjectId)
+        {
+            var aiModelProfile = AIModelCatalog.GetAllAIModels().FirstOrDefault(a => a.Name == aiModelName);
+
+            if (aiModelProfile == null)
+            {
+                throw new InvalidOperationException($"The AI model profile {aiModelName} was not found.");
+            }
+
+            aiModelProfile.EndpointObjectId = endpointObjectId;
+            var response = await managementClient.AIModels.UpsertAIModel(aiModelProfile);
+
+            return GetObjectId(response);
+        }
+
         public async Task<VectorizationRequest> GetVectorizationRequest(VectorizationRequest vectorizationRequest)
         {
             return await managementRestClient.Resources.GetResourcesAsync<VectorizationRequest>(
@@ -211,7 +240,8 @@ public async Task DeleteAgent(string agentName)
 
         /// <inheritdoc/>
         public async Task<AgentBase> CreateAgent(string agentName, string? indexingProfileName, 
-                string? textEmbeddingProfileName, string? textPartitioningProfileName)
+                string? textEmbeddingProfileName, string? textPartitioningProfileName,
+                string? apiEndpointName, string? apiEndpointUrl, string? aiModelName)
         {
             // All test agents should have a corresponding prompt in the catalog.
             // Retrieve the agent and prompt from the test catalog.
@@ -260,6 +290,15 @@ public async Task<AgentBase> CreateAgent(string agentName, string? indexingProfi
 
             agent.ObjectId = GetObjectId(response);
 
+            // Create APIEndpointConfiguration and AIModel object
+            if(!string.IsNullOrWhiteSpace(apiEndpointName) 
+                && !string.IsNullOrWhiteSpace(apiEndpointUrl)
+                && !string.IsNullOrWhiteSpace(aiModelName))
+            {
+                var endpointObjectId = await CreateAPIEndpointConfiguration(apiEndpointName, apiEndpointUrl);
+                agent.AIModelObjectId = await CreateAIModel(aiModelName, endpointObjectId);
+            }
+
             return agent;
         }