ai-functions
Version:
Core AI primitives for building intelligent applications
110 lines • 4.6 kB
JavaScript
/**
* budgetMiddleware — record token usage + cost into a {@link BudgetTracker}
*
* Replaces the post-hoc duck-typing in
* `services-as-software/src/v3/invoke/cost-estimate.ts` with a single
* AI-SDK-6 middleware: on `doGenerate` / `doStream` completion, read the
* `LanguageModelV3Usage` shape directly off the result and call
* `tracker.recordUsage(...)`. The pricing overlay is supplied via
* `customPricing` on the {@link BudgetTracker} (or we hand the tracker the
* pricing at construction time when the caller wants per-call isolation).
*
* Key V3 → BudgetTracker mapping detail: AI SDK 6 reports
* `usage.inputTokens.total` / `usage.outputTokens.total` as
* `number | undefined`. We coerce undefined → 0 so partial-streaming results
* (where the upstream provider didn't emit token counts) don't blow up the
* tracker. The `inputTokens.cacheRead` / `inputTokens.cacheWrite` breakdown
* is *not* propagated yet — round 13+ work to add prompt-cache awareness to
* BudgetTracker.
*
* Composition note: install **after** cache (so a cache hit still records
* the cost — the wrapped result is the same regardless of which layer
* served it) and **before** trace (so the trace event sees the final
* computed cost via the tracker).
*
* @packageDocumentation
*/
// ============================================================================
// Helpers
// ============================================================================
function coerceUsage(usage) {
if (!usage)
return { inputTokens: 0, outputTokens: 0 };
return {
inputTokens: usage.inputTokens?.total ?? 0,
outputTokens: usage.outputTokens?.total ?? 0,
};
}
function record(tracker, pricing, modelId, usage) {
const { inputTokens, outputTokens } = coerceUsage(usage);
if (inputTokens === 0 && outputTokens === 0)
return;
// The pricing overlay is wired in via the tracker's `customPricing`
// already (set at BudgetTracker construction time by the caller). When
// the caller wants per-call pricing override, they install
// `pricing[modelId]` ahead of time. We expose `pricing` here as a
// forward-looking hook so we can later add per-call pricing without a
// breaking change.
void pricing;
tracker.recordUsage({ inputTokens, outputTokens, model: modelId });
}
// ============================================================================
// Middleware
// ============================================================================
/**
* Build a budget middleware for `wrapLanguageModel`. Records
* {@link LanguageModelV3Usage} into the supplied {@link BudgetTracker} on
* every successful `doGenerate` / `doStream` completion. Errors from the
* downstream model propagate unchanged — the tracker is only updated on
* success.
*
* For streaming calls, we accumulate the final `usage` from the `'finish'`
* stream part (per the V3 spec, the final `'finish'` event carries the
* authoritative usage shape) and record once on stream end.
*
* @example
* ```ts
* import { wrapLanguageModel } from 'ai'
* import { BudgetTracker, budgetMiddleware } from 'ai-functions'
*
* const tracker = new BudgetTracker({ maxCost: 1.0 })
* const model = wrapLanguageModel({
* model: openai('gpt-4o'),
* middleware: budgetMiddleware({ tracker }),
* })
* ```
*/
export function budgetMiddleware(options) {
const { tracker, pricing, modelIdOverride } = options;
return {
specificationVersion: 'v3',
async wrapGenerate({ doGenerate, model }) {
const result = await doGenerate();
const modelId = modelIdOverride ?? model.modelId;
record(tracker, pricing, modelId, result.usage);
return result;
},
async wrapStream({ doStream, model }) {
const result = await doStream();
const modelId = modelIdOverride ?? model.modelId;
let finalUsage;
const transformedStream = result.stream.pipeThrough(new TransformStream({
transform(chunk, controller) {
if (chunk.type === 'finish') {
finalUsage = chunk.usage;
}
controller.enqueue(chunk);
},
flush() {
record(tracker, pricing, modelId, finalUsage);
},
}));
const wrapped = {
...result,
stream: transformedStream,
};
return wrapped;
},
};
}
//# sourceMappingURL=budget.js.map