Adjust agent

Author: intitni

Tool/Sources/LangChain/Agent.swift

@@ -64,6 +64,7 @@ public struct AgentFinish<Output: AgentOutputParsable> {
 public enum AgentNextStep<Output: AgentOutputParsable> {
     case actions([AgentAction])
     case finish(AgentFinish<Output>)
+    case thought(String)
 }
 
 public struct AgentScratchPad<Content: Equatable>: Equatable {
@@ -145,6 +146,8 @@ public extension Agent {
                 return finish
             case .actions:
                 return .init(returnValue: .unstructured(output.content ?? ""), log: output.content ?? "")
+            case let .thought(content):
+                return .init(returnValue: .unstructured(content), log: content)
             }
         }
     }

Tool/Sources/LangChain/AgentExecutor.swift

@@ -57,12 +57,14 @@ public actor AgentExecutor<InnerAgent: Agent>: Chain
         }
 
         while shouldContinue() {
+            try Task.checkCancellation()
             let nextStepOutput = try await takeNextStep(
                 input: input,
                 intermediateSteps: intermediateSteps,
                 callbackManagers: callbackManagers
             )
 
+            try Task.checkCancellation()
             switch nextStepOutput {
             case let .finish(finish):
                 return end(
@@ -82,6 +84,8 @@ public actor AgentExecutor<InnerAgent: Agent>: Chain
                         callbackManagers: callbackManagers
                     )
                 }
+            case .thought:
+                break
             }
             iterations += 1
         }
@@ -156,6 +160,7 @@ extension AgentExecutor {
                 }
                 var completedActions = [AgentAction]()
                 for try await action in taskGroup {
+                    try Task.checkCancellation()
                     completedActions.append(action)
                     callbackManagers
                         .forEach { $0.send(CallbackEvents.AgentActionDidEnd(info: action)) }
@@ -164,6 +169,15 @@ extension AgentExecutor {
             }
 
             return .actions(completedActions)
+        case let .thought(content):
+            return .actions([
+                .init(
+                    toolName: "Thought",
+                    toolInput: content,
+                    log: "Thought: \(content)",
+                    observation: nil
+                ),
+            ])
         }
     }
 
@@ -212,3 +226,4 @@ extension Double: AgentOutputParsable {
 
     public var botReadableContent: String { String(self) }
 }
+

Tool/Sources/LangChain/Agents/FunctionCallingChatAgent.swift

@@ -4,7 +4,7 @@ import OpenAIService
 
 public class FunctionCallingChatAgent<Output: AgentOutputParsable & Decodable>: Agent {
     public typealias ScratchPadContent = [ChatMessage]
-    
+
     public struct EndFunction: ChatGPTFunction {
         public typealias Argument = Output
         public typealias Result = String
@@ -79,7 +79,8 @@ public class FunctionCallingChatAgent<Output: AgentOutputParsable & Decodable>:
     public init(
         configuration: ChatGPTConfiguration = UserPreferenceChatGPTConfiguration(),
         tools: [AgentTool] = [],
-        endFunction: EndFunction
+        endFunction: EndFunction,
+        extraSystemPrompt: String = ""
     ) {
         let functions = tools.compactMap { $0 as? (any ChatGPTFunction) }
         let otherTools = tools.filter { !($0 is (any ChatGPTFunction)) }
@@ -103,13 +104,13 @@ public class FunctionCallingChatAgent<Output: AgentOutputParsable & Decodable>:
                     .init(
                         role: .system,
                         content: """
-                        Respond to the human as helpfully and accurately as possible. \
-                        Save the final answer when it's ready
-
-                        Begin!
+                        Gather information using functions, and generate a final answer to my query as helpfully and accurately as possible.
+                        You don't ask me for additional information.
+                        \(extraSystemPrompt)
+                        When you have the final answer, you MUST call `\(endFunction.name)` to save it.
                         """
                     ),
-                    .init(role: .user, content: agentInput.input)
+                    .init(role: .user, content: agentInput.input),
                 ] + agentInput.thoughts.content
             }
         )
@@ -118,49 +119,57 @@ public class FunctionCallingChatAgent<Output: AgentOutputParsable & Decodable>:
     public func extraPlan(input: AgentInput<String, ScratchPadContent>) {
         // no extra plan
     }
-
-    public func prepareForEarlyStopWithGenerate() -> String {
-        functionProvider.shouldFinish = true
-        return "(call sendFinalAnswer to finish)"
-    }
-
-    public func constructScratchpad(
+    
+    func constructBaseScratchpad(
         intermediateSteps: [AgentAction]
-    ) -> AgentScratchPad<ScratchPadContent> {
-        let baseScratchpad = intermediateSteps.flatMap {
-            [
+    ) -> ScratchPadContent {
+        return intermediateSteps.flatMap {
+            if let observation = $0.observation {
+                return [
+                    ChatMessage(
+                        role: .assistant,
+                        content: nil,
+                        functionCall: .init(name: $0.toolName, arguments: $0.toolInput)
+                    ),
+                    ChatMessage(role: .function, content: observation, name: $0.toolName),
+                ]
+            }
+            return [
+                ChatMessage(role: .assistant, content: $0.toolInput),
                 ChatMessage(
-                    role: .assistant,
-                    content: nil,
-                    functionCall: .init(name: $0.toolName, arguments: $0.toolInput)
+                    role: .user,
+                    content: "Please continue, call \(functionProvider.endFunction.name) when you are done."
                 ),
-                ChatMessage(role: .function, content: $0.observation),
             ]
         }
-        return .init(content: baseScratchpad)
     }
-    
-    public func constructFinalScratchpad(intermediateSteps: [AgentAction]) -> AgentScratchPad<ScratchPadContent> {
-        let baseScratchpad = intermediateSteps.flatMap {
-            [
-                ChatMessage(
-                    role: .assistant,
-                    content: nil,
-                    functionCall: .init(name: $0.toolName, arguments: $0.toolInput)
-                ),
-                ChatMessage(role: .function, content: $0.observation),
-            ]
-        }
+
+    public func constructScratchpad(
+        intermediateSteps: [AgentAction]
+    ) -> AgentScratchPad<ScratchPadContent> {
+        functionProvider.shouldFinish = false
+        let baseScratchpad = constructBaseScratchpad(intermediateSteps: intermediateSteps)
         return .init(content: baseScratchpad)
     }
 
+    public func constructFinalScratchpad(
+        intermediateSteps: [AgentAction]
+    ) -> AgentScratchPad<ScratchPadContent> {
+        functionProvider.shouldFinish = true
+        let baseScratchpad = constructBaseScratchpad(intermediateSteps: intermediateSteps)
+        return .init(content: baseScratchpad + [
+            ChatMessage(role: .assistant, content: "Now I need to save the final answer"),
+            ChatMessage(role: .user, content: "Please continue"),
+        ])
+    }
+
     public func validateTools(tools: [AgentTool]) throws {
         // no validation
     }
 
     public func parseOutput(_ message: ChatMessage) async -> AgentNextStep<Output> {
-        if message.role == .assistant, let functionCall = message.functionCall {
-            if let function = functionProvider.functionTools.first(where: {
+        if let functionCall = message.functionCall {
+            if let function = functionProvider.functions.first(where: {
                 $0.name == functionCall.name
             }) {
                 if function.name == functionProvider.endFunction.name {
@@ -204,6 +213,8 @@ public class FunctionCallingChatAgent<Output: AgentOutputParsable & Decodable>:
             case let .unstructured(x), let .structured(x):
                 return .finish(.init(returnValue: .unstructured(x), log: finish.log))
             }
+        case let .thought(content):
+            return .finish(.init(returnValue: .unstructured(content), log: content))
         }
     }
 }