langchain
Version:
Typescript bindings for langchain
1 lines • 8.54 kB
Source Map (JSON)
{"version":3,"file":"modelCallLimit.cjs","names":["z","createMiddleware","AIMessage"],"sources":["../../../src/agents/middleware/modelCallLimit.ts"],"sourcesContent":["import { z } from \"zod/v3\";\nimport { AIMessage } from \"@langchain/core/messages\";\nimport { InferInteropZodInput } from \"@langchain/core/utils/types\";\n\nimport { createMiddleware } from \"../middleware.js\";\n\nconst DEFAULT_EXIT_BEHAVIOR = \"end\";\n\nconst contextSchema = z.object({\n /**\n * The maximum number of model calls allowed per thread.\n */\n threadLimit: z.number().optional(),\n /**\n * The maximum number of model calls allowed per run.\n */\n runLimit: z.number().optional(),\n /**\n * The behavior to take when the limit is exceeded.\n * - \"error\" will throw an error and stop the agent.\n * - \"end\" will end the agent.\n * @default \"end\"\n */\n exitBehavior: z.enum([\"error\", \"end\"]).optional(),\n});\nexport type ModelCallLimitMiddlewareConfig = Partial<\n InferInteropZodInput<typeof contextSchema>\n>;\n\n/**\n * Middleware state schema to track the number of model calls made at the thread and run level.\n */\nconst stateSchema = z.object({\n threadModelCallCount: z.number().default(0),\n runModelCallCount: z.number().default(0),\n});\n\n/**\n * Error thrown when the model call limit is exceeded.\n *\n * @param threadLimit - The maximum number of model calls allowed per thread.\n * @param runLimit - The maximum number of model calls allowed per run.\n * @param threadCount - The number of model calls made at the thread level.\n * @param runCount - The number of model calls made at the run level.\n */\nclass ModelCallLimitMiddlewareError extends Error {\n constructor({\n threadLimit,\n runLimit,\n threadCount,\n runCount,\n }: {\n threadLimit?: number;\n runLimit?: number;\n threadCount?: number;\n runCount?: number;\n }) {\n const exceededHint: string[] = [];\n if (typeof threadLimit === \"number\" && typeof threadCount === \"number\") {\n exceededHint.push(\n `thread level call limit reached with ${threadCount} model calls`\n );\n }\n if (typeof runLimit === \"number\" && typeof runCount === \"number\") {\n exceededHint.push(\n `run level call limit reached with ${runCount} model calls`\n );\n }\n super(\n `Model call limits exceeded${\n exceededHint.length > 0 ? `: ${exceededHint.join(\", \")}` : \"\"\n }`\n );\n this.name = \"ModelCallLimitMiddlewareError\";\n }\n}\n\n/**\n * Creates a middleware to limit the number of model calls at both thread and run levels.\n *\n * This middleware helps prevent excessive model API calls by enforcing limits on how many\n * times the model can be invoked. It supports two types of limits:\n *\n * - **Thread-level limit**: Restricts the total number of model calls across an entire conversation thread\n * - **Run-level limit**: Restricts the number of model calls within a single agent run/invocation\n *\n * ## How It Works\n *\n * The middleware intercepts model requests before they are sent and checks the current call counts\n * against the configured limits. If either limit is exceeded, it throws a `ModelCallLimitMiddlewareError`\n * to stop execution and prevent further API calls.\n *\n * ## Use Cases\n *\n * - **Cost Control**: Prevent runaway costs from excessive model calls in production\n * - **Testing**: Ensure agents don't make too many calls during development/testing\n * - **Safety**: Limit potential infinite loops or recursive agent behaviors\n * - **Rate Limiting**: Enforce organizational policies on model usage per conversation\n *\n * @param middlewareOptions - Configuration options for the call limits\n * @param middlewareOptions.threadLimit - Maximum number of model calls allowed per thread (optional)\n * @param middlewareOptions.runLimit - Maximum number of model calls allowed per run (optional)\n *\n * @returns A middleware instance that can be passed to `createAgent`\n *\n * @throws {ModelCallLimitMiddlewareError} When either the thread or run limit is exceeded\n *\n * @example\n * ```typescript\n * import { createAgent, modelCallLimitMiddleware } from \"langchain\";\n *\n * // Limit to 10 calls per thread and 3 calls per run\n * const agent = createAgent({\n * model: \"openai:gpt-4o-mini\",\n * tools: [myTool],\n * middleware: [\n * modelCallLimitMiddleware({\n * threadLimit: 10,\n * runLimit: 3\n * })\n * ]\n * });\n * ```\n *\n * @example\n * ```typescript\n * // Limits can also be configured at runtime via context\n * const result = await agent.invoke(\n * { messages: [\"Hello\"] },\n * {\n * configurable: {\n * threadLimit: 5 // Override the default limit for this run\n * }\n * }\n * );\n * ```\n */\nexport function modelCallLimitMiddleware(\n middlewareOptions?: ModelCallLimitMiddlewareConfig\n) {\n return createMiddleware({\n name: \"ModelCallLimitMiddleware\",\n contextSchema,\n stateSchema,\n beforeModel: {\n canJumpTo: [\"end\"],\n hook: (state, runtime) => {\n let exitBehavior =\n runtime.context.exitBehavior ??\n middlewareOptions?.exitBehavior ??\n DEFAULT_EXIT_BEHAVIOR;\n\n // @ts-expect-error - throw is deprecated\n if (exitBehavior === \"throw\") {\n console.warn(\n \"The 'throw' exit behavior is deprecated. Please use 'error' instead.\"\n );\n exitBehavior = \"error\";\n }\n\n const threadLimit =\n runtime.context.threadLimit ?? middlewareOptions?.threadLimit;\n const runLimit =\n runtime.context.runLimit ?? middlewareOptions?.runLimit;\n\n const threadCount = state.threadModelCallCount;\n const runCount = state.runModelCallCount;\n if (typeof threadLimit === \"number\" && threadLimit <= threadCount) {\n const error = new ModelCallLimitMiddlewareError({\n threadLimit,\n threadCount,\n });\n if (exitBehavior === \"end\") {\n return {\n jumpTo: \"end\",\n messages: [new AIMessage(error.message)],\n };\n }\n\n throw error;\n }\n if (typeof runLimit === \"number\" && runLimit <= runCount) {\n const error = new ModelCallLimitMiddlewareError({\n runLimit,\n runCount,\n });\n if (exitBehavior === \"end\") {\n return {\n jumpTo: \"end\",\n messages: [new AIMessage(error.message)],\n };\n }\n\n throw error;\n }\n\n return state;\n },\n },\n afterModel: (state) => ({\n runModelCallCount: state.runModelCallCount + 1,\n threadModelCallCount: state.threadModelCallCount + 1,\n }),\n afterAgent: () => ({\n runModelCallCount: 0,\n }),\n });\n}\n"],"mappings":";;;;;AAMA,MAAM,wBAAwB;AAE9B,MAAM,gBAAgBA,OAAAA,EAAE,OAAO;;;;CAI7B,aAAaA,OAAAA,EAAE,QAAQ,CAAC,UAAU;;;;CAIlC,UAAUA,OAAAA,EAAE,QAAQ,CAAC,UAAU;;;;;;;CAO/B,cAAcA,OAAAA,EAAE,KAAK,CAAC,SAAS,MAAM,CAAC,CAAC,UAAU;CAClD,CAAC;;;;AAQF,MAAM,cAAcA,OAAAA,EAAE,OAAO;CAC3B,sBAAsBA,OAAAA,EAAE,QAAQ,CAAC,QAAQ,EAAE;CAC3C,mBAAmBA,OAAAA,EAAE,QAAQ,CAAC,QAAQ,EAAE;CACzC,CAAC;;;;;;;;;AAUF,IAAM,gCAAN,cAA4C,MAAM;CAChD,YAAY,EACV,aACA,UACA,aACA,YAMC;EACD,MAAM,eAAyB,EAAE;AACjC,MAAI,OAAO,gBAAgB,YAAY,OAAO,gBAAgB,SAC5D,cAAa,KACX,wCAAwC,YAAY,cACrD;AAEH,MAAI,OAAO,aAAa,YAAY,OAAO,aAAa,SACtD,cAAa,KACX,qCAAqC,SAAS,cAC/C;AAEH,QACE,6BACE,aAAa,SAAS,IAAI,KAAK,aAAa,KAAK,KAAK,KAAK,KAE9D;AACD,OAAK,OAAO;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;AAgEhB,SAAgB,yBACd,mBACA;AACA,QAAOC,mBAAAA,iBAAiB;EACtB,MAAM;EACN;EACA;EACA,aAAa;GACX,WAAW,CAAC,MAAM;GAClB,OAAO,OAAO,YAAY;IACxB,IAAI,eACF,QAAQ,QAAQ,gBAChB,mBAAmB,gBACnB;AAGF,QAAI,iBAAiB,SAAS;AAC5B,aAAQ,KACN,uEACD;AACD,oBAAe;;IAGjB,MAAM,cACJ,QAAQ,QAAQ,eAAe,mBAAmB;IACpD,MAAM,WACJ,QAAQ,QAAQ,YAAY,mBAAmB;IAEjD,MAAM,cAAc,MAAM;IAC1B,MAAM,WAAW,MAAM;AACvB,QAAI,OAAO,gBAAgB,YAAY,eAAe,aAAa;KACjE,MAAM,QAAQ,IAAI,8BAA8B;MAC9C;MACA;MACD,CAAC;AACF,SAAI,iBAAiB,MACnB,QAAO;MACL,QAAQ;MACR,UAAU,CAAC,IAAIC,yBAAAA,UAAU,MAAM,QAAQ,CAAC;MACzC;AAGH,WAAM;;AAER,QAAI,OAAO,aAAa,YAAY,YAAY,UAAU;KACxD,MAAM,QAAQ,IAAI,8BAA8B;MAC9C;MACA;MACD,CAAC;AACF,SAAI,iBAAiB,MACnB,QAAO;MACL,QAAQ;MACR,UAAU,CAAC,IAAIA,yBAAAA,UAAU,MAAM,QAAQ,CAAC;MACzC;AAGH,WAAM;;AAGR,WAAO;;GAEV;EACD,aAAa,WAAW;GACtB,mBAAmB,MAAM,oBAAoB;GAC7C,sBAAsB,MAAM,uBAAuB;GACpD;EACD,mBAAmB,EACjB,mBAAmB,GACpB;EACF,CAAC"}