Langgraph streaming with structured output

[markpeace]

Hi! I'm trying to set up a Langgraph agent that streams two bits of structured output (a response, and a flag that content has been updated). I can get streaming work, but when I add withStructuredOutput I get errors ... I'm clearly not adapting the output properly, but can't work out why/how! Heres the code, with the problematic bit commented

import { StateGraph, MessagesAnnotation } from "@langchain/langgraph";
import { ToolNode } from "@langchain/langgraph/prebuilt";
import { AIMessage } from "@langchain/core/messages";
import { ChatOpenAI } from "@langchain/openai";
import { z } from "zod";
import { prompts } from "./prompts/compiled";
import { zodToJsonSchema as toJson } from "zod-to-json-schema";
import { generate_tools } from "./tools";
import { ChatPromptTemplate, MessagesPlaceholder } from "@langchain/core/prompts";

interface AgentData {
  history: Message[];
  luminary: string;
  input: string;
  core_knowledge: Record<string, unknown>;
  user: string;
}

interface Message {
  role: string;
  content: string;
}

class Agent {
  luminary: string;
  compiled: any;
  history: any = [];
  input: string;
  core_knowledge: any = {};
  user: string;
  tools: any = [];

  constructor(data: AgentData) {
    this.history = data.history;
    this.luminary = data.luminary;
    this.input = data.input;
    this.core_knowledge = data.core_knowledge;
    this.user = data.user;

    const tools = generate_tools(prompts[this.luminary], data.user, this.luminary);
    this.tools = tools.concat(prompts[this.luminary]?.tools || []);

    const model = new ChatOpenAI({
      model: "gpt-4o-mini",
    })

      //This bit is causing an error. If I remove it, things are fine...
      .withStructuredOutput(z.object({
        response: z.string().describe("Your final response"),
        update_core_memory: z.boolean().describe("true if new information has been provided about the user"),
      }))

      .bind(this.tools);

    this.compiled = model;
  }

  load_history: any = async (state: typeof MessagesAnnotation.State) => {
    console.log("--- " + this.luminary + " PRELOADING MESSAGE HISTORY ---");

    const history = this.history.map(({ role, content }: any) => [role, content]);

    const prompt = await ChatPromptTemplate.fromMessages([
      ["system", `
        You are an AI chatbot with a specialism as a {role}.
        Here is today's date as an ISO string: {date}
        The user's email address is {user}
        {base_prompt}
        ### Core knowledge as a JSON object: {core_knowledge}
        ###
        Current chat history: {chat_history}
        ###
        User's input: {input}
      `],
      new MessagesPlaceholder("chat_history"),
      ["user", "{input}"],
      ["system", "Please provide your response now"]
    ]).formatMessages({
      role: prompts[this.luminary].role,
      date: new Date().toISOString(),
      user: this.user,
      base_prompt: prompts[this.luminary].base,
      core_knowledge: this.core_knowledge,
      chat_history: history,
      input: this.input
    });

    return { messages: prompt };
  }

  invoke(messages: any) {
    console.log("--- " + this.luminary + " INVOKED ---");
    return this.compiled.invoke(messages);
  }
}

// Refactored stream handling with structured output (response + core_knowledge)
const updateCoreKnowledge = (data: any) => {
  if (data.output.name === "update_core_knowledge") {
    return JSON.stringify(JSON.parse(data.input.input).query);
  }
  return "";
};

export async function POST(request: Request) {
  const data = await request.json();
  const agent = new Agent(data);

  const graph = new StateGraph(MessagesAnnotation)
    .addNode("init", agent.load_history)
    .addNode("agent", async (state: typeof MessagesAnnotation.State) => {
      const messages = state.messages;
      const response = await agent.invoke(messages);
      return { messages: [response] };
    })
    .addNode("tools", new ToolNode(agent.tools))
    .addEdge("__start__", "init")
    .addEdge("init", "agent")
    .addConditionalEdges("agent", (state: typeof MessagesAnnotation.State) => {
      const { messages } = state;
      const lastMessage: AIMessage = messages[messages.length - 1];
      return lastMessage.tool_calls?.length ? "tools" : "__end__";
    })
    .addEdge("tools", "agent")
    .compile();

  const events = graph.streamEvents({ messages: [] }, { version: "v2" });

  return doStream(events);
}

function doStream(eventStream: any) {
  console.log("--- PROCESSING FINAL STREAM ---");
  const textEncoder = new TextEncoder();
  let core_knowledge = "";

  const transformStream = new ReadableStream({
    async start(controller) {
      for await (const { event, data } of eventStream) {
        if (event === "on_tool_end") {
          core_knowledge = updateCoreKnowledge(data);
        }

        if (event === "on_chat_model_stream" && data.chunk.content) {
          controller.enqueue(textEncoder.encode(data.chunk.content));
        }
      }

      // Once the stream ends, we return the final structured JSON object.
      const response = {
        response: core_knowledge ? JSON.stringify({ chunk: core_knowledge }) : "",
        update_core_knowledge: core_knowledge || ""
      };

      controller.enqueue(textEncoder.encode(JSON.stringify(response)));
      controller.close();
    }
  });

  return new Response(transformStream, {
    headers: { "Content-Type": "application/json" }
  });
}

[hinthornw]

What errors do you get?