Put tool call responses into the tool call struct

intitni · intitni · commit 21d7ae4ce481 · 2024-03-03T22:00:32.000+08:00
diff --git a/Core/Sources/ChatContextCollectors/WebChatContextCollector/WebChatContextCollector.swift b/Core/Sources/ChatContextCollectors/WebChatContextCollector/WebChatContextCollector.swift
@@ -32,8 +32,7 @@ extension WebChatContextCollector {
     static func detectLinks(from messages: [ChatMessage]) -> [String] {
         return messages.lazy
             .compactMap {
-                $0.content ?? $0.toolCallContext?.toolCalls.map(\.function.arguments)
-                    .joined(separator: " ") ?? ""
+                $0.content ?? $0.toolCalls?.map(\.function.arguments).joined(separator: " ") ?? ""
             }
             .map(detectLinks(from:))
             .flatMap { $0 }
diff --git a/Core/Sources/ChatGPTChatTab/Chat.swift b/Core/Sources/ChatGPTChatTab/Chat.swift
@@ -327,15 +327,13 @@ struct Chat: ReducerProtocol {
                         }
                     ))
 
-                    if let responses = message.toolCallContext?.responses {
-                        for response in responses {
-                            all.append(.init(
-                                id: message.id + response.id,
-                                role: .tool,
-                                text: response.summary ?? response.content,
-                                references: []
-                            ))
-                        }
+                    for call in message.toolCalls ?? [] {
+                        all.append(.init(
+                            id: message.id + call.response.id,
+                            role: .tool,
+                            text: call.response.summary ?? call.response.content,
+                            references: []
+                        ))
                     }
 
                     return all
diff --git a/Core/Sources/ChatService/ChatService.swift b/Core/Sources/ChatService/ChatService.swift
@@ -126,7 +126,7 @@ public final class ChatService: ObservableObject {
 
         // if it's stopped before the tool calls finish, remove the message.
         await memory.mutateHistory { history in
-            if history.last?.role == .assistant, history.last?.toolCallContext?.toolCalls != nil {
+            if history.last?.role == .assistant, history.last?.toolCalls != nil {
                 history.removeLast()
             }
         }
diff --git a/Pro b/Pro
@@ -1 +1 @@
-Subproject commit a2e8aa56ff9b30bb3f3ae50a01b183cfbabb14cb
+Subproject commit ede561a8f50276d915ee672e8fad59226d349b08
diff --git a/Tool/Sources/LangChain/Chains/LLMChain.swift b/Tool/Sources/LangChain/Chains/LLMChain.swift
@@ -33,7 +33,7 @@ public class ChatModelChain<Input>: Chain {
     public func parseOutput(_ output: Output) -> String {
         if let content = output.content {
             return content
-        } else if let toolCalls = output.toolCallContext?.toolCalls {
+        } else if let toolCalls = output.toolCalls {
             return toolCalls.map { "[\($0.id)] \($0.function.name): \($0.function.arguments)" }
                 .joined(separator: "\n")
         }
diff --git a/Tool/Sources/LangChain/Chains/RefineDocumentChain.swift b/Tool/Sources/LangChain/Chains/RefineDocumentChain.swift
@@ -153,7 +153,7 @@ public final class RefineDocumentChain: Chain {
     }
 
     func extractAnswer(_ chatMessage: ChatMessage) -> IntermediateAnswer {
-        for functionCall in chatMessage.toolCallContext?.toolCalls.map(\.function) ?? [] {
+        for functionCall in chatMessage.toolCalls?.map(\.function) ?? [] {
             do {
                 let intermediateAnswer = try JSONDecoder().decode(
                     IntermediateAnswer.self,
diff --git a/Tool/Sources/LangChain/Chains/RelevantInformationExtractionChain.swift b/Tool/Sources/LangChain/Chains/RelevantInformationExtractionChain.swift
@@ -104,7 +104,7 @@ public final class RelevantInformationExtractionChain: Chain {
                             callbackManagers: callbackManagers
                         )
 
-                        if let functionCall = output.toolCallContext?.toolCalls
+                        if let functionCall = output.toolCalls?
                             .first(where: { $0.function.name == FinalAnswer().name })?.function
                         {
                             do {
diff --git a/Tool/Sources/LangChain/Chains/StructuredOutputChatModelChain.swift b/Tool/Sources/LangChain/Chains/StructuredOutputChatModelChain.swift
@@ -108,7 +108,7 @@ public class StructuredOutputChatModelChain<Output: Decodable>: Chain {
     }
 
     public func parseOutput(_ message: ChatMessage) async -> Output? {
-        if let functionCall = message.toolCallContext?.toolCalls.first?.function {
+        if let functionCall = message.toolCalls?.first?.function {
             do {
                 let result = try JSONDecoder().decode(
                     EndFunction.Arguments.self,
diff --git a/Tool/Sources/OpenAIService/APIs/GoogleAIChatCompletionsService.swift b/Tool/Sources/OpenAIService/APIs/GoogleAIChatCompletionsService.swift
@@ -208,13 +208,13 @@ extension ModelContent {
         case .user:
             return message.content ?? " "
         case .assistant:
-            if let toolCallContext = message.toolCallContext {
-                return toolCallContext.toolCalls.map { call in
-                    let response = toolCallContext.responses.first(where: { $0.id == call.id })
+            if let toolCalls = message.toolCalls {
+                return toolCalls.map { call in
+                    let response = call.response
                     return """
                     Call function: \(call.function.name)
                     Arguments: \(call.function.arguments)
-                    Result: \(response?.content ?? "N/A")
+                    Result: \(response.content)
                     """
                 }.joined(separator: "\n")
             } else {
diff --git a/Tool/Sources/OpenAIService/ChatGPTService.swift b/Tool/Sources/OpenAIService/ChatGPTService.swift
@@ -169,7 +169,7 @@ public class ChatGPTService: ChatGPTServiceType {
                 role: .user,
                 content: content,
                 name: nil,
-                toolCallContext: nil,
+                toolCalls: nil,
                 summary: summary,
                 references: []
             )
@@ -222,7 +222,7 @@ public class ChatGPTService: ChatGPTServiceType {
 
                             pendingToolCalls = await memory.history
                                 .last { $0.id == sourceMessageId }?
-                                .toolCallContext?.toolCalls ?? []
+                                .toolCalls ?? []
 
                             #if DEBUG
                             Debugger.didReceiveResponse(content: reply)
@@ -261,7 +261,7 @@ public class ChatGPTService: ChatGPTServiceType {
         return try await Debugger.$id.withValue(.init()) {
             let message = try await sendMemoryAndWait()
             var finalResult = message?.content
-            var toolCalls = message?.toolCallContext?.toolCalls
+            var toolCalls = message?.toolCalls
             while let sourceMessageId = message?.id, let calls = toolCalls, !calls.isEmpty {
                 try Task.checkCancellation()
                 if !configuration.runFunctionsAutomatically {
@@ -273,7 +273,7 @@ public class ChatGPTService: ChatGPTServiceType {
                 }
                 guard let nextMessage = try await sendMemoryAndWait() else { break }
                 finalResult = nextMessage.content
-                toolCalls = nextMessage.toolCallContext?.toolCalls
+                toolCalls = nextMessage.toolCalls
             }
 
             #if DEBUG
@@ -438,13 +438,11 @@ extension ChatGPTService {
             }(),
             content: choice.content,
             name: choice.name,
-            toolCallContext: choice.toolCalls.map {
-                .init(toolCalls: $0.map {
-                    ChatMessage.ToolCall(id: $0.id, type: $0.type, function: .init(
-                        name: $0.function.name,
-                        arguments: $0.function.arguments ?? ""
-                    ))
-                }, responses: [])
+            toolCalls: choice.toolCalls?.map {
+                ChatMessage.ToolCall(id: $0.id, type: $0.type, function: .init(
+                    name: $0.function.name,
+                    arguments: $0.function.arguments ?? ""
+                ))
             },
             references: prompt.references
         )
@@ -589,7 +587,7 @@ extension ChatGPTService {
                 }(),
                 content: chatMessage.content ?? "",
                 name: chatMessage.name,
-                toolCalls: chatMessage.toolCallContext?.toolCalls.map {
+                toolCalls: chatMessage.toolCalls?.map {
                     .init(
                         id: $0.id,
                         type: $0.type,
@@ -601,14 +599,12 @@ extension ChatGPTService {
                 }
             ))
 
-            if let responses = chatMessage.toolCallContext?.responses {
-                for response in responses {
-                    all.append(ChatCompletionsRequestBody.Message(
-                        role: .tool,
-                        content: response.content,
-                        toolCallId: response.id
-                    ))
-                }
+            for call in chatMessage.toolCalls ?? [] {
+                all.append(ChatCompletionsRequestBody.Message(
+                    role: .tool,
+                    content: call.response.content,
+                    toolCallId: call.response.id
+                ))
             }
 
             return all
diff --git a/Tool/Sources/OpenAIService/Memory/AutoManagedChatGPTMemoryStrategy/AutoManagedChatGPTMemoryOpenAIStrategy.swift b/Tool/Sources/OpenAIService/Memory/AutoManagedChatGPTMemoryStrategy/AutoManagedChatGPTMemoryOpenAIStrategy.swift
@@ -37,18 +37,15 @@ extension TokenEncoder {
             encodingContent.append(name)
             total += 1
         }
-        if let toolCallContext = message.toolCallContext {
-            for toolCall in toolCallContext.toolCalls {
+        if let toolCalls = message.toolCalls {
+            for toolCall in toolCalls {
                 encodingContent.append(toolCall.id)
                 encodingContent.append(toolCall.type)
                 encodingContent.append(toolCall.function.name)
                 encodingContent.append(toolCall.function.arguments)
-            }
-            
-            for response in toolCallContext.responses {
                 total += 4
-                encodingContent.append(response.content)
-                encodingContent.append(response.id)
+                encodingContent.append(toolCall.response.content)
+                encodingContent.append(toolCall.response.id)
             }
         }
         total += await withTaskGroup(of: Int.self, body: { group in
diff --git a/Tool/Sources/OpenAIService/Memory/ChatGPTMemory.swift b/Tool/Sources/OpenAIService/Memory/ChatGPTMemory.swift
@@ -65,21 +65,15 @@ public extension ChatGPTMemory {
         summary: String? = nil
     ) async {
         await updateMessage(id: id) { message in
-            if let index = message.toolCallContext?.responses.firstIndex(where: {
+            if let index = message.toolCalls?.firstIndex(where: {
                 $0.id == toolCallId
             }) {
                 if let content {
-                    message.toolCallContext?.responses[index].content = content
+                    message.toolCalls?[index].response.content = content
                 }
                 if let summary {
-                    message.toolCallContext?.responses[index].summary = summary
+                    message.toolCalls?[index].response.summary = summary
                 }
-            } else {
-                message.toolCallContext?.responses.append(.init(
-                    id: toolCallId,
-                    content: content ?? "",
-                    summary: summary ?? ""
-                ))
             }
         }
     }
@@ -107,7 +101,7 @@ public extension ChatGPTMemory {
                     message.role = role
                 }
                 if let toolCalls {
-                    if var existedToolCalls = message.toolCallContext?.toolCalls {
+                    if var existedToolCalls = message.toolCalls {
                         for pair in toolCalls.sorted(by: { $0.key <= $1.key }) {
                             let (proposedIndex, toolCall) = pair
                             let index = {
@@ -130,12 +124,9 @@ public extension ChatGPTMemory {
                                 existedToolCalls.append(toolCall)
                             }
                         }
-                        message.toolCallContext?.toolCalls = existedToolCalls
+                        message.toolCalls = existedToolCalls
                     } else {
-                        message.toolCallContext = .init(
-                            toolCalls: toolCalls.sorted(by: { $0.key <= $1.key }).map(\.value),
-                            responses: []
-                        )
+                        message.toolCalls = toolCalls.sorted(by: { $0.key <= $1.key }).map(\.value)
                     }
                 }
                 if let summary {
@@ -155,12 +146,7 @@ public extension ChatGPTMemory {
                     role: role ?? .system,
                     content: content,
                     name: name,
-                    toolCallContext: toolCalls.map { calls in
-                        .init(
-                            toolCalls: calls.sorted(by: { $0.key <= $1.key }).map(\.value),
-                            responses: []
-                        )
-                    },
+                    toolCalls: toolCalls?.sorted(by: { $0.key <= $1.key }).map(\.value),
                     summary: summary,
                     references: references ?? []
                 ))
diff --git a/Tool/Sources/OpenAIService/Models.swift b/Tool/Sources/OpenAIService/Models.swift
@@ -30,10 +30,12 @@ public struct ChatMessage: Equatable, Codable {
         public var id: String
         public var type: String
         public var function: FunctionCall
+        public var response: ToolCallResponse
         public init(id: String, type: String, function: FunctionCall) {
             self.id = id
             self.type = type
             self.function = function
+            response = .init(id: id, content: "", summary: nil)
         }
     }
     
@@ -47,11 +49,6 @@ public struct ChatMessage: Equatable, Codable {
             self.summary = summary
         }
     }
-    
-    public struct ToolCallContext: Codable, Equatable {
-        public var toolCalls: [ToolCall]
-        public var responses: [ToolCallResponse]
-    }
 
     public struct Reference: Codable, Equatable {
         public enum Kind: String, Codable {
@@ -109,7 +106,7 @@ public struct ChatMessage: Equatable, Codable {
     }
 
     /// A function call from the bot.
-    public var toolCallContext: ToolCallContext? {
+    public var toolCalls: [ToolCall]? {
         didSet { tokensCount = nil }
     }
 
@@ -134,7 +131,7 @@ public struct ChatMessage: Equatable, Codable {
     /// Is the message considered empty.
     var isEmpty: Bool {
         if let content, !content.isEmpty { return false }
-        if let toolCallContext, !toolCallContext.toolCalls.isEmpty { return false }
+        if let toolCalls, !toolCalls.isEmpty { return false }
         if let name, !name.isEmpty { return false }
         return true
     }
@@ -144,15 +141,15 @@ public struct ChatMessage: Equatable, Codable {
         role: Role,
         content: String?,
         name: String? = nil,
-        toolCallContext: ToolCallContext? = nil,
+        toolCalls: [ToolCall]? = nil,
         summary: String? = nil,
         tokenCount: Int? = nil,
         references: [Reference] = []
     ) {
         self.role = role
         self.content = content
         self.name = name
-        self.toolCallContext = toolCallContext
+        self.toolCalls = toolCalls
         self.summary = summary
         self.id = id
         tokensCount = tokenCount

Original file line number	Diff line number	Diff line change
`@@ -126,7 +126,7 @@ public final class ChatService: ObservableObject {`
`126`	`126`
`127`	`127`	`// if it's stopped before the tool calls finish, remove the message.`
`128`	`128`	`await memory.mutateHistory { history in`
`129`		`- if history.last?.role == .assistant, history.last?.toolCallContext?.toolCalls != nil {`
	`129`	`+ if history.last?.role == .assistant, history.last?.toolCalls != nil {`
`130`	`130`	`history.removeLast()`
`131`	`131`	`}`
`132`	`132`	`}`
Original file line number	Diff line number	Diff line change
`@@ -33,7 +33,7 @@ public class ChatModelChain<Input>: Chain {`
`33`	`33`	`public func parseOutput(_ output: Output) -> String {`
`34`	`34`	`if let content = output.content {`
`35`	`35`	`return content`
`36`		`- } else if let toolCalls = output.toolCallContext?.toolCalls {`
	`36`	`+ } else if let toolCalls = output.toolCalls {`
`37`	`37`	`return toolCalls.map { "[\($0.id)] \($0.function.name): \($0.function.arguments)" }`
`38`	`38`	`.joined(separator: "\n")`
`39`	`39`	`}`
Original file line number	Diff line number	Diff line change
`@@ -153,7 +153,7 @@ public final class RefineDocumentChain: Chain {`
`153`	`153`	`}`
`154`	`154`
`155`	`155`	`func extractAnswer(_ chatMessage: ChatMessage) -> IntermediateAnswer {`
`156`		`- for functionCall in chatMessage.toolCallContext?.toolCalls.map(\.function) ?? [] {`
	`156`	`+ for functionCall in chatMessage.toolCalls?.map(\.function) ?? [] {`
`157`	`157`	`do {`
`158`	`158`	`let intermediateAnswer = try JSONDecoder().decode(`
`159`	`159`	`IntermediateAnswer.self,`
Original file line number	Diff line number	Diff line change
`@@ -104,7 +104,7 @@ public final class RelevantInformationExtractionChain: Chain {`
`104`	`104`	`callbackManagers: callbackManagers`
`105`	`105`	`)`
`106`	`106`
`107`		`- if let functionCall = output.toolCallContext?.toolCalls`
	`107`	`+ if let functionCall = output.toolCalls?`
`108`	`108`	`.first(where: { $0.function.name == FinalAnswer().name })?.function`
`109`	`109`	`{`
`110`	`110`	`do {`
Original file line number	Diff line number	Diff line change
`@@ -108,7 +108,7 @@ public class StructuredOutputChatModelChain<Output: Decodable>: Chain {`
`108`	`108`	`}`
`109`	`109`
`110`	`110`	`public func parseOutput(_ message: ChatMessage) async -> Output? {`
`111`		`- if let functionCall = message.toolCallContext?.toolCalls.first?.function {`
	`111`	`+ if let functionCall = message.toolCalls?.first?.function {`
`112`	`112`	`do {`
`113`	`113`	`let result = try JSONDecoder().decode(`
`114`	`114`	`EndFunction.Arguments.self,`