Rename to QAInformationRetrievalChain

Author: intitni

Core/Sources/ChatContextCollectors/WebChatContextCollector/QueryWebsiteFunction.swift

@@ -64,8 +64,8 @@ struct QueryWebsiteFunction: ChatGPTFunction {
 
                         if let database = await TemporaryUSearch.view(identifier: urlString) {
                             await reportProgress("Generating answers..")
-                            let qa = RetrievalQAChain(vectorStore: database, embedding: embedding)
-                            return try await qa.call(.init(arguments.query)).answer
+                            let qa = QAInformationRetrievalChain(vectorStore: database, embedding: embedding)
+                            return try await qa.call(.init(arguments.query)).information
                         }
                         let loader = WebLoader(urls: [url])
                         let documents = try await loader.load()
@@ -83,9 +83,9 @@ struct QueryWebsiteFunction: ChatGPTFunction {
                         try await database.set(embeddedDocuments)
                         // 4. generate answer
                         await reportProgress("Generating answers..")
-                        let qa = RetrievalQAChain(vectorStore: database, embedding: embedding)
+                        let qa = QAInformationRetrievalChain(vectorStore: database, embedding: embedding)
                         let result = try await qa.call(.init(arguments.query))
-                        return result.answer
+                        return result.information
                     }
                 }
 

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

@@ -106,7 +106,7 @@ struct QAForm: View {
             }
         }()
 
-        let qa = RetrievalQAChain(
+        let qa = QAInformationRetrievalChain(
             vectorStore: store,
             embedding: embedding
         )

Tool/Sources/LangChain/Chains/RetrievalQA.swift

@@ -1,12 +1,12 @@
 import Foundation
 import OpenAIService
 
-public final class RetrievalQAChain: Chain {
+public final class QAInformationRetrievalChain: Chain {
     let vectorStore: VectorStore
     let embedding: Embeddings
 
     public struct Output {
-        public var answer: String
+        public var information: String
         public var sourceDocuments: [Document]
     }
 
@@ -26,7 +26,9 @@ public final class RetrievalQAChain: Chain {
         let documents = try await vectorStore.searchWithDistance(
             embeddings: embeddedQuestion,
             count: 5
-        )
+        ).filter { item in
+            item.distance < 0.31
+        }
 
         callbackManagers.send(CallbackEvents.RetrievalQADidExtractRelevantContent(info: documents))
 
@@ -36,24 +38,21 @@ public final class RetrievalQAChain: Chain {
             callbackManagers: callbackManagers
         )
 
-        return .init(answer: relevantInformation, sourceDocuments: documents.map(\.document))
+        return .init(information: relevantInformation, sourceDocuments: documents.map(\.document))
     }
 
     public func parseOutput(_ output: Output) -> String {
-        return output.answer
+        return output.information
     }
 }
 
 public extension CallbackEvents {
     struct RetrievalQADidExtractRelevantContent: CallbackEvent {
         public let info: [(document: Document, distance: Float)]
     }
-    
+
     var retrievalQADidExtractRelevantContent: RetrievalQADidExtractRelevantContent.Type {
         RetrievalQADidExtractRelevantContent.self
     }
 }
 
-
-
-