From 8ba65988f9ccda78a5380df5657933da0efef657 Mon Sep 17 00:00:00 2001
From: Thorsten Sommer <mail@tsommer.org>
Date: Sat, 8 Mar 2025 13:54:07 +0100
Subject: [PATCH] Improved the augmentation and generation part of RAG by
 passing the augmented data into the system prompt

---
 app/MindWork AI Studio/Chat/ChatThread.cs     | 32 ++++++++++++++++---
 .../AugmentationProcesses/AugmentationOne.cs  | 18 ++---------
 .../wwwroot/changelog/v0.9.32.md              |  1 +
 3 files changed, 30 insertions(+), 21 deletions(-)

diff --git a/app/MindWork AI Studio/Chat/ChatThread.cs b/app/MindWork AI Studio/Chat/ChatThread.cs
index d7dbf9a4..73182e7c 100644
--- a/app/MindWork AI Studio/Chat/ChatThread.cs	
+++ b/app/MindWork AI Studio/Chat/ChatThread.cs	
@@ -40,6 +40,11 @@ public sealed record ChatThread
     /// </summary>
     public IReadOnlyList<DataSourceAgentSelected> AISelectedDataSources { get; set; } = [];
 
+    /// <summary>
+    /// The augmented data for this chat thread. Will be inserted into the system prompt.
+    /// </summary>
+    public string AugmentedData { get; set; } = string.Empty;
+
     /// <summary>
     /// The name of the chat thread. Usually generated by an AI model or manually edited by the user.
     /// </summary>
@@ -74,31 +79,48 @@ public sealed record ChatThread
     /// <returns>The prepared system prompt.</returns>
     public string PrepareSystemPrompt(SettingsManager settingsManager, ChatThread chatThread, ILogger logger)
     {
+        var isAugmentedDataAvailable = !string.IsNullOrWhiteSpace(chatThread.AugmentedData);
+        var systemPromptWithAugmentedData = isAugmentedDataAvailable switch
+        {
+            true => $"""
+                     {chatThread.SystemPrompt}
+                     
+                     {chatThread.AugmentedData}
+                     """,
+
+            false => chatThread.SystemPrompt,
+        };
+        
+        if(isAugmentedDataAvailable)
+            logger.LogInformation("Augmented data is available for the chat thread.");
+        else
+            logger.LogInformation("No augmented data is available for the chat thread.");
+        
         //
         // Prepare the system prompt:
         //
         string systemPromptText;
         var logMessage = $"Using no profile for chat thread '{chatThread.Name}'.";
         if (string.IsNullOrWhiteSpace(chatThread.SelectedProfile))
-            systemPromptText = chatThread.SystemPrompt;
+            systemPromptText = systemPromptWithAugmentedData;
         else
         {
             if(!Guid.TryParse(chatThread.SelectedProfile, out var profileId))
-                systemPromptText = chatThread.SystemPrompt;
+                systemPromptText = systemPromptWithAugmentedData;
             else
             {
                 if(chatThread.SelectedProfile == Profile.NO_PROFILE.Id || profileId == Guid.Empty)
-                    systemPromptText = chatThread.SystemPrompt;
+                    systemPromptText = systemPromptWithAugmentedData;
                 else
                 {
                     var profile = settingsManager.ConfigurationData.Profiles.FirstOrDefault(x => x.Id == chatThread.SelectedProfile);
                     if(profile == default)
-                        systemPromptText = chatThread.SystemPrompt;
+                        systemPromptText = systemPromptWithAugmentedData;
                     else
                     {
                         logMessage = $"Using profile '{profile.Name}' for chat thread '{chatThread.Name}'.";
                         systemPromptText = $"""
-                                            {chatThread.SystemPrompt}
+                                            {systemPromptWithAugmentedData}
 
                                             {profile.ToSystemPrompt()}
                                             """;
diff --git a/app/MindWork AI Studio/Tools/RAG/AugmentationProcesses/AugmentationOne.cs b/app/MindWork AI Studio/Tools/RAG/AugmentationProcesses/AugmentationOne.cs
index fd6e30af..fff91251 100644
--- a/app/MindWork AI Studio/Tools/RAG/AugmentationProcesses/AugmentationOne.cs	
+++ b/app/MindWork AI Studio/Tools/RAG/AugmentationProcesses/AugmentationOne.cs	
@@ -66,22 +66,8 @@ public sealed class AugmentationOne : IAugmentationProcess
         // Let's convert all retrieval contexts to Markdown:
         await retrievalContexts.AsMarkdown(sb, token);
         
-        //
-        // Append the entire augmentation to the chat thread,
-        // just before the user prompt:
-        //
-        chatThread.Blocks.Insert(chatThread.Blocks.Count - 1, new()
-        {
-            Role = ChatRole.RAG,
-            Time = DateTimeOffset.UtcNow,
-            ContentType = ContentType.TEXT,
-            HideFromUser = true,
-            Content = new ContentText
-            {
-                Text = sb.ToString(),
-            }
-        });
-        
+        // Add the augmented data to the chat thread:
+        chatThread.AugmentedData = sb.ToString();
         return chatThread;
     }
 
diff --git a/app/MindWork AI Studio/wwwroot/changelog/v0.9.32.md b/app/MindWork AI Studio/wwwroot/changelog/v0.9.32.md
index e462100b..007b7d72 100644
--- a/app/MindWork AI Studio/wwwroot/changelog/v0.9.32.md	
+++ b/app/MindWork AI Studio/wwwroot/changelog/v0.9.32.md	
@@ -5,5 +5,6 @@
 - Improved the error handling in the ERI data source info dialog in cases where servers respond with an invalid message.
 - Improved the error handling for the entire RAG process.
 - Improved chat thread persistence after modifications through the RAG process.
+- Improved the augmentation and generation part of RAG by passing the augmented data into the system prompt.
 - Fixed the chat thread we use for the data retrieval by removing the last block, which is meant to be for the final AI answer.
 - Fixed the data source name for ERI data sources when performing data retrieval.
\ No newline at end of file