@langchain/community
Version:
Third-party integrations for LangChain.js
1 lines • 8.94 kB
Source Map (JSON)
{"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"}