improve subscriptions (#4304)

microsoft · Nov 22, 2024 · 97fd6cc · 97fd6cc
1 parent eb67e4a
commit 97fd6cc
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/dotnet/src/Microsoft.AutoGen/Agents/Services/Orleans/SubscriptionsGrain.cs b/dotnet/src/Microsoft.AutoGen/Agents/Services/Orleans/SubscriptionsGrain.cs
@@ -6,8 +6,13 @@ namespace Microsoft.AutoGen.Agents;
 internal sealed class SubscriptionsGrain([PersistentState("state", "PubSubStore")] IPersistentState<SubscriptionsState> state) : Grain, ISubscriptionsGrain
 {
     private readonly Dictionary<string, List<string>> _subscriptions = new();
-    public ValueTask<Dictionary<string, List<string>>> GetSubscriptions(string agentType)
+    public ValueTask<Dictionary<string, List<string>>> GetSubscriptions(string? agentType = null)
     {
+        //if agentType is null, return all subscriptions else filter on agentType
+        if (agentType != null)
+        {
+            return new ValueTask<Dictionary<string, List<string>>>(_subscriptions.Where(x => x.Value.Contains(agentType)).ToDictionary(x => x.Key, x => x.Value));
+        }
         return new ValueTask<Dictionary<string, List<string>>>(_subscriptions);
     }
     public ValueTask Subscribe(string agentType, string topic)