From 899f54d23da379c85ddb8d171421fafb723a5387 Mon Sep 17 00:00:00 2001
From: Rafael Azriaiev <rafael.azrv@gmail.com>
Date: Thu, 2 Jan 2025 16:37:08 +0200
Subject: [PATCH] fixed getting kernel by type returns only one pool

- added ConcurrentDictionary to maintain round robin indices for AIServiceProviderType
- update GetKernelAsync to get all configs of request type and select based on index (round robin)
- update GetKernelByName Linq method to FirstOrDefault to use our own exception throw
---
 SemanticKernelPooling/KernelPoolManager.cs | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/SemanticKernelPooling/KernelPoolManager.cs b/SemanticKernelPooling/KernelPoolManager.cs
index 5961f54..02f9c07 100644
--- a/SemanticKernelPooling/KernelPoolManager.cs
+++ b/SemanticKernelPooling/KernelPoolManager.cs
@@ -19,6 +19,8 @@ public class KernelPoolManager : IKernelPoolManager
 
     // Dictionary to maintain round-robin indices for each scope
     private readonly ConcurrentDictionary<string, int> _scopeIndices = new();
+    // Dictionary to maintain round-robin indices for each AIServiceProviderType
+    private readonly ConcurrentDictionary<AIServiceProviderType, int> _providerTypeIndices = new();
     private readonly IConfiguration _configuration;
 
     /// <summary>
@@ -109,7 +111,7 @@ private async Task<KernelWrapper> GetKernelAsync(AIServiceProviderConfiguration
     public async Task<KernelWrapper> GetKernelByNameAsync(string uniqueName)
     {
         //get the service configuration
-        var config = AIConfigurations.First(c => c.UniqueName == uniqueName);
+        var config = AIConfigurations.FirstOrDefault(c => c.UniqueName == uniqueName);
         if (config == null)
             throw new InvalidOperationException($"No configuration found for kernel pool name {uniqueName}");
 
@@ -138,12 +140,16 @@ public async Task<KernelWrapper> GetKernelByScopeAsync(string scope)
     /// <inheritdoc />
     public async Task<KernelWrapper> GetKernelAsync(AIServiceProviderType aiServiceProviderType)
     {
-        //get the service configuration by type
-        var config = AIConfigurations.First(c => c.ServiceType == aiServiceProviderType);
-        if (config == null)
+        //get all the service configurations by type
+        var configs = AIConfigurations.Where(c => c.ServiceType == aiServiceProviderType).ToList();
+        if (configs.Count == 0)
             throw new InvalidOperationException($"No configuration found for kernel pool type {aiServiceProviderType}");
 
-        var kernelWrapper = await GetKernelAsync(config);
+        // Use round-robin selection to choose a configuration
+        var index = _providerTypeIndices.AddOrUpdate(aiServiceProviderType, 0, (_, oldValue) => (oldValue + 1)% configs.Count);
+        var selectedConfig = configs[index];
+
+        var kernelWrapper = await GetKernelAsync(selectedConfig);
 
         return kernelWrapper;
     }