Reimplement CallbackManager to be more easily extensible

intitni · intitni · commit 7a37f6c49ee2 · 2023-07-10T00:31:10.000+08:00
diff --git a/Core/Sources/ChatPlugins/SearchChatPlugin/SearchQuery.swift b/Core/Sources/ChatPlugins/SearchChatPlugin/SearchQuery.swift
@@ -1,6 +1,7 @@
 import BingSearchService
 import Foundation
 import LangChain
+import OpenAIService
 
 enum SearchEvent {
     case startAction(String)
@@ -42,7 +43,10 @@ func search(_ query: String) async throws
         ),
     ]
 
-    let chatModel = OpenAIChat(temperature: 0, stream: true)
+    let chatModel = OpenAIChat(
+        configuration: UserPreferenceChatGPTConfiguration().overriding { $0.temperature = 0 },
+        stream: true
+    )
 
     let agentExecutor = AgentExecutor(
         agent: ChatAgent(
@@ -55,63 +59,37 @@ func search(_ query: String) async throws
         earlyStopHandleType: .generate
     )
 
-    class ResultCallbackManager: ChainCallbackManager {
+    return (AsyncThrowingStream<SearchEvent, Error> { continuation in
         var accumulation: String = ""
         var isGeneratingFinalAnswer = false
-        var onFinalAnswerToken: (String) -> Void
-        var onAgentActionStart: (String) -> Void
-        var onAgentActionEnd: (String) -> Void
-
-        init(
-            onFinalAnswerToken: @escaping (String) -> Void,
-            onAgentActionStart: @escaping (String) -> Void,
-            onAgentActionEnd: @escaping (String) -> Void
-        ) {
-            self.onFinalAnswerToken = onFinalAnswerToken
-            self.onAgentActionStart = onAgentActionStart
-            self.onAgentActionEnd = onAgentActionEnd
-        }
-
-        func onChainStart<T>(type: T.Type, input: T.Input) where T: LangChain.Chain {}
 
-        func onAgentFinish(output: LangChain.AgentFinish) {}
-
-        func onAgentActionStart(action: LangChain.AgentAction) {
-            onAgentActionStart("\(action.toolName): \(action.toolInput)")
-        }
-
-        func onAgentActionEnd(action: LangChain.AgentAction) {
-            onAgentActionEnd("\(action.toolName): \(action.toolInput)")
-        }
-
-        func onLLMNewToken(token: String) {
-            if isGeneratingFinalAnswer {
-                onFinalAnswerToken(token)
-                return
+        let callbackManager = CallbackManager { manager in
+            manager.on(CallbackEvents.AgentActionDidStart.self) {
+                continuation.yield(.startAction("\($0.toolName): \($0.toolInput)"))
             }
-            accumulation.append(token)
-            if accumulation.hasSuffix("Final Answer: ") {
-                isGeneratingFinalAnswer = true
-                accumulation = ""
+
+            manager.on(CallbackEvents.AgentActionDidEnd.self) {
+                continuation.yield(.endAction("\($0.toolName): \($0.toolInput)"))
             }
-        }
-    }
 
-    return (AsyncThrowingStream<SearchEvent, Error> { continuation in
-        let callback = ResultCallbackManager(
-            onFinalAnswerToken: {
-                continuation.yield(.answerToken($0))
-            },
-            onAgentActionStart: {
-                continuation.yield(.startAction($0))
-            },
-            onAgentActionEnd: {
-                continuation.yield(.endAction($0))
+            manager.on(CallbackEvents.LLMDidProduceNewToken.self) {
+                if isGeneratingFinalAnswer {
+                    continuation.yield(.answerToken($0))
+                    return
+                }
+                accumulation.append($0)
+                if accumulation.hasSuffix("Final Answer: ") {
+                    isGeneratingFinalAnswer = true
+                    accumulation = ""
+                }
             }
-        )
+        }
         Task {
             do {
-                let finalAnswer = try await agentExecutor.run(query, callbackManagers: [callback])
+                let finalAnswer = try await agentExecutor.run(
+                    query,
+                    callbackManagers: [callbackManager]
+                )
                 continuation.yield(.finishAnswer(finalAnswer, linkStorage.links))
                 continuation.finish()
             } catch {
diff --git a/Tool/Sources/LangChain/Agent.swift b/Tool/Sources/LangChain/Agent.swift
@@ -86,7 +86,7 @@ public extension Agent {
     func plan(
         input: Input,
         intermediateSteps: [AgentAction],
-        callbackManagers: [ChainCallbackManager]
+        callbackManagers: [CallbackManager]
     ) async throws -> AgentNextStep {
         let input = getFullInputs(input: input, intermediateSteps: intermediateSteps)
         let output = try await chatModelChain.call(input, callbackManagers: callbackManagers)
@@ -97,7 +97,7 @@ public extension Agent {
         input: Input,
         earlyStoppedHandleType: AgentEarlyStopHandleType,
         intermediateSteps: [AgentAction],
-        callbackManagers: [ChainCallbackManager]
+        callbackManagers: [CallbackManager]
     ) async throws -> AgentFinish {
         switch earlyStoppedHandleType {
         case .force:
diff --git a/Tool/Sources/LangChain/AgentExecutor.swift b/Tool/Sources/LangChain/AgentExecutor.swift
@@ -1,13 +1,5 @@
 import Foundation
 
-public protocol ChainCallbackManager {
-    func onChainStart<T: Chain>(type: T.Type, input: T.Input)
-    func onAgentFinish(output: AgentFinish)
-    func onAgentActionStart(action: AgentAction)
-    func onAgentActionEnd(action: AgentAction)
-    func onLLMNewToken(token: String)
-}
-
 public actor AgentExecutor<InnerAgent: Agent>: Chain where InnerAgent.Input == String {
     public typealias Input = String
     public struct Output {
@@ -39,7 +31,7 @@ public actor AgentExecutor<InnerAgent: Agent>: Chain where InnerAgent.Input == S
 
     public func callLogic(
         _ input: Input,
-        callbackManagers: [ChainCallbackManager]
+        callbackManagers: [CallbackManager]
     ) async throws -> Output {
         try agent.validateTools(tools: Array(tools.values))
 
@@ -89,7 +81,7 @@ public actor AgentExecutor<InnerAgent: Agent>: Chain where InnerAgent.Input == S
             }
             iterations += 1
         }
-        
+
         let output = try await agent.returnStoppedResponse(
             input: input,
             earlyStoppedHandleType: earlyStopHandleType,
@@ -106,7 +98,7 @@ public actor AgentExecutor<InnerAgent: Agent>: Chain where InnerAgent.Input == S
     public nonisolated func parseOutput(_ output: Output) -> String {
         output.finalOutput
     }
-    
+
     public func cancel() {
         isCancelled = true
         earlyStopHandleType = .force
@@ -119,10 +111,10 @@ extension AgentExecutor {
     func end(
         output: AgentFinish,
         intermediateSteps: [AgentAction],
-        callbackManagers: [ChainCallbackManager]
+        callbackManagers: [CallbackManager]
     ) -> Output {
         for callbackManager in callbackManagers {
-            callbackManager.onAgentFinish(output: output)
+            callbackManager.send(CallbackEvents.AgentDidFinish(info: output))
         }
         let finalOutput = output.returnValue
         return .init(finalOutput: finalOutput, intermediateSteps: intermediateSteps)
@@ -131,7 +123,7 @@ extension AgentExecutor {
     func takeNextStep(
         input: Input,
         intermediateSteps: [AgentAction],
-        callbackManagers: [ChainCallbackManager]
+        callbackManagers: [CallbackManager]
     ) async throws -> AgentNextStep {
         let output = try await agent.plan(
             input: input,
@@ -144,7 +136,8 @@ extension AgentExecutor {
             let completedActions = try await withThrowingTaskGroup(of: AgentAction.self) {
                 taskGroup in
                 for action in actions {
-                    callbackManagers.forEach { $0.onAgentActionStart(action: action) }
+                    callbackManagers
+                        .forEach { $0.send(CallbackEvents.AgentActionDidStart(info: action)) }
                     guard let tool = tools[action.toolName] else { throw InvalidToolError() }
                     taskGroup.addTask {
                         let observation = try await tool.run(input: action.toolInput)
@@ -154,7 +147,8 @@ extension AgentExecutor {
                 var completedActions = [AgentAction]()
                 for try await action in taskGroup {
                     completedActions.append(action)
-                    callbackManagers.forEach { $0.onAgentActionEnd(action: action) }
+                    callbackManagers
+                        .forEach { $0.send(CallbackEvents.AgentActionDidEnd(info: action)) }
                 }
                 return completedActions
             }
diff --git a/Tool/Sources/LangChain/Callback.swift b/Tool/Sources/LangChain/Callback.swift
@@ -0,0 +1,56 @@
+import Foundation
+
+public protocol CallbackEvent {
+    associatedtype Info
+    var info: Info { get }
+}
+
+public enum CallbackEvents {}
+
+public struct CallbackManager {
+    private var observers = [Any]()
+
+    public init() {}
+
+    public init(observers: (inout CallbackManager) -> Void) {
+        var manager = CallbackManager()
+        observers(&manager)
+        self = manager
+    }
+
+    public mutating func on<Event: CallbackEvent>(
+        _: Event.Type = Event.self,
+        _ 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)
+        }
+    }
+}
+
+public extension CallbackEvents {
+    struct AgentDidFinish: CallbackEvent {
+        public let info: AgentFinish
+    }
+
+    struct AgentActionDidStart: CallbackEvent {
+        public let info: AgentAction
+    }
+
+    struct AgentActionDidEnd: CallbackEvent {
+        public let info: AgentAction
+    }
+
+    struct LLMDidProduceNewToken: CallbackEvent {
+        public let info: String
+    }
+
+    struct ChainDidStart<T: Chain>: CallbackEvent {
+        public let info: (type: T.Type, input: T.Input)
+    }
+}
+
diff --git a/Tool/Sources/LangChain/Chain.swift b/Tool/Sources/LangChain/Chain.swift
@@ -3,19 +3,20 @@ import Foundation
 public protocol Chain {
     associatedtype Input
     associatedtype Output
-    func callLogic(_ input: Input, callbackManagers: [ChainCallbackManager]) async throws -> Output
+    func callLogic(_ input: Input, callbackManagers: [CallbackManager]) async throws -> Output
     func parseOutput(_ output: Output) -> String
 }
 
 public extension Chain {
-    func run(_ input: Input, callbackManagers: [ChainCallbackManager] = []) async throws -> String {
+    func run(_ input: Input, callbackManagers: [CallbackManager] = []) async throws -> String {
         let output = try await call(input, callbackManagers: callbackManagers)
         return parseOutput(output)
     }
 
-    func call(_ input: Input, callbackManagers: [ChainCallbackManager] = []) async throws -> Output {
+    func call(_ input: Input, callbackManagers: [CallbackManager] = []) async throws -> Output {
         for callbackManager in callbackManagers {
-            callbackManager.onChainStart(type: Self.self, input: input)
+            callbackManager
+                .send(CallbackEvents.ChainDidStart(info: (type: Self.self, input: input)))
         }
         return try await callLogic(input, callbackManagers: callbackManagers)
     }
@@ -35,7 +36,7 @@ public struct SimpleChain<Input, Output>: Chain {
 
     public func callLogic(
         _ input: Input,
-        callbackManagers: [ChainCallbackManager]
+        callbackManagers: [CallbackManager]
     ) async throws -> Output {
         return try await block(input)
     }
@@ -54,7 +55,7 @@ public struct ConnectedChain<A: Chain, B: Chain>: Chain where B.Input == A.Outpu
 
     public func callLogic(
         _ input: Input,
-        callbackManagers: [ChainCallbackManager] = []
+        callbackManagers: [CallbackManager] = []
     ) async throws -> Output {
         let a = try await chainA.call(input, callbackManagers: callbackManagers)
         let b = try await chainB.call(a, callbackManagers: callbackManagers)
@@ -75,7 +76,7 @@ public struct PairedChain<A: Chain, B: Chain>: Chain {
 
     public func callLogic(
         _ input: Input,
-        callbackManagers: [ChainCallbackManager] = []
+        callbackManagers: [CallbackManager] = []
     ) async throws -> Output {
         async let a = chainA.call(input.0, callbackManagers: callbackManagers)
         async let b = chainB.call(input.1, callbackManagers: callbackManagers)
@@ -96,7 +97,7 @@ public struct MappedChain<A: Chain, NewOutput>: Chain {
 
     public func callLogic(
         _ input: Input,
-        callbackManagers: [ChainCallbackManager]
+        callbackManagers: [CallbackManager]
     ) async throws -> Output {
         let output = try await chain.call(input, callbackManagers: callbackManagers)
         return map(output)
diff --git a/Tool/Sources/LangChain/Chains/LLMChain.swift b/Tool/Sources/LangChain/Chains/LLMChain.swift
@@ -19,7 +19,7 @@ public class ChatModelChain<Input>: Chain {
 
     public func callLogic(
         _ input: Input,
-        callbackManagers: [ChainCallbackManager]
+        callbackManagers: [CallbackManager]
     ) async throws -> Output {
         let prompt = promptTemplate(input)
         let output = try await chatModel.generate(
diff --git a/Tool/Sources/LangChain/ChatModel/ChatModel.swift b/Tool/Sources/LangChain/ChatModel/ChatModel.swift
@@ -4,7 +4,7 @@ public protocol ChatModel {
     func generate(
         prompt: [ChatMessage],
         stops: [String],
-        callbackManagers: [ChainCallbackManager]
+        callbackManagers: [CallbackManager]
     ) async throws -> String
 }
 
diff --git a/Tool/Sources/LangChain/ChatModel/OpenAIChat.swift b/Tool/Sources/LangChain/ChatModel/OpenAIChat.swift
@@ -2,26 +2,22 @@ import Foundation
 import OpenAIService
 
 public struct OpenAIChat: ChatModel {
-    public var temperature: Double
+    public var configuration: ChatGPTConfiguration
     public var stream: Bool
 
     public init(
-        temperature: Double = 0.7,
-        stream: Bool = false
+        configuration: ChatGPTConfiguration,
+        stream: Bool
     ) {
-        self.temperature = temperature
+        self.configuration = configuration
         self.stream = stream
     }
 
     public func generate(
         prompt: [ChatMessage],
         stops: [String],
-        callbackManagers: [ChainCallbackManager]
+        callbackManagers: [CallbackManager]
     ) async throws -> String {
-        let configuration = UserPreferenceChatGPTConfiguration().overriding(.init(
-            temperature: temperature,
-            stop: stops
-        ))
         let memory = AutoManagedChatGPTMemory(
             systemPrompt: "",
             configuration: configuration,
@@ -47,7 +43,8 @@ public struct OpenAIChat: ChatModel {
             var message = ""
             for try await trunk in stream {
                 message.append(trunk)
-                callbackManagers.forEach { $0.onLLMNewToken(token: trunk) }
+                callbackManagers
+                    .forEach { $0.send(CallbackEvents.LLMDidProduceNewToken(info: trunk)) }
             }
             return message
         } else {

Original file line number	Diff line number	Diff line change
`@@ -4,7 +4,7 @@ public protocol ChatModel {`
`4`	`4`	`func generate(`
`5`	`5`	`prompt: [ChatMessage],`
`6`	`6`	`stops: [String],`
`7`		`- callbackManagers: [ChainCallbackManager]`
	`7`	`+ callbackManagers: [CallbackManager]`
`8`	`8`	`) async throws -> String`
`9`	`9`	`}`
`10`	`10`