Move tool call responses into its source assistant message

Author: intitni

Core/Sources/ChatContextCollectors/WebChatContextCollector/WebChatContextCollector.swift

@@ -32,7 +32,8 @@ extension WebChatContextCollector {
     static func detectLinks(from messages: [ChatMessage]) -> [String] {
         return messages.lazy
             .compactMap {
-                $0.content ?? $0.toolCalls?.map(\.function.arguments).joined(separator: " ") ?? ""
+                $0.content ?? $0.toolCallContext?.toolCalls.map(\.function.arguments)
+                    .joined(separator: " ") ?? ""
             }
             .map(detectLinks(from:))
             .flatMap { $0 }

Core/Sources/ChatGPTChatTab/Chat.swift

@@ -15,7 +15,7 @@ public struct DisplayedChatMessage: Equatable {
 
     public struct Reference: Equatable {
         public typealias Kind = ChatMessage.Reference.Kind
-        
+
         public var title: String
         public var subtitle: String
         public var uri: String
@@ -135,7 +135,7 @@ struct Chat: ReducerProtocol {
                     await send(.focusOnTextField)
                     await send(.refresh)
                 }
-                
+
             case .refresh:
                 return .run { send in
                     await send(.chatMenu(.refresh))
@@ -298,8 +298,9 @@ struct Chat: ReducerProtocol {
                 }.cancellable(id: CancelID.observeDefaultScopesChange(id), cancelInFlight: true)
 
             case .historyChanged:
-                state.history = service.chatHistory.map { message in
-                    .init(
+                state.history = service.chatHistory.flatMap { message in
+                    var all = [DisplayedChatMessage]()
+                    all.append(.init(
                         id: message.id,
                         role: {
                             switch message.role {
@@ -312,7 +313,6 @@ struct Chat: ReducerProtocol {
                                     return .assistant
                                 }
                                 return .ignored
-                            case .tool: return .tool
                             }
                         }(),
                         text: message.summary ?? message.content ?? "",
@@ -325,7 +325,20 @@ struct Chat: ReducerProtocol {
                                 kind: $0.kind
                             )
                         }
-                    )
+                    ))
+
+                    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: []
+                            ))
+                        }
+                    }
+
+                    return all
                 }
 
                 state.title = {
@@ -401,7 +414,7 @@ struct ChatMenu: ReducerProtocol {
                 return .run {
                     await $0(.refresh)
                 }
-                
+
             case .refresh:
                 state.temperatureOverride = service.configuration.overriding.temperature
                 state.chatModelIdOverride = service.configuration.overriding.modelId

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?.toolCalls != nil {
+            if history.last?.role == .assistant, history.last?.toolCallContext?.toolCalls != nil {
                 history.removeLast()
             }
         }

Core/Sources/ChatService/ContextAwareAutoManagedChatGPTMemory.swift

@@ -39,7 +39,7 @@ public final class ContextAwareAutoManagedChatGPTMemory: ChatGPTMemory {
 
     public func generatePrompt() async -> ChatGPTPrompt {
         let content = (await memory.history)
-            .last(where: { $0.role == .user || $0.role == .tool })?.content
+            .last(where: { $0.role == .user  })?.content
         try? await contextController.collectContextInformation(
             systemPrompt: """
             \(chatService?.systemPrompt ?? "")

Pro

@@ -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.toolCalls {
+        } else if let toolCalls = output.toolCallContext?.toolCalls {
             return toolCalls.map { "[\($0.id)] \($0.function.name): \($0.function.arguments)" }
                 .joined(separator: "\n")
         }

Tool/Sources/LangChain/Chains/LLMChain.swift

@@ -153,7 +153,7 @@ public final class RefineDocumentChain: Chain {
     }
 
     func extractAnswer(_ chatMessage: ChatMessage) -> IntermediateAnswer {
-        for functionCall in chatMessage.toolCalls?.map(\.function) ?? [] {
+        for functionCall in chatMessage.toolCallContext?.toolCalls.map(\.function) ?? [] {
             do {
                 let intermediateAnswer = try JSONDecoder().decode(
                     IntermediateAnswer.self,

Tool/Sources/LangChain/Chains/RefineDocumentChain.swift

@@ -104,7 +104,7 @@ public final class RelevantInformationExtractionChain: Chain {
                             callbackManagers: callbackManagers
                         )
 
-                        if let functionCall = output.toolCalls?
+                        if let functionCall = output.toolCallContext?.toolCalls
                             .first(where: { $0.function.name == FinalAnswer().name })?.function
                         {
                             do {

Tool/Sources/LangChain/Chains/RelevantInformationExtractionChain.swift

@@ -108,7 +108,7 @@ public class StructuredOutputChatModelChain<Output: Decodable>: Chain {
     }
 
     public func parseOutput(_ message: ChatMessage) async -> Output? {
-        if let functionCall = message.toolCalls?.first?.function {
+        if let functionCall = message.toolCallContext?.toolCalls.first?.function {
             do {
                 let result = try JSONDecoder().decode(
                     EndFunction.Arguments.self,

Tool/Sources/LangChain/Chains/StructuredOutputChatModelChain.swift

@@ -5,8 +5,28 @@ import Preferences
 
 struct ChatCompletionsRequestBody: Codable, Equatable {
     struct Message: Codable, Equatable {
+        enum Role: String, Codable, Equatable {
+            case system
+            case user
+            case assistant
+            case tool
+            
+            var asChatMessageRole: ChatMessage.Role {
+                switch self {
+                case .system:
+                    return .system
+                case .user:
+                    return .user
+                case .assistant:
+                    return .assistant
+                case .tool:
+                    return .user
+                }
+            }
+        }
+
         /// The role of the message.
-        var role: ChatMessage.Role
+        var role: Role
         /// The content of the message.
         var content: String
         /// When we want to reply to a function call with the result, we have to provide the
@@ -149,12 +169,13 @@ struct ChatCompletionsStreamDataChunk {
         }
 
         struct ToolCall {
+            var index: Int?
             var id: String?
             var type: String?
             var function: FunctionCall?
         }
 
-        var role: ChatMessage.Role?
+        var role: ChatCompletionsRequestBody.Message.Role?
         var content: String?
         var toolCalls: [ToolCall]?
     }
@@ -174,7 +195,7 @@ protocol ChatCompletionsAPI {
 
 struct ChatCompletionResponseBody: Codable, Equatable {
     typealias Message = ChatCompletionsRequestBody.Message
-    
+
     var id: String?
     var object: String
     var model: String