Merge branch 'feature/retrieval-qa-improvement' into develop

Author: intitni

Core/Sources/ChatContextCollectors/WebChatContextCollector/QueryWebsiteFunction.swift

@@ -64,13 +64,7 @@ struct QueryWebsiteFunction: ChatGPTFunction {
 
                         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
-                                )
-                            }
+                            let qa = RetrievalQAChain(vectorStore: database, embedding: embedding)
                             return try await qa.call(.init(arguments.query)).answer
                         }
                         let loader = WebLoader(urls: [url])
@@ -89,13 +83,7 @@ struct QueryWebsiteFunction: ChatGPTFunction {
                         try await database.set(embeddedDocuments)
                         // 4. generate answer
                         await reportProgress("Generating answers..")
-                        let qa = RetrievalQAChain(vectorStore: database, embedding: embedding) {
-                            OpenAIChat(
-                                configuration: UserPreferenceChatGPTConfiguration()
-                                    .overriding(.init(temperature: 0)),
-                                stream: true
-                            )
-                        }
+                        let qa = RetrievalQAChain(vectorStore: database, embedding: embedding)
                         let result = try await qa.call(.init(arguments.query))
                         return result.answer
                     }

Core/Sources/ChatService/ChatFunctionProvider.swift

@@ -15,5 +15,9 @@ final class ChatFunctionProvider {
     }
 }
 
-extension ChatFunctionProvider: ChatGPTFunctionProvider {}
+extension ChatFunctionProvider: ChatGPTFunctionProvider {
+    var functionCallStrategy: OpenAIService.FunctionCallStrategy? {
+        nil
+    }
+}
 

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

@@ -1,55 +1,97 @@
 import AppKit
+import Foundation
 import LangChain
 import OpenAIService
 import PlaygroundSupport
 import SwiftUI
 
 struct QAForm: View {
-    @State var intermediateAnswers = [String]()
+    @State var intermediateAnswers = [RefineDocumentChain.IntermediateAnswer]()
+    @State var relevantDocuments = [(document: Document, distance: Float)]()
+    @State var duration: TimeInterval = 0
     @State var answer: String = ""
     @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"
 
     var body: some View {
-        Form {
-            Section(header: Text("Input")) {
-                TextField("URL", text: $url)
-                TextField("Question", text: $question)
-                Button("Ask") {
-                    Task {
-                        do {
-                            try await ask()
-                        } catch {
-                            answer = error.localizedDescription
+        HStack(spacing: 0) {
+            ScrollView {
+                Form {
+                    Section(header: Text("Input")) {
+                        TextField("URL", text: $url)
+                        TextField("Question", text: $question)
+                        HStack {
+                            Button("Ask") {
+                                Task {
+                                    do {
+                                        try await ask()
+                                    } catch {
+                                        answer = error.localizedDescription
+                                    }
+                                }
+                            }
+                            .disabled(isProcessing)
+                            
+                            Text("\(duration) seconds")
+                        }
+                    }
+                    Section(header: Text("Answer")) {
+                        Text(answer)
+                    }
+                    Section(header: Text("Intermediate Answers")) {
+                        ForEach(0..<intermediateAnswers.endIndex, id: \.self) { index in
+                            let answer = intermediateAnswers[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))
+                                }
+                                Divider()
+                            }
+                            .textSelection(.enabled)
                         }
                     }
                 }
-                .disabled(isProcessing)
+                .formStyle(.grouped)
             }
-            Section(header: Text("Answer")) {
-                Text(answer)
-            }
-            Section(header: Text("Intermediate Answers")) {
-                ForEach(intermediateAnswers, id: \.self) { answer in
-                    Text(answer)
-                    Divider()
-                }
+
+            ScrollView {
+                Form {
+                    Section(header: Text("Relevant Documents")) {
+                        ForEach(0..<relevantDocuments.endIndex, id: \.self) { index in
+                            let document = relevantDocuments[index]
+                            VStack(alignment: .leading) {
+                                Text("\(document.distance)")
+                                Text(document.document.pageContent)
+                                Divider()
+                            }
+                            .textSelection(.enabled)
+                        }
+                    }
+                }.formStyle(.grouped)
             }
         }
-        .formStyle(.grouped)
     }
 
     func ask() async throws {
+        let start = Date().timeIntervalSince1970
+        answer = ""
+        relevantDocuments = []
         intermediateAnswers = []
+        duration = 0
         isProcessing = true
         defer { isProcessing = false }
         guard let url = URL(string: url) else {
             answer = "Invalid URL"
             return
         }
-        let chatGPTConfiguration = UserPreferenceChatGPTConfiguration()
-            .overriding { $0.temperature = 0 }
         let embeddingConfiguration = UserPreferenceEmbeddingConfiguration().overriding()
         let embedding = OpenAIEmbedding(configuration: embeddingConfiguration)
         let store: VectorStore = try await {
@@ -72,8 +114,7 @@ struct QAForm: View {
 
         let qa = RetrievalQAChain(
             vectorStore: store,
-            embedding: embedding,
-            chatModelFactory: { OpenAIChat(configuration: chatGPTConfiguration, stream: false) }
+            embedding: embedding
         )
         answer = try await qa.run(
             question,
@@ -82,15 +123,19 @@ struct QAForm: View {
                     $0.on(CallbackEvents.RetrievalQADidGenerateIntermediateAnswer.self) {
                         intermediateAnswers.append($0)
                     }
+                    $0.on(CallbackEvents.RetrievalQADidExtractRelevantContent.self) {
+                        relevantDocuments = $0
+                    }
                 },
             ]
         )
+        duration = Date().timeIntervalSince1970 - start
     }
 }
 
 let hostingView = NSHostingController(
     rootView: QAForm()
-        .frame(width: 600, height: 800)
+        .frame(width: 800, height: 800)
 )
 
 PlaygroundPage.current.needsIndefiniteExecution = true

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

@@ -104,7 +104,7 @@ public extension Agent {
     ) async throws -> AgentNextStep {
         let input = getFullInputs(input: input, intermediateSteps: intermediateSteps)
         let output = try await chatModelChain.call(input, callbackManagers: callbackManagers)
-        return parseOutput(output)
+        return parseOutput(output.content ?? "")
     }
 
     func returnStoppedResponse(
@@ -128,12 +128,13 @@ public extension Agent {
             """
             let input = AgentInput(input: input, thoughts: .text(thoughts))
             let output = try await chatModelChain.call(input, callbackManagers: callbackManagers)
-            let nextAction = parseOutput(output)
+            let reply = output.content ?? ""
+            let nextAction = parseOutput(reply)
             switch nextAction {
             case let .finish(finish):
                 return finish
             case .actions:
-                return AgentFinish(returnValue: output, log: output)
+                return AgentFinish(returnValue: reply, log: reply)
             }
         }
     }

Tool/Sources/LangChain/Agent.swift

@@ -1,7 +1,7 @@
 import Foundation
 
 public class ChatModelChain<Input>: Chain {
-    public typealias Output = String
+    public typealias Output = ChatMessage
 
     var chatModel: ChatModel
     var promptTemplate: (Input) -> [ChatMessage]
@@ -31,7 +31,13 @@ public class ChatModelChain<Input>: Chain {
     }
 
     public func parseOutput(_ output: Output) -> String {
-        output
+        if let content = output.content {
+            return content
+        } else if let functionCall = output.functionCall {
+            return "\(functionCall.name): \(functionCall.arguments)"
+        }
+        
+        return ""
     }
 }