Update QAInformationRetrievalChain to support multiple vector store

Author: intitni

Tool/Sources/LangChain/Chains/RetrievalQA.swift

@@ -2,8 +2,9 @@ import Foundation
 import OpenAIService
 
 public final class QAInformationRetrievalChain: Chain {
-    let vectorStore: VectorStore
+    let vectorStores: [VectorStore]
     let embedding: Embeddings
+    let maxCount: Int
 
     public struct Output {
         public var information: String
@@ -12,24 +13,51 @@ public final class QAInformationRetrievalChain: Chain {
 
     public init(
         vectorStore: VectorStore,
-        embedding: Embeddings
+        embedding: Embeddings,
+        maxCount: Int = 5
     ) {
-        self.vectorStore = vectorStore
+        vectorStores = [vectorStore]
         self.embedding = embedding
+        self.maxCount = maxCount
+    }
+
+    public init(
+        vectorStores: [VectorStore],
+        embedding: Embeddings,
+        maxCount: Int = 5
+    ) {
+        self.vectorStores = vectorStores
+        self.embedding = embedding
+        self.maxCount = maxCount
     }
 
     public func callLogic(
         _ input: String,
         callbackManagers: [CallbackManager]
     ) async throws -> Output {
         let embeddedQuestion = try await embedding.embed(query: input)
-        let documents = try await vectorStore.searchWithDistance(
-            embeddings: embeddedQuestion,
-            count: 5
-        ).filter { item in
-            item.distance < 0.31
-        }
+        let documentsSlice = await withTaskGroup(
+            of: [(document: Document, distance: Float)].self
+        ) { group in
+            for vectorStore in vectorStores {
+                group.addTask {
+                    (try? await vectorStore.searchWithDistance(
+                        embeddings: embeddedQuestion,
+                        count: 5
+                    ).filter { item in
+                        item.distance < 0.31
+                    }) ?? []
+                }
+            }
+            var result = [(document: Document, distance: Float)]()
+            for await items in group {
+                result.append(contentsOf: items)
+            }
+            return result
+        }.sorted { $0.distance < $1.distance }.prefix(maxCount)
 
+        let documents = Array(documentsSlice)
+        
         callbackManagers.send(CallbackEvents.RetrievalQADidExtractRelevantContent(info: documents))
 
         let relevantInformationChain = RelevantInformationExtractionChain()