Agent streaming

[vercel-labs]

Agent streaming

[RobertHH-IS]

Yea parsing the stream for all sorts of conditions is a bit of a struggle. I was also interested in doing streaming with agents that involved all sorts of tools and outputs, but when it came to showing notifications of tool use, steps, images, links, user input etc... I have found that simply parsing the JSON on completion, instead of trying to implement it during the stream, makes life 10x easier. Completing the project "onCompletion" and then backtracking later to implement streaming where it makes sense - once langchain and vercel ai library evolve better support and more solid examples - is the route I went.

[aranlucas]

There's a couple of things that probably need to be in an FAQ or something

Here the code that worked for me.

import { LangChainStream } from "@/lib/LangChainStream";
import { PineconeClient } from "@pinecone-database/pinecone";
import { Message, StreamingTextResponse } from "ai";
import { initializeAgentExecutorWithOptions } from "langchain/agents";
import { ChatOpenAI } from "langchain/chat_models/openai";
import { Calculator } from "langchain/tools/calculator";

export const runtime = "edge";

export async function POST(req: Request) {
  const { messages } = await req.json();

  // My own LangChainStream implementation in PR201
  const { stream, handlers } = LangChainStream();

  const model = new ChatOpenAI({
    streaming: true,
  });

  const tools = [new Calculator()];

  const executor = await initializeAgentExecutorWithOptions(tools, model, {
    agentType: "chat-zero-shot-react-description",
    returnIntermediateSteps: true,
  });

  const input = messages[messages.length - 1].content;
  executor.call({ input, verbose: true }, [handlers]).catch(console.error);

  return new StreamingTextResponse(stream);
}

You'll get something like

Thought: We can use the calculator tool to add the numbers. Action: ``` { "action": "calculator", "action_input": "2 + 2" } ``` Final Answer: 4

Things to change

1. You can't await a streaming response. Once you put that await, it no longer tries to stream.
2. Change the handler function to be request scoped, not on constructor
3. you should only use the last message to send to executor.
4. Current langchain handlers closes stream too early. Here's my PR attempt to fix vercel/ai#201. I created my own handler until that's fixed.

[RobertHH-IS]

Thanks, I´ll see if we dont get an update to the LangChainStream soon. Would you mind pasting your file in lib so i can see if this is working?

[aranlucas]

@lib/LangChainStream

import { type AIStreamCallbacks, createCallbacksTransformer } from 'ai'

export function LangChainStream(callbacks?: AIStreamCallbacks) {
  const stream = new TransformStream()
  const writer = stream.writable.getWriter()

  const runs = new Set()

  const handleError = async (e: Error, runId: string) => {
    runs.delete(runId)
    await writer.ready
    await writer.abort(e)
  }

  const handleStart = (runId: string) => {
    runs.add(runId)
  }

  const handleEnd = async (runId: string) => {
    runs.delete(runId)

    if (runs.size === 0) {
      await writer.ready
      await writer.close()
    }
  }

  return {
    stream: stream.readable.pipeThrough(createCallbacksTransformer(callbacks)),
    handlers: {
      handleLLMNewToken: async (token: string) => {
        await writer.ready
        await writer.write(token)
      },
      handleLLMStart: async (_llm: any, _prompts: string[], runId: string) => {
        handleStart(runId)
      },
      handleLLMEnd: async (_output: any, runId: string) => {
        await handleEnd(runId)
      },
      handleLLMError: async (e: Error, runId: string) => {
        await handleError(e, runId)
      },
      handleChatModelStart: async (
        _llm: any,
        _messages: any,
        runId: string
      ) => {
        handleStart(runId)
      },
      handleChatModelEnd: async (_output: any, runId: string) => {
        await handleEnd(runId)
      },
      handleChainStart: async (_chain: any, _inputs: any, runId: string) => {
        handleStart(runId)
      },
      handleChainEnd: async (_outputs: any, runId: string) => {
        await handleEnd(runId)
      },
      handleChainError: async (e: Error, runId: string) => {
        await handleError(e, runId)
      },
      handleToolStart: async (_tool: any, _input: string, runId: string) => {
        handleStart(runId)
      },
      handleToolEnd: async (_output: string, runId: string) => {
        await handleEnd(runId)
      },
      handleToolError: async (e: Error, runId: string) => {
        await handleError(e, runId)
      },
      handleAgentAction: async (_action: any, runId: string) => {
        handleStart(runId)
      },
      handleAgentEnd: async (_output: any, runId: string) => {
        await handleEnd(runId)
      }
    }
  }
}

[RobertHH-IS]

Do you have a good way to accomplish this with agents:
"you should only use the last message to send to executor"

I have tried to set returnIntermediateSteps to false, but it still outputs thoughts and actions. The python version has a FinalStreamingStdOutCallbackHandler but I cannot find anything similar in JS version.

[aranlucas]

It's all in the handlers

From a brief look at https://api.python.langchain.com/en/stable/_modules/langchain/callbacks/streaming_stdout_final_only.html, you'd modify the handleNewToken

    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        """Run on new LLM token. Only available when streaming is enabled."""

        # Remember the last n tokens, where n = len(answer_prefix_tokens)
        self.append_to_last_tokens(token)

        # Check if the last n tokens match the answer_prefix_tokens list ...
        if self.check_if_answer_reached():
            self.answer_reached = True
            if self.stream_prefix:
                for t in self.last_tokens:
                    sys.stdout.write(t)
                sys.stdout.flush()
            return

        # ... if yes, then print tokens from now on
        if self.answer_reached:
            sys.stdout.write(token)
            sys.stdout.flush()

To keep checking if answer_prefix_tokens =["Final", "Answer", ":"] is equal to the last 3 tokens you have received, and only await writer.write(token) at that point

[DanielhCarranza]

I'm wondering how you did that implementation to stream the final response of an Agent in typescript.
Here is mine, but doesn't seem to work, any idea why?

...
    // Stream Final Response from Agent 
    const answerPrefixTokens: string[] = ["Final", "Answer",":"]
    let appendNewToken: string[] = []
    let lastTokens: string[] = []
    let answerReached: boolean = false

    const appendToLastToken = (token: string): void => {
        lastTokens.push(token)
        if (lastTokens.length > answerPrefixTokens.length) {
            lastTokens.shift()
        }
    }
    const streamFunc = new TransformStream()
    const writer = streamFunc.writable.getWriter()
    
    const { stream, handlers } = LangChainStream({
      async onCompletion(completion) {
        const title = body.messages[0].content.substring(0, 100)
        const userId = session?.user?.id
        if (userId) {
          const id = body.id ?? nanoid()
          const createdAt = Date.now()
          const path = `/chat/${id}`
          const payload = {
            id,
            title,
            userId,
            createdAt,
            path,
            messages: [
              ...messages,
              {
                content: completion,
                role: 'assistant'
              }
            ]
          }
          await kv.hmset(`chat:${id}`, payload)
          await kv.zadd(`user:chat:${userId}`, {
            score: createdAt,
            member: `chat:${id}`
          })
        }
        
    },
      async onToken(token) {
        handlers.handleLLMNewToken = async (token) => {
            appendToLastToken(token)
            if (lastTokens.join(" ") === answerPrefixTokens.join(" ")) {
                answerReached = true
            }
            if (answerReached) {
                await writer.write(token)
            }
        }
    },
    });

[RobertHH-IS]

@DanielhCarranza What does not work - what are you trying to acheive in the code?

[DanielhCarranza]

I'm trying to stream only the final output from the agent

[MaxLeiter]

@RobertHH-IS and @DanielhCarranza please share any feedback in the https://github.com/vercel-labs/ai repo, as we're actively investigating building the SDK to support more complex use cases.

[leerob]

The AI SDK now supports OpenAI Function calling: https://sdk.vercel.ai/docs/guides/providers/openai-functions