@@ -8,55 +8,65 @@ public enum ChatReducerTriggerEvent
88 BeforeMessageAdded , AfterMessagesRetrieval
99}
1010
11+ public sealed class State
12+ {
13+ public State ( string sessionDbKey )
14+ {
15+ this . SessionDbKey = sessionDbKey ?? throw new ArgumentNullException ( nameof ( sessionDbKey ) ) ;
16+ }
17+
18+ public string SessionDbKey { get ; }
19+ }
20+
1121public sealed class VectorChatHistoryProvider : ChatHistoryProvider
1222{
23+ private readonly ProviderSessionState < State > _sessionState ;
24+ private IReadOnlyList < string > ? _stateKeys ;
1325 private readonly VectorStore _vectorStore ;
1426 private readonly IChatReducer ? _chatReducer ;
1527 private readonly ChatReducerTriggerEvent _reducerTriggerEvent ;
1628
1729 public VectorChatHistoryProvider (
1830 VectorStore vectorStore ,
19- JsonElement serializedStoreState ,
20- JsonSerializerOptions ? jsonSerializerOptions = null )
21- : this ( vectorStore , serializedStoreState , jsonSerializerOptions , null , ChatReducerTriggerEvent . AfterMessagesRetrieval )
31+ Func < AgentSession ? , State > ? stateInitializer = null ,
32+ string ? stateKey = null )
33+ : this ( vectorStore , stateInitializer , stateKey , null , ChatReducerTriggerEvent . AfterMessagesRetrieval )
2234 {
2335
2436 }
2537
2638 public VectorChatHistoryProvider (
2739 VectorStore vectorStore ,
28- JsonElement serializedStoreState ,
29- JsonSerializerOptions ? jsonSerializerOptions = null ,
40+ Func < AgentSession ? , State > ? stateInitializer = null ,
41+ string ? stateKey = null ,
3042 IChatReducer ? chatReducer = null ,
3143 ChatReducerTriggerEvent reducerTriggerEvent = ChatReducerTriggerEvent . AfterMessagesRetrieval )
3244 {
3345 Console . WriteLine ( $ "VectorStore: { vectorStore . GetType ( ) } ") ;
3446 _vectorStore = vectorStore ?? throw new ArgumentNullException ( nameof ( vectorStore ) ) ;
3547 _chatReducer = chatReducer ;
3648 _reducerTriggerEvent = reducerTriggerEvent ;
37- if ( serializedStoreState . ValueKind is JsonValueKind . String )
38- {
39- SessionDbKey = serializedStoreState . Deserialize < string > ( ) ;
40- Console . WriteLine ( $ "SessionDbKey: { SessionDbKey } ") ;
41- }
49+ this . _sessionState = new ProviderSessionState < State > (
50+ stateInitializer ?? ( _ => new State ( Guid . NewGuid ( ) . ToString ( "N" ) ) ) ,
51+ stateKey ?? this . GetType ( ) . Name ) ;
4252 }
4353
44- public string ? SessionDbKey { get ; private set ; }
54+ public string ? SessionDbKey => _sessionState . StateKey ;
4555
46- public override async ValueTask InvokedAsync ( InvokedContext context , CancellationToken cancellationToken = default )
56+
57+ protected override async ValueTask StoreChatHistoryAsync ( InvokedContext context , CancellationToken cancellationToken = default )
4758 {
4859 if ( context . InvokeException is not null )
4960 {
5061 return ;
5162 }
5263
53- var messages = context . RequestMessages . Concat ( context . AIContextProviderMessages ?? [ ] ) . Concat ( context . ResponseMessages ?? [ ] ) ;
64+ var messages = context . RequestMessages . Concat ( context . RequestMessages ?? [ ] ) . Concat ( context . ResponseMessages ?? [ ] ) ;
5465 if ( _reducerTriggerEvent is ChatReducerTriggerEvent . BeforeMessageAdded && _chatReducer is not null )
5566 {
5667 messages = await _chatReducer . ReduceAsync ( messages , cancellationToken ) . ConfigureAwait ( false ) ;
5768 }
5869
59- SessionDbKey ??= Guid . NewGuid ( ) . ToString ( "N" ) ;
6070 var collection = _vectorStore . GetCollection < string , ChatHistoryItem > ( "agent_chat_history" ) ;
6171 await collection . EnsureCollectionExistsAsync ( cancellationToken ) ;
6272 await collection . UpsertAsync ( messages . Select ( x => new ChatHistoryItem ( )
@@ -69,7 +79,7 @@ public override async ValueTask InvokedAsync(InvokedContext context, Cancellatio
6979 } ) , cancellationToken ) ;
7080 }
7181
72- public override async ValueTask < IEnumerable < ChatMessage > > InvokingAsync ( InvokingContext context , CancellationToken cancellationToken = default )
82+ protected override async ValueTask < IEnumerable < ChatMessage > > ProvideChatHistoryAsync ( InvokingContext context , CancellationToken cancellationToken = default )
7383 {
7484 var collection = _vectorStore . GetCollection < string , ChatHistoryItem > ( "agent_chat_history" ) ;
7585 await collection . EnsureCollectionExistsAsync ( cancellationToken ) ;
@@ -94,10 +104,6 @@ public override async ValueTask<IEnumerable<ChatMessage>> InvokingAsync(Invoking
94104 return messages ;
95105 }
96106
97- public override JsonElement Serialize ( JsonSerializerOptions ? jsonSerializerOptions = null ) =>
98- // We have to serialize the session id, so that on deserialization you can retrieve the messages using the same session id.
99- JsonSerializer . SerializeToElement ( SessionDbKey ) ;
100-
101107 private sealed class ChatHistoryItem
102108 {
103109 [ VectorStoreKey ]
0 commit comments