Make process message method more flexible

LLama.Examples/Examples/ChatSessionWithHistory.cs

@@ -48,6 +48,10 @@ public class ChatSessionWithHistory
 
         Console.ForegroundColor = ConsoleColor.Yellow;
         Console.WriteLine("The chat session has started.");
+        Console.WriteLine("Type 'exit' to end the chat session.");
+        Console.WriteLine("Type 'save' to save the chat session to disk.");
+        Console.WriteLine("Type 'load' to load the chat session from disk.");
+        Console.WriteLine("Type 'regenerate' to regenerate the last response.");
 
         // show the prompt
         Console.ForegroundColor = ConsoleColor.Green;
@@ -55,12 +59,14 @@ public class ChatSessionWithHistory
 
         while (userInput != "exit")
         {
+            // Save the chat state to disk
             if (userInput == "save")
             {
                 session.SaveSession("Assets/chat-with-bob");
                 Console.ForegroundColor = ConsoleColor.Yellow;
                 Console.WriteLine("Session saved.");
             }
+            // Load the chat state from disk
             else if (userInput == "load")
             {
                 session.LoadSession("Assets/chat-with-bob");

LLama.Examples/Examples/ChatSessionWithRestart.cs

@@ -37,8 +37,11 @@ public class ChatSessionWithRestart
         };
 
         Console.ForegroundColor = ConsoleColor.Yellow;
-        Console.WriteLine("The chat session has started. Write `save` to save session in memory."
-            + " Write `reset` to start from the last saved checkpoint");
+        Console.WriteLine("The chat session has started. Starting point saved.");
+        Console.WriteLine("Type 'exit' to end the chat session.");
+        Console.WriteLine("Type 'save' to save chat session state in memory.");
+        Console.WriteLine("Type 'reset' to reset the chat session to its saved state.");
+        Console.WriteLine("Type 'answer for assistant' to add and process provided user and assistant messages.");
 
         // show the prompt
         Console.ForegroundColor = ConsoleColor.Green;
@@ -46,6 +49,7 @@ public class ChatSessionWithRestart
 
         while (userInput != "exit")
         {
+            // Load the session state from the reset state
             if(userInput == "reset")
             {
                 session.LoadSession(resetState);
@@ -53,25 +57,33 @@ public class ChatSessionWithRestart
                 Console.ForegroundColor = ConsoleColor.Yellow;
                 Console.WriteLine("Session reset.");
             }
+            // Assign new reset state.
             else if (userInput == "save")
             {
                 resetState = session.GetSessionState();
                 Console.ForegroundColor = ConsoleColor.Yellow;
                 Console.WriteLine("Session saved.");
             }
-            else if (userInput == "regenerate")
+            // Provide user and override assistant answer with your own.
+            else if (userInput == "answer for assistant")
             {
                 Console.ForegroundColor = ConsoleColor.Yellow;
-                Console.WriteLine("Regenerating last response ...");
+                Console.WriteLine("Provide user input: ");
 
-                await foreach (
-                    var text
-                    in session.RegenerateAssistantMessageAsync(
-                        inferenceParams))
-                {
-                    Console.ForegroundColor = ConsoleColor.White;
-                    Console.Write(text);
-                }
+                Console.ForegroundColor = ConsoleColor.Green;
+                string userInputOverride = Console.ReadLine() ?? "";
+
+                Console.ForegroundColor = ConsoleColor.Yellow;
+                Console.WriteLine("Provide assistant input: ");
+                
+                Console.ForegroundColor = ConsoleColor.Green;
+                string assistantInputOverride = Console.ReadLine() ?? "";
+                
+                await session.AddAndProcessUserMessage(userInputOverride);
+                await session.AddAndProcessAssistantMessage(assistantInputOverride);
+
+                Console.ForegroundColor = ConsoleColor.Yellow;
+                Console.WriteLine("User and assistant messages processed. Provide next user message:");
             }
             else
             {

LLama/ChatSession.cs

@@ -262,33 +262,6 @@ public class ChatSession
         return this;
     }
 
-    
-    /// <summary>
-    /// Compute KV cache for the system message and add it to the chat history.
-    /// </summary>
-    /// <param name="content"></param>
-    /// <returns></returns>
-    public async Task<ChatSession> ProcessSystemMessage(string content)
-    {
-        if (Executor is not StatefulExecutorBase statefulExecutor)
-        {
-            throw new InvalidOperationException("Executor must be a StatefulExecutorBase to support pre-processing of system messages.");
-        }
-        if (History.Messages.Count > 0)
-        {
-            throw new ArgumentException("Cannot add a system message after another message", nameof(content));
-        }        
-        foreach (var inputTransform in InputTransformPipeline)
-        {
-            content = inputTransform.Transform(content);
-        }
-
-        await statefulExecutor.PrefillPromptAsync(content);
-
-        History.AddMessage(AuthorRole.System, content);
-        return this;
-    }
-
     /// <summary>
     /// Add a system message to the chat history.
     /// </summary>
@@ -323,6 +296,49 @@ public class ChatSession
         return this;
     }
 
+    /// <summary>
+    /// Compute KV cache for the message and add it to the chat history.
+    /// </summary>
+    /// <param name="message"></param>
+    /// <returns></returns>
+    public async Task<ChatSession> AddAndProcessMessage(ChatHistory.Message message)
+    {
+        if (Executor is not StatefulExecutorBase statefulExecutor)
+        {
+            throw new InvalidOperationException("Executor must be a StatefulExecutorBase to support pre-processing of system messages.");
+        }
+        AddMessage(message);
+        var content = message.Content;
+        if (message.AuthorRole != AuthorRole.Assistant)
+        {
+            foreach (var inputTransform in InputTransformPipeline)
+            {
+                content = inputTransform.Transform(content);
+            }
+        }
+
+        await statefulExecutor.PrefillPromptAsync(content);
+        return this;
+    }
+
+    /// <summary>
+    /// Compute KV cache for the system message and add it to the chat history.
+    /// </summary>
+    public Task<ChatSession> AddAndProcessSystemMessage(string content)
+        => AddAndProcessMessage(new ChatHistory.Message(AuthorRole.System, content));
+
+    /// <summary>
+    /// Compute KV cache for the user message and add it to the chat history.
+    /// </summary>
+    public Task<ChatSession> AddAndProcessUserMessage(string content)
+        => AddAndProcessMessage(new ChatHistory.Message(AuthorRole.User, content));
+
+    /// <summary>
+    /// Compute KV cache for the assistant message and add it to the chat history.
+    /// </summary>
+    public Task<ChatSession> AddAndProcessAssistantMessage(string content)
+        => AddAndProcessMessage(new ChatHistory.Message(AuthorRole.Assistant, content));
+
     /// <summary>
     /// Replace a user message with a new message and remove all messages after the new message.
     /// This is useful when the user wants to edit a message. And regenerate the response.