Merge branch 'feature/retrieval-qa-to-get-information-only' into develop

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

@@ -4,12 +4,14 @@ import LangChain
 import OpenAIService
 import PlaygroundSupport
 import SwiftUI
+import TokenEncoder
 
 struct QAForm: View {
-    @State var intermediateAnswers = [RefineDocumentChain.IntermediateAnswer]()
+    @State var relevantInformation = [String]()
     @State var relevantDocuments = [(document: Document, distance: Float)]()
     @State var duration: TimeInterval = 0
     @State var answer: String = ""
+    @State var tokenCount: Int = 0
     @State var question: String = "What is Swift macros?"
     @State var isProcessing: Bool = false
     @State var url: String = "https://developer.apple.com/documentation/swift/applying-macros"
@@ -36,23 +38,14 @@ struct QAForm: View {
                             Text("\(duration) seconds")
                         }
                     }
-                    Section(header: Text("Answer")) {
+                    Section(header: Text("All Relevant Information (\(tokenCount) words)")) {
                         Text(answer)
                     }
-                    Section(header: Text("Intermediate Answers")) {
-                        ForEach(0..<intermediateAnswers.endIndex, id: \.self) { index in
-                            let answer = intermediateAnswers[index]
+                    Section(header: Text("Relevant Information")) {
+                        ForEach(0..<relevantInformation.endIndex, id: \.self) { index in
+                            let information = relevantInformation[index]
                             VStack(alignment: .leading) {
-                                Text(answer.answer)
-                                VStack(alignment: .leading) {
-                                    Text("Usefulness: \(answer.usefulness)")
-                                    Text("Needs more context: \(answer.more ? "Yes" : "No")")
-                                }
-                                .padding()
-                                .background {
-                                    RoundedRectangle(cornerRadius: 8)
-                                        .fill(Color(NSColor.textBackgroundColor))
-                                }
+                                Text(information)
                                 Divider()
                             }
                             .textSelection(.enabled)
@@ -84,8 +77,9 @@ struct QAForm: View {
         let start = Date().timeIntervalSince1970
         answer = ""
         relevantDocuments = []
-        intermediateAnswers = []
+        relevantInformation = []
         duration = 0
+        tokenCount = 0
         isProcessing = true
         defer { isProcessing = false }
         guard let url = URL(string: url) else {
@@ -112,23 +106,24 @@ struct QAForm: View {
             }
         }()
 
-        let qa = RetrievalQAChain(
+        let qa = QAInformationRetrievalChain(
             vectorStore: store,
             embedding: embedding
         )
         answer = try await qa.run(
             question,
             callbackManagers: [
                 .init {
-                    $0.on(CallbackEvents.RetrievalQADidGenerateIntermediateAnswer.self) {
-                        intermediateAnswers.append($0)
+                    $0.on(\.relevantInformationExtractionChainDidExtractPartialRelevantContent) {
+                        relevantInformation.append($0)
                     }
-                    $0.on(CallbackEvents.RetrievalQADidExtractRelevantContent.self) {
+                    $0.on(\.retrievalQADidExtractRelevantContent) {
                         relevantDocuments = $0
                     }
                 },
             ]
         )
+        tokenCount = answer.split(separator: " ").count
         duration = Date().timeIntervalSince1970 - start
     }
 }

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

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<Timeline
+   version = "3.0">
+   <TimelineItems>
+   </TimelineItems>
+</Timeline>

Tool/Sources/LangChain/Agent.swift

@@ -24,14 +24,26 @@ public extension CallbackEvents {
     struct AgentDidFinish: CallbackEvent {
         public let info: AgentFinish
     }
+    
+    var agentDidFinish: AgentDidFinish.Type {
+        AgentDidFinish.self
+    }
 
     struct AgentActionDidStart: CallbackEvent {
         public let info: AgentAction
     }
+    
+    var agentActionDidStart: AgentActionDidStart.Type {
+        AgentActionDidStart.self
+    }
 
     struct AgentActionDidEnd: CallbackEvent {
         public let info: AgentAction
     }
+    
+    var agentActionDidEnd: AgentActionDidEnd.Type {
+        AgentActionDidEnd.self
+    }
 }
 
 public struct AgentFinish: Equatable {

Tool/Sources/LangChain/Callback.swift

@@ -5,7 +5,9 @@ public protocol CallbackEvent {
     var info: Info { get }
 }
 
-public enum CallbackEvents {}
+public struct CallbackEvents {
+    private init() {}
+}
 
 public struct CallbackManager {
     fileprivate var observers = [Any]()
@@ -25,19 +27,39 @@ public struct CallbackManager {
         observers.append(handler)
     }
 
+    public mutating func on<Event: CallbackEvent>(
+        _: KeyPath<CallbackEvents, Event.Type>,
+        _ handler: @escaping (Event.Info) -> Void
+    ) {
+        observers.append(handler)
+    }
+
     public func send<Event: CallbackEvent>(_ event: Event) {
         for case let observer as ((Event.Info) -> Void) in observers {
             observer(event.info)
         }
     }
+
+    func send<Event: CallbackEvent>(
+        _: KeyPath<CallbackEvents, Event.Type>,
+        _ info: Event.Info
+    ) {
+        for case let observer as ((Event.Info) -> Void) in observers {
+            observer(info)
+        }
+    }
 }
 
 public extension [CallbackManager] {
     func send<Event: CallbackEvent>(_ event: Event) {
-        for cb in self {
-            for case let observer as ((Event.Info) -> Void) in cb.observers {
-                observer(event.info)
-            }
-        }
+        for cb in self { cb.send(event) }
+    }
+
+    func send<Event: CallbackEvent>(
+        _ keyPath: KeyPath<CallbackEvents, Event.Type>,
+        _ info: Event.Info
+    ) {
+        for cb in self { cb.send(keyPath, info) }
     }
 }
+

Tool/Sources/LangChain/Chain.swift

@@ -10,7 +10,7 @@ public protocol Chain {
 public extension Chain {
     typealias ChainDidStart = CallbackEvents.ChainDidStart<Self>
     typealias ChainDidEnd = CallbackEvents.ChainDidEnd<Self>
-    
+
     func run(_ input: Input, callbackManagers: [CallbackManager] = []) async throws -> String {
         let output = try await call(input, callbackManagers: callbackManagers)
         return parseOutput(output)