.NET: API to manage AgentThreads in hosting scenarios

[DeagleGross]

This is a prototype / design proposal for the API which allows managing AgentThread in relationship with AIAgent.

PR showcases a simple scenario where Hosting.A2A package now obtains an ability to load the thread from in-app-memory based on contextId (input of the user), and agentId (the id of a specific AIAgent). It then saves the agent thread content back via serialization.

Default behavior is to not have any thread store - the NoContextAgentThreadStore simply never saves or returns a thread. I think this is useful to not force any default behavior here, because storing all agent thread content in memory can quickly drain all of memory.

The API added: IAgentThreadStore which regulates how agent thread data is being stored:

+ public interface IAgentThreadStore
{
+    ValueTask SaveThreadAsync( string conversationId, string agentId, AgentThread thread, CancellationToken cancellationToken = default);
+    ValueTask<JsonElement?> GetThreadAsync( string conversationId, string agentId, CancellationToken cancellationToken = default);
}

and AIHostAgent which wraps the AIAgent and a specific IAgentThreadStore implementation.

+ public class AIHostAgent : AIAgent
{
+ public async ValueTask<AgentThread> GetOrCreateThreadAsync(string conversationId, CancellationToken cancellationToken = default)

+ public async ValueTask SaveThreadAsync(string conversationId, AgentThread thread, CancellationToken cancellationToken = default)
}

Schema of flow

This is a flow which shows the implementation as in the current PR. It leverages the separate IAgentThreadStore which is used to . It definitely works, uses the existing AIAgent abstraction and uses a separate interface to save/restore thread via the conversationId.

sequenceDiagram
  actor User as User
  participant code as Server Endpoint Handler
  participant store as IAgentThreadStore
  participant db as IDatabase
  participant agent as AIAgent
  participant msgStore as IMessageStore
  User ->> code: POST /v1/responses (conversationId)
    activate code
  code ->> store: GetThreadAsync(agent, conversationId)
    activate store
  store -->> db: JsonElement GetState(conversationId)
    activate db
  db -->> store: 
    deactivate db  
  store -->> code: state not null ? AIAgent.DeserializeThread(state, options)
  store -->> code: state null ? AIAgent.GetNewThread()
    deactivate store
  code ->> agent: AIAgent.RunAsync(messages, thread)
    activate agent
  agent ->> msgStore: GetMessagesAsync()
    activate msgStore
  msgStore -->> agent:
    deactivate msgStore
  agent -->> code:
    deactivate agent
  code ->> store: SaveThreadAsync(agent, thread)
    activate store
  store ->> db: SaveState(thread) as JsonElement
    activate db
  db -->> store:
    deactivate db
  store -->> code: 
    deactivate store
  code ->> User: { Response as JSON }
    deactivate code

Loading

This is a "direct" flow, where no "extra" store is used, but the restoration of thread happens inside of the same MessageStore.

sequenceDiagram
  actor User as User
  participant code as Server Endpoint Handler
  participant agent as AIHostAgent
  participant someStore as ? (maybe MessageStore)
  participant msgStore as IMessageStore
  User ->> code: POST /v1/responses (conversationId)
    activate code
  code ->> agent: AIHostAgent.RestoreThread(conversationId)
    activate agent
    agent ->> someStore:
    someStore -->> agent:
  agent -->> code:
    deactivate agent
  code ->> agent: AIAgent.RunAsync(messages, thread)
    activate agent
  agent ->> msgStore: GetMessagesAsync()
    activate msgStore
  msgStore -->> agent:
    deactivate msgStore
  agent -->> code:
    deactivate agent
  code ->> User: { Response as JSON }
    deactivate code

Loading

[SergeyMenshykh]

I wonder if the store should be responsible for obtaining a new thread from the agent, in addition to its existing "store" responsibilities. Shouldn't the store return null to indicate to the calling code that there is no thread available, allowing the calling/consumer code to decide whether to create a new thread or not? This would provide a clear separation of responsibilities and additional flexibility: the store would be responsible for storage, while the calling code would handle the absence of a thread.

[DeagleGross]

It's a bit confusing at the moment, I agree. Current implementation is what we decided before via the call with Reuben, Weslie and Stephen: see #1520 (comment).

The point is that we have a "better" API on the store interface: otherwise I would expose JsonElement? which is basically exposing inner implementation details. Maybe not everyone expects that type to be a representation of an AgentThread.

also whatever mechanism under the hood has to return a thread, otherwise you would not pass anything into AIAgent.RunAsync() and thread will not be returned to you. So we have to create a thread (that's why return type is not nullable)