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/@langchain+community@0.3.16_@aws-crypto+sha256-js@5.2.0_@aws-sdk+client-dynamodb@3.699.0_@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/@langchain+core@0.3.18_openai@4.73.0_encoding@0.1.13_zod@3.23.8_/node_modules/@langchain/core/dist/callbacks/manager.js:559:21)
    at async eval (webpack-internal:///(rsc)/./node_modules/.pnpm/@langchain+core@0.3.18_openai@4.73.0_encoding@0.1.13_zod@3.23.8_/node_modules/@langchain/core/dist/singletons/callbacks.js:54:17)
    at async run (webpack-internal:///(rsc)/./node_modules/.pnpm/p-queue@6.6.2/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;
  }
}