diff --git a/nodejs/src/ingestion/pipelines/ai/costs/cost-model-matching.test.ts b/nodejs/src/ingestion/pipelines/ai/costs/cost-model-matching.test.ts index 527fe146da3a..97cd30eb0bfd 100644 --- a/nodejs/src/ingestion/pipelines/ai/costs/cost-model-matching.test.ts +++ b/nodejs/src/ingestion/pipelines/ai/costs/cost-model-matching.test.ts @@ -33,6 +33,43 @@ jest.mock('./providers', () => { }, }, }, + 'anthropic/claude-bedrock-profile-test': { + model: 'anthropic/claude-bedrock-profile-test', + cost: { + default: { prompt_token: 0.000001, completion_token: 0.000002 }, + 'amazon-bedrock-claude-on-aws': { + prompt_token: 0.000003, + completion_token: 0.000004, + }, + 'amazon-bedrock-global': { + prompt_token: 0.000005, + completion_token: 0.000006, + }, + 'amazon-bedrock-us-east-1': { + prompt_token: 0.000007, + completion_token: 0.000008, + }, + anthropic: { prompt_token: 0.000009, completion_token: 0.00001 }, + }, + }, + 'anthropic/claude-bedrock-arn-test': { + model: 'anthropic/claude-bedrock-arn-test', + cost: { + default: { prompt_token: 0.000011, completion_token: 0.000012 }, + 'amazon-bedrock-claude-on-aws': { + prompt_token: 0.000013, + completion_token: 0.000014, + }, + 'amazon-bedrock-us-east-1': { + prompt_token: 0.000015, + completion_token: 0.000016, + }, + 'amazon-bedrock-us-west-2': { + prompt_token: 0.000017, + completion_token: 0.000018, + }, + }, + }, 'google/gemini-2.5-pro-preview': { model: 'google/gemini-2.5-pro-preview', cost: { @@ -259,6 +296,64 @@ describe('findCostFromModel()', () => { }) }) + describe('Bedrock inference profiles', () => { + test.each([ + ['us', 'bedrock', 'amazon-bedrock-us-east-1'], + ['us', 'amazon_bedrock', 'amazon-bedrock-us-east-1'], + ['global', 'bedrock', 'amazon-bedrock-global'], + ['global', 'amazon_bedrock', 'amazon-bedrock-global'], + ])('uses the %s profile and %s provider alias to select %s pricing', (profile, provider, expectedProvider) => { + const result = findCostFromModel(`${profile}.anthropic.claude-bedrock-profile-test`, { + $ai_provider: provider, + }) + + expect(result).toBeDefined() + expect(result!.cost.model).toBe('anthropic/claude-bedrock-profile-test') + expect(result!.cost.provider).toBe(expectedProvider) + }) + + it('keeps an explicit Bedrock provider authoritative over the model profile', () => { + const result = findCostFromModel('us.anthropic.claude-bedrock-profile-test', { + $ai_provider: 'amazon-bedrock-global', + }) + + expect(result).toBeDefined() + expect(result!.cost.provider).toBe('amazon-bedrock-global') + }) + + it('uses the ARN region to select pricing when a US profile spans multiple regions', () => { + const result = findCostFromModel( + 'arn:aws:bedrock:us-west-2:123456789012:inference-profile/us.anthropic.claude-bedrock-arn-test', + { $ai_provider: 'bedrock' } + ) + + expect(result).toBeDefined() + expect(result!.cost.model).toBe('anthropic/claude-bedrock-arn-test') + expect(result!.cost.provider).toBe('amazon-bedrock-us-west-2') + }) + + test.each(['au', 'unknown'])( + 'preserves generic Bedrock fallback when the %s profile has no catalog provider', + (profile) => { + const result = findCostFromModel(`${profile}.anthropic.claude-bedrock-profile-test`, { + $ai_provider: 'amazon_bedrock', + }) + + expect(result).toBeDefined() + expect(result!.cost.provider).toBe('amazon-bedrock-claude-on-aws') + } + ) + + it('does not infer a Bedrock provider for another provider', () => { + const result = findCostFromModel('us.anthropic.claude-bedrock-profile-test', { + $ai_provider: 'anthropic', + }) + + expect(result).toBeDefined() + expect(result!.cost.provider).toBe('anthropic') + }) + }) + describe('gateway provider scenarios', () => { it('finds manual cost when model has provider prefix, even with gateway provider', () => { const result = findCostFromModel('anthropic/claude-3-5-sonnet', { $ai_provider: 'gateway' }) diff --git a/nodejs/src/ingestion/pipelines/ai/costs/cost-model-matching.ts b/nodejs/src/ingestion/pipelines/ai/costs/cost-model-matching.ts index 47ecbb26211e..74ae60acc194 100644 --- a/nodejs/src/ingestion/pipelines/ai/costs/cost-model-matching.ts +++ b/nodejs/src/ingestion/pipelines/ai/costs/cost-model-matching.ts @@ -1,7 +1,7 @@ import { logger } from '~/common/utils/logger' import { Properties } from '~/plugin-scaffold' -import { resolveModelCostForProvider } from './provider-matching' +import { resolveModelCostForProvider, resolveProviderAliases } from './provider-matching' import { manualCostsByModel, openRouterCostsByModel } from './providers' import type { ModelCostRow, ResolvedModelCost } from './providers/types' @@ -38,6 +38,39 @@ const findManualCost = (model: string): ModelCostRow | undefined => { return undefined } +const resolveBedrockInferenceProfileProvider = ( + model: string, + providerCosts: ModelCostRow['cost'], + provider: string | undefined +): string | undefined => { + if (!provider || resolveProviderAliases(provider) !== 'amazon-bedrock') { + return provider + } + + const lowerCaseModel = model.toLowerCase() + const inferenceProfileArn = + /^arn:(?:aws|aws-cn|aws-us-gov):bedrock:([a-z0-9-]+):\d{12}:inference-profile\/[^/]+$/.exec(lowerCaseModel) + const arnProvider = inferenceProfileArn ? `amazon-bedrock-${inferenceProfileArn[1]}` : undefined + + if (arnProvider && providerCosts[arnProvider]) { + return arnProvider + } + + const modelId: string = lowerCaseModel.split('/').pop() ?? lowerCaseModel + const profilePrefix: string = modelId.split('.')[0] + const profileProviderPrefix = `amazon-bedrock-${profilePrefix}` + + if (providerCosts[profileProviderPrefix]) { + return profileProviderPrefix + } + + const regionalProviders = Object.keys(providerCosts).filter( + (providerKey) => providerKey.startsWith(`${profileProviderPrefix}-`) && providerCosts[providerKey] + ) + + return regionalProviders.length === 1 ? regionalProviders[0] : provider +} + export const findCostFromModel = (model: string, properties: Properties): CostModelResult | undefined => { const providerProperty: unknown = properties['$ai_provider'] @@ -46,7 +79,11 @@ export const findCostFromModel = (model: string, properties: Properties): CostMo const manualMatch: ModelCostRow | undefined = findManualCost(model) const resolvedManualMatch: ResolvedModelCost | undefined = manualMatch - ? resolveModelCostForProvider(manualMatch.cost, provider, manualMatch.model) + ? resolveModelCostForProvider( + manualMatch.cost, + resolveBedrockInferenceProfileProvider(model, manualMatch.cost, provider), + manualMatch.model + ) : undefined if (resolvedManualMatch) { @@ -56,7 +93,11 @@ export const findCostFromModel = (model: string, properties: Properties): CostMo const openRouterMatch: ModelCostRow | undefined = searchModelInCosts(model, openRouterCostsByModel) const resolvedOpenRouterMatch: ResolvedModelCost | undefined = openRouterMatch - ? resolveModelCostForProvider(openRouterMatch.cost, provider, openRouterMatch.model) + ? resolveModelCostForProvider( + openRouterMatch.cost, + resolveBedrockInferenceProfileProvider(model, openRouterMatch.cost, provider), + openRouterMatch.model + ) : undefined if (resolvedOpenRouterMatch) {