From 97fd6cc1e091fe6c0e5226e747b5b0fb6ecce2d6 Mon Sep 17 00:00:00 2001
From: Ryan Sweet <rysweet@microsoft.com>
Date: Fri, 22 Nov 2024 05:57:11 -0800
Subject: [PATCH] improve subscriptions (#4304)

---
 .../Agents/Services/Orleans/SubscriptionsGrain.cs          | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/dotnet/src/Microsoft.AutoGen/Agents/Services/Orleans/SubscriptionsGrain.cs b/dotnet/src/Microsoft.AutoGen/Agents/Services/Orleans/SubscriptionsGrain.cs
index 905dc8e914ac..682073f0b97c 100644
--- 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)