Move CallbackEvents to the right places

Author: intitni

Tool/Sources/LangChain/Agent.swift

@@ -20,6 +20,20 @@ public struct AgentAction: Equatable {
     }
 }
 
+public extension CallbackEvents {
+    struct AgentDidFinish: CallbackEvent {
+        public let info: AgentFinish
+    }
+
+    struct AgentActionDidStart: CallbackEvent {
+        public let info: AgentAction
+    }
+
+    struct AgentActionDidEnd: CallbackEvent {
+        public let info: AgentAction
+    }
+}
+
 public struct AgentFinish: Equatable {
     public var returnValue: String
     public var log: String
@@ -38,12 +52,12 @@ public enum AgentNextStep: Equatable {
 public enum AgentScratchPad: Equatable {
     case text(String)
     case messages([String])
-    
+
     var isEmpty: Bool {
         switch self {
-        case .text(let text):
+        case let .text(text):
             return text.isEmpty
-        case .messages(let messages):
+        case let .messages(messages):
             return messages.isEmpty
         }
     }
@@ -108,7 +122,7 @@ public extension Agent {
         case .generate:
             var thoughts = constructBaseScratchpad(intermediateSteps: intermediateSteps)
             thoughts += """
-            
+
             \(llmPrefix)I now need to return a final answer based on the previous steps:
             (Please continue with `Final Answer:`)
             """

Tool/Sources/LangChain/Callback.swift

@@ -31,26 +31,3 @@ public struct CallbackManager {
         }
     }
 }
-
-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)
-    }
-}
-

Tool/Sources/LangChain/Chain.swift

@@ -8,6 +8,9 @@ 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)
@@ -18,10 +21,26 @@ public extension Chain {
             callbackManager
                 .send(CallbackEvents.ChainDidStart(info: (type: Self.self, input: input)))
         }
+        defer {
+            for callbackManager in callbackManagers {
+                callbackManager
+                    .send(CallbackEvents.ChainDidEnd(info: (type: Self.self, input: input)))
+            }
+        }
         return try await callLogic(input, callbackManagers: callbackManagers)
     }
 }
 
+public extension CallbackEvents {
+    struct ChainDidStart<T: Chain>: CallbackEvent {
+        public let info: (type: T.Type, input: T.Input)
+    }
+
+    struct ChainDidEnd<T: Chain>: CallbackEvent {
+        public let info: (type: T.Type, input: T.Input)
+    }
+}
+
 public struct SimpleChain<Input, Output>: Chain {
     let block: (Input) async throws -> Output
     let parseOutputBlock: (Output) -> String

Tool/Sources/LangChain/ChatModel/ChatModel.swift

@@ -23,3 +23,9 @@ public struct ChatMessage {
         self.content = content
     }
 }
+
+public extension CallbackEvents {
+    struct LLMDidProduceNewToken: CallbackEvent {
+        public let info: String
+    }
+}