Update QueryWebsiteFunction to use RetrievalQAChain

Author: intitni

Core/Sources/ChatContextCollectors/WebChatContextCollector/QueryWebsiteFunction.swift

@@ -10,14 +10,10 @@ struct QueryWebsiteFunction: ChatGPTFunction {
     }
 
     struct Result: ChatGPTFunctionResult {
-        var relevantDocuments: [Document]
+        var answers: [String]
 
         var botReadableContent: String {
-            // don't forget to remove overlaps
-            if relevantDocuments.isEmpty {
-                return "No relevant information found"
-            }
-            return relevantDocuments.map(\.pageContent).joined(separator: "\n\n")
+            return answers.joined(separator: "\n")
         }
     }
 
@@ -57,66 +53,71 @@ struct QueryWebsiteFunction: ChatGPTFunction {
 
     func call(arguments: Arguments) async throws -> Result {
         do {
-            let embedding = OpenAIEmbedding(
-                configuration: UserPreferenceEmbeddingConfiguration()
-            )
+            let embedding = OpenAIEmbedding(configuration: UserPreferenceEmbeddingConfiguration())
 
-            let queryEmbeddings = try await embedding.embed(query: arguments.query)
-            let searchCount = UserDefaults.shared.value(for: \.chatGPTMaxToken) > 5000 ? 3 : 20
-            
-            let result = try await withThrowingTaskGroup(
-                of: [(document: Document, distance: Float)].self
-            ) { group in
+            let result = try await withThrowingTaskGroup(of: String.self) { group in
                 for urlString in arguments.urls {
                     guard let url = URL(string: urlString) else { continue }
                     group.addTask {
-                        if let database = await TemporaryUSearch.view(identifier: urlString) {
-                            return try await database.searchWithDistance(
-                                embeddings: queryEmbeddings,
-                                count: searchCount
-                            )
-                        }
                         // 1. grab the website content
                         await reportProgress("Loading \(url)..")
-                        print("== load \(url)")
+                        
+                        if let database = await TemporaryUSearch.view(identifier: urlString) {
+                            await reportProgress("Generating answers..")
+                            let qa = RetrievalQAChain(vectorStore: database, embedding: embedding) {
+                                OpenAIChat(
+                                    configuration: UserPreferenceChatGPTConfiguration()
+                                        .overriding(.init(temperature: 0)),
+                                    stream: true
+                                )
+                            }
+                            return try await qa.call(.init(arguments.query)).answer
+                        }
                         let loader = WebLoader(urls: [url])
                         let documents = try await loader.load()
                         await reportProgress("Processing \(url)..")
-                        print("== loaded \(url), documents: \(documents.count)")
                         // 2. split the content
                         let splitter = RecursiveCharacterTextSplitter(
                             chunkSize: 1000,
                             chunkOverlap: 100
                         )
                         let splitDocuments = try await splitter.transformDocuments(documents)
-                        print("== split \(url), documents: \(splitDocuments.count)")
                         // 3. embedding and store in db
                         await reportProgress("Embedding \(url)..")
                         let embeddedDocuments = try await embedding.embed(documents: splitDocuments)
-                        print("== embedded \(url)")
                         let database = TemporaryUSearch(identifier: urlString)
                         try await database.set(embeddedDocuments)
-                        print("== save to database \(url)")
-                        let result = try await database.searchWithDistance(
-                            embeddings: queryEmbeddings,
-                            count: searchCount
-                        )
-                        print("== result of \(url): \(result)")
-                        return result
+                        // 4. generate answer
+                        await reportProgress("Generating answers..")
+                        let qa = RetrievalQAChain(vectorStore: database, embedding: embedding) {
+                            OpenAIChat(
+                                configuration: UserPreferenceChatGPTConfiguration()
+                                    .overriding(.init(temperature: 0)),
+                                stream: true
+                            )
+                        }
+                        let result = try await qa.call(.init(arguments.query))
+                        return result.answer
                     }
                 }
 
-                var all = [(document: Document, distance: Float)]()
+                var all = [String]()
                 for try await result in group {
-                    all.append(contentsOf: result)
+                    all.append(result)
                 }
-                await reportProgress("Finish reading websites.")
+                await reportProgress("""
+                Finish reading websites.
+                \(
+                    arguments.urls
+                        .map { "- [\($0)](\($0))" }
+                        .joined(separator: "\n")
+                )
+                """)
+                                     
                 return all
-                    .sorted { $0.distance < $1.distance }
-                    .prefix(searchCount)
             }
 
-            return .init(relevantDocuments: result.map(\.document))
+            return .init(answers: result)
         } catch {
             await reportProgress("Failed reading websites.")
             throw error

Tool/Sources/LangChain/Chains/RetrievalQA.swift

@@ -6,8 +6,8 @@ public final class RetrievalQAChain: Chain {
     let chatModelFactory: () -> ChatModel
 
     public struct Output {
-        var answer: String
-        var sourceDocuments: [Document]
+        public var answer: String
+        public var sourceDocuments: [Document]
     }
 
     public init(