Upstash Rate Limiter with callback not working correctly

[kunalhexus]

Checklist

• I added a very descriptive title to this issue.
• I included a link to the documentation page I am referring to (if applicable).

Issue with current documentation:

I am following the instructions on this link for setting up rate limiting with upstash.

I am using chain.streamEvents for streaming, passing the rate limit handler to callbacks.

On server side, I am able to see the rate limit error, but the chain still continues to run and return the correct response.

⨯ unhandledRejection: UpstashRatelimitError: Request limit reached!
    at UpstashRatelimitHandler.handleChainStart (webpack-internal:///(rsc)/./node_modules/.pnpm/@[email protected]_@[email protected]_@[email protected]_@aws_nzbkjz35sxoqfmioyf4dm75tau/node_modules/@langchain/community/dist/callbacks/handlers/upstash_ratelimit.js:152:23)
    at process.processTicksAndRejections (node:internal/process/task_queues:105:5)
    at async eval (webpack-internal:///(rsc)/./node_modules/.pnpm/@[email protected][email protected][email protected][email protected]_/node_modules/@langchain/core/dist/callbacks/manager.js:559:21)
    at async eval (webpack-internal:///(rsc)/./node_modules/.pnpm/@[email protected][email protected][email protected][email protected]_/node_modules/@langchain/core/dist/singletons/callbacks.js:54:17)
    at async run (webpack-internal:///(rsc)/./node_modules/.pnpm/[email protected]/node_modules/p-queue/dist/index.js:163:29) {
  type: 'request',
  limit: 1,
  reset: 1732665900000
}

Doing something like below is not catching the exception

try {
        const stream = chain.streamEvents({
                messages: messages
            },
            { version: "v2", encoding: "text/event-stream", callbacks: [handler] },
            { includeNames: includeTypes: ["parser"] }
        );
        return new Response(stream, { headers: { "content-type": "text/event-stream" } });   
    } catch (e) {
        console.error(e);
        return NextResponse.json({error: true}, {status: 429});
    }

Idea or request for content:

Looking for correct way to use rate limiter with langchain

[dosubot]

Hey there, @kunalhexus! I'm here to help you with any bugs, questions, or contributions you have. Let's tackle this together while we wait for a human maintainer.

The issue you're encountering with the UpstashRatelimitHandler not stopping the chain when a rate limit is exceeded might be due to how the error is being handled in the streaming context. The UpstashRatelimitError should be thrown when the rate limit is reached, but it seems like the error is not being caught properly in your try-catch block.

Here are a few suggestions to address this:

1. Ensure Proper Error Propagation: Make sure that the UpstashRatelimitError is being propagated correctly from the handler. The error should be thrown and not just logged. Verify that the handler is correctly integrated with the streamEvents method.

2. Check Error Handling in Streaming: The streamEvents method might handle errors differently compared to synchronous methods. Ensure that the error is not being swallowed by the streaming mechanism. You might need to handle the error within the stream processing logic itself.

3. Debugging: Add additional logging inside the streamEvents method to see if the error is being caught at any point before your try-catch block. This can help identify where the error handling might be failing.

4. Review Documentation and Tests: Check if there are any specific examples or test cases in the Langchain codebase that demonstrate handling UpstashRatelimitError in a streaming context. This might provide insights into the correct implementation.

If these steps do not resolve the issue, you might need to look into the implementation details of streamEvents and how it interacts with the callback mechanism to ensure that errors are handled as expected.

To continue talking to Dosu, mention @dosu.

[CahidArda]

Hi @kunalhexus,

I just tried this myself and reproduced the behavior. The chain doesn't stop. It seems something was changed after implementing Upstash Ratelimit callback.

I have found this workaround for now:

const handler = new UpstashRatelimitHandler(user_id, {
  requestRatelimit: ratelimit,
});
handler.awaitHandlers = true

In the meantime, I will open a PR to fix this.

[Zamoca42]

In my case, I was able to reproduce the issue when I added the ratelimitHandler to the callbacks in invoke. However, it works properly when the ratelimitHandler is added to the model’s callbacks.

export async function getArchitecture(input: string) {
  const cache = new UpstashRedisCache({
    client: redis,
    ttl: process.env.NODE_ENV === "production" ? CACHE_TTL : DEV_CACHE_TTL,
  });

  const ratelimitHandler = new UpstashRatelimitHandler(userId, {
    requestRatelimit,
  });

  const { success } = await requestRatelimit.limit(userId);

  try {
      const model = new ChatGoogleGenerativeAI({
        model: "gemini-1.5-flash",
        temperature: 1.0,
        cache,
        callbacks: [ratelimitHandler], // it works
      });

      const architecture = await model.pipe(parser).invoke(formattedPrompt, {
        callbacks: [ratelimitHandler], // unhandledRejection: UpstashRatelimitError: Request limit reached!
      });

    return { architecture, success };
  } catch (error) {
    throw error;
  }
}

[CahidArda]

Hi @jacoblee93,

I think Upstash Ratelimit should be able to stop the chain whether it's passed to the model or in invoke. The callback is written so that it would be used in invoke.

Currently it doesn't seem to work when passed to invoke. I was able to find this workaround:

const handler = new UpstashRatelimitHandler(user_id, {
  requestRatelimit: ratelimit,
});
handler.awaitHandlers = true

Should we update UpstashRatelimitHandler code by adding awaitHandlers = true by default?

[jacoblee93]

Yeah that would make sense here

[CahidArda]

I opened the PR, can you take a look when you have time? @jacoblee93

Best,