How to stream LangChain, LangGraph's final generation

[Mohan-Sharma]

How do we stream LangGraph's final generation using vercel's AI sdk? It works pretty well with LangChain LCEL as explain in this blog. But how do we do this using LangGraph and vercel ai?

I am getting error TypeError: stream.pipeThrough is not a function while doing LangChainAdapter.toDataStreamResponse(final_generation) in handling the graph's stream

When user clicks submit the below POST is invoked, which in turn invokes the graph flow. This is where the error occurs.

//src/routes/api/chat
import { LangChainAdapter } from 'ai';
import type { RequestHandler } from './$types';
import type { Message } from 'ai/svelte';
import { Workflow } from '$lib/server/graph/workflow';

export const POST: RequestHandler = async ({ request }) => {

    const { messages }: { messages: Message[] } = await request.json();
    let final_generation =  null;
    const eventStream = await Workflow.getCompiledStateGraph().streamEvents({'question': messages.pop(), 'chat_history': messages}, { version: "v2"});
    return LangChainAdapter.toDataStreamResponse(eventStream);
};

A simple CompiledStateGraph

export class Workflow {
    // @ts-ignore
    private static COMPILED_STATE_GRAPH: CompiledStateGraph | null = null;
    
    private constructor() {}
    
    public static getCompiledStateGraph() {
        if (!Workflow.COMPILED_STATE_GRAPH) {
            const graph = new StateGraph(State)
            .addNode("retrieve", retrieveDocuments)
            .addNode("llm_search", generate)
            .addConditionalEdges(START, routeQuestion)
            .addEdge("llm_search", END)
            .addEdge("retrieve", END);
            Workflow.COMPILED_STATE_GRAPH  = graph.compile();
        }
        return Workflow.COMPILED_STATE_GRAPH;
    }
}

generate node

import { State } from '$lib/server/graph/state';
import { LLMClient } from '$lib/server/llm-client';
import { StringOutputParser } from '@langchain/core/output_parsers';
import { ChatPromptTemplate } from '@langchain/core/prompts';

export const generate = async (state: typeof State.State): Promise<Partial<typeof State.State>> => {
    console.log("---LLM Inference---");
    const PROMPT_TEMPLATE = 'You are a helpful assistant!';
    const prompt = ChatPromptTemplate.fromMessages([
        ['system', PROMPT_TEMPLATE],
        ['human', "{question}"],
    ]);
    const chain = prompt.pipe(LLMClient.getClient()).pipe(new StringOutputParser());
    const stream = await chain.stream({question: state.question});
    return { generation: stream };
};

State

import type { Document } from "@langchain/core/documents";
import { Annotation } from "@langchain/langgraph";
import type { Message } from "ai/svelte";

export const State = Annotation.Root({
    chat_history: Annotation<{messages:Message[]}>,
    question: Annotation<string>,
    generation: Annotation<string>,
    documents: Annotation<Document[]>({
        reducer: (_, y) => y,
        default: () => [],
    })
})

The input binding, handling the user message is done by useChat() of vercel ai

Rendering part, here the HumanInput component binds to the input and submits the user query to handleSubmit which in turn invokes our previous POST server function.

<script lang="ts">
    import HumanInput from "$lib/components/HumanInput.svelte";
    import MaxWidthWrapper from '$lib/components/MaxWidthWrapper.svelte';
    import DisplayMessages from "$lib/components/DisplayMessages.svelte";
    import {useChat} from '@ai-sdk/svelte';
    const { input, handleSubmit, messages } = useChat();
</script>

<div class="flex flex-col h-screen">
    <div class="flex-grow overflow-hidden">
        <MaxWidthWrapper class_="h-full flex flex-col">
            <DisplayMessages {messages} />
            <HumanInput {input} {handleSubmit}/>
        </MaxWidthWrapper>
    </div>
</div>