vercel / ai

Build AI-powered applications with React, Svelte, Vue, and Solid
https://sdk.vercel.ai/docs
Other
8.72k stars 1.23k forks source link

Unable to only stream back the last chain of my langchain chain #788

Open ElectricCodeGuy opened 7 months ago

ElectricCodeGuy commented 7 months ago

Description

I am encountering an issue with the AI SDK where I am unable to stream back only the last output from my LangChain chain. The SDK seems to be streaming all outputs, including intermediate ones, which is not the intended behavior. I am using the LangChainStream function provided by the SDK in a Next.js environment with TypeScript.

Set up a LangChain using the SDK. Implement handlers within the LangChainStream, specifically focusing on the handleChainEnd handler. Override the handleChainEnd handler to conditionally write to the stream only when the output_text is present. Call the chain with these modified handlers. Observe that all outputs, including intermediate ones, are streamed back instead of only the final output.

Code example

import OpenAI from 'openai';
import { LangChainStream, StreamingTextResponse } from 'ai';
// ... other necessary imports

export async function processChain(req: Request) {
  // ... setup and initialization code

  const { stream, writer, handlers } = LangChainStream();

  handlers.handleChainEnd = async (_outputs: any) => {
    if (_outputs && 'output_text' in _outputs) {
      writer.write(_outputs.output_text); // Expecting only this to be streamed
    }
  };

  const combineDocsChain = loadSummarizationChain(OpenAI, {
    // ... configuration
  });

  combineDocsChain.call({ input_documents: docsSummary }, [handlers]);

  return new StreamingTextResponse(stream, {
    status: 200,
    headers: { 'Content-Type': 'application/json' }
  });
}

Also Tried:

    const { stream, writer, handlers } = LangChainStream({
      onFinal: (completion: string) => {
        // Stream only the final output
        writer.write(completion);
      }
    });

Additional context

Expected Behavior: The stream should only return the output_text from the final output of the LangChain chain.

Actual Behavior: The stream is returning all outputs, including intermediate ones, not just the output_text.

If I await the .call then we block the output stream. Im not sure it is supported, but is it possible to await the intermedia chains in the .call and then only stream the final output_text ?

The only other solution that seems to be working is calling a new chain and applying the handlers to that one and awaiting the chain with the intermedia steps. But This adds additional delays to the output stream.

lgrammel commented 7 months ago

@ElectricCodeGuy thanks for reporting the issue. A reproduction scenario would be tremendously helpful for debugging this bug.

Would you mind sharing such a scenario, e.g. by

  1. forking the ai sdk
  2. adding your reproduction to examples/next-langchain
  3. share the branch on your fork with the reproduction