The main concept we need to understand here is how Vercel AI and LangChain handles the messages. While AI SDK understands Message from ai package, LangChain deals with subtypes of BaseMessage from @langchain/core/messages package.

The trick to solving this issue was to translate the message between these two formats.

//src/api/chat/+server.ts
import { LangChainAdapter } from 'ai';
import type { Message } from 'ai/svelte';
import { Workflow } from '$lib/server/graph/workflow';
import { convertLangChainMessageToVercelMessage, convertVercelMessageToLangChainMessage } from '$lib/utils/utility';

export const POST = async ({ request, params }) => {
    const config = { configurable: { thread_id: params.id}, version: "v2" };
    const messages: { messages: Message[] } = await request.json();
    const userQuery = messages.messages[messages.messages.length - 1].content;
    
    let history = (messages.messages ?? [])
    .slice(0, -1)
    .filter(
        (message: Message) =>
            message.role === 'user' || message.role === 'assistant'
    )
    .map(convertVercelMessageToLangChainMessage);
    
    let compiledStateGraph = Workflow.getCompiledStateGraph();
    const stream = await compiledStateGraph.streamEvents({question: userQuery,  messages: history}, config);
    const transformStream = new ReadableStream({
        async start(controller) {
            for await (const { event, data, tags } of stream) {
                if (event === 'on_chat_model_stream') {
                    if (!!data.chunk.content  && tags.includes("llm_inference")) {
                        const aiMessage = convertLangChainMessageToVercelMessage(data.chunk);
                        controller.enqueue(aiMessage);
                    }
                }
            }
            controller.close();
        }
    });
    return LangChainAdapter.toDataStreamResponse(transformStream);
};

Before we invoke the graph flow, I am extracting the current user query from the message array of useChat() and then converting the rest of the messages to langchain understandable format using convertVercelMessageToLangChainMessage function. Similarly once i receive the AIMesssageChunks from LangChain streams, I am converting them back to vercel ai understandable format using convertLangChainMessageToVercelMessage function before returning the stream back for useChat() to handle the new messages.

import type { Message } from 'ai/svelte';
import { AIMessage, BaseMessage, ChatMessage, HumanMessage } from '@langchain/core/messages';

/**
 * Converts a Vercel message to a LangChain message.
 * @param message - The message to convert.
 * @returns The converted LangChain message.
 */
export const convertVercelMessageToLangChainMessage = (message: Message): BaseMessage => {
  switch (message.role) {
    case 'user':
      return new HumanMessage({ content: message.content });
    case 'assistant':
      return new AIMessage({ content: message.content });
    default:
      return new ChatMessage({ content: message.content, role: message.role });
  }
};

/**
 * Converts a LangChain message to a Vercel message.
 * @param message - The message to convert.
 * @returns The converted Vercel message.
 */
export const convertLangChainMessageToVercelMessage = (message: BaseMessage) => {
  switch (message.getType()) {
    case 'human':
      return { content: message.content, role: 'user' };
    case 'ai':
      return {
        content: message.content,
        role: 'assistant',
        tool_calls: (message as AIMessage).tool_calls
      };
    default:
      return { content: message.content, role: message._getType() };
  }
};

Also notice if (!!data.chunk.content && tags.includes("llm_inference")) this is how we can filter the last generation. As during the last node execution of graph we can tag the LLM with configs(tags) which we can later use to get the result of that node's execution.

export const llmInference = async (state: typeof GraphState.State) => {
    console.log("---LLM Inference---");
    
    const PROMPT_TEMPLATE = 'You are a helpful assistant! Please answer the user query. Use the chat history to provide context.';
    
    const prompt = ChatPromptTemplate.fromMessages([
        ['system', PROMPT_TEMPLATE],
        ['human', "{question}"],
        ['human', "Chat History: {messages}"],
    ]);
    
    const inferenceChain = prompt.pipe(LLMClient.getClient().withConfig({ tags: ["llm_inference"]}));
    const generation = await inferenceChain.invoke(state);
    return { messages: [generation] };
};