From 8330640eb57dfcac8bebd2b71ddaf8ac2bf4d2b2 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 9 Nov 2023 15:49:53 +0800
Subject: [PATCH] Move context system prompt to right next to the latest
 message

---
 .../DynamicContextController.swift            |  5 ++--
 .../Memory/AutoManagedChatGPTMemory.swift     | 28 ++++++++++++++++---
 2 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/Core/Sources/ChatService/DynamicContextController.swift b/Core/Sources/ChatService/DynamicContextController.swift
index 2f905999e..2f7cf0149 100644
--- a/Core/Sources/ChatService/DynamicContextController.swift
+++ b/Core/Sources/ChatService/DynamicContextController.swift
@@ -92,7 +92,7 @@ final class DynamicContextController {
             return contexts
         }
 
-        let extraSystemPrompt = contexts
+        let contextSystemPrompt = contexts
             .map(\.systemPrompt)
             .filter { !$0.isEmpty }
             .joined(separator: "\n\n")
@@ -104,9 +104,10 @@ final class DynamicContextController {
 
         let contextualSystemPrompt = """
         \(language.isEmpty ? "" : "You must always reply in \(language)")
-        \(systemPrompt)\(extraSystemPrompt.isEmpty ? "" : "\n\(extraSystemPrompt)")
+        \(systemPrompt)
         """
         await memory.mutateSystemPrompt(contextualSystemPrompt)
+        await memory.mutateContextSystemPrompt(contextSystemPrompt)
         await memory.mutateRetrievedContent(contextPrompts.map(\.content))
         functionProvider.append(functions: contexts.flatMap(\.functions))
     }
diff --git a/Tool/Sources/OpenAIService/Memory/AutoManagedChatGPTMemory.swift b/Tool/Sources/OpenAIService/Memory/AutoManagedChatGPTMemory.swift
index 9a117d632..0a4f36e60 100644
--- a/Tool/Sources/OpenAIService/Memory/AutoManagedChatGPTMemory.swift
+++ b/Tool/Sources/OpenAIService/Memory/AutoManagedChatGPTMemory.swift
@@ -9,6 +9,7 @@ public actor AutoManagedChatGPTMemory: ChatGPTMemory {
     public private(set) var remainingTokens: Int?
 
     public var systemPrompt: String
+    public var contextSystemPrompt: String
     public var retrievedContent: [String] = []
     public var history: [ChatMessage] = [] {
         didSet { onHistoryChange() }
@@ -27,6 +28,7 @@ public actor AutoManagedChatGPTMemory: ChatGPTMemory {
         functionProvider: ChatGPTFunctionProvider
     ) {
         self.systemPrompt = systemPrompt
+        contextSystemPrompt = ""
         self.configuration = configuration
         self.functionProvider = functionProvider
         _ = Self.encoder // force pre-initialize
@@ -40,6 +42,10 @@ public actor AutoManagedChatGPTMemory: ChatGPTMemory {
         systemPrompt = newPrompt
     }
 
+    public func mutateContextSystemPrompt(_ newPrompt: String) {
+        contextSystemPrompt = newPrompt
+    }
+
     public func mutateRetrievedContent(_ newContent: [String]) {
         retrievedContent = newContent
     }
@@ -67,6 +73,8 @@ public actor AutoManagedChatGPTMemory: ChatGPTMemory {
     ///     [Retrieved Content B]
     /// [Functions] priority: high
     /// [Message History] priority: medium
+    /// [Context System Prompt] priority: high
+    /// [Latest Message] priority: high
     /// ```
     func generateSendingHistory(
         maxNumberOfMessages: Int = UserDefaults.shared.value(for: \.chatGPTMaxMessageCount),
@@ -80,7 +88,12 @@ public actor AutoManagedChatGPTMemory: ChatGPTMemory {
         }
 
         var smallestSystemPromptMessage = ChatMessage(role: .system, content: systemPrompt)
+        var contextSystemPromptMessage = ChatMessage(role: .system, content: contextSystemPrompt)
         let smallestSystemMessageTokenCount = countToken(&smallestSystemPromptMessage)
+        let contextSystemPromptTokenCount = !contextSystemPrompt.isEmpty
+            ? countToken(&contextSystemPromptMessage)
+            : 0
+
         let functionTokenCount = functionProvider.functions.reduce(into: 0) { partial, function in
             var count = encoder.countToken(text: function.name)
                 + encoder.countToken(text: function.description)
@@ -92,6 +105,7 @@ public actor AutoManagedChatGPTMemory: ChatGPTMemory {
             partial += count
         }
         let mandatoryContentTokensCount = smallestSystemMessageTokenCount
+            + contextSystemPromptTokenCount
             + functionTokenCount
             + 3 // every reply is primed with <|start|>assistant<|message|>
 
@@ -135,13 +149,13 @@ public actor AutoManagedChatGPTMemory: ChatGPTMemory {
         for (index, content) in retrievedContent.filter({ !$0.isEmpty }).enumerated() {
             if index == 0 {
                 if !appendToSystemPrompt("""
-                
-                
+
+
                 ## Relevant Content
-                
+
                 Below are information related to the conversation, separated by \(separator)
 
-                
+
                 """) { break }
             } else {
                 if !appendToSystemPrompt("\n\(separator)\n") { break }
@@ -154,16 +168,22 @@ public actor AutoManagedChatGPTMemory: ChatGPTMemory {
             let message = ChatMessage(role: .system, content: systemPrompt)
             allMessages.append(message)
         }
+        
+        if !contextSystemPrompt.isEmpty {
+            allMessages.insert(contextSystemPromptMessage, at: 1)
+        }
 
         #if DEBUG
         Logger.service.info("""
         Sending tokens count
         - system prompt: \(smallestSystemMessageTokenCount)
+        - context system prompt: \(contextSystemPromptTokenCount)
         - functions: \(functionTokenCount)
         - messages: \(messageTokenCount)
         - retrieved content: \(retrievedContentTokenCount)
         - total: \(
             smallestSystemMessageTokenCount
+                + contextSystemPromptTokenCount
                 + functionTokenCount
                 + messageTokenCount
                 + retrievedContentTokenCount