Langchain Agent Streaming LLM & Function Call & Tools

[nikohann]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

import { NextRequest, NextResponse } from "next/server";
import { StreamingTextResponse, LangChainStream } from "ai";
import { zodToJsonSchema } from "zod-to-json-schema";
import { z } from "zod";
import {
    type BaseMessage,
    AIMessage,
    FunctionMessage,
} from "@langchain/core/messages";
import { RunnableSequence } from "@langchain/core/runnables";
import { ChatPromptTemplate, MessagesPlaceholder } from "@langchain/core/prompts";
import { ChatOpenAI } from "@langchain/openai";
import {
    AgentExecutor,
    type AgentFinish,
    type AgentStep,
} from "langchain/agents";
import { DynamicTool } from "@langchain/core/tools";
import type { FunctionsAgentAction } from "langchain/agents/openai/output_parser";
import { convertToOpenAIFunction } from "@langchain/core/utils/function_calling";

import { TavilySearchAPIRetriever } from "@langchain/community/retrievers/tavily_search_api";

export const runtime = "edge";

export async function POST(req: NextRequest) {
    try {

        const json = await req.json()
        const { messages } = json;

        const llm = new ChatOpenAI({
            modelName: "gpt-4-1106-preview",
            streaming: true
        });

        const searchTool = new DynamicTool({
            name: "web-search-tool",
            description: "Tool for getting the latest information from the web",
            func: async (searchQuery: string, runManager) => {
                const retriever = new TavilySearchAPIRetriever();
                const docs = await retriever.getRelevantDocuments(searchQuery, runManager?.getChild());
                return docs.map((doc) => doc.pageContent).join("\n-----\n");
            },
        });

        const prompt = ChatPromptTemplate.fromMessages([
            [
                "system",
                "You are a helpful assistant. You must always call one of the provided tools.",
            ],
            ["user", "{input}"],
            new MessagesPlaceholder("agent_scratchpad"),
        ]);

        const responseSchema = z.object({
            answer: z.string().describe("The final answer to return to the user"),
            sources: z
                .array(z.string())
                .describe(
                    "List of page chunks that contain answer to the question. Only include a page chunk if it contains relevant information"
                ),
        });

        const responseOpenAIFunction = {
            name: "response",
            description: "Return the response to the user",
            parameters: zodToJsonSchema(responseSchema),
        };

        const structuredOutputParser = (
            message: AIMessage
        ): FunctionsAgentAction | AgentFinish => {

            if (message.content && typeof message.content !== "string") {
                throw new Error("This agent cannot parse non-string model responses.");
            }

            if (message.additional_kwargs.function_call) {
                const { function_call } = message.additional_kwargs;
                try {
                    const toolInput = function_call.arguments
                        ? JSON.parse(function_call.arguments)
                        : {};
                    // If the function call name is `response` then we know it's used our final
                    // response function and can return an instance of `AgentFinish`
                    if (function_call.name === "response") {
                        return { returnValues: { ...toolInput }, log: message.content };
                    }

                    return {
                        tool: function_call.name,
                        toolInput,
                        log: `Invoking "${function_call.name}" with ${function_call.arguments ?? "{}"
                            }\n${message.content}`,
                        messageLog: [message],
                    };
                } catch (error) {
                    throw new Error(
                        `Failed to parse function arguments from chat model response. Text: "${function_call.arguments}". ${error}`
                    );
                }
            } else {
                return {
                    returnValues: { output: message.content },
                    log: message.content,
                };
            }
        };

        const formatAgentSteps = (steps: AgentStep[]): BaseMessage[] =>
            steps.flatMap(({ action, observation }) => {
                if ("messageLog" in action && action.messageLog !== undefined) {
                    const log = action.messageLog as BaseMessage[];
                    return log.concat(new FunctionMessage(observation, action.tool));
                } else {
                    return [new AIMessage(action.log)];
                }
            });

        const llmWithTools = llm.bind({
            functions: [convertToOpenAIFunction(searchTool), responseOpenAIFunction],
        });
        /** Create the runnable */
        const runnableAgent = RunnableSequence.from<{
            input: string;
            steps: Array<AgentStep>;
        }>([
            {
                input: (i) => i.input,
                agent_scratchpad: (i) => formatAgentSteps(i.steps),
            },
            prompt,
            llmWithTools,
            structuredOutputParser,
        ]);

        const executor = AgentExecutor.fromAgentAndTools({
            agent: runnableAgent,
            tools: [searchTool],
        });

        const { handlers, stream } = LangChainStream();

        executor.invoke({
            input: messages[messages.length - 1].content,
        }, {
            callbacks: [handlers]
        });

        return new StreamingTextResponse(stream);

    } catch (e: any) {
        console.log((e as Error).stack);
        return NextResponse.json({ error: e.message }, { status: e.status ?? 500 });
    }
}

Description

https://js.langchain.com/docs/modules/agents/how_to/agent_structured

For some reason agent executor does not want to stream response to LangChainStream and StreamingTextResponse.

Without streaming it works:

export async function POST(req: NextRequest) {
    try {

        const json = await req.json()
        const { messages } = json;

        const llm = new ChatOpenAI({
            modelName: "gpt-4-1106-preview",
        });

        const searchTool = new DynamicTool({
            name: "web-search-tool",
            description: "Tool for getting the latest information from the web",
            func: async (searchQuery: string, runManager) => {
                const retriever = new TavilySearchAPIRetriever();
                const docs = await retriever.getRelevantDocuments(searchQuery, runManager?.getChild());
                return docs.map((doc) => doc.pageContent).join("\n-----\n");
            },
        });

        const prompt = ChatPromptTemplate.fromMessages([
            [
                "system",
                "You are a helpful assistant. You must always call one of the provided tools.",
            ],
            ["user", "{input}"],
            new MessagesPlaceholder("agent_scratchpad"),
        ]);

        const responseSchema = z.object({
            answer: z.string().describe("The final answer to return to the user"),
            sources: z
                .array(z.string())
                .describe(
                    "List of page chunks that contain answer to the question. Only include a page chunk if it contains relevant information"
                ),
        });

        const responseOpenAIFunction = {
            name: "response",
            description: "Return the response to the user",
            parameters: zodToJsonSchema(responseSchema),
        };

        const structuredOutputParser = (
            message: AIMessage
        ): FunctionsAgentAction | AgentFinish => {

            if (message.content && typeof message.content !== "string") {
                throw new Error("This agent cannot parse non-string model responses.");
            }

            if (message.additional_kwargs.function_call) {
                const { function_call } = message.additional_kwargs;
                try {
                    const toolInput = function_call.arguments
                        ? JSON.parse(function_call.arguments)
                        : {};
                    // If the function call name is `response` then we know it's used our final
                    // response function and can return an instance of `AgentFinish`
                    if (function_call.name === "response") {
                        return { returnValues: { ...toolInput }, log: message.content };
                    }

                    return {
                        tool: function_call.name,
                        toolInput,
                        log: `Invoking "${function_call.name}" with ${function_call.arguments ?? "{}"
                            }\n${message.content}`,
                        messageLog: [message],
                    };
                } catch (error) {
                    throw new Error(
                        `Failed to parse function arguments from chat model response. Text: "${function_call.arguments}". ${error}`
                    );
                }
            } else {
                return {
                    returnValues: { output: message.content },
                    log: message.content,
                };
            }
        };

        const formatAgentSteps = (steps: AgentStep[]): BaseMessage[] =>
            steps.flatMap(({ action, observation }) => {
                if ("messageLog" in action && action.messageLog !== undefined) {
                    const log = action.messageLog as BaseMessage[];
                    return log.concat(new FunctionMessage(observation, action.tool));
                } else {
                    return [new AIMessage(action.log)];
                }
            });

        const llmWithTools = llm.bind({
            functions: [convertToOpenAIFunction(searchTool), responseOpenAIFunction],
        });
        /** Create the runnable */
        const runnableAgent = RunnableSequence.from<{
            input: string;
            steps: Array<AgentStep>;
        }>([
            {
                input: (i) => i.input,
                agent_scratchpad: (i) => formatAgentSteps(i.steps),
            },
            prompt,
            llmWithTools,
            structuredOutputParser,
        ]);

        const executor = AgentExecutor.fromAgentAndTools({
            agent: runnableAgent,
            tools: [searchTool],
        });

        const res = await executor.invoke({
            input: messages[messages.length - 1].content,
        });

        return NextResponse.json(res.answer);

    } catch (e: any) {
        console.log((e as Error).stack);
        return NextResponse.json({ error: e.message }, { status: e.status ?? 500 });
    }
}
``



### System Info

Langchain v0.1.27
Macbook Air 2020 M1
Node v18.19.1
NPM v10.2.4

[dosubot[bot]]

Based on the code snippet you've provided and the context from the LangChain JS repository, it seems that the streaming functionality should be working as expected. The LangChainStream and StreamingTextResponse are designed to handle streaming of tokens in real-time, handle each new token event, concatenate these tokens to form the complete response, and finally log the complete response.

However, there have been similar issues reported in the repository, such as Streaming OpenAI Function Call output in LangChainStreaming and LangchainJs Streaming with OpenAIFunctions Agent. These discussions might provide some insights into the problem you're facing.

One closed issue, Streaming stopped working since the latest version (0.0.30), was resolved by a pull request: handleLLMNewToken should not use the overall verbose flag of the model, as it is only fired when asked for by the user. This suggests that there might have been some changes in the streaming functionality in the past.

Unfortunately, I couldn't find specific changes in the LangChain JS version 0.1.28 that could potentially affect the streaming functionality, or the key differences between LangChain JS version 0.1.27 and 0.1.28. Also, there are no known issues or bugs related to the streaming functionality in LangChain JS version 0.1.27 in the repository.

To better understand the issue, could you provide more details about the behavior you're observing? Specifically, are there any error messages or unusual logs when you try to stream? Also, could you share the version of the @langchain/core and @langchain/openai packages you're using? This information could help identify potential compatibility issues or bugs.

Sources

• Streaming OpenAI Function Call output in LangChainStreaming
• LangchainJs Streaming with OpenAIFunctions Agent
• Streaming stopped working since the latest version (0.0.30)
• examples/src/agents/openai_runnable_stream_log.ts
• examples/src/models/llm/llm_streaming.ts

About Dosu

This response is meant to be useful and save you time. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

If you want to continue the conversation, start your reply with @dosu-bot.