fix: langchain stream not closing

[aranlucas]

fixes: #97

The problem

1. Streaming handlers should be used in Request callbacks - https://js.langchain.com/docs/production/callbacks/#when-do-you-want-to-use-each-of-these
2. You can have multiple chains. You need to know when all the chains before closing the stream. Currently, the stream is being closed after the first chain End

Solution

1. Update docs to reflect request callbaks
2. Add internal state that keeps track of all of the runs happening. If there are no runs, then the stream is safe to close.

[changeset-bot]

🦋 Changeset detected

Latest commit: 0a26a5a

The changes in this PR will be included in the next version bump.

This PR includes changesets to release 1 package

Name	Type
ai	Patch

Not sure what this means? Click here to learn what changesets are.

Click here if you're a maintainer who wants to add another changeset to this PR

[MaxLeiter]

Thank you! Feel free to open an issue here for tracking the langchain discussion

[MaxLeiter]

Actually, I just remembered why @jaredpalmer made the change. Here's the initial issue: #63

It seems like we may want to let the consumer choose which callback to end on, as it's model dependent

[aranlucas]

Yeah I figured that was the reasoning, thanks for confirming.

This change will probably break SequentialChain use case.

It seems like your suggestion might be the best one. Or langchainjs can expose have something like handleEnd where it doesn't depend on a model.

[aranlucas]

Last one here before I create an issue: https://js.langchain.com/docs/production/callbacks/#multiple-handlers

Before

  const { stream, handlers } = LangChainStream()
 
  const llm = new ChatOpenAI({
    streaming: true,
    callbackManager: CallbackManager.fromHandlers(handlers)
  })
 
  llm
    .call(
      (messages as Message[]).map(m =>
        m.role == 'user'
          ? new HumanChatMessage(m.content)
          : new AIChatMessage(m.content)
      )
    )
    .catch(console.error)
 
  return new StreamingTextResponse(stream)

After

  // Different handlers for llm, agents, chains
  const { stream, llmHandlers, chainHandlers } = LangChainStream()
 
  // Do not pass handlers here
  const llm = new ChatOpenAI({
    streaming: true,
  })
 
  // Pass handlers here instead
  llm
    .call(
      (messages as Message[]).map((m) =>
        m.role == "user"
          ? new HumanChatMessage(m.content)
          : new AIChatMessage(m.content)
      ),
      [llmHandlers]
    )
    .catch(console.error);
 
  return new StreamingTextResponse(stream)

[aranlucas]

Created #205 for further discussions

This PR will probably break the linked usecase - can be closed or left until we find something.

[jaredpalmer]

Big fan of const { stream, llmHandlers, chainHandlers } = LangChainStream() lets do that

[aranlucas]

@jaredpalmer that actually won't work after doing some research #205

As an exaggerated example, a chain can spawn a chain that spawns a chain, meaning that on the first chainEnd, the stream will close. You need to be able to count how many things are running otherwise only the first chain to finish will write to the stream, locking out the other chain results.

I'm leaning towards #205 (comment) as the solution now. It should cover all use-cases, and keep the API the same. (Still need to update docs though to use on "Request callbacks" rather than "Constructor callbacks")

[aranlucas]

Updated with my latest findings 👀

It may be oversimplified compared to what they have done with tracing but I think it covers all the use-cases.

[aranlucas]

Take #205 (comment) as an example. Let's say we wanted to stream both the responses.

in first chain

Tragedy at Sunset on the Beach is a story of love, loss, and redemption. It follows the story of two young lovers, Jack and Jill, who meet on a beach at sunset. They quickly fall in love and plan to spend the rest of their lives together. \n\nHowever, tragedy strikes when Jack is killed in a car accident. Jill is left devastated and unable to cope with the loss of her beloved. She spirals into a deep depression and begins to lose hope. \n\nJust when all seems lost, Jill discovers that Jack had left her a letter before his death. In the letter, he tells her that he will always love her and that she should never give up hope. With this newfound strength, Jill is able to find the courage to move on with her life and find happiness again. \n\nTragedy at Sunset on the Beach is a story of love, loss, and redemption. It is a story of hope and courage in the face of tragedy. It is a story of finding strength in the darkest of times and of never giving up.",     

handleChainEnd gets called, stream is closed. Second chain can't write the second response

Tragedy at Sunset on the Beach is a powerful and moving story of love, loss, and redemption. The play follows the story of two young lovers, Jack and Jill, whose plans for a future together are tragically cut short when Jack is killed in a car accident. \n\nThe play is beautifully written and the performances are outstanding. The actors bring a depth of emotion to their characters that is both heartbreaking and inspiring. The story is full of unexpected twists and turns, and the audience is taken on an emotional rollercoaster as Jill struggles to come to terms with her loss. \n\nThe play is ultimately a story of hope and resilience, and it is sure to leave audiences with a newfound appreciation for life and love. Tragedy at Sunset on the Beach is a must-see for anyone looking for an emotionally charged and thought-provoking experience."

There's probably many ways to fix this, but the way I suggested in this PR seems to work for most use cases I've tried.

[nfcampos]

You might want to be careful with this one, in some chains that eg run multiple things in parallel there could be other events arriving after an error event, it doesn't look like you're handling those

[aranlucas]

How should it be handled? If there's an error, can it recover gracefully from it?

Right now this is definitely giving up on first error, but I'll be honest I haven't encountered a scenario where the error was recovered and it could continue the chain.

[nfcampos]

So it will overall be cancelled when the first error occurs. The issue is if it's a chain with multiple LLM calls in parallel you may still receive new tokens from call #2 after call #1 has failed. In which case currently what would happen is it would probably throw an exception in handleLLMNewToken as you try to write to an already aborted writer?

[aranlucas]

Yeah I noticed that in #205 (comment).

Error in handler Handler, handleChainEnd: TypeError: The stream (in closed state) is not in the writable state and cannot be closed      

I wonder what would be the best behavior, if to abort the call (ignoring the second one) or allow call 2 to keep writing to stream.

[aranlucas]

Note: the exception of writing to a stream in closed state seems to be handled somewhere that doesn't seem to cause catastrophic errors.

[nfcampos]

Yes any exceptions inside a callback handler are caught by langchain to avoid an error in a callback handler interrupting a run

[aranlucas]

Ok I see your point.

How does this suggestion sound.

Keep an error state handler

let errors = []

Whenever we encounter an error on any handle*Error push it errors

Then handleEnd

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

    if (runs.size === 0) {
      if (error.length !=0) {
          // for loop over errors or join errors?
          await writer.ready
          await writer.abort(errors)
          return;
      }
      await writer.ready
      await writer.close()
    }
  }

[aranlucas]

Or I can just re-throw the error similar to https://github.com/vercel-labs/ai/blob/main/packages/core/streams/ai-stream.ts#L152-L154

I think an issue I had yesterday working with this is that even though the stream had ended, the background processing had not stopped. This led me to have to restart dev server to clear out running processes

[MaxLeiter]

I'm a fan of rethrowing the error

[jaredpalmer]

Are we good to merge @nfcampos @MaxLeiter ?

[nfcampos]

looks good to me

[jaredpalmer]

@nfcampos I added this [] at all call sites in order to satisfy typescript. If these aren't necessary then, this may be an upstream issue with LC, likely missing an overload.

[nfcampos]

I think that should be undefined

[jaredpalmer]

stylistic suggestion: y'all should have only a two arguments to .call, passing undefined is very odd.

[nfcampos]

Yep agreed, I've improved this here langchain-ai/langchainjs#1868 once that is merged you'll be able to do instead llm.call(messages, { callbacks: handlers })

[aranlucas]

Thank you 🙏