.NET: [BREAKING] Subclass AgentThread so that different agents have their own threads with their own typed settings. (#798)

* Subclass AgentThread so that different agents have their own threads with their own typed settings.

* Address PR comment.

* Add unit tests for base abstract threads

* Fix style warning

* Fix stlying

* FIx and suppress warnings as needed.

* Remove covariant thread response types and fix some styling.

* Remove unecessary json property name attributes and make OrchestratingAgentThread private

* Fix break from merge from main.

* Fix formatting

* Fix deserialization bug in Memory sample

* Remove thread deletion from basic samples.

* Remove public constructors for thread subclasses and add more factory methods to concrete agent types.

* Update AgentProxy thread constructors to be internal as well.

* Revert AgentProxyThread to internal

* Change AIContextProvider to internal set

* Change conversation id and message store properties to internal set

* Update styling.

* Seal various thread types.

* Add thread type check for thread deletion

* Fix tests after latest merge from main

* Add thread type checks for thread deletion.

---------

Co-authored-by: Chris <66376200+crickman@users.noreply.github.com>

dotnet/samples/GettingStarted/AgentOpenTelemetry/Program.cs

@@ -123,7 +123,7 @@ appLogger.LogInformation("Agent created successfully with ID: {AgentId}", agent.
 
 // Create a parent span for the entire agent session
 using var sessionActivity = activitySource.StartActivity("Agent Session");
-var sessionId = thread.ConversationId ?? Guid.NewGuid().ToString();
+var sessionId = Guid.NewGuid().ToString();
 sessionActivity?
     .SetTag("agent.name", "OpenTelemetryDemoAgent")
     .SetTag("session.id", sessionId)

dotnet/samples/GettingStarted/AgentOrchestration/AgentSample.cs

@@ -153,9 +153,9 @@ public class AgentSample(ITestOutputHelper output) : BaseSample(output)
         await persistentAgentsClient.Administration.DeleteAgentAsync(agent.Id, cancellationToken);
 
         // If a thread is provided, delete it as well.
-        if (thread is not null)
+        if (thread is ChatClientAgentThread chatThread)
         {
-            await persistentAgentsClient.Threads.DeleteThreadAsync(thread.ConversationId, cancellationToken);
+            await persistentAgentsClient.Threads.DeleteThreadAsync(chatThread.ConversationId, cancellationToken);
         }
     }
 
@@ -169,9 +169,9 @@ public class AgentSample(ITestOutputHelper output) : BaseSample(output)
         await assistantClient.DeleteAssistantAsync(agent.Id, cancellationToken);
 
         // If a thread is provided, delete it as well.
-        if (thread is not null)
+        if (thread is ChatClientAgentThread chatThread)
         {
-            await assistantClient.DeleteThreadAsync(thread.ConversationId, cancellationToken);
+            await assistantClient.DeleteThreadAsync(chatThread.ConversationId, cancellationToken);
         }
     }
 

dotnet/samples/GettingStarted/AgentProviders/Agent_With_AzureFoundry/Program.cs

@@ -36,6 +36,5 @@ AgentThread thread = agent1.GetNewThread();
 Console.WriteLine(await agent1.RunAsync("Tell me a joke about a pirate.", thread));
 
 // Cleanup for sample purposes.
-await persistentAgentsClient.Threads.DeleteThreadAsync(thread.ConversationId);
 await persistentAgentsClient.Administration.DeleteAgentAsync(agent1.Id);
 await persistentAgentsClient.Administration.DeleteAgentAsync(agent2.Id);

dotnet/samples/GettingStarted/AgentProviders/Agent_With_CustomImplementation/Program.cs

@@ -8,6 +8,7 @@ using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Runtime.CompilerServices;
+using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
@@ -30,6 +31,12 @@ namespace SampleApp
     // Custom agent that parrot's the user input back in upper case.
     internal sealed class UpperCaseParrotAgent : AIAgent
     {
+        public override AgentThread GetNewThread()
+            => new CustomAgentThread();
+
+        public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
+            => new CustomAgentThread(serializedThread, jsonSerializerOptions);
+
         public override async Task<AgentRunResponse> RunAsync(IEnumerable<ChatMessage> messages, AgentThread? thread = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)
         {
             // Create a thread if the user didn't supply one.
@@ -96,5 +103,17 @@ namespace SampleApp
 
                 return messageClone;
             });
+
+        /// <summary>
+        /// A thread type for our custom agent that only supports in memory storage of messages.
+        /// </summary>
+        internal sealed class CustomAgentThread : InMemoryAgentThread
+        {
+            internal CustomAgentThread()
+                : base() { }
+
+            internal CustomAgentThread(JsonElement serializedThreadState, JsonSerializerOptions? jsonSerializerOptions = null)
+                : base(serializedThreadState, jsonSerializerOptions) { }
+        }
     }
 }

dotnet/samples/GettingStarted/AgentProviders/Agent_With_OpenAIAssistants/Program.cs

@@ -37,6 +37,5 @@ AgentThread thread = agent1.GetNewThread();
 Console.WriteLine(await agent1.RunAsync("Tell me a joke about a pirate.", thread));
 
 // Cleanup for sample purposes.
-await assistantClient.DeleteThreadAsync(thread.ConversationId);
 await assistantClient.DeleteAssistantAsync(agent1.Id);
 await assistantClient.DeleteAssistantAsync(agent2.Id);

dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyThreadStorage/Program.cs

@@ -38,12 +38,12 @@ AIAgent agent = new AzureOpenAIClient(
      {
          Name = JokerName,
          Instructions = JokerInstructions,
-         ChatMessageStoreFactory = (jsonElement, jso) =>
+         ChatMessageStoreFactory = ctx =>
          {
              // Create a new chat message store for this agent that stores the messages in a vector store.
              // Each thread must get its own copy of the VectorChatMessageStore, since the store
              // also contains the id that the thread is stored under.
-             return new VectorChatMessageStore(vectorStore, jsonElement, jso);
+             return new VectorChatMessageStore(vectorStore, ctx.SerializedState, ctx.JsonSerializerOptions);
          }
      });
 

dotnet/samples/GettingStarted/Agents/Agent_Step13_Memory/Program.cs

@@ -38,7 +38,7 @@ ChatClient chatClient = new AzureOpenAIClient(
 AIAgent agent = chatClient.CreateAIAgent(new ChatClientAgentOptions()
 {
     Instructions = "You are a friendly assistant. Always address the user by their name.",
-    AIContextProviderFactory = (jse, jso) => new UserInfoMemory(chatClient.AsIChatClient(), jse, jso)
+    AIContextProviderFactory = ctx => new UserInfoMemory(chatClient.AsIChatClient(), ctx.SerializedState, ctx.JsonSerializerOptions)
 });
 
 // Create a new thread for the conversation.
@@ -63,7 +63,7 @@ Console.WriteLine(await agent.RunAsync("What is my name and age?", deserializedT
 Console.WriteLine("\n>> Read memories from memory component\n");
 
 // It's possible to access the memory component via the thread's AIContextProvider property.
-var userInfo = ((UserInfoMemory)deserializedThread.AIContextProvider!).UserInfo;
+var userInfo = ((deserializedThread as ChatClientAgentThread)!.AIContextProvider as UserInfoMemory)!.UserInfo;
 
 // Output the user info that was captured by the memory component.
 Console.WriteLine($"MEMORY - User Name: {userInfo.UserName}");
@@ -71,18 +71,15 @@ Console.WriteLine($"MEMORY - User Age: {userInfo.UserAge}");
 
 Console.WriteLine("\n>> Use new thread with previously created memories\n");
 
-// Create a new thread.
-thread = agent.GetNewThread();
-
-// It is also possible to add the memory component to an individual thread only instead of all
-// threads via the factory above.
-// In this case we will also use the same user info object, so this thread will share the same
-// memories as the previous thread.
-thread.AIContextProvider = new UserInfoMemory(chatClient.AsIChatClient(), userInfo);
+// It is also possible to set the memories in a memory component on an individual thread.
+// This is useful if we want to start a new thread, but have it share the same memories as a previous thread.
+// For this scenario, we have to know the underlying agent thread type and the memory component type.
+var newThread = agent.GetNewThread();
+((newThread as ChatClientAgentThread)!.AIContextProvider as UserInfoMemory)!.UserInfo = userInfo;
 
 // Invoke the agent and output the text result.
 // This time the agent should remember the user's name and use it in the response.
-Console.WriteLine(await agent.RunAsync("What is my name and age?", thread));
+Console.WriteLine(await agent.RunAsync("What is my name and age?", newThread));
 
 namespace SampleApp
 {
@@ -102,7 +99,10 @@ namespace SampleApp
         public UserInfoMemory(IChatClient chatClient, JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null)
         {
             this._chatClient = chatClient;
-            this.UserInfo = serializedState.Deserialize<UserInfo>(jsonSerializerOptions) ?? new UserInfo();
+
+            this.UserInfo = serializedState.ValueKind == JsonValueKind.Object ?
+                serializedState.Deserialize<UserInfo>(jsonSerializerOptions)! :
+                new UserInfo();
         }
 
         public UserInfo UserInfo { get; set; }