diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.metering.test.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.metering.test.ts index 9609f08ea..b579b8576 100644 --- a/src/backend/drivers/ai-chat/ChatCompletionDriver.metering.test.ts +++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.metering.test.ts @@ -434,6 +434,119 @@ describe('ChatCompletionDriver credit gate against real metering', () => { }); }); +describe('ChatCompletionDriver credit gate with an AI cost factor', () => { + const FACTOR = 1.3; + let seenKeys: string[]; + const listener = (key: string, event: { factor: number }) => { + seenKeys.push(key); + event.factor = FACTOR; + }; + + beforeEach(() => { + seenKeys = []; + server.clients.event.on('ai.cost.factor.*', listener as never); + }); + + afterEach(() => { + server.clients.event.off('ai.cost.factor.*', listener as never); + }); + + const okResult = { + message: { role: 'assistant', content: [{ type: 'text', text: 'ok' }] }, + usage: { input_tokens: 1, output_tokens: 1 }, + finish_reason: 'stop', + } as never; + + // Usage is recorded at cost × factor, so a cap sized at the bare cost + // lets a completion spend past what's left. + it('caps output at the factored price', async () => { + const actor = freeUser() as never; + const metering = server.services.metering; + const allowance = (await metering.getActorSubscription(actor)) + .monthUsageAllowance; + await metering.incrementUsage( + actor, + 'test:prior-spend', + 1, + Math.floor(allowance * 0.9), + ); + const remaining = await metering.getRemainingUsage(actor); + const completeSpy = vi + .spyOn(FakeChatProvider.prototype, 'complete') + .mockResolvedValueOnce(okResult); + + await withTestActor( + () => + driver.complete({ + model: 'priced', + messages: [{ role: 'user', content: 'go' }], + max_tokens: 100_000, + }), + actor, + ); + + const passed = completeSpy.mock.calls[0]![0] as { max_tokens: number }; + expect(passed.max_tokens * 2000 * FACTOR).toBeLessThanOrEqual( + remaining, + ); + expect(passed.max_tokens).toBeGreaterThan( + remaining / (2000 * FACTOR) - 10, + ); + }); + + it('reserves the factored worst case', async () => { + const actor = freeUser() as never; + const reserve = vi.spyOn(server.services.metering, 'reserveCredits'); + const completeSpy = vi + .spyOn(FakeChatProvider.prototype, 'complete') + .mockResolvedValueOnce(okResult); + + await withTestActor( + () => + driver.complete({ + model: 'priced', + messages: [{ role: 'user', content: 'go' }], + max_tokens: 10, + }), + actor, + ); + + const passed = completeSpy.mock.calls[0]![0] as { max_tokens: number }; + expect(passed.max_tokens).toBe(10); + const [, held] = reserve.mock.calls[0]!; + // At least the 10 capped output tokens at 2000 × factor. + expect(held).toBeGreaterThanOrEqual(10 * 2000 * FACTOR); + }); + + // The factor is looked up by the key the provider records under, or a + // per-model factor would price the gate and the charge differently. + it('asks for the factor under the metering key the provider records with', async () => { + const proto = FakeChatProvider.prototype as unknown as { + meteringModelKey?: (id: string) => string; + }; + proto.meteringModelKey = (id) => `fake-recorded:${id}`; + try { + vi.spyOn( + FakeChatProvider.prototype, + 'complete', + ).mockResolvedValueOnce(okResult); + await withTestActor( + () => + driver.complete({ + model: 'priced', + messages: [{ role: 'user', content: 'go' }], + }), + freeUser() as never, + ); + expect(seenKeys).toContain( + 'ai.cost.factor.ai-chat.fake-recorded:priced', + ); + } finally { + delete proto.meteringModelKey; + } + }); +}); + describe('ChatCompletionDriver credit gate on multimodal prompts', () => { it('prices attachments into the pre-flight affordability check', async () => { const actor = freeUser() as never; diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.ts index 65f939b2b..005f8c4f6 100644 --- a/src/backend/drivers/ai-chat/ChatCompletionDriver.ts +++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.ts @@ -948,11 +948,21 @@ export class ChatCompletionDriver extends PuterDriver { // A prompt estimated past a long-context threshold pays the raised // rates on input and output alike. const multipliers = longContextMultipliers(model, promptTokenEstimate); + // Recorded costs pass through the AI cost factor; the gate prices at + // the same rate or it under-reserves on every request. + const costFactor = await this.#aiMetering.costFactor( + actor, + this.#meteringModelKey(model), + ); // `|| 0` also catches NaN from a malformed cost table. const inputTokenCost = - (Number(model.costs?.[inputKey] ?? 0) || 0) * multipliers.input; + (Number(model.costs?.[inputKey] ?? 0) || 0) * + multipliers.input * + costFactor; const outputTokenCost = - (Number(model.costs?.[outputKey] ?? 0) || 0) * multipliers.output; + (Number(model.costs?.[outputKey] ?? 0) || 0) * + multipliers.output * + costFactor; const approximateInputCost = promptTokenEstimate * inputTokenCost; const minimumCredits = Number(model.minimumCredits || 1); @@ -1021,7 +1031,8 @@ export class ChatCompletionDriver extends PuterDriver { // What this attempt can cost at worst: the prompt, plus output run to // the cap just set. Capped output is what makes the number finite — // for a model with no output price the output term is zero and the - // prompt estimate stands alone. + // prompt estimate stands alone. Already factored, so it's held through + // the plain service rather than scaled a second time. const worstCaseCost = approximateInputCost + (args.max_tokens ?? 0) * outputTokenCost; return this.services.metering.reserveCredits( @@ -1030,6 +1041,14 @@ export class ChatCompletionDriver extends PuterDriver { ); } + /** The key `model`'s usage is recorded and cost-factored under. */ + #meteringModelKey(model: IChatModel): string { + return ( + this.#providers[model.provider!]?.meteringModelKey?.(model.id) ?? + `${model.provider}:${model.id}` + ); + } + /** * Charge a stream that produced output but never reported usage. * @@ -1085,7 +1104,7 @@ export class ChatCompletionDriver extends PuterDriver { [`estimated_${outputKey}`]: outputTokens, }, actor, - `${model.provider}:${model.id}`, + this.#meteringModelKey(model), { // Undefined when the model has no cost table: the entry is // recorded unpriced rather than free. diff --git a/src/backend/drivers/ai-chat/providers/alibaba/AlibabaProvider.ts b/src/backend/drivers/ai-chat/providers/alibaba/AlibabaProvider.ts index 0ab886afb..61b0999f8 100644 --- a/src/backend/drivers/ai-chat/providers/alibaba/AlibabaProvider.ts +++ b/src/backend/drivers/ai-chat/providers/alibaba/AlibabaProvider.ts @@ -58,6 +58,11 @@ export class AlibabaProvider implements IChatProvider { return modelLookupNames(this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `alibaba:${modelId}`; + } + async complete({ messages, stream, @@ -100,7 +105,7 @@ export class AlibabaProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor!, - `alibaba:${modelUsed.id}`, + this.meteringModelKey(modelUsed.id), costsOverride, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/azure/AzureChatProvider.ts b/src/backend/drivers/ai-chat/providers/azure/AzureChatProvider.ts index bbc92c4af..9e55e0345 100644 --- a/src/backend/drivers/ai-chat/providers/azure/AzureChatProvider.ts +++ b/src/backend/drivers/ai-chat/providers/azure/AzureChatProvider.ts @@ -112,6 +112,11 @@ export class AzureChatProvider implements IChatProvider { return this.#defaultModel; } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string | undefined): string { + return `azure-openai:${modelId}`; + } + async complete( params: ICompleteArguments, ): ReturnType { @@ -275,7 +280,7 @@ export class AzureChatProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor!, - `azure-openai:${modelUsed?.id}`, + this.meteringModelKey(modelUsed?.id), costsOverrideFromModel, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/azure/AzureResponsesProvider.ts b/src/backend/drivers/ai-chat/providers/azure/AzureResponsesProvider.ts index 2d24a4c92..ee97209fd 100644 --- a/src/backend/drivers/ai-chat/providers/azure/AzureResponsesProvider.ts +++ b/src/backend/drivers/ai-chat/providers/azure/AzureResponsesProvider.ts @@ -94,6 +94,11 @@ export class AzureResponsesProvider implements IChatProvider { return this.#defaultModel; } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string | undefined): string { + return `azure-openai:${modelId}`; + } + async complete({ messages, model, @@ -262,7 +267,7 @@ export class AzureResponsesProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor, - `azure-openai:${modelUsed?.id}`, + this.meteringModelKey(modelUsed?.id), costsOverrideFromModel, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/byteplus/BytePlusProvider.ts b/src/backend/drivers/ai-chat/providers/byteplus/BytePlusProvider.ts index b6f2a9921..31bd1f2d1 100644 --- a/src/backend/drivers/ai-chat/providers/byteplus/BytePlusProvider.ts +++ b/src/backend/drivers/ai-chat/providers/byteplus/BytePlusProvider.ts @@ -80,6 +80,11 @@ export class BytePlusProvider implements IChatProvider { return modelLookupNames(this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `byteplus:${modelId}`; + } + async complete( params: ICompleteArguments, ): ReturnType { @@ -151,7 +156,7 @@ export class BytePlusProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor!, - `byteplus:${modelUsed.id}`, + this.meteringModelKey(modelUsed.id), costsOverride, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/claude/ClaudeProvider.ts b/src/backend/drivers/ai-chat/providers/claude/ClaudeProvider.ts index ff88f533c..8dd181278 100644 --- a/src/backend/drivers/ai-chat/providers/claude/ClaudeProvider.ts +++ b/src/backend/drivers/ai-chat/providers/claude/ClaudeProvider.ts @@ -135,6 +135,11 @@ export class ClaudeProvider implements IChatProvider { return modelLookupNames(this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `claude:${modelId}`; + } + async complete({ messages, stream, @@ -669,7 +674,7 @@ export class ClaudeProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( usageSum, actor, - `claude:${modelUsed.id}`, + this.meteringModelKey(modelUsed.id), costsOverrideFromModel, ); chatStream.end(usageSum); @@ -696,7 +701,7 @@ export class ClaudeProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( usage, actor, - `claude:${modelUsed.id}`, + this.meteringModelKey(modelUsed.id), costsOverrideFromModel, ); diff --git a/src/backend/drivers/ai-chat/providers/deepseek/DeepSeekProvider.ts b/src/backend/drivers/ai-chat/providers/deepseek/DeepSeekProvider.ts index c750b2321..e188cca7d 100644 --- a/src/backend/drivers/ai-chat/providers/deepseek/DeepSeekProvider.ts +++ b/src/backend/drivers/ai-chat/providers/deepseek/DeepSeekProvider.ts @@ -52,6 +52,11 @@ export class DeepSeekProvider implements IChatProvider { return modelLookupNames(this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `deepseek:${modelId}`; + } + async complete({ messages, stream, @@ -131,7 +136,7 @@ export class DeepSeekProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor!, - `deepseek:${modelUsed.id}`, + this.meteringModelKey(modelUsed.id), costsOverrideFromModel, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.ts b/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.ts index a70a5c813..572890fe6 100644 --- a/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.ts +++ b/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.ts @@ -58,6 +58,11 @@ export class GeminiChatProvider implements IChatProvider { return modelLookupNames(await this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string | undefined): string { + return `gemini:${modelId}`; + } + async complete({ messages, stream, @@ -156,7 +161,7 @@ export class GeminiChatProvider implements IChatProvider { this.meteringService.utilRecordUsageObject( trackedUsage, actor!, - `gemini:${modelUsed?.id}`, + this.meteringModelKey(modelUsed?.id), costsOverrideFromModel, ); diff --git a/src/backend/drivers/ai-chat/providers/groq/GroqAIProvider.ts b/src/backend/drivers/ai-chat/providers/groq/GroqAIProvider.ts index 690ca2e82..8f1fd8e7b 100644 --- a/src/backend/drivers/ai-chat/providers/groq/GroqAIProvider.ts +++ b/src/backend/drivers/ai-chat/providers/groq/GroqAIProvider.ts @@ -51,6 +51,11 @@ export class GroqAIProvider implements IChatProvider { return modelLookupNames(this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `groq:${modelId}`; + } + async complete({ messages, model, @@ -99,7 +104,7 @@ export class GroqAIProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor, - `groq:${modelUsed.id}`, + this.meteringModelKey(modelUsed.id), costsOverride, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/hoonify/HoonifyProvider.ts b/src/backend/drivers/ai-chat/providers/hoonify/HoonifyProvider.ts index e07d799b6..9325b76ff 100644 --- a/src/backend/drivers/ai-chat/providers/hoonify/HoonifyProvider.ts +++ b/src/backend/drivers/ai-chat/providers/hoonify/HoonifyProvider.ts @@ -77,6 +77,11 @@ export class HoonifyProvider implements IChatProvider { return modelIds; } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return modelId; + } + async complete( params: ICompleteArguments, ): ReturnType { @@ -149,7 +154,7 @@ export class HoonifyProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor, - modelUsed.id, + this.meteringModelKey(modelUsed.id), costsOverrideFromModel, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/infron/InfronProvider.ts b/src/backend/drivers/ai-chat/providers/infron/InfronProvider.ts index 90f43123f..ab6360074 100644 --- a/src/backend/drivers/ai-chat/providers/infron/InfronProvider.ts +++ b/src/backend/drivers/ai-chat/providers/infron/InfronProvider.ts @@ -223,6 +223,11 @@ export class InfronProvider implements IChatProvider { } /** AI Chat completion method. See AIChatService for more details. */ + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return modelId; + } + async complete({ messages, stream, @@ -311,7 +316,7 @@ export class InfronProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( billedTrackedUsage, actor, - modelUsed.id, + this.meteringModelKey(modelUsed.id), costOverwrites, ); (billedTrackedUsage as Record).usd_cents = @@ -331,7 +336,7 @@ export class InfronProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor, - modelUsed.id, + this.meteringModelKey(modelUsed.id), costOverwrites, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/meta/MetaProvider.ts b/src/backend/drivers/ai-chat/providers/meta/MetaProvider.ts index c5d0e36ae..b4a21825e 100644 --- a/src/backend/drivers/ai-chat/providers/meta/MetaProvider.ts +++ b/src/backend/drivers/ai-chat/providers/meta/MetaProvider.ts @@ -99,6 +99,11 @@ export class MetaProvider implements IChatProvider { return modelLookupNames(this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `meta:${modelId}`; + } + async complete( params: ICompleteArguments, ): ReturnType { @@ -227,7 +232,7 @@ export class MetaProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor!, - `meta:${modelUsed.id}`, + this.meteringModelKey(modelUsed.id), costsOverride, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/minimax/MiniMaxProvider.ts b/src/backend/drivers/ai-chat/providers/minimax/MiniMaxProvider.ts index bdf0871c8..53939f43d 100644 --- a/src/backend/drivers/ai-chat/providers/minimax/MiniMaxProvider.ts +++ b/src/backend/drivers/ai-chat/providers/minimax/MiniMaxProvider.ts @@ -58,6 +58,11 @@ export class MiniMaxProvider implements IChatProvider { return modelLookupNames(this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `minimax:${modelId}`; + } + async complete({ messages, stream, @@ -111,7 +116,7 @@ export class MiniMaxProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor!, - `minimax:${modelUsed.id}`, + this.meteringModelKey(modelUsed.id), costsOverride, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/mistral/MistralAiProvider.ts b/src/backend/drivers/ai-chat/providers/mistral/MistralAiProvider.ts index 83d6256cb..c9addfce8 100644 --- a/src/backend/drivers/ai-chat/providers/mistral/MistralAiProvider.ts +++ b/src/backend/drivers/ai-chat/providers/mistral/MistralAiProvider.ts @@ -165,6 +165,11 @@ export class MistralAIProvider implements IChatProvider { }); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `mistral:${modelId}`; + } + async complete({ messages, stream, @@ -396,7 +401,7 @@ export class MistralAIProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor, - `mistral:${selectedModel.id}`, + this.meteringModelKey(selectedModel.id), costsOverrideFromModel, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/moonshot/MoonshotProvider.ts b/src/backend/drivers/ai-chat/providers/moonshot/MoonshotProvider.ts index ebb26de62..1d656bdb7 100644 --- a/src/backend/drivers/ai-chat/providers/moonshot/MoonshotProvider.ts +++ b/src/backend/drivers/ai-chat/providers/moonshot/MoonshotProvider.ts @@ -56,6 +56,11 @@ export class MoonshotProvider implements IChatProvider { return modelLookupNames(this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `moonshotai:${modelId}`; + } + async complete({ messages, stream, @@ -107,7 +112,7 @@ export class MoonshotProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor, - `moonshotai:${modelUsed.id}`, + this.meteringModelKey(modelUsed.id), costsOverride, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/neuralwatt/NeuralwattProvider.ts b/src/backend/drivers/ai-chat/providers/neuralwatt/NeuralwattProvider.ts index f72dedde3..a548e994b 100644 --- a/src/backend/drivers/ai-chat/providers/neuralwatt/NeuralwattProvider.ts +++ b/src/backend/drivers/ai-chat/providers/neuralwatt/NeuralwattProvider.ts @@ -157,6 +157,11 @@ export class NeuralwattProvider implements IChatProvider { return undefined; } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return modelId; + } + async complete({ messages, stream, @@ -305,7 +310,7 @@ export class NeuralwattProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( billedTrackedUsage, actor!, - modelUsed.id, + this.meteringModelKey(modelUsed.id), costOverwrites, ); const result = annotate(billedTrackedUsage); @@ -327,7 +332,7 @@ export class NeuralwattProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor!, - modelUsed.id, + this.meteringModelKey(modelUsed.id), costOverwrites, ); return annotate(trackedUsage); diff --git a/src/backend/drivers/ai-chat/providers/ollama/OllamaProvider.ts b/src/backend/drivers/ai-chat/providers/ollama/OllamaProvider.ts index 4a2affda0..3e6a060b5 100644 --- a/src/backend/drivers/ai-chat/providers/ollama/OllamaProvider.ts +++ b/src/backend/drivers/ai-chat/providers/ollama/OllamaProvider.ts @@ -108,6 +108,11 @@ export class OllamaChatProvider implements IChatProvider { } return model_names; } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return modelId; + } + async complete({ messages, stream, diff --git a/src/backend/drivers/ai-chat/providers/openai/OpenAiChatCompletionsProvider.ts b/src/backend/drivers/ai-chat/providers/openai/OpenAiChatCompletionsProvider.ts index 9237df7db..de9da5349 100644 --- a/src/backend/drivers/ai-chat/providers/openai/OpenAiChatCompletionsProvider.ts +++ b/src/backend/drivers/ai-chat/providers/openai/OpenAiChatCompletionsProvider.ts @@ -96,6 +96,11 @@ export class OpenAiChatProvider implements IChatProvider { return this.#defaultModel; } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string | undefined): string { + return `openai:${modelId}`; + } + async complete( params: ICompleteArguments, ): ReturnType { @@ -247,7 +252,7 @@ export class OpenAiChatProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor, - `openai:${modelUsed?.id}`, + this.meteringModelKey(modelUsed?.id), costsOverrideFromModel, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/openai/OpenAiChatResponsesProvider.ts b/src/backend/drivers/ai-chat/providers/openai/OpenAiChatResponsesProvider.ts index 7a4f444c7..57396f41c 100644 --- a/src/backend/drivers/ai-chat/providers/openai/OpenAiChatResponsesProvider.ts +++ b/src/backend/drivers/ai-chat/providers/openai/OpenAiChatResponsesProvider.ts @@ -88,6 +88,11 @@ export class OpenAiResponsesChatProvider implements IChatProvider { return this.#defaultModel; } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string | undefined): string { + return `openai:${modelId}`; + } + async complete({ messages, model, @@ -280,7 +285,7 @@ export class OpenAiResponsesChatProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor!, - `openai:${modelUsed?.id}`, + this.meteringModelKey(modelUsed?.id), costsOverrideFromModel, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/openrouter/OpenRouterProvider.ts b/src/backend/drivers/ai-chat/providers/openrouter/OpenRouterProvider.ts index be20cb48c..ce2299ce8 100644 --- a/src/backend/drivers/ai-chat/providers/openrouter/OpenRouterProvider.ts +++ b/src/backend/drivers/ai-chat/providers/openrouter/OpenRouterProvider.ts @@ -93,6 +93,11 @@ export class OpenRouterProvider implements IChatProvider { } /** AI Chat completion method. See AIChatService for more details. */ + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return modelId; + } + async complete({ messages, stream, @@ -194,7 +199,7 @@ export class OpenRouterProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor, - modelUsed.id, + this.meteringModelKey(modelUsed.id), costOverwrites, ); (trackedUsage as Record).usd_cents = @@ -221,7 +226,7 @@ export class OpenRouterProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor, - modelUsed.id, + this.meteringModelKey(modelUsed.id), costOverwrites, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.ts b/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.ts index 5278b4376..617edcf2f 100644 --- a/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.ts +++ b/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.ts @@ -110,6 +110,11 @@ export class TogetherAIProvider implements IChatProvider { return modelLookupNames(await this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `togetherai:${modelId.replace(/^togetherai:/, '')}`; + } + async complete({ messages, stream, @@ -174,7 +179,7 @@ export class TogetherAIProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor, - `togetherai:${modelIdForParams}`, + this.meteringModelKey(modelUsed.id), costsOverride, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/xai/XAIProvider.ts b/src/backend/drivers/ai-chat/providers/xai/XAIProvider.ts index e09b2ffe1..2a39b6d5c 100644 --- a/src/backend/drivers/ai-chat/providers/xai/XAIProvider.ts +++ b/src/backend/drivers/ai-chat/providers/xai/XAIProvider.ts @@ -56,6 +56,11 @@ export class XAIProvider implements IChatProvider { return modelLookupNames(this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `xai:${modelId}`; + } + async complete({ messages, stream, @@ -105,7 +110,7 @@ export class XAIProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor, - `xai:${modelUsed.id}`, + this.meteringModelKey(modelUsed.id), costsOverride, ); return metered; diff --git a/src/backend/drivers/ai-chat/providers/zai/ZAIProvider.ts b/src/backend/drivers/ai-chat/providers/zai/ZAIProvider.ts index c1eac5006..eea785d06 100644 --- a/src/backend/drivers/ai-chat/providers/zai/ZAIProvider.ts +++ b/src/backend/drivers/ai-chat/providers/zai/ZAIProvider.ts @@ -79,6 +79,11 @@ export class ZAIProvider implements IChatProvider { return modelLookupNames(this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `zai:${modelId}`; + } + async complete( params: ICompleteArguments, ): ReturnType { @@ -162,7 +167,7 @@ export class ZAIProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor!, - `zai:${modelUsed.id}`, + this.meteringModelKey(modelUsed.id), costsOverrideFromModel, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/types.ts b/src/backend/drivers/ai-chat/types.ts index 0d8bf5c3f..937abc256 100644 --- a/src/backend/drivers/ai-chat/types.ts +++ b/src/backend/drivers/ai-chat/types.ts @@ -173,4 +173,10 @@ export interface IChatProvider { * attempt. */ readonly resolvesPuterPaths?: boolean; + /** + * The model key usage is recorded under, which is also what the AI cost + * factor is looked up by. The driver assumes `:` when + * absent. + */ + meteringModelKey?(modelId: string): string; }