@stackfactor/agent-utils 1.0.19 → 1.0.20

package/dist/cjs/const.d.ts

@@ -1,6 +1,7 @@
 declare const _default: {
     HTTP_CODES: {
         BAD_REQUEST: number;
+        PAYMENT_REQUIRED: number;
         UNPROCESSABLE_ENTITY: number;
         INTERNAL_SERVER_ERROR: number;
         BAD_GATEWAY: number;
@@ -9,6 +10,7 @@ declare const _default: {
         UNABLE_TO_GENERATE_CONTENT: string;
         UNEXPECTED_ERROR: string;
         UNSUPPORTED_MODEL: string;
+        QUOTA_EXHAUSTED: string;
     };
 };
 export default _default;

package/dist/cjs/const.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"const.d.ts","sourceRoot":"","sources":["../../src/const.ts"],"names":[],"mappings":";;;;;;;;;;;;;AAAA,wBAaE"}
+{"version":3,"file":"const.d.ts","sourceRoot":"","sources":["../../src/const.ts"],"names":[],"mappings":";;;;;;;;;;;;;;;AAAA,wBAgBE"}

package/dist/cjs/const.js

@@ -3,6 +3,7 @@ Object.defineProperty(exports, "__esModule", { value: true });
 exports.default = {
     HTTP_CODES: {
         BAD_REQUEST: 400,
+        PAYMENT_REQUIRED: 402,
         UNPROCESSABLE_ENTITY: 422,
         INTERNAL_SERVER_ERROR: 500,
         BAD_GATEWAY: 502,
@@ -11,5 +12,6 @@ exports.default = {
         UNABLE_TO_GENERATE_CONTENT: "Unable to generate content",
         UNEXPECTED_ERROR: "An unexpected error occured. If the issue persists please contact the StackFactor support team at support@stackfactor.ai",
         UNSUPPORTED_MODEL: "The specified model is not supported",
+        QUOTA_EXHAUSTED: "Agent session quota exhausted: no remaining budget for additional LLM calls.",
     },
 };

package/dist/cjs/langChain.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"langChain.d.ts","sourceRoot":"","sources":["../../src/langChain.ts"],"names":[],"mappings":";wBAgWQ,MAAM,aACD,MAAM,gBACH,MAAM,SACb,GAAG,EAAE,kBACI,GAAG,UACX,GAAG,KACV,GAAG;sBAyBG,GAAG,UACF,MAAM,UACN,GAAG,eACC,QAAQ,GAAG,IAAI,KAC1B,OAAO,CAAC,GAAG,CAAC;wCAmfF,MAAM,UACT,GAAG,UACH,GAAG,oBACO,GAAG,KACpB,GAAG;oCAxYO,MAAM,UACT,GAAG,UACH,GAAG,oBACO,GAAG,eACT,MAAM,eACN,MAAM,wBACG,OAAO,WACpB,GAAG,cACA,MAAM,UACV,GAAG,EAAE,KACX,OAAO,CAAC,GAAG,CAAC;sDA6rBF,MAAM,UACT,GAAG,UACH,MAAM,YACL,GAAG,KACX,OAAO,CAAC,GAAG,CAAC;0CA7vB8B,GAAG,KAAG,MAAM;;AAwzBzD,wBAOE"}
+{"version":3,"file":"langChain.d.ts","sourceRoot":"","sources":["../../src/langChain.ts"],"names":[],"mappings":";wBAogBQ,MAAM,aACD,MAAM,gBACH,MAAM,SACb,GAAG,EAAE,kBACI,GAAG,UACX,GAAG,KACV,GAAG;sBAyBG,GAAG,UACF,MAAM,UACN,GAAG,eACC,QAAQ,GAAG,IAAI,KAC1B,OAAO,CAAC,GAAG,CAAC;wCAmgBF,MAAM,UACT,GAAG,UACH,GAAG,oBACO,GAAG,KACpB,GAAG;oCAhZO,MAAM,UACT,GAAG,UACH,GAAG,oBACO,GAAG,eACT,MAAM,eACN,MAAM,wBACG,OAAO,WACpB,GAAG,cACA,MAAM,UACV,GAAG,EAAE,KACX,OAAO,CAAC,GAAG,CAAC;sDA4sBF,MAAM,UACT,GAAG,UACH,MAAM,YACL,GAAG,KACX,OAAO,CAAC,GAAG,CAAC;0CA5wB8B,GAAG,KAAG,MAAM;;AAu0BzD,wBAOE"}

package/dist/cjs/langChain.js

@@ -16,12 +16,157 @@ const zod_to_json_schema_1 = require("zod-to-json-schema");
 const JSON_ESCAPE_INSTRUCTION = `
 CRITICAL - Your response must be valid JSON. Escape ALL special characters in string values:
 - Newlines → \\n
-- Tabs → \\t  
+- Tabs → \\t
 - Carriage returns → \\r
 - Double quotes inside strings → \\"
 - Backslashes → \\\\
 Do NOT include raw newlines, tabs, or unescaped quotes inside JSON string values.
 `.trim();
+/**
+ * Reads `global.quota` (set by the StackFactor host before invoking `main()`) and
+ * throws a `PAYMENT_REQUIRED` error when the session has no remaining budget. No-op
+ * when `global.quota` is absent (e.g. local tests) or `remaining` is not a number,
+ * which keeps the library usable outside the StackFactor runtime.
+ */
+const assertQuotaAvailable = () => {
+    const quota = globalThis.quota;
+    if (!quota || typeof quota.remaining !== "number")
+        return;
+    if (quota.remaining <= 0) {
+        throw errorHandling_js_1.default.create(const_js_1.default.HTTP_CODES.PAYMENT_REQUIRED, const_js_1.default.ERROR.QUOTA_EXHAUSTED);
+    }
+};
+/**
+ * Subtracts the given USD cost from `global.quota.remaining` and adds it to
+ * `global.quota.usedThisSession`. Silently no-ops when the quota globals are not
+ * present so callers do not need to branch. Negative, zero, or non-finite costs
+ * are ignored to keep counters monotonic.
+ */
+const recordCost = (cost) => {
+    if (!cost || cost <= 0 || !Number.isFinite(cost))
+        return;
+    const quota = globalThis.quota;
+    if (!quota)
+        return;
+    if (typeof quota.remaining === "number") {
+        quota.remaining -= cost;
+    }
+    if (typeof quota.usedThisSession === "number") {
+        quota.usedThisSession += cost;
+    }
+    else {
+        quota.usedThisSession = cost;
+    }
+};
+/**
+ * Looks up the per-model pricing entry from `config.modelPricing`. Returns `null`
+ * when pricing is not configured for the model, in which case cost recording
+ * becomes a no-op (the call still succeeds — pricing data is the host's
+ * responsibility, not the agent's).
+ */
+const getModelPrice = (modelName, config) => {
+    const pricing = config?.modelPricing;
+    if (!pricing)
+        return null;
+    return pricing[modelName] || null;
+};
+/**
+ * Computes USD cost for a text LLM call. `config.modelPricing[modelName]` is
+ * expected to provide `input` and `output` rates in dollars per million tokens.
+ */
+const calculateTextCost = (modelName, usage, config) => {
+    if (!usage)
+        return 0;
+    const price = getModelPrice(modelName, config);
+    if (!price)
+        return 0;
+    const inputCost = ((usage.input_tokens || 0) / 1_000_000) * (price.input || 0);
+    const outputCost = ((usage.output_tokens || 0) / 1_000_000) * (price.output || 0);
+    return inputCost + outputCost;
+};
+/**
+ * Computes USD cost for an image generation call. `config.modelPricing[modelName]`
+ * is expected to provide a `perImage` rate in dollars.
+ */
+const calculateImageCost = (modelName, numImages, config) => {
+    const price = getModelPrice(modelName, config);
+    if (!price)
+        return 0;
+    return (numImages || 0) * (price.perImage || 0);
+};
+/**
+ * Extracts a normalized token-usage object from a single LangChain `invoke()`
+ * response. Reads from `usage_metadata` first (standardized in LangChain v1),
+ * then falls back to provider-specific shapes under `response_metadata`.
+ */
+const extractUsageFromInvoke = (response) => {
+    if (!response)
+        return null;
+    const um = response.usage_metadata;
+    if (um) {
+        return {
+            input_tokens: um.input_tokens || 0,
+            output_tokens: um.output_tokens || 0,
+        };
+    }
+    const rm = response.response_metadata;
+    if (rm?.usage) {
+        return {
+            input_tokens: rm.usage.input_tokens || rm.usage.prompt_tokens || 0,
+            output_tokens: rm.usage.output_tokens || rm.usage.completion_tokens || 0,
+        };
+    }
+    if (rm?.tokenUsage) {
+        return {
+            input_tokens: rm.tokenUsage.promptTokens || 0,
+            output_tokens: rm.tokenUsage.completionTokens || 0,
+        };
+    }
+    return null;
+};
+/**
+ * Accumulates token usage from a streaming chunk into a running total. LangChain
+ * typically attaches `usage_metadata` to the final chunk; earlier chunks carry no
+ * usage and are no-ops here.
+ */
+const accumulateChunkUsage = (acc, chunk) => {
+    if (!chunk)
+        return acc;
+    const um = chunk.usage_metadata;
+    if (um) {
+        acc.input_tokens += um.input_tokens || 0;
+        acc.output_tokens += um.output_tokens || 0;
+        return acc;
+    }
+    const rm = chunk.response_metadata;
+    if (rm?.usage) {
+        acc.input_tokens += rm.usage.input_tokens || rm.usage.prompt_tokens || 0;
+        acc.output_tokens +=
+            rm.usage.output_tokens || rm.usage.completion_tokens || 0;
+    }
+    return acc;
+};
+/**
+ * Sums token usage across every message in an agent invocation response. Each
+ * AIMessage in `response.messages` may carry its own `usage_metadata` (one per
+ * LLM round-trip the agent made).
+ */
+const sumAgentResponseUsage = (response) => {
+    const messages = response?.messages;
+    if (!Array.isArray(messages) || messages.length === 0)
+        return null;
+    const total = { input_tokens: 0, output_tokens: 0 };
+    let found = false;
+    for (const msg of messages) {
+        const um = msg?.usage_metadata;
+        if (um) {
+            total.input_tokens += um.input_tokens || 0;
+            total.output_tokens += um.output_tokens || 0;
+            found = true;
+        }
+    }
+    return found ? total : null;
+};
 /**
  * Converts a Zod validation error object into a human-readable multi-line string.
  * Each failing field is described with its dot-notation path and a contextual message
@@ -377,12 +522,15 @@ const runAgent = async (agent, prompt, config, onProgress = null) => {
             },
         ]
         : undefined;
+    assertQuotaAvailable();
     const response = await agent.invoke({
         messages: [{ role: "user", content: prompt }],
     }, {
         recursionLimit: config.recursionLimit || 25,
         ...(callbacks ? { callbacks } : {}),
     });
+    const modelName = agent.options?.model?.modelName || agent.options?.model?.model || "";
+    recordCost(calculateTextCost(modelName, sumAgentResponseUsage(response), config));
     const endTime = Date.now();
     const duration = endTime - startTime;
     logger_js_1.default.log(null, logger_js_1.default.levels.info, `Agent "${agent.options?.name}" completed in ${Math.round(duration / 1000)} seconds.`);
@@ -589,8 +737,11 @@ const runPromptWithModel = async (modelName, config, prompt, onProgressReport, m
         // overshoots. This avoids the magic "expected length" constant — longer responses
         // simply slow the curve down rather than exceeding the range.
         const expectedDurationMs = 15_000; // Tune: expected typical response time
+        const streamUsage = { input_tokens: 0, output_tokens: 0 };
+        assertQuotaAvailable();
         const stream = await llm.stream(messagesToSend);
         for await (const chunk of stream) {
+            accumulateChunkUsage(streamUsage, chunk);
             const content = chunk?.content || chunk;
             if (typeof content === "string") {
                 rawContent += content;
@@ -608,6 +759,7 @@ const runPromptWithModel = async (modelName, config, prompt, onProgressReport, m
                 }
             }
         }
+        recordCost(calculateTextCost(modelName, streamUsage, config));
         await onProgressReport({
             message: "Processing complete",
             progress: maxPercent,
@@ -694,7 +846,9 @@ const runPromptWithModel = async (modelName, config, prompt, onProgressReport, m
             messagesToSend = messages;
         }
         // Simply invoke without streaming
+        assertQuotaAvailable();
         const response = await llm.invoke(messagesToSend);
+        recordCost(calculateTextCost(modelName, extractUsageFromInvoke(response), config));
         const rawContent = response?.content || response;
         // If not expecting JSON, return raw content directly
         if (!expectsJsonResponse) {
@@ -877,7 +1031,9 @@ const generateImageWithOpenAI = async (modelName, config, prompt, options) => {
             requestParams.style = style;
         }
     }
+    assertQuotaAvailable();
     const response = await openai.images.generate(requestParams);
+    recordCost(calculateImageCost(modelName, response.data?.length || n, config));
     // Format response based on number of images
     if (n === 1) {
         const imageData = response.data[0];
@@ -975,6 +1131,7 @@ const generateImageWithGoogle = async (modelName, config, prompt, options) => {
             safetySettings: safetySettings,
         },
     };
+    assertQuotaAvailable();
     const response = await ai.models.generateContent(req);
     // Extract images from response
     const images = [];
@@ -993,6 +1150,7 @@ const generateImageWithGoogle = async (modelName, config, prompt, options) => {
     if (images.length === 0) {
         throw errorHandling_js_1.default.create(const_js_1.default.HTTP_CODES.INTERNAL_SERVER_ERROR, `No images were generated by ${modelName}`);
     }
+    recordCost(calculateImageCost(modelName, images.length, config));
     if (numberOfImages === 1 || images.length === 1) {
         return images[0];
     }

package/dist/esm/const.d.ts

@@ -1,6 +1,7 @@
 declare const _default: {
     HTTP_CODES: {
         BAD_REQUEST: number;
+        PAYMENT_REQUIRED: number;
         UNPROCESSABLE_ENTITY: number;
         INTERNAL_SERVER_ERROR: number;
         BAD_GATEWAY: number;
@@ -9,6 +10,7 @@ declare const _default: {
         UNABLE_TO_GENERATE_CONTENT: string;
         UNEXPECTED_ERROR: string;
         UNSUPPORTED_MODEL: string;
+        QUOTA_EXHAUSTED: string;
     };
 };
 export default _default;

package/dist/esm/const.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"const.d.ts","sourceRoot":"","sources":["../../src/const.ts"],"names":[],"mappings":";;;;;;;;;;;;;AAAA,wBAaE"}
+{"version":3,"file":"const.d.ts","sourceRoot":"","sources":["../../src/const.ts"],"names":[],"mappings":";;;;;;;;;;;;;;;AAAA,wBAgBE"}

package/dist/esm/const.js

@@ -1,6 +1,7 @@
 export default {
     HTTP_CODES: {
         BAD_REQUEST: 400,
+        PAYMENT_REQUIRED: 402,
         UNPROCESSABLE_ENTITY: 422,
         INTERNAL_SERVER_ERROR: 500,
         BAD_GATEWAY: 502,
@@ -9,5 +10,6 @@ export default {
         UNABLE_TO_GENERATE_CONTENT: "Unable to generate content",
         UNEXPECTED_ERROR: "An unexpected error occured. If the issue persists please contact the StackFactor support team at support@stackfactor.ai",
         UNSUPPORTED_MODEL: "The specified model is not supported",
+        QUOTA_EXHAUSTED: "Agent session quota exhausted: no remaining budget for additional LLM calls.",
     },
 };

package/dist/esm/langChain.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"langChain.d.ts","sourceRoot":"","sources":["../../src/langChain.ts"],"names":[],"mappings":";wBAgWQ,MAAM,aACD,MAAM,gBACH,MAAM,SACb,GAAG,EAAE,kBACI,GAAG,UACX,GAAG,KACV,GAAG;sBAyBG,GAAG,UACF,MAAM,UACN,GAAG,eACC,QAAQ,GAAG,IAAI,KAC1B,OAAO,CAAC,GAAG,CAAC;wCAmfF,MAAM,UACT,GAAG,UACH,GAAG,oBACO,GAAG,KACpB,GAAG;oCAxYO,MAAM,UACT,GAAG,UACH,GAAG,oBACO,GAAG,eACT,MAAM,eACN,MAAM,wBACG,OAAO,WACpB,GAAG,cACA,MAAM,UACV,GAAG,EAAE,KACX,OAAO,CAAC,GAAG,CAAC;sDA6rBF,MAAM,UACT,GAAG,UACH,MAAM,YACL,GAAG,KACX,OAAO,CAAC,GAAG,CAAC;0CA7vB8B,GAAG,KAAG,MAAM;;AAwzBzD,wBAOE"}
+{"version":3,"file":"langChain.d.ts","sourceRoot":"","sources":["../../src/langChain.ts"],"names":[],"mappings":";wBAogBQ,MAAM,aACD,MAAM,gBACH,MAAM,SACb,GAAG,EAAE,kBACI,GAAG,UACX,GAAG,KACV,GAAG;sBAyBG,GAAG,UACF,MAAM,UACN,GAAG,eACC,QAAQ,GAAG,IAAI,KAC1B,OAAO,CAAC,GAAG,CAAC;wCAmgBF,MAAM,UACT,GAAG,UACH,GAAG,oBACO,GAAG,KACpB,GAAG;oCAhZO,MAAM,UACT,GAAG,UACH,GAAG,oBACO,GAAG,eACT,MAAM,eACN,MAAM,wBACG,OAAO,WACpB,GAAG,cACA,MAAM,UACV,GAAG,EAAE,KACX,OAAO,CAAC,GAAG,CAAC;sDA4sBF,MAAM,UACT,GAAG,UACH,MAAM,YACL,GAAG,KACX,OAAO,CAAC,GAAG,CAAC;0CA5wB8B,GAAG,KAAG,MAAM;;AAu0BzD,wBAOE"}

package/dist/esm/langChain.js

@@ -11,12 +11,157 @@ import { zodToJsonSchema } from "zod-to-json-schema";
 const JSON_ESCAPE_INSTRUCTION = `
 CRITICAL - Your response must be valid JSON. Escape ALL special characters in string values:
 - Newlines → \\n
-- Tabs → \\t  
+- Tabs → \\t
 - Carriage returns → \\r
 - Double quotes inside strings → \\"
 - Backslashes → \\\\
 Do NOT include raw newlines, tabs, or unescaped quotes inside JSON string values.
 `.trim();
+/**
+ * Reads `global.quota` (set by the StackFactor host before invoking `main()`) and
+ * throws a `PAYMENT_REQUIRED` error when the session has no remaining budget. No-op
+ * when `global.quota` is absent (e.g. local tests) or `remaining` is not a number,
+ * which keeps the library usable outside the StackFactor runtime.
+ */
+const assertQuotaAvailable = () => {
+    const quota = globalThis.quota;
+    if (!quota || typeof quota.remaining !== "number")
+        return;
+    if (quota.remaining <= 0) {
+        throw errorHandlingHelper.create(constants.HTTP_CODES.PAYMENT_REQUIRED, constants.ERROR.QUOTA_EXHAUSTED);
+    }
+};
+/**
+ * Subtracts the given USD cost from `global.quota.remaining` and adds it to
+ * `global.quota.usedThisSession`. Silently no-ops when the quota globals are not
+ * present so callers do not need to branch. Negative, zero, or non-finite costs
+ * are ignored to keep counters monotonic.
+ */
+const recordCost = (cost) => {
+    if (!cost || cost <= 0 || !Number.isFinite(cost))
+        return;
+    const quota = globalThis.quota;
+    if (!quota)
+        return;
+    if (typeof quota.remaining === "number") {
+        quota.remaining -= cost;
+    }
+    if (typeof quota.usedThisSession === "number") {
+        quota.usedThisSession += cost;
+    }
+    else {
+        quota.usedThisSession = cost;
+    }
+};
+/**
+ * Looks up the per-model pricing entry from `config.modelPricing`. Returns `null`
+ * when pricing is not configured for the model, in which case cost recording
+ * becomes a no-op (the call still succeeds — pricing data is the host's
+ * responsibility, not the agent's).
+ */
+const getModelPrice = (modelName, config) => {
+    const pricing = config?.modelPricing;
+    if (!pricing)
+        return null;
+    return pricing[modelName] || null;
+};
+/**
+ * Computes USD cost for a text LLM call. `config.modelPricing[modelName]` is
+ * expected to provide `input` and `output` rates in dollars per million tokens.
+ */
+const calculateTextCost = (modelName, usage, config) => {
+    if (!usage)
+        return 0;
+    const price = getModelPrice(modelName, config);
+    if (!price)
+        return 0;
+    const inputCost = ((usage.input_tokens || 0) / 1_000_000) * (price.input || 0);
+    const outputCost = ((usage.output_tokens || 0) / 1_000_000) * (price.output || 0);
+    return inputCost + outputCost;
+};
+/**
+ * Computes USD cost for an image generation call. `config.modelPricing[modelName]`
+ * is expected to provide a `perImage` rate in dollars.
+ */
+const calculateImageCost = (modelName, numImages, config) => {
+    const price = getModelPrice(modelName, config);
+    if (!price)
+        return 0;
+    return (numImages || 0) * (price.perImage || 0);
+};
+/**
+ * Extracts a normalized token-usage object from a single LangChain `invoke()`
+ * response. Reads from `usage_metadata` first (standardized in LangChain v1),
+ * then falls back to provider-specific shapes under `response_metadata`.
+ */
+const extractUsageFromInvoke = (response) => {
+    if (!response)
+        return null;
+    const um = response.usage_metadata;
+    if (um) {
+        return {
+            input_tokens: um.input_tokens || 0,
+            output_tokens: um.output_tokens || 0,
+        };
+    }
+    const rm = response.response_metadata;
+    if (rm?.usage) {
+        return {
+            input_tokens: rm.usage.input_tokens || rm.usage.prompt_tokens || 0,
+            output_tokens: rm.usage.output_tokens || rm.usage.completion_tokens || 0,
+        };
+    }
+    if (rm?.tokenUsage) {
+        return {
+            input_tokens: rm.tokenUsage.promptTokens || 0,
+            output_tokens: rm.tokenUsage.completionTokens || 0,
+        };
+    }
+    return null;
+};
+/**
+ * Accumulates token usage from a streaming chunk into a running total. LangChain
+ * typically attaches `usage_metadata` to the final chunk; earlier chunks carry no
+ * usage and are no-ops here.
+ */
+const accumulateChunkUsage = (acc, chunk) => {
+    if (!chunk)
+        return acc;
+    const um = chunk.usage_metadata;
+    if (um) {
+        acc.input_tokens += um.input_tokens || 0;
+        acc.output_tokens += um.output_tokens || 0;
+        return acc;
+    }
+    const rm = chunk.response_metadata;
+    if (rm?.usage) {
+        acc.input_tokens += rm.usage.input_tokens || rm.usage.prompt_tokens || 0;
+        acc.output_tokens +=
+            rm.usage.output_tokens || rm.usage.completion_tokens || 0;
+    }
+    return acc;
+};
+/**
+ * Sums token usage across every message in an agent invocation response. Each
+ * AIMessage in `response.messages` may carry its own `usage_metadata` (one per
+ * LLM round-trip the agent made).
+ */
+const sumAgentResponseUsage = (response) => {
+    const messages = response?.messages;
+    if (!Array.isArray(messages) || messages.length === 0)
+        return null;
+    const total = { input_tokens: 0, output_tokens: 0 };
+    let found = false;
+    for (const msg of messages) {
+        const um = msg?.usage_metadata;
+        if (um) {
+            total.input_tokens += um.input_tokens || 0;
+            total.output_tokens += um.output_tokens || 0;
+            found = true;
+        }
+    }
+    return found ? total : null;
+};
 /**
  * Converts a Zod validation error object into a human-readable multi-line string.
  * Each failing field is described with its dot-notation path and a contextual message
@@ -372,12 +517,15 @@ const runAgent = async (agent, prompt, config, onProgress = null) => {
             },
         ]
         : undefined;
+    assertQuotaAvailable();
     const response = await agent.invoke({
         messages: [{ role: "user", content: prompt }],
     }, {
         recursionLimit: config.recursionLimit || 25,
         ...(callbacks ? { callbacks } : {}),
     });
+    const modelName = agent.options?.model?.modelName || agent.options?.model?.model || "";
+    recordCost(calculateTextCost(modelName, sumAgentResponseUsage(response), config));
     const endTime = Date.now();
     const duration = endTime - startTime;
     logger.log(null, logger.levels.info, `Agent "${agent.options?.name}" completed in ${Math.round(duration / 1000)} seconds.`);
@@ -584,8 +732,11 @@ const runPromptWithModel = async (modelName, config, prompt, onProgressReport, m
         // overshoots. This avoids the magic "expected length" constant — longer responses
         // simply slow the curve down rather than exceeding the range.
         const expectedDurationMs = 15_000; // Tune: expected typical response time
+        const streamUsage = { input_tokens: 0, output_tokens: 0 };
+        assertQuotaAvailable();
         const stream = await llm.stream(messagesToSend);
         for await (const chunk of stream) {
+            accumulateChunkUsage(streamUsage, chunk);
             const content = chunk?.content || chunk;
             if (typeof content === "string") {
                 rawContent += content;
@@ -603,6 +754,7 @@ const runPromptWithModel = async (modelName, config, prompt, onProgressReport, m
                 }
             }
         }
+        recordCost(calculateTextCost(modelName, streamUsage, config));
         await onProgressReport({
             message: "Processing complete",
             progress: maxPercent,
@@ -689,7 +841,9 @@ const runPromptWithModel = async (modelName, config, prompt, onProgressReport, m
             messagesToSend = messages;
         }
         // Simply invoke without streaming
+        assertQuotaAvailable();
         const response = await llm.invoke(messagesToSend);
+        recordCost(calculateTextCost(modelName, extractUsageFromInvoke(response), config));
         const rawContent = response?.content || response;
         // If not expecting JSON, return raw content directly
         if (!expectsJsonResponse) {
@@ -872,7 +1026,9 @@ const generateImageWithOpenAI = async (modelName, config, prompt, options) => {
             requestParams.style = style;
         }
     }
+    assertQuotaAvailable();
     const response = await openai.images.generate(requestParams);
+    recordCost(calculateImageCost(modelName, response.data?.length || n, config));
     // Format response based on number of images
     if (n === 1) {
         const imageData = response.data[0];
@@ -970,6 +1126,7 @@ const generateImageWithGoogle = async (modelName, config, prompt, options) => {
             safetySettings: safetySettings,
         },
     };
+    assertQuotaAvailable();
     const response = await ai.models.generateContent(req);
     // Extract images from response
     const images = [];
@@ -988,6 +1145,7 @@ const generateImageWithGoogle = async (modelName, config, prompt, options) => {
     if (images.length === 0) {
         throw errorHandlingHelper.create(constants.HTTP_CODES.INTERNAL_SERVER_ERROR, `No images were generated by ${modelName}`);
     }
+    recordCost(calculateImageCost(modelName, images.length, config));
     if (numberOfImages === 1 || images.length === 1) {
         return images[0];
     }

package/package.json

@@ -3,7 +3,7 @@
   "publishConfig": {
     "access": "restricted"
   },
-  "version": "1.0.19",
+  "version": "1.0.20",
   "description": "",
   "main": "dist/cjs/index.js",
   "module": "dist/esm/index.js",