Improved how the AI remembers its earlier tool calls (#1185)

Co-authored-by: Thorsten Sommer <SommerEngineering@users.noreply.github.com>

app/MindWork AI Studio/Chat/ChatThread.cs

@@ -280,8 +280,11 @@ public sealed record ChatThread
     /// </summary>
     /// <remarks>
     /// Addresses the model wrote into the call itself are left out, see ToolCallOutcome.ReturnedWebAddresses.
-    /// Not stored, like the results of the tools themselves: after a restart, a chat knows fewer
-    /// addresses, never more.
+    /// Not stored: each answer stores those addresses its own tools returned which a later message
+    /// still shows the model, see ContentText.WebAddressesFromTools, so they outlast a restart. This
+    /// set covers the runs without an answer in the chat, such as those of the assistants, and the
+    /// addresses no later message shows, such as the links of a page read. It is never emptied, so
+    /// the addresses of an answer which was deleted or generated anew count until the chat is closed.
     /// </remarks>
     [JsonIgnore]
     public HashSet<string> RuntimeWebAddressesFromTools { get; } = new(StringComparer.Ordinal);
@@ -303,7 +306,8 @@ public sealed record ChatThread
     /// <remarks>
     /// Given means: it stands in the system prompt, see RuntimeSystemPromptWebAddresses, in a
     /// message of the user or a document attached to one, see ContentText.RuntimeAttachmentWebAddresses,
-    /// or a tool returned it, see RuntimeWebAddressesFromTools. What the model wrote itself never
+    /// or a tool returned it, see RuntimeWebAddressesFromTools and ContentText.WebAddressesFromTools
+    /// of the answers in this chat. What the model wrote itself never
     /// counts, its answers included, because the model could have put anything of the chat into it.
     /// Neither does what any model wrote and what reached this chat as a message of the user, such
     /// as an answer sent here from another chat, see WebAddressesWrittenByModels. The comparison is
@@ -320,6 +324,10 @@ public sealed record ChatThread
         if (this.RuntimeWebAddressesFromTools.Contains(requestKey) || this.RuntimeSystemPromptWebAddresses.Contains(requestKey))
             return true;
 
+        // A tool returned it for one of the answers, possibly before a restart:
+        if (this.Blocks.Any(block => block.Role is ChatRole.AI && block.Content is ContentText answer && answer.WebAddressesFromTools.Contains(requestKey)))
+            return true;
+
         var messagesOfTheUser = this.Blocks
             .Where(block => block.Role is ChatRole.USER)
             .Select(block => block.Content)
@@ -488,6 +496,22 @@ public sealed record ChatThread
             parts.Add(new(SystemPromptPartKind.TOOL_INSTRUCTIONS, toolPolicy) { Tools = instructedTools });
         }
 
+        //
+        // Explain the notes in front of the earlier answers, see ToolCallHistoryNote. Also without
+        // tools: the note then tells the model what it read, and this tells it that it no longer
+        // has the content.
+        //
+        if (ToolCallHistoryNote.IsInAnyOf(this.Blocks))
+        {
+            systemPromptText = $"""
+                                {systemPromptText}
+
+                                {ToolCallHistoryNote.SYSTEM_PROMPT_EXPLANATION}
+                                """;
+
+            parts.Add(new(SystemPromptPartKind.HISTORY_NOTES, ToolCallHistoryNote.SYSTEM_PROMPT_EXPLANATION));
+        }
+
         var explanation = $"{logMessage} {profileNote}";
         if(!this.IncludeDateTime)
             return new(systemPromptText, systemPromptTextWithChatTemplate, explanation, parts, switchedOffProfile);
@@ -804,7 +828,10 @@ public sealed record ChatThread
     /// What goes unnoticed is a change beside the messages: another system prompt, profile, or
     /// selection of tools. That shows only with the next answer. Noticing it would take a
     /// fingerprint of the system prompt as it was sent, after the data sources added to it, which
-    /// is a lot of machinery for a number which corrects itself one answer later.
+    /// is a lot of machinery for a number which corrects itself one answer later. The same holds for
+    /// the explanation of the notes on earlier answers, which the system prompt gets once the first
+    /// answer has a note, see ToolCallHistoryNote. The note of the last answer itself is counted: it
+    /// changes with every answer, and the answer goes out with it, see ToolCallHistoryNote.TextForModel.
     /// </remarks>
     /// <param name="model">The model the next request would go to.</param>
     /// <returns>
@@ -825,7 +852,7 @@ public sealed record ChatThread
         if (!string.Equals(reported.ModelId, model.Id, StringComparison.Ordinal))
             return ReportedHistory.UNKNOWN;
 
-        return ReportedHistory.Of(reported.ToTokenUsage(), answer.Text);
+        return ReportedHistory.Of(reported.ToTokenUsage(), ToolCallHistoryNote.TextForModel(this.Blocks[^1], answer.Text));
     }
 
     private static void DeleteManagedAttachments(ContentBlock block)

app/MindWork AI Studio/Chat/ContentText.cs

@@ -57,6 +57,45 @@ public sealed class ContentText : IContent
 
     public List<ToolInvocationTrace> ToolInvocations { get; set; } = [];
 
+    /// <summary>
+    /// The web addresses the tools returned while this answer came about and which a later message
+    /// still shows the model, as request keys, see WebAddresses.CreateRequestKey.
+    /// </summary>
+    /// <remarks>
+    /// Stored with the answer, so that the model may read them again in a later message, also after
+    /// a restart, while its free address choice is off, see ChatThread.IsWebAddressGivenToTheModel.
+    /// A tool returned them, so the model neither made them up nor wrote anything of the chat into
+    /// them; addresses it wrote into the call itself are left out, see ToolCallOutcome.ReturnedWebAddresses.<br/><br/>
+    ///
+    /// Only those the model can still see later are kept: in the text of the answer, among its
+    /// sources, or in the arguments of a call which ran, see EndToolRun. A page a tool read returns
+    /// all of its links, often hundreds. Once the result is gone, the model has no way of knowing
+    /// them, short of reading the page again, which returns them once more. Storing them all would
+    /// only make the chat grow.<br/><br/>
+    ///
+    /// Kept per answer rather than per chat: an answer which is generated anew takes its addresses
+    /// along, and so does one which is deleted, once the chat is opened again, see
+    /// ChatThread.RuntimeWebAddressesFromTools.
+    /// </remarks>
+    public HashSet<string> WebAddressesFromTools { get; set; } = new(StringComparer.Ordinal);
+
+    /// <summary>
+    /// The web addresses the tools of the running request returned so far, as request keys.
+    /// </summary>
+    /// <remarks>
+    /// Never stored: a chat is saved while its answer is still being written, and would hold every
+    /// link of every page read until then. EndToolRun keeps those which a later message shows, in
+    /// WebAddressesFromTools.
+    /// </remarks>
+    [JsonIgnore]
+    public HashSet<string> PendingWebAddressesFromTools { get; private set; } = new(StringComparer.Ordinal);
+
+    /// <summary>
+    /// Notes the web addresses a tool returned for this answer, see PendingWebAddressesFromTools.
+    /// </summary>
+    /// <param name="requestKeys">The addresses as request keys, see ToolCallOutcome.ReturnedWebAddresses.</param>
+    public void KeepWebAddressesFromTools(IEnumerable<string> requestKeys) => this.PendingWebAddressesFromTools.UnionWith(requestKeys);
+
     /// <summary>
     /// Optional structured data for a specialized visual presentation of this text.
     /// </summary>
@@ -231,13 +270,15 @@ public sealed class ContentText : IContent
     /// Clears what the previous run of the tools left behind.
     /// </summary>
     /// <remarks>
-    /// Both parts at once, because both belong to one request: the traces the user reads and the
-    /// payload the counting needs. They were cleared separately for exactly as long as there was
-    /// only one of them.
+    /// All parts at once, because all belong to one request: the traces the user reads, the
+    /// payload the counting needs, and the web addresses the tools returned. They were cleared
+    /// separately for exactly as long as there was only one of them.
     /// </remarks>
     public void BeginToolRun()
     {
         this.ToolInvocations.Clear();
+        this.WebAddressesFromTools = new(StringComparer.Ordinal);
+        this.PendingWebAddressesFromTools = new(StringComparer.Ordinal);
         this.PendingToolConversation = [];
     }
 
@@ -246,9 +287,50 @@ public sealed class ContentText : IContent
     /// </summary>
     /// <remarks>
     /// The traces stay -- they are what the user reads afterwards to see how the answer came
-    /// about. What goes is the payload, which belonged to a request that is over.
+    /// about. What goes is the payload, which belonged to a request that is over.<br/><br/>
+    ///
+    /// Of the web addresses the tools returned, the answer keeps those which a later message still
+    /// shows the model, see WebAddressesFromTools. The set is replaced as a whole, since the chat
+    /// may be saved meanwhile. Whoever ends a request may call this once more, which keeps what the
+    /// first call kept.
     /// </remarks>
-    public void EndToolRun() => this.PendingToolConversation = [];
+    public void EndToolRun()
+    {
+        this.PendingToolConversation = [];
+
+        var shownLater = this.FindRequestKeysShownLater();
+        var kept = new HashSet<string>(StringComparer.Ordinal);
+        kept.UnionWith(this.WebAddressesFromTools.Where(shownLater.Contains));
+        kept.UnionWith(this.PendingWebAddressesFromTools.Where(shownLater.Contains));
+
+        this.WebAddressesFromTools = kept;
+        this.PendingWebAddressesFromTools = new(StringComparer.Ordinal);
+    }
+
+    /// <summary>
+    /// The request keys of the web addresses a later message shows the model with this answer.
+    /// </summary>
+    /// <remarks>
+    /// Those in its text, those of its sources, and those in the arguments of its calls which ran,
+    /// as the note in front of the answer names them, see ToolCallHistoryNote. The values of
+    /// sensitive arguments do not count, since the note hides them.
+    /// </remarks>
+    private HashSet<string> FindRequestKeysShownLater()
+    {
+        var texts = new List<string> { this.Text };
+        texts.AddRange(this.Sources.Select(source => source.URL));
+        foreach (var trace in this.ToolInvocations.Where(trace => trace is { Status: ToolInvocationTraceStatus.SUCCESS, WasExecuted: true }))
+            texts.AddRange(trace.FindKeptTexts());
+
+        var requestKeys = new HashSet<string>(StringComparer.Ordinal);
+        foreach (var address in texts.SelectMany(WebAddresses.Find))
+        {
+            if (WebAddresses.TryCreateRequestKey(address, out var requestKey))
+                requestKeys.Add(requestKey);
+        }
+
+        return requestKeys;
+    }
 
     /// <summary>
     /// Keeps what the provider says the request behind this answer cost.
@@ -550,6 +632,7 @@ public sealed class ContentText : IContent
         Sources = [..this.Sources],
         SourcePassages = [..this.SourcePassages],
         FileAttachments = [..this.FileAttachments],
+        WebAddressesFromTools = new HashSet<string>(this.WebAddressesFromTools, StringComparer.Ordinal),
         Presentation = this.Presentation?.DeepClone(),
         ToolInvocations = [..this.ToolInvocations.Select(x => new ToolInvocationTrace
         {
@@ -562,6 +645,7 @@ public sealed class ContentText : IContent
             WasExecuted = x.WasExecuted,
             StatusMessage = x.StatusMessage,
             Arguments = new Dictionary<string, string>(x.Arguments, StringComparer.Ordinal),
+            ArgumentsJson = x.ArgumentsJson,
             Result = x.Result,
             JsonResult = x.JsonResult?.DeepClone(),
         })],

app/MindWork AI Studio/Chat/ConversationParts.cs

@@ -155,10 +155,15 @@ public sealed record ConversationParts
                 if (string.IsNullOrWhiteSpace(text.Text))
                     continue;
 
+                //
+                // An earlier answer travels with the note of the tools and sources it rested on, see
+                // ToolCallHistoryNote.TextForModel, so it is counted with it. The note of an answer does
+                // not change anymore, so the text stays the same and its count stays in the cache.
+                //
                 if (text.IsStreaming)
                     growing.Add(text.Text);
                 else
-                    texts.Add(text.Text);
+                    texts.Add(ToolCallHistoryNote.TextForModel(block, text.Text));
 
                 Sort(text.FileAttachments, documents, ref images);
             }

app/MindWork AI Studio/Chat/ListContentBlockExtensions.cs

@@ -4,6 +4,7 @@ using AIStudio.Settings;
 using AIStudio.Tools.Rust;
 using AIStudio.Tools.Security;
 using AIStudio.Tools.Services;
+using AIStudio.Tools.ToolCallingSystem;
 
 namespace AIStudio.Chat;
 
@@ -81,7 +82,7 @@ public static class ListContentBlockExtensions
                 return Task.Run(async () => new TextMessage
                 {
                     Role = roleTransformer(block.Role),
-                    Content = await text.PrepareTextContentForAI(reportingScope),
+                    Content = await PrepareTextAsync(block, text),
                 } as IMessageBase);
             }
         
@@ -100,13 +101,19 @@ public static class ListContentBlockExtensions
                         Role = roleTransformer(block.Role),
                         Content =
                         [
-                            textSubContentFactory(await text.PrepareTextContentForAI(reportingScope)),
+                            textSubContentFactory(await PrepareTextAsync(block, text)),
                             ..images.OfType<ISubContent>(),
                         ]
                     } as IMessageBase;
                 });
             }
 
+            //
+            // Local function to prepare the text of one message. An earlier answer of the model gets the
+            // note of the tools and sources it rested on in front, see ToolCallHistoryNote.TextForModel.
+            //
+            async Task<string> PrepareTextAsync(ContentBlock block, ContentText text) => ToolCallHistoryNote.TextForModel(block, await text.PrepareTextContentForAI(reportingScope));
+
             //
             // Local function to prepare one image for the model. An image which cannot be prepared
             // stays out of the message, and the user learns why. It used to go along as an image

app/MindWork AI Studio/Chat/SystemPromptPartKind.cs

@@ -38,4 +38,9 @@ public enum SystemPromptPartKind
     /// How to use the tools the request offers.
     /// </summary>
     TOOL_INSTRUCTIONS,
+
+    /// <summary>
+    /// What the notes in front of the earlier answers mean, see ToolCallHistoryNote.
+    /// </summary>
+    HISTORY_NOTES,
 }