feat: return structured object from llm.exec (#2198)

.changeset/five-peaches-brush.md

@@ -0,0 +1,5 @@
+---
+"@llamaindex/core": patch
+---
+
+feat: return structured object from llm.exec

.changeset/red-bushes-brake.md

@@ -0,0 +1,5 @@
+---
+"@llamaindex/workflow": patch
+---
+
+feat: return structured data when using agent.run()

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

@@ -37,6 +37,58 @@ console.log(result.data.result); // Baby Llama is called cria
 console.log(result.data.message); // { role: 'assistant', content: 'Baby Llama is called cria' }
 ```
 
+### Structured Output
+
+You can extract structured data from agent responses by providing a `responseFormat` with a Zod schema. This is useful when you need the agent's response in a specific format for further processing:
+
+```typescript
+import { z } from "zod";
+import { tool } from "llamaindex";
+import { agent } from "@llamaindex/workflow";
+import { openai } from "@llamaindex/openai";
+
+// Define a weather tool
+const weatherTool = tool({
+  name: "weatherTool",
+  description: "Get weather information",
+  parameters: z.object({
+    location: z.string(),
+  }),
+  execute: ({ location }) => {
+    return `The weather in ${location} is sunny. The temperature is 72 degrees. The humidity is 50%. The wind speed is 10 mph.`;
+  },
+});
+
+// Define the structure you want for the response
+const responseSchema = z.object({
+  temperature: z.number(),
+  humidity: z.number(),
+  windSpeed: z.number(),
+});
+
+// Create the agent
+const weatherAgent = agent({
+  name: "weatherAgent",
+  tools: [weatherTool],
+  llm: openai({ model: "gpt-4.1-mini" }),
+});
+
+// Run with structured output
+const result = await weatherAgent.run("What's the weather in Tokyo?", {
+  responseFormat: responseSchema,
+});
+
+console.log("Natural language result:", result.data.result);
+console.log("Structured data:", result.data.object);
+// Output: { temperature: 72, humidity: 50, windSpeed: 10 }
+```
+
+The agent will:
+1. Use the weather tool to get the raw weather information
+2. Process that information through the LLM
+3. Extract structured data according to your schema
+4. Return both the natural language response and the structured object
+
 ### Event Streaming
 
 Agent Workflows provide a unified interface for event streaming, making it easy to track and respond to different events during execution:

apps/next/src/content/docs/llamaindex/modules/agents/low-level.mdx

@@ -9,6 +9,7 @@ Sometimes your need more control over LLM interactions than what high-level agen
 Use `llm.exec` when you need to:
 - Build custom agent logic in [workflow](/docs/llamaindex/modules/agents/workflows) steps
 - Have precise control over message handling and tool execution
+- Extract structured data from LLM responses
 
 ## Basic Usage
 
@@ -51,6 +52,38 @@ messages.push(...newMessages);
 
 > `newMessages` is an array as each tool call generates two messages: a tool call message and the tool call result message.
 
+## Structured Output
+
+You can use `responseFormat` with a Zod schema to get structured data from the LLM response:
+
+```ts
+import { openai } from "@llamaindex/openai";
+import { ChatMessage } from "llamaindex";
+import z from "zod";
+
+const llm = openai({ model: "gpt-4.1-mini" });
+
+const schema = z.object({
+  title: z.string(),
+  author: z.string(),
+  year: z.number(),
+});
+
+const messages = [
+  {
+    role: "user",
+    content: "I have been reading La Divina Commedia by Dante Alighieri, published in 1321",
+  } as ChatMessage,
+];
+
+const { newMessages, toolCalls, object } = await llm.exec({
+  messages,
+  responseFormat: schema,
+});
+
+console.log(object); // { title: "La Divina Commedia", author: "Dante Alighieri", year: 1321 }
+```
+
 ## Agent Loop Pattern
 
 A common pattern is to use `llm.exec` in a loop until the LLM stops making tool calls:
@@ -102,7 +135,7 @@ For real-time responses, use the `stream` option to get the assistant's response
 
 ```ts
 import { openai } from "@llamaindex/openai";
-import { tool } from "llamaindex";
+import { ChatMessage, tool } from "llamaindex";
 import z from "zod";
 
 async function streamingAgentLoop() {
@@ -153,6 +186,7 @@ async function streamingAgentLoop() {
 
 - **`newMessages`**: Array of new chat messages including the LLM response and any tool call messages (call or result). This is a function return the array when streaming.
 - **`toolCalls`**: Array of tool calls made by the LLM
+- **`object`**: The structured object when using `responseFormat` with a Zod schema (undefined if no schema is provided)
 - **`stream`**: Async iterable for streaming responses (only when `stream: true`)
 
 ## Best Practices
@@ -161,4 +195,4 @@ For using `llm.exec` in an agent loop, take care to:
 
 1. **Maintain message history**: Always add `newMessages` to your conversation history
 2. **Set exit conditions**: Implement proper logic to avoid infinite loops
-
+3. **Handle structured output**: When using `responseFormat`, the `object` property contains your parsed data

examples/agents/agent/structured-object.ts

@@ -0,0 +1,39 @@
+import { z } from "zod";
+
+import { openai } from "@llamaindex/openai";
+import { agent } from "@llamaindex/workflow";
+import { tool } from "llamaindex";
+
+const weatherTool = tool({
+  name: "weatherTool",
+  description: "Get weather information",
+  parameters: z.object({
+    location: z.string(),
+  }),
+  execute: ({ location }) => {
+    return `The weather in ${location} is sunny. The temperature is 72 degrees. The humidity is 50%. The wind speed is 10 mph.`;
+  },
+});
+
+const responseSchema = z.object({
+  temperature: z.number(),
+  humidity: z.number(),
+  windSpeed: z.number(),
+});
+
+const myAgent = agent({
+  name: "myAgent",
+  tools: [weatherTool],
+  llm: openai({ model: "gpt-4.1-mini" }),
+});
+
+async function main() {
+  const result = await myAgent.run("What's the weather in Tokyo?", {
+    responseFormat: responseSchema,
+  });
+
+  console.log("result.data.result: ", result.data.result);
+  console.log("result.data.object: ", result.data.object);
+}
+
+main().catch(console.error);

examples/agents/tools/exec-object.ts

@@ -0,0 +1,48 @@
+import { openai } from "@llamaindex/openai";
+import { ChatMessage } from "llamaindex";
+import z from "zod";
+
+const llm = openai({ model: "gpt-4.1-mini" });
+
+const schema = z.object({
+  title: z.string(),
+  author: z.string(),
+  year: z.number(),
+});
+
+const messages: ChatMessage[] = [
+  {
+    role: "user",
+    content: `I have been reading La Divina Commedia by Dante Alighieri, published in 1321`,
+  },
+];
+
+async function main() {
+  {
+    // Non-streaming
+    const { object } = await llm.exec({ messages, responseFormat: schema });
+    console.log("Non-streaming object:", object);
+  }
+
+  {
+    // Streaming
+    let exit = false;
+    do {
+      const { stream, newMessages, toolCalls, object } = await llm.exec({
+        messages,
+        stream: true,
+        responseFormat: schema,
+      });
+
+      for await (const chunk of stream) {
+        console.log(chunk.delta);
+      }
+      console.log("Streaming object:", object);
+
+      messages.push(...newMessages());
+      exit = toolCalls.length === 0;
+    } while (!exit);
+  }
+}
+
+main().catch(console.error);

examples/agents/tools/response-format-exec.ts

@@ -13,6 +13,7 @@ const responseSchema = z.object({
 async function main() {
   const messages: ChatMessage[] = [];
   let toolCalls: ToolCall[] = [];
+  let object: z.infer<typeof responseSchema> | undefined;
   do {
     const result = await llm.exec({
       messages: [
@@ -27,13 +28,14 @@ async function main() {
       ],
       responseFormat: responseSchema,
     });
-    console.log(result.newMessages[0].content);
+    object = result.object;
     messages.push(...result.newMessages);
     toolCalls = result.toolCalls;
   } while (toolCalls.length == 0);
 
-  console.log(messages[1].content);
+  console.log(messages);
   console.log(toolCalls);
+  console.log(object);
 }
 
 main().catch(console.error);

packages/core/src/llms/base.ts

@@ -3,8 +3,17 @@ import type { JSONObject } from "../global";
 import { tool } from "../tools/";
 import { extractText } from "../utils/llms";
 import { streamConverter } from "../utils/stream";
-import { isZodSchema, safeParseSchema } from "../zod";
-import { callToolToMessage, getToolCallsFromResponse } from "./tool-call";
+import {
+  isZodSchema,
+  safeParseSchema,
+  type ZodInfer,
+  type ZodSchema,
+} from "../zod";
+import {
+  callToolToMessage,
+  getToolCallsFromResponse,
+  type CallToolToMessageResult,
+} from "./tool-call";
 import type {
   ChatMessage,
   ChatResponse,
@@ -22,6 +31,8 @@ import type {
   ToolCallLLMMessageOptions,
 } from "./type";
 
+const STRUCTURED_OUTPUT_TOOL_NAME = "format_output";
+
 export abstract class BaseLLM<
   AdditionalChatOptions extends object = object,
   AdditionalMessageOptions extends object = object,
@@ -77,33 +88,40 @@ export abstract class BaseLLM<
     >,
   ): Promise<ChatResponse<AdditionalMessageOptions>>;
 
-  exec(
+  exec<Z extends ZodSchema>(
     params: LLMChatParamsStreaming<
       AdditionalChatOptions,
-      AdditionalMessageOptions
+      AdditionalMessageOptions,
+      Z
     >,
-  ): Promise<ExecStreamResponse<AdditionalMessageOptions>>;
-  exec(
+  ): Promise<ExecStreamResponse<AdditionalMessageOptions, ZodInfer<Z>>>;
+  exec<Z extends ZodSchema>(
     params: LLMChatParamsNonStreaming<
       AdditionalChatOptions,
-      AdditionalMessageOptions
+      AdditionalMessageOptions,
+      Z
     >,
-  ): Promise<ExecResponse<AdditionalMessageOptions>>;
-  async exec(
+  ): Promise<ExecResponse<AdditionalMessageOptions, ZodInfer<Z>>>;
+  async exec<Z extends ZodSchema>(
     params:
-      | LLMChatParamsStreaming<AdditionalChatOptions, AdditionalMessageOptions>
+      | LLMChatParamsStreaming<
+          AdditionalChatOptions,
+          AdditionalMessageOptions,
+          Z
+        >
       | LLMChatParamsNonStreaming<
           AdditionalChatOptions,
-          AdditionalMessageOptions
+          AdditionalMessageOptions,
+          Z
         >,
   ): Promise<
-    | ExecResponse<AdditionalMessageOptions>
-    | ExecStreamResponse<AdditionalMessageOptions>
+    | ExecResponse<AdditionalMessageOptions, ZodInfer<Z>>
+    | ExecStreamResponse<AdditionalMessageOptions, ZodInfer<Z>>
   > {
     const responseFormat = params.responseFormat;
     if (typeof responseFormat != "undefined" && isZodSchema(responseFormat)) {
       const structuredTool = tool({
-        name: "format_output",
+        name: STRUCTURED_OUTPUT_TOOL_NAME,
         description: "Respond with a JSON object",
         parameters: responseFormat,
         execute: (args) => {
@@ -133,31 +151,40 @@ export abstract class BaseLLM<
     const response = await this.chat(params);
     newMessages.push(response.message);
     const toolCalls = getToolCallsFromResponse(response);
+
+    let structuredOutput: ZodInfer<Z> | undefined = undefined;
     if (params.tools && toolCalls.length > 0) {
       for (const toolCall of toolCalls) {
-        const toolResultMessage =
-          await callToolToMessage<AdditionalMessageOptions>(
-            params.tools,
-            toolCall,
-            logger,
-          );
+        const toolResultMessage = await callToolToMessage(
+          params.tools,
+          toolCall,
+          logger,
+        );
         if (toolResultMessage) {
-          newMessages.push(toolResultMessage);
+          newMessages.push(
+            toolResultMessage as ChatMessage<AdditionalMessageOptions>,
+          );
+        }
+        if (toolCall.name === STRUCTURED_OUTPUT_TOOL_NAME) {
+          structuredOutput = toolResultMessage?.options?.toolResult
+            ?.result as ZodInfer<Z>;
         }
       }
     }
     return {
       newMessages,
       toolCalls,
+      object: structuredOutput,
     };
   }
 
-  async streamExec(
+  async streamExec<Z extends ZodSchema>(
     params: LLMChatParamsStreaming<
       AdditionalChatOptions,
-      AdditionalMessageOptions
+      AdditionalMessageOptions,
+      Z
     >,
-  ): Promise<ExecStreamResponse<AdditionalMessageOptions>> {
+  ): Promise<ExecStreamResponse<AdditionalMessageOptions, ZodInfer<Z>>> {
     const logger = params.logger ?? emptyLogger;
     const responseStream = await this.chat(params);
     const iterator = responseStream[Symbol.asyncIterator]();
@@ -170,6 +197,16 @@ export abstract class BaseLLM<
       firstChunk?.options && "toolCall" in firstChunk.options;
 
     if (!hasToolCallsInFirst) {
+      // extract structured output from the last message
+      const lastMessage = params.messages[params.messages.length - 1];
+      const toolResult = (
+        lastMessage?.options as { toolResult: CallToolToMessageResult }
+      )?.toolResult;
+      const structuredOutput =
+        toolResult?.name === STRUCTURED_OUTPUT_TOOL_NAME
+          ? (toolResult.result as ZodInfer<Z>)
+          : undefined;
+
       let content = firstChunk?.delta ?? "";
       let finished = false;
       return {
@@ -201,6 +238,7 @@ export abstract class BaseLLM<
               ]
             : [];
         },
+        object: structuredOutput,
       };
     }
     // Helper function to process a chunk
@@ -252,15 +290,21 @@ export abstract class BaseLLM<
           toolCall: toolCalls,
         } as AdditionalMessageOptions,
       });
+      let structuredOutput: ZodInfer<Z> | undefined = undefined;
       for (const toolCall of toolCalls) {
-        const toolResultMessage =
-          await callToolToMessage<AdditionalMessageOptions>(
-            params.tools,
-            toolCall,
-            logger,
-          );
+        const toolResultMessage = await callToolToMessage(
+          params.tools,
+          toolCall,
+          logger,
+        );
         if (toolResultMessage) {
-          messages.push(toolResultMessage);
+          messages.push(
+            toolResultMessage as ChatMessage<AdditionalMessageOptions>,
+          );
+        }
+        if (toolCall.name === STRUCTURED_OUTPUT_TOOL_NAME) {
+          structuredOutput = toolResultMessage?.options?.toolResult
+            ?.result as ZodInfer<Z>;
         }
       }
       return {
@@ -269,6 +313,7 @@ export abstract class BaseLLM<
           return messages;
         },
         toolCalls,
+        object: structuredOutput,
       };
     } else {
       throw new Error("Cannot get tool calls from response");

packages/core/src/llms/mock.ts

@@ -8,6 +8,8 @@ import type {
   LLMCompletionParamsNonStreaming,
   LLMCompletionParamsStreaming,
   LLMMetadata,
+  ToolCall,
+  ToolCallLLMMessageOptions,
 } from "./type";
 
 export class MockLLM extends ToolCallLLM {
@@ -15,18 +17,29 @@ export class MockLLM extends ToolCallLLM {
   options: {
     timeBetweenToken: number;
     responseMessage: string;
+    mockToolCallResponse?: {
+      toolCalls: ToolCall[];
+      responseMessage?: string;
+    };
   };
-  supportToolCall: boolean = false;
+  supportToolCall: boolean = true;
 
   constructor(options?: {
     timeBetweenToken?: number;
     responseMessage?: string;
     metadata?: LLMMetadata;
+    mockToolCallResponse?: {
+      toolCalls: ToolCall[];
+      responseMessage?: string;
+    };
   }) {
     super();
     this.options = {
       timeBetweenToken: options?.timeBetweenToken ?? 20,
       responseMessage: options?.responseMessage ?? "This is a mock response",
+      ...(options?.mockToolCallResponse && {
+        mockToolCallResponse: options.mockToolCallResponse,
+      }),
     };
     this.metadata = options?.metadata ?? {
       model: "MockLLM",
@@ -34,7 +47,7 @@ export class MockLLM extends ToolCallLLM {
       topP: 0.5,
       contextWindow: 1024,
       tokenizer: undefined,
-      structuredOutput: false,
+      structuredOutput: true,
     };
   }
 
@@ -51,14 +64,46 @@ export class MockLLM extends ToolCallLLM {
   ): Promise<AsyncIterable<ChatResponseChunk> | ChatResponse<object>> {
     const responseMessage = this.options.responseMessage;
     const timeBetweenToken = this.options.timeBetweenToken;
+    const mockToolCallResponse = this.options.mockToolCallResponse;
+
+    // Check if we have tools and should simulate tool calls
+    const shouldSimulateToolCalls = params.tools && mockToolCallResponse;
 
     if (params.stream) {
-      return (async function* () {
-        for (const char of responseMessage) {
-          yield { delta: char, raw: {} };
-          await new Promise((resolve) => setTimeout(resolve, timeBetweenToken));
-        }
-      })();
+      if (shouldSimulateToolCalls) {
+        return (async function* () {
+          // First yield the tool call
+          yield {
+            delta: "",
+            raw: {},
+            options: {
+              toolCall: mockToolCallResponse.toolCalls,
+            } as ToolCallLLMMessageOptions,
+          };
+        })();
+      } else {
+        return (async function* () {
+          for (const char of responseMessage) {
+            yield { delta: char, raw: {} };
+            await new Promise((resolve) =>
+              setTimeout(resolve, timeBetweenToken),
+            );
+          }
+        })();
+      }
+    }
+
+    if (shouldSimulateToolCalls) {
+      return {
+        message: {
+          content: mockToolCallResponse.responseMessage || "",
+          role: "assistant",
+          options: {
+            toolCall: mockToolCallResponse.toolCalls,
+          } as ToolCallLLMMessageOptions,
+        },
+        raw: {},
+      };
     }
 
     return {

packages/core/src/llms/tool-call.ts

@@ -1,5 +1,6 @@
 import { type Logger } from "@llamaindex/env";
 import { callTool } from "../agent/utils.js";
+import type { JSONValue } from "../global";
 import { stringifyJSONToMessageContent } from "../utils";
 import type {
   BaseTool,
@@ -37,28 +38,30 @@ export const getToolCallsFromResponse = (
   return [];
 };
 
-export const callToolToMessage = async <
-  AdditionalMessageOptions extends object = object,
->(
+export type CallToolToMessageResult = {
+  id: string;
+  name: string;
+  result: JSONValue;
+  isError: boolean;
+};
+
+export const callToolToMessage = async (
   tools: BaseTool[],
   toolCall: ToolCall,
   logger: Logger,
-): Promise<ChatMessage<AdditionalMessageOptions> | null> => {
+): Promise<ChatMessage<{ toolResult: CallToolToMessageResult }>> => {
   const tool = tools?.find((t) => t.metadata.name === toolCall.name);
-
   const toolOutput = await callTool(tool, toolCall, logger);
-
-  const toolResultMessage: ChatMessage<AdditionalMessageOptions> = {
+  return {
     role: "user",
     content: stringifyJSONToMessageContent(toolOutput.output),
     options: {
       toolResult: {
         id: toolCall.id,
+        name: toolCall.name,
         result: toolOutput.output,
         isError: toolOutput.isError,
       },
-    } as AdditionalMessageOptions,
+    },
   };
-
-  return toolResultMessage;
 };

packages/core/src/llms/type.ts

@@ -99,18 +99,22 @@ export type ChatResponseChunk<
 
 export interface ExecResponse<
   AdditionalMessageOptions extends object = object,
+  O = JSONObject,
 > {
   newMessages: ChatMessage<AdditionalMessageOptions>[];
   toolCalls: ToolCall[];
+  object?: O | undefined;
 }
 
 export interface ExecStreamResponse<
   AdditionalMessageOptions extends object = object,
+  O = JSONObject,
 > {
   stream: AsyncIterable<ChatResponseChunk<AdditionalMessageOptions>>;
   // this is a function as while streaming, the assistant message is not ready yet - can be called after the stream is done
   newMessages(): ChatMessage<AdditionalMessageOptions>[];
   toolCalls: ToolCall[];
+  object?: O | undefined;
 }
 
 export interface CompletionResponse {
@@ -136,25 +140,36 @@ export type LLMMetadata = {
 export interface LLMChatParamsBase<
   AdditionalChatOptions extends object = object,
   AdditionalMessageOptions extends object = object,
+  Schema extends ZodSchema = ZodSchema,
 > {
   messages: ChatMessage<AdditionalMessageOptions>[];
   additionalChatOptions?: AdditionalChatOptions | undefined;
   tools?: BaseTool[] | undefined;
-  responseFormat?: ZodSchema | object | undefined;
+  responseFormat?: Schema | object | undefined;
   logger?: Logger | undefined;
 }
 
 export interface LLMChatParamsStreaming<
   AdditionalChatOptions extends object = object,
   AdditionalMessageOptions extends object = object,
-> extends LLMChatParamsBase<AdditionalChatOptions, AdditionalMessageOptions> {
+  Schema extends ZodSchema = ZodSchema,
+> extends LLMChatParamsBase<
+    AdditionalChatOptions,
+    AdditionalMessageOptions,
+    Schema
+  > {
   stream: true;
 }
 
 export interface LLMChatParamsNonStreaming<
   AdditionalChatOptions extends object = object,
   AdditionalMessageOptions extends object = object,
-> extends LLMChatParamsBase<AdditionalChatOptions, AdditionalMessageOptions> {
+  Schema extends ZodSchema = ZodSchema,
+> extends LLMChatParamsBase<
+    AdditionalChatOptions,
+    AdditionalMessageOptions,
+    Schema
+  > {
   stream?: false;
 }
 

packages/core/tests/llms/base.test.ts

@@ -1,7 +1,8 @@
 import { MockLLM } from "@llamaindex/core/llms/mock";
+import { tool } from "@llamaindex/core/tools";
 import { describe, expect, it } from "vitest";
+import { z } from "zod/v3";
 
-// TODO: add tests for tool calls
 describe("BaseLLM exec", () => {
   it("should stream text response when no tool call is made", async () => {
     const responseMessage = "This is a response message while streaming";
@@ -40,4 +41,89 @@ describe("BaseLLM exec", () => {
     ]);
     expect(toolCalls).toEqual([]);
   });
+
+  it("should handle tool calls with weather tool", async () => {
+    const weatherTool = tool({
+      name: "get_weather",
+      description: "Get the current weather for a location",
+      parameters: z.object({
+        address: z.string().describe("The address"),
+      }),
+      execute: ({ address }) => {
+        return `It's sunny in ${address}!`;
+      },
+    });
+
+    const mockToolCallInput = { address: "San Francisco" };
+    const expectedWeatherResult = "It's sunny in San Francisco!";
+
+    const llm = new MockLLM({
+      mockToolCallResponse: {
+        toolCalls: [
+          {
+            id: "weather-tool-call-id",
+            name: "get_weather",
+            input: mockToolCallInput,
+          },
+        ],
+        responseMessage: "",
+      },
+    });
+
+    const { newMessages, toolCalls } = await llm.exec({
+      messages: [
+        { content: "What's the weather in San Francisco?", role: "user" },
+      ],
+      tools: [weatherTool],
+    });
+
+    expect(toolCalls).toHaveLength(1);
+    expect(toolCalls[0]!.name).toBe("get_weather");
+    expect(toolCalls[0]!.input).toEqual(mockToolCallInput);
+    expect(newMessages).toHaveLength(2); // assistant message + tool result message
+
+    // Check that the tool was executed and result is in the tool result message
+    const toolResultMessage = newMessages[1];
+    expect(toolResultMessage!.content).toBe(expectedWeatherResult);
+    expect(toolResultMessage!.role).toBe("user");
+    expect(toolResultMessage!.options).toHaveProperty("toolResult");
+  });
+
+  it("should return structured output with responseFormat", async () => {
+    const schema = z.object({
+      title: z.string(),
+      author: z.string(),
+      year: z.number(),
+    });
+
+    const mockObject = {
+      title: "La Divina Commedia",
+      author: "Dante Alighieri",
+      year: 1321,
+    };
+
+    const llm = new MockLLM({
+      mockToolCallResponse: {
+        toolCalls: [
+          {
+            id: "test-tool-call-id",
+            name: "format_output",
+            input: mockObject,
+          },
+        ],
+        responseMessage: "",
+      },
+    });
+
+    const { newMessages, toolCalls, object } = await llm.exec({
+      messages: [{ content: "Extract book info", role: "user" }],
+      responseFormat: schema,
+    });
+
+    expect(toolCalls).toHaveLength(1);
+    expect(toolCalls[0]!.name).toBe("format_output");
+    expect(object).toBeDefined();
+    expect(object).toEqual(mockObject);
+    expect(newMessages).toHaveLength(2); // assistant message + tool result message
+  });
 });

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

@@ -1,15 +1,16 @@
 import { callTool } from "@llamaindex/core/agent";
-import type { JSONValue } from "@llamaindex/core/global";
+import type { JSONObject, 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 { z } from "@llamaindex/core/zod";
+import { z, type ZodInfer, type ZodSchema } from "@llamaindex/core/zod";
 import { consoleLogger, emptyLogger, type Logger } from "@llamaindex/env";
 import {
   createWorkflow,
   workflowEvent,
+  WorkflowStream,
   type Handler,
   type Workflow,
   type WorkflowContext,
@@ -58,10 +59,11 @@ export const startAgentEvent = workflowEvent<
   debugLabel: "llamaindex-start",
 });
 
-export type AgentResultData = {
+export type AgentResultData<O = JSONObject> = {
   result: MessageContent;
   message: ChatMessage;
   state?: AgentWorkflowState | undefined;
+  object?: O | undefined;
 };
 export const stopAgentEvent = workflowEvent<AgentResultData, "llamaindex-stop">(
   {
@@ -456,10 +458,20 @@ export class AgentWorkflow implements Workflow {
       };
       const content = await agent.finalize(context.state, agentOutput);
 
+      // if responseFormat is set, use the structured tool to parse the response
+      let object: JSONObject | undefined = undefined;
+      if (context.state.responseFormat) {
+        object = await agent.getStructuredOutput(
+          context.state.responseFormat,
+          content.response,
+        );
+      }
+
       return stopAgentEvent.with({
         message: content.response,
         result: content.response.content,
         state: context.state,
+        object,
       });
     }
 
@@ -636,19 +648,23 @@ export class AgentWorkflow implements Workflow {
     };
   }
 
-  runStream(
+  runStream<Z extends ZodSchema>(
     userInput: MessageContent,
     params?: {
       chatHistory?: ChatMessage[];
       state?: AgentWorkflowState;
+      responseFormat?: Z;
     },
-  ) {
+  ): WorkflowStream<WorkflowEventData<AgentResultData<ZodInfer<Z>>>> {
     if (this.agents.size === 0) {
       throw new Error("No agents added to workflow");
     }
     const state = params?.state ?? this.createInitialState();
 
-    const { sendEvent, stream } = this.workflow.createContext(state);
+    const { sendEvent, stream } = this.workflow.createContext({
+      ...state,
+      responseFormat: params?.responseFormat,
+    });
     sendEvent(
       startAgentEvent.with({
         userInput: userInput,
@@ -658,13 +674,14 @@ export class AgentWorkflow implements Workflow {
     return stream.until(stopAgentEvent);
   }
 
-  async run(
+  async run<Z extends ZodSchema>(
     userInput: MessageContent,
     params?: {
       chatHistory?: ChatMessage[];
       state?: AgentWorkflowState;
+      responseFormat?: Z;
     },
-  ): Promise<WorkflowEventData<AgentResultData>> {
+  ): Promise<WorkflowEventData<AgentResultData<ZodInfer<Z>>>> {
     const finalEvent = (await this.runStream(userInput, params).toArray()).at(
       -1,
     );

packages/workflow/src/agent/base.ts

@@ -1,5 +1,7 @@
+import type { JSONObject } from "@llamaindex/core/global";
 import type { BaseToolWithCall, ChatMessage, LLM } from "@llamaindex/core/llms";
 import { Memory } from "@llamaindex/core/memory";
+import type { ZodSchema } from "@llamaindex/core/zod";
 import type { WorkflowContext } from "@llamaindex/workflow-core";
 import type { AgentOutput, AgentToolCallResult } from "./events";
 
@@ -9,6 +11,7 @@ export type AgentWorkflowState = {
   agents: string[];
   currentAgentName: string;
   nextAgentName?: string | null;
+  responseFormat?: ZodSchema | undefined;
 };
 
 /**
@@ -22,6 +25,14 @@ export interface BaseWorkflowAgent {
   readonly llm: LLM;
   readonly canHandoffTo: string[];
 
+  /**
+   * Take the final response and convert it to a structured output
+   */
+  getStructuredOutput(
+    responseFormat: ZodSchema,
+    response: ChatMessage,
+  ): Promise<JSONObject>;
+
   /**
    * Take a single step with the agent
    * Using memory directly to get messages instead of requiring them to be passed in

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

@@ -7,7 +7,12 @@ import {
   type ChatResponseChunk,
 } from "@llamaindex/core/llms";
 import { tool } from "@llamaindex/core/tools";
-import { isZodV3Schema, z, zodToJsonSchema } from "@llamaindex/core/zod";
+import {
+  isZodV3Schema,
+  z,
+  zodToJsonSchema,
+  type ZodSchema,
+} from "@llamaindex/core/zod";
 import {
   type WorkflowContext,
   type WorkflowEvent,
@@ -131,6 +136,14 @@ export class FunctionAgent implements BaseWorkflowAgent {
     this.canHandoffTo = this.initHandOffNames(canHandoffTo ?? []);
   }
 
+  async getStructuredOutput(responseFormat: ZodSchema, response: ChatMessage) {
+    const { object } = await this.llm.exec({
+      messages: [response],
+      responseFormat,
+    });
+    return object ?? {};
+  }
+
   private initHandOffNames(
     handoffTo: string[] | BaseWorkflowAgent[] | AgentWorkflow[],
   ): string[] {

packages/workflow/test/agent-workflow.test.ts

@@ -249,6 +249,58 @@ describe("agent", () => {
     expect(fileMessage.mimeType).toEqual("application/pdf");
     expect(fileMessage.data).toEqual(pdfBuffer);
   });
+
+  test("agent with responseFormat returns structured output", async () => {
+    const weatherTool = FunctionTool.from(
+      ({ location }: { location: string }) => {
+        return `The weather in ${location} is sunny. The temperature is 72 degrees. The humidity is 50%. The wind speed is 10 mph.`;
+      },
+      {
+        name: "weatherTool",
+        description: "Get weather information",
+        parameters: z.object({
+          location: z.string(),
+        }),
+      },
+    );
+
+    const responseSchema = z.object({
+      temperature: z.number(),
+      humidity: z.number(),
+      windSpeed: z.number(),
+    });
+
+    // Define the expected structured output
+    const expectedObject = {
+      temperature: 72,
+      humidity: 50,
+      windSpeed: 10,
+    };
+
+    // Create a mock LLM with structured output support
+    const mockLLM = new MockLLM({
+      responseMessage: "Weather information retrieved and formatted.",
+    });
+
+    mockLLM.exec = vi.fn().mockResolvedValue({
+      object: expectedObject,
+    });
+
+    // Spy on the weather tool to verify it gets called
+    vi.spyOn(weatherTool, "call");
+
+    const myAgent = agent({
+      name: "myAgent",
+      tools: [weatherTool],
+      llm: mockLLM,
+    });
+
+    const result = await myAgent.run("What's the weather in Tokyo?", {
+      responseFormat: responseSchema,
+    });
+
+    expect(result.data.object).toMatchObject(expectedObject);
+  });
 });
 
 describe("Multiple agents", () => {