Add cancellation support to agent executor

Author: intitni

Core/Sources/ChatPlugins/SearchChatPlugin/SearchChatPlugin.swift

@@ -28,14 +28,17 @@ public actor SearchChatPlugin: ChatPlugin {
         }
 
         do {
-            let eventStream = try await search(content)
+            let (eventStream, cancelAgent) = try await search(content)
 
             var actions = [String]()
             var finishedActions = Set<String>()
             var message = ""
 
             for try await event in eventStream {
-                guard !isCancelled else { return }
+                guard !isCancelled else {
+                    await cancelAgent()
+                    break
+                }
                 switch event {
                 case let .startAction(content):
                     actions.append(content)

Core/Sources/ChatPlugins/SearchChatPlugin/SearchQuery.swift

@@ -9,7 +9,9 @@ enum SearchEvent {
     case finishAnswer(String, [(title: String, link: String)])
 }
 
-func search(_ query: String) async throws -> AsyncThrowingStream<SearchEvent, Error> {
+func search(_ query: String) async throws
+    -> (stream: AsyncThrowingStream<SearchEvent, Error>, cancel: () async -> Void)
+{
     let bingSearch = BingSearchService(
         subscriptionKey: UserDefaults.shared.value(for: \.bingSearchSubscriptionKey),
         searchURL: UserDefaults.shared.value(for: \.bingSearchEndpoint)
@@ -91,7 +93,7 @@ func search(_ query: String) async throws -> AsyncThrowingStream<SearchEvent, Er
         }
     }
 
-    return AsyncThrowingStream<SearchEvent, Error> { continuation in
+    return (AsyncThrowingStream<SearchEvent, Error> { continuation in
         let callback = ResultCallbackManager(
             onFinalAnswerToken: {
                 continuation.yield(.answerToken($0))
@@ -112,6 +114,8 @@ func search(_ query: String) async throws -> AsyncThrowingStream<SearchEvent, Er
                 continuation.finish(throwing: error)
             }
         }
-    }
+    }, {
+        await agentExecutor.cancel()
+    })
 }
 

Tool/Sources/LangChain/AgentExecutor.swift

@@ -19,8 +19,9 @@ public actor AgentExecutor<InnerAgent: Agent>: Chain where InnerAgent.Input == S
     let tools: [String: AgentTool]
     let maxIteration: Int?
     let maxExecutionTime: Double?
-    let earlyStopHandleType: AgentEarlyStopHandleType
+    var earlyStopHandleType: AgentEarlyStopHandleType
     var now: () -> Date = { Date() }
+    var isCancelled = false
 
     public init(
         agent: InnerAgent,
@@ -47,6 +48,7 @@ public actor AgentExecutor<InnerAgent: Agent>: Chain where InnerAgent.Input == S
         var intermediateSteps: [AgentAction] = []
 
         func shouldContinue() -> Bool {
+            if isCancelled { return false }
             if let maxIteration = maxIteration, iterations >= maxIteration {
                 return false
             }
@@ -87,7 +89,7 @@ public actor AgentExecutor<InnerAgent: Agent>: Chain where InnerAgent.Input == S
             }
             iterations += 1
         }
-
+        
         let output = try await agent.returnStoppedResponse(
             input: input,
             earlyStoppedHandleType: earlyStopHandleType,
@@ -104,6 +106,11 @@ 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
+    }
 }
 
 struct InvalidToolError: Error {}