withStructuredOutput streaming return only one chunk (inconsistent with Langchain Python)

azhong-git commented 1 month ago

Checked other resources

[X] I added a very descriptive title to this issue.
[X] I searched the LangChain.js documentation with the integrated search.
[X] I used the GitHub search to find a similar question and didn't find it.
[X] I am sure that this is a bug in LangChain.js rather than my code.
[X] The bug is not resolved by updating to the latest stable version of LangChain (or the specific integration package).

Example Code

I noticed that withStructuredOutput only return one chunk in stream mode, unlike in Langchain Python.

  const model = new ChatOpenAI({ streaming: true });
  const schema = z.object({
    items: z.array(
      z.object({
        title: z.string(),
        description: z.string().optional(),
      })
    ),
  });
  const llm = model.withStructuredOutput(schema);
  const chain = prompt.pipe(llm);
  if (streaming) {
    const stream = await chain.stream({ text: text });
    return stream;
  } else {
    const answerText = chain.invoke({ text: text });
    return answerText;
  }

Error Message and Stack Trace (if applicable)

No response

Description

I expect to see JSON output being streamed in more than one chunks like in the Langchain Python doc

In my example, I am extracting a list of objects and convert to JSON and stream to user once the first few objects in the list are available can be really helpful.

  const schema = z.object({
    items: z.array(
      z.object({
        title: z.string(),
        description: z.string().optional(),
      })
    ),
  });

System Info

mac Node v20.5.1 langchain@0.2.13

dosubot[bot] commented 1 month ago

Hey @azhong-git! 👋 I'm here to help you with any bugs, questions, or contributions you have for the LangChain.js repository. While you're waiting for a human maintainer, feel free to ask me anything!

To achieve streaming JSON output in multiple chunks with withStructuredOutput in LangChain.js, you can use the JsonOutputParser class. This class is designed to handle streaming JSON data in chunks, similar to the behavior in LangChain Python.

Here's an example of how you can modify your code to use JsonOutputParser:

import { z } from "zod";
import { zodToJsonSchema } from "zod-to-json-schema";
import { ChatOpenAI } from "@langchain/openai";
import { ChatPromptTemplate } from "@langchain/core/prompts";
import { JsonOutputParser } from "@langchain/core/output_parsers";
import { applyPatch } from "@langchain/core/utils/json_patch";

const schema = z.object({
  items: z.array(
    z.object({
      title: z.string(),
      description: z.string().optional(),
    })
  ),
});

const modelParams = {
  functions: [
    {
      name: "items",
      description: "A list of items",
      parameters: zodToJsonSchema(schema),
    },
  ],
  function_call: { name: "items" },
};

async function streamJsonOutput(text: string) {
  const prompt = ChatPromptTemplate.fromTemplate(`Extract items from: {text}`);
  const model = new ChatOpenAI({
    temperature: 0,
    streaming: true,
  }).bind(modelParams);

  const parser = new JsonOutputParser({ diff: true });
  const chain = prompt.pipe(model).pipe(parser);

  const stream = await chain.stream({ text });

  const chunks = [];
  let aggregate: any = {};
  for await (const chunk of stream) {
    console.log(chunk);
    chunks.push(chunk);
    aggregate = applyPatch(aggregate, chunk).newDocument;
  }

  console.log(aggregate);
  return chunks;
}

// Usage
streamJsonOutput("Your input text here");

In this example, the JsonOutputParser class is used to read from a stream and process the JSON data in chunks. The applyPatch method is called in a loop to aggregate each chunk until the stream is done. Each chunk can then be processed as needed.

This should help you achieve the desired behavior of streaming JSON output in multiple chunks, similar to LangChain Python ^[1]^[2]^[3]^[4]^[5].

To continue talking to Dosu, mention @dosu.

diksipav commented 1 month ago

@azhong-git did this work for you? I have the same problem and I also only get one chunk. If I try stream_tool_client also get one chunk only. Suggestion here by the bot gives me a stream that is completed right away, I don't get any chunks.

azhong-git commented 1 month ago

No it does not work for me in streaming mode. I ended up directly using string as output.

diksipav commented 1 month ago

Thank you. I lost a day on this. Also seems that langchain doesn't have discord help anymore.

jacoblee93 commented 2 days ago

Hey folks, sorry about missing this - should be fixed in the latest version of core 0.3 and 0.2 for OpenAI and a few others.

langchain-ai / langchainjs