feat: return raw output for agent toolcall result (#1806)

.changeset/three-seahorses-know.md

@@ -0,0 +1,6 @@
+---
+"@llamaindex/server": patch
+"@llamaindex/workflow": patch
+---
+
+feat: return raw output for agent toolcall result

packages/server/next/app/page.tsx

@@ -14,7 +14,7 @@ const ChatSection = dynamic(() => import("./components/chat-section"), {
 
 export default function Home() {
   return (
-    <main className="background-gradient flex h-screen w-screen items-center justify-center">
+    <main className="background-gradient flex h-screen w-screen items-center justify-center overflow-hidden">
       <div className="w-[90%] space-y-2 lg:w-[60rem] lg:space-y-10">
         <ChatSection />
       </div>

packages/server/src/events.ts

@@ -55,7 +55,7 @@ export function toSourceEventNode(node: NodeWithScore<Metadata>) {
   const { file_name, pipeline_id } = node.node.metadata;
 
   const filePath = pipeline_id
-    ? `output/llamacloud/${pipeline_id}${file_name}`
+    ? `output/llamacloud/${pipeline_id}$${file_name}`
     : `data/${file_name}`;
 
   return {

packages/server/src/utils/workflow.ts

@@ -3,17 +3,27 @@ import type {
   AgentInputData,
   ChatResponseChunk,
   EngineResponse,
+  Metadata,
+  NodeWithScore,
   WorkflowEvent,
 } from "llamaindex";
 import {
   AgentStream,
+  AgentToolCall,
+  AgentToolCallResult,
   AgentWorkflow,
+  LLamaCloudFileService,
   StopEvent,
   Workflow,
   type AgentWorkflowContext,
 } from "llamaindex";
 import { ReadableStream } from "stream/web";
-import { SourceEvent, type SourceEventNode } from "../events";
+import {
+  SourceEvent,
+  toAgentRunEvent,
+  toSourceEvent,
+  type SourceEventNode,
+} from "../events";
 import type { ServerWorkflow } from "../types";
 import { downloadFile } from "./file";
 import { sendSuggestedQuestionsEvent } from "./suggestion";
@@ -124,13 +134,49 @@ function appendEventDataToAnnotations(
   dataStream: StreamData,
   event: WorkflowEvent<unknown>,
 ) {
+  const transformedEvent = transformWorkflowEvent(event);
+
   // for SourceEvent, we need to trigger download files from LlamaCloud (if having)
-  if (event instanceof SourceEvent) {
-    const sourceNodes = event.data.data.nodes;
+  if (transformedEvent instanceof SourceEvent) {
+    const sourceNodes = transformedEvent.data.data.nodes;
     downloadLlamaCloudFilesFromNodes(sourceNodes); // download files in background
   }
 
-  dataStream.appendMessageAnnotation(event.data as JSONValue);
+  dataStream.appendMessageAnnotation(transformedEvent.data as JSONValue);
+}
+
+// transform WorkflowEvent to another WorkflowEvent for annotations display purpose
+// this useful for handling AgentWorkflow events, because we cannot easily append custom events like custom workflows
+function transformWorkflowEvent(
+  event: WorkflowEvent<unknown>,
+): WorkflowEvent<unknown> {
+  // convert AgentToolCall event to AgentRunEvent
+  if (event instanceof AgentToolCall) {
+    const inputString = JSON.stringify(event.data.toolKwargs);
+    return toAgentRunEvent({
+      agent: event.data.agentName,
+      text: `Using tool: '${event.data.toolName}' with inputs: '${inputString}'`,
+      type: "text",
+    });
+  }
+
+  // modify AgentToolCallResult event
+  if (event instanceof AgentToolCallResult) {
+    const rawOutput = event.data.raw;
+
+    // if AgentToolCallResult contains sourceNodes, convert it to SourceEvent
+    if (
+      rawOutput &&
+      typeof rawOutput === "object" &&
+      "sourceNodes" in rawOutput // TODO: better use Zod to validate and extract sourceNodes from toolCallResult
+    ) {
+      return toSourceEvent(
+        rawOutput.sourceNodes as unknown as NodeWithScore<Metadata>[],
+      );
+    }
+  }
+
+  return event;
 }
 
 async function downloadLlamaCloudFilesFromNodes(nodes: SourceEventNode[]) {
@@ -141,7 +187,13 @@ async function downloadLlamaCloudFilesFromNodes(nodes: SourceEventNode[]) {
     if (downloadedFiles.includes(node.filePath)) continue; // skip if file already downloaded
     if (!node.metadata.pipeline_id) continue; // only download files from LlamaCloud
 
-    await downloadFile(node.url, node.filePath);
+    const downloadUrl = await LLamaCloudFileService.getFileUrl(
+      node.metadata.pipeline_id,
+      node.fileName,
+    );
+    if (!downloadUrl) continue;
+
+    await downloadFile(downloadUrl, node.filePath);
 
     downloadedFiles.push(node.filePath);
   }

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

@@ -399,9 +399,11 @@ export class AgentWorkflow {
           isError: false,
         },
         returnDirect: false,
+        raw: {},
       });
       try {
         const output = await this.callTool(toolCall, ctx);
+        toolResult.data.raw = output;
         toolResult.data.toolOutput.result =
           stringifyJSONToMessageContent(output);
         toolResult.data.returnDirect = toolCall.data.toolName === "handOff";

packages/workflow/src/agent/events.ts

@@ -9,13 +9,13 @@ export class AgentToolCall extends WorkflowEvent<{
   toolId: string;
 }> {}
 
-// TODO: Check for if we need a raw tool output
 export class AgentToolCallResult extends WorkflowEvent<{
   toolName: string;
   toolKwargs: Record<string, JSONValue>;
   toolId: string;
   toolOutput: ToolResult;
   returnDirect: boolean;
+  raw: JSONValue;
 }> {}
 
 export class AgentInput extends WorkflowEvent<{