fix: ensure correct message content in agent workflow (#2114)

Co-authored-by: Marcus Schiesser <mail@marcusschiesser.de>

.changeset/pink-planets-tan.md

@@ -0,0 +1,7 @@
+---
+"@llamaindex/doc": patch
+"@llamaindex/core": patch
+"@llamaindex/workflow": patch
+---
+
+fix: ensure correct message content in agent workflow

apps/next/src/content/docs/llamaindex/modules/agents/agent_workflow.mdx

@@ -34,6 +34,7 @@ const jokeAgent = agent({
 // Run the workflow
 const result = await jokeAgent.run("Tell me something funny");
 console.log(result.data.result); // Baby Llama is called cria
+console.log(result.data.message); // { role: 'assistant', content: 'Baby Llama is called cria' }
 ```
 
 ### Event Streaming

examples/agents/agent/single-agent.ts

@@ -24,6 +24,7 @@ async function main() {
     state: result.data.state,
   });
   console.log(`${JSON.stringify(caResult, null, 2)}`);
+  console.log("assistant message:", result.data.message);
 }
 
 main().catch((error) => {

packages/core/src/utils/index.ts

@@ -56,10 +56,45 @@ export function prettifyError(error: unknown): string {
   }
 }
 
+/**
+ * Returns a stringfied JSON with double quotes removed.
+ *
+ * @param value - The JSON value to stringify
+ * @returns The stringified JSON with no double quotes
+ */
 export function stringifyJSONToMessageContent(value: JSONValue): string {
   return JSON.stringify(value, null, 2).replace(/"([^"]*)"/g, "$1");
 }
 
+export function assertIsJSONValue(value: unknown): asserts value is JSONValue {
+  if (
+    typeof value === "string" ||
+    typeof value === "number" ||
+    typeof value === "boolean"
+  ) {
+    return;
+  }
+
+  if (Array.isArray(value)) {
+    for (const item of value) {
+      assertIsJSONValue(item);
+    }
+    return;
+  }
+
+  if (typeof value === "object" && value !== null) {
+    for (const [key, val] of Object.entries(value)) {
+      if (typeof key !== "string") {
+        throw new Error(`Invalid object key: ${key}`);
+      }
+      assertIsJSONValue(val);
+    }
+    return;
+  }
+
+  throw new Error(`Value is not a valid JSONValue: ${String(value)}`);
+}
+
 export {
   extractDataUrlComponents,
   extractImage,

packages/workflow/src/agent/agent-workflow.ts

@@ -1,8 +1,12 @@
+import type { JSONValue } from "@llamaindex/core/global";
 import type { ChatMessage, MessageContent } from "@llamaindex/core/llms";
 import { createMemory, Memory } from "@llamaindex/core/memory";
 import { PromptTemplate } from "@llamaindex/core/prompts";
 import { tool } from "@llamaindex/core/tools";
-import { stringifyJSONToMessageContent } from "@llamaindex/core/utils";
+import {
+  assertIsJSONValue,
+  stringifyJSONToMessageContent,
+} from "@llamaindex/core/utils";
 import {
   createWorkflow,
   getContext,
@@ -55,6 +59,7 @@ export const startAgentEvent = workflowEvent<
 
 export type AgentResultData = {
   result: MessageContent;
+  message: ChatMessage;
   state?: AgentWorkflowState | undefined;
 };
 export const stopAgentEvent = workflowEvent<AgentResultData, "llamaindex-stop">(
@@ -423,6 +428,7 @@ export class AgentWorkflow implements Workflow {
       );
 
       return stopAgentEvent.with({
+        message: content.response,
         result: content.response.content,
         state: this.stateful.getContext().state,
       });
@@ -502,18 +508,17 @@ export class AgentWorkflow implements Workflow {
     if (directResult) {
       const isHandoff = directResult.toolName === "handOff";
 
-      const output =
-        typeof directResult.toolOutput.result === "string"
-          ? directResult.toolOutput.result
-          : JSON.stringify(directResult.toolOutput.result);
+      const raw = directResult.raw;
+      const output = typeof raw === "string" ? raw : JSON.stringify(raw);
+      const responseMessage: ChatMessage = {
+        role: "assistant" as const,
+        content: output, // use stringified tool output for assistant message
+      };
 
       const agentOutput = {
-        response: {
-          role: "assistant" as const,
-          content: output,
-        },
+        response: responseMessage,
         toolCalls: [],
-        raw: output,
+        raw,
         currentAgentName: agent.name,
       };
 
@@ -542,6 +547,7 @@ export class AgentWorkflow implements Workflow {
       }
 
       return stopAgentEvent.with({
+        message: responseMessage,
         result: output,
         state: this.stateful.getContext().state,
       });
@@ -566,14 +572,16 @@ export class AgentWorkflow implements Workflow {
     this.workflow.handle([toolResultsEvent], this.processToolResults);
   }
 
-  private callTool(toolCall: AgentToolCall) {
+  private async callTool(toolCall: AgentToolCall): Promise<JSONValue> {
     const tool = this.agents
       .get(toolCall.agentName)
       ?.tools.find((t) => t.metadata.name === toolCall.toolName);
     if (!tool) {
       throw new Error(`Tool ${toolCall.toolName} not found`);
     }
-    return tool.call(toolCall.toolKwargs);
+    const output = await tool.call(toolCall.toolKwargs);
+    assertIsJSONValue(output);
+    return output;
   }
 
   private createInitialState(): AgentWorkflowState {