Fix retrieved content number in rag when context window is not big enough

Author: intitni

Tool/Sources/OpenAIService/Memory/AutoManagedChatGPTMemory.swift

@@ -286,8 +286,6 @@ extension AutoManagedChatGPTMemory {
                     text += """
                     Here are the information you know about the system and the project, \
                     separated by \(separator)
-
-                    
                     """
                 }
 
@@ -302,7 +300,7 @@ extension AutoManagedChatGPTMemory {
         {
             var right = retrievedContent.count
             var left = 0
-            var retrievedContent = retrievedContent
+            var gappedRetrievedContent = retrievedContent
             var tokenCount: Int?
             var proposedMessage = buildMessage(retrievedContent: [])
 
@@ -340,7 +338,7 @@ extension AutoManagedChatGPTMemory {
                 let (isValid, _tokenCount) = await checkValid(proposedMessage: _proposedMessage)
                 if isValid {
                     proposedMessage = _proposedMessage
-                    retrievedContent = _retrievedContent
+                    gappedRetrievedContent = _retrievedContent
                     tokenCount = _tokenCount
                     left = count + 1
                 } else {
@@ -355,7 +353,7 @@ extension AutoManagedChatGPTMemory {
             } else {
                 await strategy.countToken(proposedMessage)
             }
-            return (proposedMessage, retrievedContent, finalCount)
+            return (proposedMessage, gappedRetrievedContent, finalCount)
         }
 
         let (message, references, tokensCount) = await buildMessageThatFits()

Tool/Tests/OpenAIServiceTests/AutoManagedChatGPTMemoryRetrievedContentTests.swift

@@ -0,0 +1,178 @@
+import Foundation
+import XCTest
+
+@testable import OpenAIService
+
+class AutoManagedChatGPTMemoryRetrievedContentTests: XCTestCase {
+    let separator = String(repeating: "=", count: 32)
+
+    func ref(_ text: String) -> ChatMessage.Reference {
+        .init(
+            title: "",
+            subTitle: "",
+            content: text,
+            uri: "",
+            startLine: nil,
+            endLine: nil,
+            kind: .text
+        )
+    }
+
+    func test_retrieved_content_when_the_context_window_is_large_enough() async {
+        let strategy = Strategy()
+        
+        let memory = AutoManagedChatGPTMemory(
+            systemPrompt: "",
+            configuration: UserPreferenceChatGPTConfiguration(),
+            functionProvider: EmptyFunctionProvider()
+        )
+
+        await memory.mutateRetrievedContent([
+            ref("A"), ref("B"), ref("C"), ref("D"), ref("E"),
+        ])
+
+        let fullContent = """
+        Here are the information you know about the system and the project, \
+        separated by \(separator)
+
+        \(separator)[DOCUMENT 0]
+
+        A
+
+        \(separator)[DOCUMENT 1]
+
+        B
+
+        \(separator)[DOCUMENT 2]
+
+        C
+
+        \(separator)[DOCUMENT 3]
+
+        D
+
+        \(separator)[DOCUMENT 4]
+
+        E
+        """
+
+        let maxTokenCount = await strategy.countToken(.init(role: .user, content: fullContent))
+        
+        let result = await memory.generateRetrievedContentMessage(
+            maxTokenCount: maxTokenCount,
+            strategy: strategy
+        )
+
+        XCTAssertEqual(result.references.count, 5)
+        XCTAssertEqual(result.retrievedContent.role, .user)
+        XCTAssertEqual(result.retrievedContent.content, """
+        Here are the information you know about the system and the project, \
+        separated by \(separator)
+
+        \(separator)[DOCUMENT 0]
+
+        A
+
+        \(separator)[DOCUMENT 1]
+
+        B
+
+        \(separator)[DOCUMENT 2]
+
+        C
+
+        \(separator)[DOCUMENT 3]
+
+        D
+
+        \(separator)[DOCUMENT 4]
+
+        E
+        """)
+    }
+
+    func test_retrieved_content_when_the_context_window_is_just_not_large_enough() async {
+        let strategy = Strategy()
+
+        let memory = AutoManagedChatGPTMemory(
+            systemPrompt: "",
+            configuration: UserPreferenceChatGPTConfiguration(),
+            functionProvider: EmptyFunctionProvider()
+        )
+
+        await memory.mutateRetrievedContent([
+            ref("A"), ref("B"), ref("C"), ref("D"), ref("E"),
+        ])
+
+        let fullContent = """
+        Here are the information you know about the system and the project, \
+        separated by \(separator)
+
+        \(separator)[DOCUMENT 0]
+
+        A
+
+        \(separator)[DOCUMENT 1]
+
+        B
+
+        \(separator)[DOCUMENT 2]
+
+        C
+
+        \(separator)[DOCUMENT 3]
+
+        D
+
+        \(separator)[DOCUMENT 4]
+
+        E
+        """
+
+        let maxTokenCount = await strategy.countToken(.init(role: .user, content: fullContent))
+
+        let result = await memory.generateRetrievedContentMessage(
+            maxTokenCount: maxTokenCount - 1,
+            strategy: strategy
+        )
+
+        XCTAssertEqual(result.references.count, 4)
+        XCTAssertEqual(result.retrievedContent.role, .user)
+        XCTAssertEqual(result.retrievedContent.content, """
+        Here are the information you know about the system and the project, \
+        separated by \(separator)
+
+        \(separator)[DOCUMENT 0]
+
+        A
+
+        \(separator)[DOCUMENT 1]
+
+        B
+
+        \(separator)[DOCUMENT 2]
+
+        C
+
+        \(separator)[DOCUMENT 3]
+
+        D
+        """)
+    }
+}
+
+private struct EmptyFunctionProvider: ChatGPTFunctionProvider {
+    var functions: [any ChatGPTFunction] { [] }
+    var functionCallStrategy: FunctionCallStrategy? { nil }
+}
+
+private struct Strategy: AutoManagedChatGPTMemoryStrategy {
+    func countToken(_ message: OpenAIService.ChatMessage) async -> Int {
+        message.content?.count ?? 0
+    }
+
+    func countToken<F>(_: F) async -> Int where F: ChatGPTFunction {
+        0
+    }
+}
+

Tool/Tests/OpenAIServiceTests/LimitMessagesTests.swift

@@ -4,7 +4,7 @@ import XCTest
 
 @testable import OpenAIService
 
-final class AutoManagedChatGPTMemoryTests: XCTestCase {
+final class AutoManagedChatGPTMemoryLimitTests: XCTestCase {
     func test_send_all_messages_if_not_reached_token_limit() async {
         let (messages, memory) = await runService(
             systemPrompt: "system",