Adjust implementation of QAInformationRetrievalChain

Author: intitni

Tool/Sources/LangChain/Chains/RelevantInformationExtractionChain.swift

@@ -15,44 +15,74 @@ public final class RelevantInformationExtractionChain: Chain {
     public typealias Output = String
 
     class FunctionProvider: ChatGPTFunctionProvider {
-        var functionCallStrategy: FunctionCallStrategy? = .auto
-        var functions: [any ChatGPTFunction] = [NoneFunction()]
+        var functionCallStrategy: FunctionCallStrategy? = .name("saveFinalAnswer")
+        var functions: [any ChatGPTFunction] = [FinalAnswer()]
     }
 
-    struct NoneFunction: ChatGPTArgumentsCollectingFunction {
-        typealias Arguments = NoArguments
-        var name: String = "noInformationFound"
-        var description: String = "Call when you can't find any relevant information from the document, or the question was not mentioned in the document"
+    struct FinalAnswer: ChatGPTArgumentsCollectingFunction {
+        struct Arguments: Decodable {
+            var relevantInformation: String
+            var noRelevantInformationFound: Bool?
+        }
+
+        var name: String = "saveFinalAnswer"
+        var description: String =
+            "save the relevant information"
+        var argumentSchema: JSONSchemaValue {
+            [
+                .type: "object",
+                .properties: [
+                    "relevantInformation": [.type: "string"],
+                    "noRelevantInformationFound": [.type: "boolean"],
+                ],
+                .required: ["relevantInformation", "noRelevantInformationFound"],
+            ]
+        }
+    }
+
+    let filterMetadata: (String) -> Bool
+    let hint: String
+
+    init(filterMetadata: @escaping (String) -> Bool = { _ in true }, hint: String) {
+        self.filterMetadata = filterMetadata
+        self.hint = hint
     }
 
     func buildChatModel() -> ChatModelChain<TaskInput> {
         .init(
             chatModel: OpenAIChat(
                 configuration: UserPreferenceChatGPTConfiguration().overriding {
-                    $0.temperature = 0
+                    $0.temperature = 0.5
                     $0.runFunctionsAutomatically = false
                 },
                 memory: EmptyChatGPTMemory(),
                 functionProvider: FunctionProvider(),
                 stream: false
             )
-        ) { input in [
+        ) { [filterMetadata, hint] input in [
             .init(
                 role: .system,
                 content: """
                 Extract the relevant information from the Document according to the Question.
+                The information may not directly answer the question, but it should be relevant to the question, \
+                please think carefully and make you decision.
                 Make the information clear, concise and short.
                 If found code, wrap it in markdown code block.
+                \(hint)
                 """
             ),
             .init(
                 role: .user,
                 content: """
                 Question:###
+                (how, when, what or why)
                 \(input.question)
                 ###
                 Document:###
-                \(input.document)
+                \(input.document.metadata.filter { key, _ in
+                    filterMetadata(key)
+                })
+                \(input.document.pageContent)
                 ###
                 """
             ),
@@ -73,6 +103,22 @@ public final class RelevantInformationExtractionChain: Chain {
                             taskInput,
                             callbackManagers: callbackManagers
                         )
+                        
+                        if let functionCall = output.functionCall {
+                            do {
+                                let arguments = try JSONDecoder().decode(
+                                    FinalAnswer.Arguments.self,
+                                    from: functionCall.arguments.data(using: .utf8) ?? Data()
+                                )
+                                if arguments.noRelevantInformationFound ?? false {
+                                    return ""
+                                }
+                                return arguments.relevantInformation
+                            } catch {
+                                return output.content ?? ""
+                            }
+                        }
+                        
                         return output.content ?? ""
                     }
 
@@ -119,3 +165,4 @@ public extension CallbackEvents {
         RelevantInformationExtractionChainDidExtractPartialRelevantContent.self
     }
 }
+

Tool/Sources/LangChain/Chains/RetrievalQA.swift

@@ -5,6 +5,8 @@ public final class QAInformationRetrievalChain: Chain {
     let vectorStores: [VectorStore]
     let embedding: Embeddings
     let maxCount: Int
+    let filterMetadata: (String) -> Bool
+    let hint: String
 
     public struct Output {
         public var information: String
@@ -14,21 +16,29 @@ public final class QAInformationRetrievalChain: Chain {
     public init(
         vectorStore: VectorStore,
         embedding: Embeddings,
-        maxCount: Int = 5
+        maxCount: Int = 5,
+        filterMetadata: @escaping (String) -> Bool = { _ in true },
+        hint: String = ""
     ) {
         vectorStores = [vectorStore]
         self.embedding = embedding
         self.maxCount = maxCount
+        self.filterMetadata = filterMetadata
+        self.hint = hint
     }
 
     public init(
         vectorStores: [VectorStore],
         embedding: Embeddings,
-        maxCount: Int = 5
+        maxCount: Int = 5,
+        filterMetadata: @escaping (String) -> Bool = { _ in true },
+        hint: String = ""
     ) {
         self.vectorStores = vectorStores
         self.embedding = embedding
         self.maxCount = maxCount
+        self.filterMetadata = filterMetadata
+        self.hint = hint
     }
 
     public func callLogic(
@@ -57,10 +67,13 @@ public final class QAInformationRetrievalChain: Chain {
         }.sorted { $0.distance < $1.distance }.prefix(maxCount)
 
         let documents = Array(documentsSlice)
-        
+
         callbackManagers.send(CallbackEvents.RetrievalQADidExtractRelevantContent(info: documents))
 
-        let relevantInformationChain = RelevantInformationExtractionChain()
+        let relevantInformationChain = RelevantInformationExtractionChain(
+            filterMetadata: filterMetadata,
+            hint: hint
+        )
         let relevantInformation = try await relevantInformationChain.run(
             .init(question: input, documents: documents),
             callbackManagers: callbackManagers