Update ChatModel to return ChatMessage

Author: intitni

Tool/Sources/LangChain/Agent.swift

@@ -104,7 +104,7 @@ public extension Agent {
     ) async throws -> AgentNextStep {
         let input = getFullInputs(input: input, intermediateSteps: intermediateSteps)
         let output = try await chatModelChain.call(input, callbackManagers: callbackManagers)
-        return parseOutput(output)
+        return parseOutput(output.content ?? "")
     }
 
     func returnStoppedResponse(
@@ -128,12 +128,13 @@ public extension Agent {
             """
             let input = AgentInput(input: input, thoughts: .text(thoughts))
             let output = try await chatModelChain.call(input, callbackManagers: callbackManagers)
-            let nextAction = parseOutput(output)
+            let reply = output.content ?? ""
+            let nextAction = parseOutput(reply)
             switch nextAction {
             case let .finish(finish):
                 return finish
             case .actions:
-                return AgentFinish(returnValue: output, log: output)
+                return AgentFinish(returnValue: reply, log: reply)
             }
         }
     }

Tool/Sources/LangChain/Chains/RetrievalQA.swift

@@ -183,16 +183,17 @@ public final class RefineDocumentChain: Chain {
             output = try await refinementChatModel.call(
                 .init(
                     question: input.question,
-                    previousAnswer: output,
+                    previousAnswer: content,
                     document: document.document.pageContent,
                     distance: document.distance
                 ),
                 callbackManagers: callbackManagers
             )
+            content = output.content ?? ""
             callbackManagers
-                .send(CallbackEvents.RetrievalQADidGenerateIntermediateAnswer(info: output))
+                .send(CallbackEvents.RetrievalQADidGenerateIntermediateAnswer(info: content))
         }
-        return output
+        return content
     }
 
     public func parseOutput(_ output: String) -> String {