Make AgentScratchPad generic

intitni · intitni · commit 59e9a042bfb6 · 2023-09-18T14:05:05.000+08:00
diff --git a/Tool/Sources/LangChain/Agent.swift b/Tool/Sources/LangChain/Agent.swift
@@ -66,25 +66,15 @@ public enum AgentNextStep<Output: AgentOutputParsable> {
     case finish(AgentFinish<Output>)
 }
 
-public enum AgentScratchPad: Equatable {
-    case text(String)
-    case messages([String])
-
-    var isEmpty: Bool {
-        switch self {
-        case let .text(text):
-            return text.isEmpty
-        case let .messages(messages):
-            return messages.isEmpty
-        }
-    }
+public struct AgentScratchPad<Content: Equatable>: Equatable {
+    var content: Content
 }
 
-public struct AgentInput<T> {
+public struct AgentInput<T, ScratchPadContent: Equatable> {
     var input: T
-    var thoughts: AgentScratchPad
+    var thoughts: AgentScratchPad<ScratchPadContent>
 
-    public init(input: T, thoughts: AgentScratchPad) {
+    public init(input: T, thoughts: AgentScratchPad<ScratchPadContent>) {
         self.input = input
         self.thoughts = thoughts
     }
@@ -100,20 +90,24 @@ public enum AgentEarlyStopHandleType: Equatable {
 public protocol Agent {
     associatedtype Input
     associatedtype Output: AgentOutputParsable
-    var chatModelChain: ChatModelChain<AgentInput<Input>> { get }
+    associatedtype ScratchPadContent: Equatable
+    var chatModelChain: ChatModelChain<AgentInput<Input, ScratchPadContent>> { get }
     var observationPrefix: String { get }
     var llmPrefix: String { get }
 
     func validateTools(tools: [AgentTool]) throws
-    func constructScratchpad(intermediateSteps: [AgentAction]) -> AgentScratchPad
-    func extraPlan(input: AgentInput<Input>)
-    func prepareForEarlyStopWithGenerate() -> String
-    func parseOutput(_ output: ChatModelChain<AgentInput<Input>>.Output) async
+    func constructScratchpad(intermediateSteps: [AgentAction]) -> AgentScratchPad<ScratchPadContent>
+    func constructFinalScratchpad(intermediateSteps: [AgentAction]) -> AgentScratchPad<ScratchPadContent>
+    func extraPlan(input: AgentInput<Input, ScratchPadContent>)
+    func parseOutput(_ output: ChatModelChain<AgentInput<Input, ScratchPadContent>>.Output) async
         -> AgentNextStep<Output>
 }
 
 public extension Agent {
-    func getFullInputs(input: Input, intermediateSteps: [AgentAction]) -> AgentInput<Input> {
+    func getFullInputs(
+        input: Input,
+        intermediateSteps: [AgentAction]
+    ) -> AgentInput<Input, ScratchPadContent> {
         let thoughts = constructScratchpad(intermediateSteps: intermediateSteps)
         return AgentInput(input: input, thoughts: thoughts)
     }
@@ -142,14 +136,8 @@ public extension Agent {
                 log: ""
             )
         case .generate:
-            var thoughts = constructBaseScratchpad(intermediateSteps: intermediateSteps)
-            thoughts += """
-
-            \(llmPrefix)I now need to return a final answer based on the previous steps:
-            \(prepareForEarlyStopWithGenerate())
-            """
-            let input = AgentInput(input: input, thoughts: .text(thoughts))
-            
+            let thoughts = constructFinalScratchpad(intermediateSteps: intermediateSteps)
+            let input = AgentInput(input: input, thoughts: thoughts)
             let output = try await chatModelChain.call(input, callbackManagers: callbackManagers)
             let nextAction = await parseOutput(output)
             switch nextAction {
@@ -160,16 +148,4 @@ public extension Agent {
             }
         }
     }
-
-    func constructBaseScratchpad(intermediateSteps: [AgentAction]) -> String {
-        var thoughts = ""
-        for step in intermediateSteps {
-            thoughts += """
-            \(step.log)
-            \(observationPrefix)\(step.observation ?? "")
-            """
-        }
-        return thoughts
-    }
 }
-
diff --git a/Tool/Sources/LangChain/AgentTool.swift b/Tool/Sources/LangChain/AgentTool.swift
@@ -31,18 +31,37 @@ public struct SimpleAgentTool: AgentTool {
     }
 }
 
-public struct FunctionCallingAgentTool: AgentTool {
-    public let function: any ChatGPTFunction
-    public var name: String { function.name }
-    public var description: String { function.description }
-    public let returnDirectly: Bool
-    
-    public init(function: any ChatGPTFunction, returnDirectly: Bool = false) {
+public struct FunctionCallingAgentTool<F: ChatGPTFunction>: AgentTool, ChatGPTFunction {
+    public func call(arguments: F.Arguments) async throws -> F.Result {
+        try await function.call(arguments: arguments)
+    }
+
+    public var argumentSchema: OpenAIService.JSONSchemaValue { function.argumentSchema }
+
+    public func prepare() async { await function.prepare() }
+
+    public var reportProgress: (String) async -> Void {
+        get { function.reportProgress }
+        set { function.reportProgress = newValue }
+    }
+
+    public typealias Arguments = F.Arguments
+    public typealias Result = F.Result
+
+    public var function: F
+    public var name: String
+    public var description: String
+    public var returnDirectly: Bool
+
+    public init(function: F, returnDirectly: Bool = false) {
         self.function = function
+        name = function.name
+        description = function.description
         self.returnDirectly = returnDirectly
     }
-    
+
     public func run(input: String) async throws -> String {
         try await function.call(argumentsJsonString: input).botReadableContent
     }
 }
+
diff --git a/Tool/Sources/LangChain/Agents/ChatAgent.swift b/Tool/Sources/LangChain/Agents/ChatAgent.swift
@@ -36,9 +36,10 @@ private func formatInstruction(toolsNames: String, preferredLanguage: String) ->
 public class ChatAgent: Agent {
     public typealias Input = String
     public typealias Output = String
+    public typealias ScratchPadContent = String
     public var observationPrefix: String { "Observation: " }
     public var llmPrefix: String { "Thought: " }
-    public let chatModelChain: ChatModelChain<AgentInput<String>>
+    public let chatModelChain: ChatModelChain<AgentInput<String, String>>
     let tools: [AgentTool]
 
     public init(chatModel: ChatModel, tools: [AgentTool], preferredLanguage: String) {
@@ -68,54 +69,61 @@ public class ChatAgent: Agent {
                         Begin! Reminder to always use the exact characters `Final Answer` when responding.
                         """
                     ),
-                    agentInput.thoughts.isEmpty
+                    agentInput.thoughts.content.isEmpty
                         ? .init(role: .user, content: agentInput.input)
                         : .init(
                             role: .user,
                             content: """
                             \(agentInput.input)
 
-                            \({
-                                switch agentInput.thoughts {
-                                case let .text(text):
-                                    return text
-                                case let .messages(messages):
-                                    return messages.map { message in
-                                        """
-                                        \(message)
-                                        """
-                                    }.joined(separator: "\n")
-                                }
-                            }())
+                            \(agentInput.thoughts.content)
                             """
                         ),
                 ]
             }
         )
     }
+    
+    func constructBaseScratchpad(intermediateSteps: [AgentAction]) -> String {
+        var thoughts = ""
+        for step in intermediateSteps {
+            thoughts += """
+            \(step.log)
+            \(observationPrefix)\(step.observation ?? "")
+            """
+        }
+        return thoughts
+    }
 
-    public func constructScratchpad(intermediateSteps: [AgentAction]) -> AgentScratchPad {
+    public func constructScratchpad(intermediateSteps: [AgentAction]) -> AgentScratchPad<String> {
         let baseScratchpad = constructBaseScratchpad(intermediateSteps: intermediateSteps)
-        if baseScratchpad.isEmpty { return .text("") }
-        return .text("""
+        if baseScratchpad.isEmpty { return .init(content: "") }
+        return .init(content: """
         This was your previous work (but I haven't seen any of it! I only see what you return as `Final Answer`):
         \(baseScratchpad)
         (Please continue with `Thought:` or `Final Answer:`)
         """)
     }
+    
+    public func constructFinalScratchpad(intermediateSteps: [AgentAction]) -> AgentScratchPad<String> {
+        let baseScratchpad = constructBaseScratchpad(intermediateSteps: intermediateSteps)
+        if baseScratchpad.isEmpty { return .init(content: "") }
+        return .init(content: """
+        This was your previous work (but I haven't seen any of it! I only see what you return as `Final Answer`):
+        \(baseScratchpad)
+        \(llmPrefix)I now need to return a final answer based on the previous steps:
+        "(Please continue with `Final Answer:`)"
+        """)
+    }
 
     public func validateTools(tools: [AgentTool]) throws {
         // no validation
     }
 
-    public func extraPlan(input: AgentInput<String>) {
+    public func extraPlan(input: AgentInput<String, String>) {
         // do nothing
     }
 
-    public func prepareForEarlyStopWithGenerate() -> String {
-        "(Please continue with `Final Answer:`)"
-    }
-
     public func parseOutput(_ output: ChatMessage) async -> AgentNextStep<Output> {
         let text = output.content ?? ""
 
diff --git a/Tool/Sources/LangChain/Agents/FunctionCallingChatAgent.swift b/Tool/Sources/LangChain/Agents/FunctionCallingChatAgent.swift
@@ -3,11 +3,13 @@ import Logger
 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
-        public var name: String { "sendFinalAnswer" }
-        public var description: String { "Send the final answer to user" }
+        public var name: String { "saveFinalAnswer" }
+        public var description: String { "Save the final answer when it's ready" }
         public let argumentSchema: JSONSchemaValue
         public var reportProgress: (String) async -> Void = { _ in }
         public func prepare() async {}
@@ -71,16 +73,16 @@ public class FunctionCallingChatAgent<Output: AgentOutputParsable & Decodable>:
     public typealias Input = String
     public var observationPrefix: String { "Observation: " }
     public var llmPrefix: String { "Thought: " }
-    public let chatModelChain: ChatModelChain<AgentInput<String>>
+    public let chatModelChain: ChatModelChain<AgentInput<String, ScratchPadContent>>
     var functionProvider: FunctionProvider
 
     public init(
         configuration: ChatGPTConfiguration = UserPreferenceChatGPTConfiguration(),
         tools: [AgentTool] = [],
         endFunction: EndFunction
     ) {
-        let functions = tools.compactMap { $0 as? FunctionCallingAgentTool }.map(\.function)
-        let otherTools = tools.filter { !($0 is FunctionCallingAgentTool) }
+        let functions = tools.compactMap { $0 as? (any ChatGPTFunction) }
+        let otherTools = tools.filter { !($0 is (any ChatGPTFunction)) }
         functionProvider = .init(
             tools: otherTools,
             functionTools: functions,
@@ -102,38 +104,18 @@ public class FunctionCallingChatAgent<Output: AgentOutputParsable & Decodable>:
                         role: .system,
                         content: """
                         Respond to the human as helpfully and accurately as possible. \
-                        Format final answer to be more readable, in a ordered list if possible. \
+                        Save the final answer when it's ready
 
                         Begin!
                         """
                     ),
-                    agentInput.thoughts.isEmpty
-                        ? .init(role: .user, content: agentInput.input)
-                        : .init(
-                            role: .user,
-                            content: """
-                            \(agentInput.input)
-
-                            \({
-                                switch agentInput.thoughts {
-                                case let .text(text):
-                                    return text
-                                case let .messages(messages):
-                                    return messages.map { message in
-                                        """
-                                        \(message)
-                                        """
-                                    }.joined(separator: "\n")
-                                }
-                            }())
-                            """
-                        ),
-                ]
+                    .init(role: .user, content: agentInput.input)
+                ] + agentInput.thoughts.content
             }
         )
     }
 
-    public func extraPlan(input: AgentInput<String>) {
+    public func extraPlan(input: AgentInput<String, ScratchPadContent>) {
         // no extra plan
     }
 
@@ -142,14 +124,34 @@ public class FunctionCallingChatAgent<Output: AgentOutputParsable & Decodable>:
         return "(call sendFinalAnswer to finish)"
     }
 
-    public func constructScratchpad(intermediateSteps: [AgentAction]) -> AgentScratchPad {
-        let baseScratchpad = constructBaseScratchpad(intermediateSteps: intermediateSteps)
-        if baseScratchpad.isEmpty { return .text("") }
-        return .text("""
-        This was your previous work (but I haven't seen any of it! I only see what you return as `Final Answer`):
-        \(baseScratchpad)
-        (Please continue with `Thought:` or call a function)
-        """)
+    public func constructScratchpad(
+        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),
+            ]
+        }
+        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),
+            ]
+        }
+        return .init(content: baseScratchpad)
     }
 
     public func validateTools(tools: [AgentTool]) throws {
@@ -185,7 +187,7 @@ public class FunctionCallingChatAgent<Output: AgentOutputParsable & Decodable>:
                 }
             }
         }
-        
+
         // fallback to normal agent.
 
         let stringBaseOutput = await ChatAgent(
diff --git a/Tool/Sources/OpenAIService/Models.swift b/Tool/Sources/OpenAIService/Models.swift
@@ -18,6 +18,10 @@ public struct ChatMessage: Equatable, Codable {
     public struct FunctionCall: Codable, Equatable {
         public var name: String
         public var arguments: String
+        public init(name: String, arguments: String) {
+            self.name = name
+            self.arguments = arguments
+        }
     }
 
     /// The role of a message.
