.Net: [fix] StepwisePlanner ChatHistory token calculation refactoring (#2788)

Refactor chat history token calculation and clipping. Leaving the
existing `ChatHistory` as-is to ensure a complete transcript is
available. Extended class `ChatHistory` to support insert.

Fixes #2773 

### Contribution Checklist

<!-- Before submitting this PR, please make sure: -->

- [x] The code builds clean without any errors or warnings
- [x] The PR follows the [SK Contribution
Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md)
and the [pre-submission formatting
script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts)
raises no violations
- [x] All unit tests pass, and I have added new tests where possible
- [x] I didn't break anyone 😄

Author: lemillermicrosoft

dotnet/src/Extensions/Planning.StepwisePlanner/ChatHistoryExtensions.cs

@@ -0,0 +1,41 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Linq;
+using Microsoft.SemanticKernel.AI.ChatCompletion;
+using static Microsoft.SemanticKernel.Text.TextChunker;
+
+namespace Microsoft.SemanticKernel.Planning.Stepwise;
+
+/// <summary>
+/// Extension methods for <see cref="ChatHistory"/> class.
+/// </summary>
+public static class ChatHistoryExtensions
+{
+    /// <summary>
+    /// Returns the number of tokens in the chat history.
+    /// </summary>
+    // <param name="chatHistory">The chat history.</param>
+    // <param name="additionalMessage">An additional message to include in the token count.</param>
+    // <param name="skipStart">The index to start skipping messages.</param>
+    // <param name="skipCount">The number of messages to skip.</param>
+    // <param name="tokenCounter">The token counter to use.</param>
+    internal static int GetTokenCount(this ChatHistory chatHistory, string? additionalMessage = null, int skipStart = 0, int skipCount = 0, TokenCounter? tokenCounter = null)
+    {
+        tokenCounter ??= DefaultTokenCounter;
+
+        var messages = string.Join("\n", chatHistory.Where((m, i) => i < skipStart || i >= skipStart + skipCount).Select(m => m.Content));
+
+        if (!string.IsNullOrEmpty(additionalMessage))
+        {
+            messages = $"{messages}\n{additionalMessage}";
+        }
+
+        var tokenCount = tokenCounter(messages);
+        return tokenCount;
+    }
+
+    private static int DefaultTokenCounter(string input)
+    {
+        return input.Length / 4;
+    }
+}

dotnet/src/Extensions/Planning.StepwisePlanner/StepwisePlanner.cs

@@ -128,7 +128,7 @@ public async Task<SKContext> ExecutePlanAsync(
             throw new SKException("ChatHistory is null.");
         }
 
-        var startingMessageCount = chatHistory.Messages.Count;
+        var startingMessageCount = chatHistory.Count;
 
         var stepsTaken = new List<SystemStep>();
         SystemStep? lastStep = null;
@@ -205,9 +205,9 @@ public async Task<SKContext> ExecutePlanAsync(
 
                 lastStep.OriginalResponse += step.OriginalResponse;
                 step = lastStep;
-                if (chatHistory.Messages.Count > startingMessageCount)
+                if (chatHistory.Count > startingMessageCount)
                 {
-                    chatHistory.Messages.RemoveAt(chatHistory.Messages.Count - 1);
+                    chatHistory.RemoveAt(chatHistory.Count - 1);
                 }
             }
             else
@@ -382,34 +382,28 @@ private Task<string> GetSystemMessage(SKContext context)
 
     private Task<string> GetNextStepCompletion(List<SystemStep> stepsTaken, ChatHistory chatHistory, IAIService aiService, int startingMessageCount, CancellationToken token)
     {
-        var tokenCount = this.GetChatHistoryTokens(chatHistory);
+        var skipStart = startingMessageCount;
+        var skipCount = 0;
 
-        var preserveFirstNSteps = 0;
-        var removalIndex = (startingMessageCount) + preserveFirstNSteps;
-        var messagesRemoved = 0;
         string? originalThought = null;
-        while (tokenCount >= this.Config.MaxTokens && chatHistory.Messages.Count > removalIndex)
-        {
-            // something needs to be removed.
-            if (string.IsNullOrEmpty(originalThought))
-            {
-                originalThought = stepsTaken[0].Thought;
-            }
 
-            // Update message history
-            chatHistory.AddAssistantMessage($"{Thought} {originalThought}");
-            preserveFirstNSteps++;
-            chatHistory.AddAssistantMessage("... I've removed some of my previous work to make room for the new stuff ...");
-            preserveFirstNSteps++;
+        var tokenCount = chatHistory.GetTokenCount();
+        while (tokenCount >= this.Config.MaxTokens && chatHistory.Count > skipStart)
+        {
+            originalThought = $"{Thought} {stepsTaken.FirstOrDefault()?.Thought}";
+            tokenCount = chatHistory.GetTokenCount($"{originalThought}\n{TrimMessage}", skipStart, ++skipCount);
+        }
 
-            removalIndex = (startingMessageCount) + preserveFirstNSteps;
+        var reducedChatHistory = new ChatHistory();
+        reducedChatHistory.AddRange(chatHistory.Where((m, i) => i < skipStart || i >= skipStart + skipCount));
 
-            chatHistory.Messages.RemoveAt(removalIndex);
-            tokenCount = this.GetChatHistoryTokens(chatHistory);
-            messagesRemoved++;
+        if (skipCount > 0 && originalThought is not null)
+        {
+            reducedChatHistory.InsertMessage(skipStart, AuthorRole.Assistant, TrimMessage);
+            reducedChatHistory.InsertMessage(skipStart, AuthorRole.Assistant, originalThought);
         }
 
-        return this.GetCompletionAsync(aiService, chatHistory, stepsTaken.Count == 0, token);
+        return this.GetCompletionAsync(aiService, reducedChatHistory, stepsTaken.Count == 0, token);
     }
 
     private async Task<string> GetCompletionAsync(IAIService aiService, ChatHistory chatHistory, bool addThought, CancellationToken token)
@@ -421,7 +415,7 @@ private async Task<string> GetCompletionAsync(IAIService aiService, ChatHistory
         }
         else if (aiService is ITextCompletion textCompletion)
         {
-            var thoughtProcess = string.Join("\n", chatHistory.Messages.Select(m => m.Content));
+            var thoughtProcess = string.Join("\n", chatHistory.Select(m => m.Content));
 
             // Add Thought to the thought process at the start of the first iteration
             if (addThought)
@@ -444,13 +438,6 @@ private async Task<string> GetCompletionAsync(IAIService aiService, ChatHistory
         throw new SKException("No AIService available for getting completions.");
     }
 
-    private int GetChatHistoryTokens(ChatHistory chatHistory)
-    {
-        var messages = string.Join("\n", chatHistory.Messages);
-        var tokenCount = messages.Length / 4;
-        return tokenCount;
-    }
-
     /// <summary>
     /// Parse LLM response into a SystemStep during execution
     /// </summary>
@@ -749,6 +736,11 @@ private static string ToFullyQualifiedName(FunctionView function)
     /// </summary>
     private const string Observation = "[OBSERVATION]";
 
+    /// <summary>
+    /// The chat message to include when trimming thought process history
+    /// </summary>
+    private const string TrimMessage = "... I've removed some of my previous work to make room for the new stuff ...";
+
     /// <summary>
     /// The regex for parsing the thought response
     /// </summary>

dotnet/src/Extensions/Planning.StepwisePlanner/StepwisePlannerConfig.cs

@@ -51,9 +51,12 @@ public sealed class StepwisePlannerConfig
     #region Execution configuration
 
     /// <summary>
-    /// The maximum number of tokens to allow in a plan.
+    /// The maximum number of tokens to allow in a request and for completion.
     /// </summary>
-    public int MaxTokens { get; set; } = 1024;
+    /// <remarks>
+    /// Default value is 2000.
+    /// </remarks>
+    public int MaxTokens { get; set; } = 2000;
 
     /// <summary>
     /// The maximum number of iterations to allow in a plan.

dotnet/src/SemanticKernel.Abstractions/AI/ChatCompletion/ChatHistory.cs

@@ -33,6 +33,17 @@ public void AddMessage(AuthorRole authorRole, string content)
         this.Add(new ChatMessage(authorRole, content));
     }
 
+    /// <summary>
+    /// Insert a message into the chat history
+    /// </summary>
+    /// <param name="index">Index of the message to insert</param>
+    /// <param name="authorRole">Role of the message author</param>
+    /// <param name="content">Message content</param>
+    public void InsertMessage(int index, AuthorRole authorRole, string content)
+    {
+        this.Insert(index, new ChatMessage(authorRole, content));
+    }
+
     /// <summary>
     /// Add a user message to the chat history
     /// </summary>