{"version":3,"file":"upstash_ratelimit.cjs","names":["BaseCallbackHandler"],"sources":["../../../src/callbacks/handlers/upstash_ratelimit.ts"],"sourcesContent":["import { Ratelimit } from \"@upstash/ratelimit\";\nimport { Serialized } from \"@langchain/core/load/serializable\";\nimport { LLMResult } from \"@langchain/core/outputs\";\nimport { ChainValues } from \"@langchain/core/utils/types\";\nimport { BaseCallbackHandler } from \"@langchain/core/callbacks/base\";\n\n/**\n * Upstash Ratelimit Error\n *\n * Raised when the rate limit is reached in `UpstashRatelimitHandler`.\n */\nclass UpstashRatelimitError extends Error {\n  type: \"token\" | \"request\";\n\n  limit?: number;\n\n  reset?: number;\n\n  /**\n   * @param message - Error message\n   * @param type - The kind of limit which was reached. One of \"token\" or \"request\"\n   * @param limit - The limit which was reached. Passed when type is request\n   * @param reset - Unix timestamp in milliseconds when the limits are reset. Passed when type is request\n   */\n  constructor(\n    message: string,\n    type: \"token\" | \"request\",\n    limit?: number,\n    reset?: number\n  ) {\n    super(message);\n    this.type = type;\n    this.limit = limit;\n    this.reset = reset;\n  }\n}\n\ninterface UpstashRatelimitHandlerOptions {\n  tokenRatelimit?: Ratelimit;\n  requestRatelimit?: Ratelimit;\n  includeOutputTokens?: boolean;\n\n  llmOutputTokenUsageField?: string;\n  llmOutputTotalTokenField?: string;\n  llmOutputPromptTokenField?: string;\n}\n\n/**\n * Callback to handle rate limiting based on the number of requests\n * or the number of tokens in the input.\n *\n * It uses Upstash Ratelimit to track the rate limit which utilizes\n * Upstash Redis to track the state.\n *\n * Should not be passed to the chain when initializing the chain.\n * This is because the handler has a state which should be fresh\n * every time invoke is called. Instead, initialize and pass a handler\n * every time you invoke.\n */\nclass UpstashRatelimitHandler extends BaseCallbackHandler {\n  name = \"UpstashRatelimit\";\n\n  raiseError = true;\n\n  private _checked = false;\n\n  identifier: string;\n\n  tokenRatelimit?: Ratelimit;\n\n  requestRatelimit?: Ratelimit;\n\n  includeOutputTokens: boolean;\n\n  llmOutputTokenUsageField: string;\n\n  llmOutputTotalTokenField: string;\n\n  llmOutputPromptTokenField: string;\n\n  /**\n   * @param identifier - The identifier to rate limit, like a user ID or an IP address\n   * @param options - Ratelimit options\n   */\n  constructor(identifier: string, options: UpstashRatelimitHandlerOptions) {\n    super();\n    if (!options.tokenRatelimit && !options.requestRatelimit) {\n      throw new Error(\n        \"You must pass at least one of tokenRatelimit or requestRatelimit.\"\n      );\n    }\n    this.identifier = identifier;\n    this.tokenRatelimit = options.tokenRatelimit;\n    this.requestRatelimit = options.requestRatelimit;\n    this.includeOutputTokens = options.includeOutputTokens ?? false;\n\n    this.llmOutputTokenUsageField =\n      options.llmOutputTokenUsageField ?? \"tokenUsage\";\n    this.llmOutputTotalTokenField =\n      options.llmOutputTotalTokenField ?? \"totalTokens\";\n    this.llmOutputPromptTokenField =\n      options.llmOutputPromptTokenField ?? \"promptTokens\";\n\n    this.awaitHandlers = true;\n  }\n\n  /**\n   * Run when the chain starts running.\n   *\n   * This method is called multiple times during a chain execution.\n   * To ensure it only runs once, it checks and updates a `_checked` state.\n   *\n   * @param _chain - Serialized chain\n   * @param _inputs - Chain input values\n   * @throws UpstashRatelimitError - If the request rate limit is reached\n   */\n  async handleChainStart(\n    _chain: Serialized,\n    _inputs: ChainValues\n  ): Promise<void> {\n    if (this.requestRatelimit && !this._checked) {\n      const response = await this.requestRatelimit.limit(this.identifier);\n      if (!response.success) {\n        throw new UpstashRatelimitError(\n          \"Request limit reached!\",\n          \"request\",\n          response.limit,\n          response.reset\n        );\n      }\n      this._checked = true;\n    }\n  }\n\n  /**\n   * Run when the LLM starts running.\n   *\n   * @param _llm - Serialized LLM\n   * @param _prompts - Prompts passed to the LLM\n   * @throws UpstashRatelimitError - If the token rate limit is reached\n   */\n  async handleLLMStart(\n    _llm: Serialized,\n    _prompts: string[],\n    _runId: string,\n    _parentRunId?: string,\n    _extraParams?: Record<string, unknown>,\n    _tags?: string[],\n    _metadata?: Record<string, unknown>,\n    _name?: string\n  ): Promise<void> {\n    if (this.tokenRatelimit) {\n      const result = await this.tokenRatelimit.getRemaining(this.identifier);\n\n      // result of getRemaining was changed from a number to an object in v2.0.0.\n      // we check to make sure that it works with versions before & after:\n      const remaining = typeof result === \"number\" ? result : result.remaining;\n\n      if (remaining <= 0) {\n        throw new UpstashRatelimitError(\"Token limit reached!\", \"token\");\n      }\n    }\n  }\n\n  /**\n   * Run when the LLM ends running.\n   *\n   * If the `includeOutputTokens` is set to true, the number of tokens\n   * in the LLM completion are counted for rate limiting.\n   *\n   * @param output - LLM result output\n   * @throws Error - If the LLM response does not include required token usage information\n   */\n  async handleLLMEnd(\n    output: LLMResult,\n    _runId: string,\n    _parentRunId?: string,\n    _tags?: string[]\n  ): Promise<void> {\n    if (this.tokenRatelimit) {\n      const llmOutput = output.llmOutput || {};\n      try {\n        const tokenUsage = llmOutput[this.llmOutputTokenUsageField];\n        const tokenCount = this.includeOutputTokens\n          ? tokenUsage[this.llmOutputTotalTokenField]\n          : tokenUsage[this.llmOutputPromptTokenField];\n\n        if (tokenCount !== undefined) {\n          await this.tokenRatelimit.limit(this.identifier, {\n            rate: tokenCount,\n          });\n        } else {\n          throw new Error(\"tokenCount not found in llm output\");\n        }\n      } catch (error) {\n        // eslint-disable-next-line no-instanceof/no-instanceof\n        if (error instanceof UpstashRatelimitError) {\n          throw error;\n        }\n        console.error(\n          `Failed to log token usage for Upstash rate limit. It could be because the LLM returns the token usage in a different format than expected. See UpstashRatelimitHandler parameters. Got error: ${error}`\n        );\n      }\n    }\n  }\n\n  /**\n   * Creates a new UpstashRatelimitHandler object with the same\n   * ratelimit configurations but with a new identifier if it's\n   * provided.\n   *\n   * Also resets the state of the handler.\n   *\n   * @param identifier - Optional new identifier to use for the new handler instance\n   * @returns New UpstashRatelimitHandler instance\n   */\n  reset(identifier?: string): UpstashRatelimitHandler {\n    return new UpstashRatelimitHandler(identifier ?? this.identifier, {\n      tokenRatelimit: this.tokenRatelimit,\n      requestRatelimit: this.requestRatelimit,\n      includeOutputTokens: this.includeOutputTokens,\n    });\n  }\n}\n\nexport {\n  UpstashRatelimitHandler,\n  UpstashRatelimitError,\n  UpstashRatelimitHandlerOptions,\n};\n"],"mappings":";;;;;;;;;;;;;AAWA,IAAM,wBAAN,cAAoC,MAAM;CACxC;CAEA;CAEA;;;;;;;CAQA,YACE,SACA,MACA,OACA,OACA;AACA,QAAM,QAAQ;AACd,OAAK,OAAO;AACZ,OAAK,QAAQ;AACb,OAAK,QAAQ;;;;;;;;;;;;;;;AA0BjB,IAAM,0BAAN,MAAM,gCAAgCA,+BAAAA,oBAAoB;CACxD,OAAO;CAEP,aAAa;CAEb,WAAmB;CAEnB;CAEA;CAEA;CAEA;CAEA;CAEA;CAEA;;;;;CAMA,YAAY,YAAoB,SAAyC;AACvE,SAAO;AACP,MAAI,CAAC,QAAQ,kBAAkB,CAAC,QAAQ,iBACtC,OAAM,IAAI,MACR,oEACD;AAEH,OAAK,aAAa;AAClB,OAAK,iBAAiB,QAAQ;AAC9B,OAAK,mBAAmB,QAAQ;AAChC,OAAK,sBAAsB,QAAQ,uBAAuB;AAE1D,OAAK,2BACH,QAAQ,4BAA4B;AACtC,OAAK,2BACH,QAAQ,4BAA4B;AACtC,OAAK,4BACH,QAAQ,6BAA6B;AAEvC,OAAK,gBAAgB;;;;;;;;;;;;CAavB,MAAM,iBACJ,QACA,SACe;AACf,MAAI,KAAK,oBAAoB,CAAC,KAAK,UAAU;GAC3C,MAAM,WAAW,MAAM,KAAK,iBAAiB,MAAM,KAAK,WAAW;AACnE,OAAI,CAAC,SAAS,QACZ,OAAM,IAAI,sBACR,0BACA,WACA,SAAS,OACT,SAAS,MACV;AAEH,QAAK,WAAW;;;;;;;;;;CAWpB,MAAM,eACJ,MACA,UACA,QACA,cACA,cACA,OACA,WACA,OACe;AACf,MAAI,KAAK,gBAAgB;GACvB,MAAM,SAAS,MAAM,KAAK,eAAe,aAAa,KAAK,WAAW;AAMtE,QAFkB,OAAO,WAAW,WAAW,SAAS,OAAO,cAE9C,EACf,OAAM,IAAI,sBAAsB,wBAAwB,QAAQ;;;;;;;;;;;;CActE,MAAM,aACJ,QACA,QACA,cACA,OACe;AACf,MAAI,KAAK,gBAAgB;GACvB,MAAM,YAAY,OAAO,aAAa,EAAE;AACxC,OAAI;IACF,MAAM,aAAa,UAAU,KAAK;IAClC,MAAM,aAAa,KAAK,sBACpB,WAAW,KAAK,4BAChB,WAAW,KAAK;AAEpB,QAAI,eAAe,KAAA,EACjB,OAAM,KAAK,eAAe,MAAM,KAAK,YAAY,EAC/C,MAAM,YACP,CAAC;QAEF,OAAM,IAAI,MAAM,qCAAqC;YAEhD,OAAO;AAEd,QAAI,iBAAiB,sBACnB,OAAM;AAER,YAAQ,MACN,iMAAiM,QAClM;;;;;;;;;;;;;;CAeP,MAAM,YAA8C;AAClD,SAAO,IAAI,wBAAwB,cAAc,KAAK,YAAY;GAChE,gBAAgB,KAAK;GACrB,kBAAkB,KAAK;GACvB,qBAAqB,KAAK;GAC3B,CAAC"}