Update RetrievalQA to support early return when the answer is good enough

Author: intitni

Core/Sources/ChatService/ChatFunctionProvider.swift

@@ -15,5 +15,9 @@ final class ChatFunctionProvider {
     }
 }
 
-extension ChatFunctionProvider: ChatGPTFunctionProvider {}
+extension ChatFunctionProvider: ChatGPTFunctionProvider {
+    var functionCallStrategy: OpenAIService.FunctionCallStrategy? {
+        nil
+    }
+}
 

Playground.playground/Pages/RetrievalQAChain.xcplaygroundpage/Contents.swift

@@ -5,7 +5,7 @@ import PlaygroundSupport
 import SwiftUI
 
 struct QAForm: View {
-    @State var intermediateAnswers = [String]()
+    @State var intermediateAnswers = [RefineDocumentChain.IntermediateAnswer]()
     @State var answer: String = ""
     @State var question: String = "What is Swift macros?"
     @State var isProcessing: Bool = false
@@ -31,9 +31,14 @@ struct QAForm: View {
                 Text(answer)
             }
             Section(header: Text("Intermediate Answers")) {
-                ForEach(intermediateAnswers, id: \.self) { answer in
-                    Text(answer)
-                    Divider()
+                ForEach(0..<intermediateAnswers.endIndex, id: \.self) { index in
+                    let answer = intermediateAnswers[index]
+                    VStack {
+                        Text(answer.answer)
+                        Text("Score: \(answer.score)")
+                        Text("Needs more context: \(answer.more ? "Yes" : "No")")
+                        Divider()
+                    }
                 }
             }
         }
@@ -48,8 +53,6 @@ struct QAForm: View {
             answer = "Invalid URL"
             return
         }
-        let chatGPTConfiguration = UserPreferenceChatGPTConfiguration()
-            .overriding { $0.temperature = 0 }
         let embeddingConfiguration = UserPreferenceEmbeddingConfiguration().overriding()
         let embedding = OpenAIEmbedding(configuration: embeddingConfiguration)
         let store: VectorStore = try await {
@@ -72,8 +75,7 @@ struct QAForm: View {
 
         let qa = RetrievalQAChain(
             vectorStore: store,
-            embedding: embedding,
-            chatModelFactory: { OpenAIChat(configuration: chatGPTConfiguration, stream: false) }
+            embedding: embedding
         )
         answer = try await qa.run(
             question,

Playground.playground/Pages/RetrievalQAChain.xcplaygroundpage/timeline.xctimeline

@@ -3,6 +3,8 @@
    version = "3.0">
    <TimelineItems>
    </TimelineItems>
+   <TimelineItems>
+   </TimelineItems>
    <TimelineItems>
       <LoggerValueHistoryTimelineItem
          documentLocation = "documentLocation"

Tool/Sources/LangChain/Chains/RetrievalQA.swift

@@ -43,7 +43,7 @@ public final class RetrievalQAChain: Chain {
 
 public extension CallbackEvents {
     struct RetrievalQADidGenerateIntermediateAnswer: CallbackEvent {
-        public let info: String
+        public let info: RefineDocumentChain.IntermediateAnswer
     }
 }
 
@@ -66,16 +66,38 @@ public final class RefineDocumentChain: Chain {
         var distance: Float
     }
 
+    public struct IntermediateAnswer: Decodable {
+        public var answer: String
+        public var score: Double
+        public var more: Bool
+
+        public enum CodingKeys: String, CodingKey {
+            case answer
+            case score
+            case more
+        }
+
+        init(answer: String, score: Double, more: Bool) {
+            self.answer = answer
+            self.score = score
+            self.more = more
+        }
+
+        public init(from decoder: Decoder) throws {
+            let container = try decoder.container(keyedBy: CodingKeys.self)
+            answer = try container.decode(String.self, forKey: .answer)
+            score = (try? container.decode(Double.self, forKey: .score)) ?? 0
+            more = (try? container.decode(Bool.self, forKey: .more)) ?? (score < 6)
+        }
+    }
+
     class FunctionProvider: ChatGPTFunctionProvider {
-        var functions: [any ChatGPTFunction] = []
+        var functionCallStrategy: FunctionCallStrategy? = .name("respond")
+        var functions: [any ChatGPTFunction] = [RespondFunction()]
     }
 
     struct RespondFunction: ChatGPTFunction {
-        struct Arguments: Codable {
-            var answer: String
-            var score: Double
-            var more: Bool
-        }
+        typealias Arguments = IntermediateAnswer
 
         struct Result: ChatGPTFunctionResult {
             var botReadableContent: String { "" }
@@ -91,17 +113,18 @@ public final class RefineDocumentChain: Chain {
                 .properties: [
                     "answer": [
                         .type: "string",
-                        .description: "The answer",
+                        .description: "The refined answer",
                     ],
                     "score": [
                         .type: "number",
-                        .description: "The score of the answer, the higher the better",
+                        .description: "The score of the answer, the higher the better. 0 to 10.",
                     ],
                     "more": [
                         .type: "boolean",
                         .description: "Whether more information is needed to complete the answer",
                     ],
                 ],
+                .required: ["answer", "score", "more"],
             ]
         }
 
@@ -114,18 +137,16 @@ public final class RefineDocumentChain: Chain {
 
     let initialChatModel: ChatModelChain<InitialInput>
     let refinementChatModel: ChatModelChain<RefinementInput>
-    let initialChatMemory: ChatGPTMemory
-    let refinementChatMemory: ChatGPTMemory
 
     public init() {
-        initialChatMemory = ConversationChatGPTMemory(systemPrompt: "")
-        refinementChatMemory = ConversationChatGPTMemory(systemPrompt: "")
-
         initialChatModel = .init(
             chatModel: OpenAIChat(
-                configuration: UserPreferenceChatGPTConfiguration()
-                    .overriding(.init(temperature: 0)),
-                memory: initialChatMemory,
+                configuration: UserPreferenceChatGPTConfiguration().overriding {
+                    $0.temperature = 0
+                    $0.runFunctionsAutomatically = false
+                },
+                memory: EmptyChatGPTMemory(),
+                functionProvider: FunctionProvider(),
                 stream: false
             ),
             promptTemplate: { input in [
@@ -140,9 +161,12 @@ public final class RefineDocumentChain: Chain {
         )
         refinementChatModel = .init(
             chatModel: OpenAIChat(
-                configuration: UserPreferenceChatGPTConfiguration()
-                    .overriding(.init(temperature: 0)),
-                memory: refinementChatMemory,
+                configuration: UserPreferenceChatGPTConfiguration().overriding {
+                    $0.temperature = 0
+                    $0.runFunctionsAutomatically = false
+                },
+                memory: EmptyChatGPTMemory(),
+                functionProvider: FunctionProvider(),
                 stream: false
             ),
             promptTemplate: { input in [
@@ -168,6 +192,26 @@ public final class RefineDocumentChain: Chain {
         guard let firstDocument = input.documents.first else {
             return ""
         }
+
+        func extractAnswer(_ chatMessage: ChatMessage) -> IntermediateAnswer {
+            if let functionCall = chatMessage.functionCall {
+                do {
+                    let intermediateAnswer = try JSONDecoder().decode(
+                        IntermediateAnswer.self,
+                        from: functionCall.arguments.data(using: .utf8) ?? Data()
+                    )
+                    return intermediateAnswer
+                } catch {
+                    let intermediateAnswer = IntermediateAnswer(
+                        answer: functionCall.arguments,
+                        score: 0,
+                        more: true
+                    )
+                    return intermediateAnswer
+                }
+            }
+            return .init(answer: chatMessage.content ?? "", score: 0, more: true)
+        }
         var output = try await initialChatModel.call(
             .init(
                 question: input.question,
@@ -176,24 +220,27 @@ public final class RefineDocumentChain: Chain {
             ),
             callbackManagers: callbackManagers
         )
-        guard var content = output.content else { return "" }
-        callbackManagers
-            .send(CallbackEvents.RetrievalQADidGenerateIntermediateAnswer(info: content))
-        for document in input.documents.dropFirst(1) {
+        var intermediateAnswer = extractAnswer(output)
+        callbackManagers.send(
+            CallbackEvents.RetrievalQADidGenerateIntermediateAnswer(info: intermediateAnswer)
+        )
+
+        for document in input.documents.dropFirst(1) where intermediateAnswer.more {
             output = try await refinementChatModel.call(
                 .init(
                     question: input.question,
-                    previousAnswer: content,
+                    previousAnswer: intermediateAnswer.answer,
                     document: document.document.pageContent,
                     distance: document.distance
                 ),
                 callbackManagers: callbackManagers
             )
-            content = output.content ?? ""
-            callbackManagers
-                .send(CallbackEvents.RetrievalQADidGenerateIntermediateAnswer(info: content))
+            intermediateAnswer = extractAnswer(output)
+            callbackManagers.send(
+                CallbackEvents.RetrievalQADidGenerateIntermediateAnswer(info: intermediateAnswer)
+            )
         }
-        return content
+        return intermediateAnswer.answer
     }
 
     public func parseOutput(_ output: String) -> String {

Tool/Sources/LangChain/ChatModel/OpenAIChat.swift

@@ -38,8 +38,7 @@ public struct OpenAIChat: ChatModel {
             var message = ""
             for try await trunk in stream {
                 message.append(trunk)
-                callbackManagers
-                    .forEach { $0.send(CallbackEvents.LLMDidProduceNewToken(info: trunk)) }
+                callbackManagers.send(CallbackEvents.LLMDidProduceNewToken(info: trunk))
             }
             return await memory.messages.last ?? .init(role: .assistant, content: "")
         } else {

Tool/Sources/OpenAIService/ChatGPTService.swift

@@ -205,7 +205,7 @@ extension ChatGPTService {
                 model: configuration.model,
                 remainingTokens: remainingTokens
             ),
-            function_call: nil,
+            function_call: functionProvider.functionCallStrategy,
             functions: functionProvider.functions.map {
                 ChatGPTFunctionSchema(
                     name: $0.name,
@@ -302,7 +302,7 @@ extension ChatGPTService {
                 model: configuration.model,
                 remainingTokens: remainingTokens
             ),
-            function_call: nil,
+            function_call: functionProvider.functionCallStrategy,
             functions: functionProvider.functions.map {
                 ChatGPTFunctionSchema(
                     name: $0.name,
@@ -318,6 +318,7 @@ extension ChatGPTService {
             url,
             requestBody
         )
+
         let response = try await api()
 
         guard let choice = response.choices.first else { return nil }

Tool/Sources/OpenAIService/CompletionStreamAPI.swift

@@ -12,6 +12,31 @@ protocol CompletionStreamAPI {
     )
 }
 
+public enum FunctionCallStrategy: Encodable, Equatable {
+    /// Forbid the bot to call any function.
+    case none
+    /// Let the bot choose what function to call.
+    case auto
+    /// Force the bot to call a function with the given name.
+    case name(String)
+
+    struct CallFunctionNamed: Codable {
+        var name: String
+    }
+
+    public func encode(to encoder: Encoder) throws {
+        var container = encoder.singleValueContainer()
+        switch self {
+        case .none:
+            try container.encode("none")
+        case .auto:
+            try container.encode("auto")
+        case let .name(name):
+            try container.encode(CallFunctionNamed(name: name))
+        }
+    }
+}
+
 /// https://platform.openai.com/docs/api-reference/chat/create
 struct CompletionRequestBody: Encodable, Equatable {
     struct Message: Codable, Equatable {
@@ -31,7 +56,7 @@ struct CompletionRequestBody: Encodable, Equatable {
         ///   "arguments": "{ \"location\": \"earth\" }"
         /// }
         /// ```
-        var function_call: MessageFunctionCall?
+        var function_call: CompletionRequestBody.MessageFunctionCall?
     }
 
     struct MessageFunctionCall: Codable, Equatable {
@@ -41,31 +66,6 @@ struct CompletionRequestBody: Encodable, Equatable {
         var arguments: String?
     }
 
-    enum FunctionCallStrategy: Encodable, Equatable {
-        /// Forbid the bot to call any function.
-        case none
-        /// Let the bot choose what function to call.
-        case auto
-        /// Force the bot to call a function with the given name.
-        case name(String)
-
-        struct CallFunctionNamed: Codable {
-            var name: String
-        }
-
-        func encode(to encoder: Encoder) throws {
-            var container = encoder.singleValueContainer()
-            switch self {
-            case .none:
-                try container.encode("none")
-            case .auto:
-                try container.encode("auto")
-            case let .name(name):
-                try container.encode(CallFunctionNamed(name: name))
-            }
-        }
-    }
-
     struct Function: Codable {
         var name: String
         var description: String

Tool/Sources/OpenAIService/FucntionCall/ChatGPTFuntionProvider.swift

@@ -2,6 +2,7 @@ import Foundation
 
 public protocol ChatGPTFunctionProvider {
     var functions: [any ChatGPTFunction] { get }
+    var functionCallStrategy: FunctionCallStrategy? { get }
 }
 
 extension ChatGPTFunctionProvider {
@@ -11,6 +12,7 @@ extension ChatGPTFunctionProvider {
 }
 
 public struct NoChatGPTFunctionProvider: ChatGPTFunctionProvider {
+    public var functionCallStrategy: FunctionCallStrategy?
     public var functions: [any ChatGPTFunction] { [] }
     public init() {}
 }