Replace hardcoded subtraction chain with registry-driven decomposition

.pre-commit-config.yaml

@@ -20,6 +20,7 @@ repos:
     rev: v2.3.0
     hooks:
       - id: codespell
+        exclude: "^tests/cassettes/"
         additional_dependencies:
           - tomli
 

packages/js/src/__tests__/calcPrice.test.ts

@@ -64,14 +64,18 @@ describe('Core Price Calculation Function', () => {
       }
       const modelPrice: ModelPrice = {
         input_audio_mtok: 10.0,
+        input_mtok: 5.0,
         output_audio_mtok: 20.0,
+        output_mtok: 10.0,
       }
 
       const result = calcPrice(usage, modelPrice)
 
+      // input leaf: 100 - 100 = 0 non-audio text, 100 audio
+      // output leaf: 50 - 50 = 0 non-audio text, 50 audio
       expect(result).toMatchObject({
-        input_price: 0.001, // 100 * 10.0 / 1_000_000
-        output_price: 0.001, // 50 * 20.0 / 1_000_000
+        input_price: 0.001, // (0 * 5.0 + 100 * 10.0) / 1_000_000
+        output_price: 0.001, // (0 * 10.0 + 50 * 20.0) / 1_000_000
         total_price: 0.002,
       })
     })

packages/js/src/__tests__/decompose.test.ts

@@ -0,0 +1,176 @@
+import { describe, expect, it } from 'vitest'
+
+import type { Usage } from '../types'
+
+import { computeLeafValues, isDescendantOrSelf, validateAncestorCoverage } from '../decompose'
+import { getUnit, TOKENS_FAMILY } from '../units'
+
+describe('Containment', () => {
+  it('self is descendant-or-self', () => {
+    const unit = getUnit('input_mtok')
+    expect(isDescendantOrSelf(unit, unit)).toBe(true)
+  })
+
+  it('child is descendant', () => {
+    expect(isDescendantOrSelf(getUnit('input_mtok'), getUnit('cache_read_mtok'))).toBe(true)
+  })
+
+  it('parent is NOT descendant of child', () => {
+    expect(isDescendantOrSelf(getUnit('cache_read_mtok'), getUnit('input_mtok'))).toBe(false)
+  })
+
+  it('grandchild is descendant', () => {
+    expect(isDescendantOrSelf(getUnit('input_mtok'), getUnit('cache_audio_read_mtok'))).toBe(true)
+  })
+
+  it('sibling not descendant', () => {
+    expect(isDescendantOrSelf(getUnit('cache_read_mtok'), getUnit('cache_write_mtok'))).toBe(false)
+  })
+
+  it('different direction not descendant', () => {
+    expect(isDescendantOrSelf(getUnit('input_mtok'), getUnit('output_mtok'))).toBe(false)
+  })
+})
+
+describe('Leaf Values', () => {
+  const family = TOKENS_FAMILY
+
+  it('simple text model', () => {
+    const priced = new Set(['input_mtok', 'output_mtok'])
+    const usage = { input_tokens: 1000, output_tokens: 500 }
+    expect(computeLeafValues(priced, usage, family)).toEqual({ input_mtok: 1000, output_mtok: 500 })
+  })
+
+  it('with cache', () => {
+    const priced = new Set(['cache_read_mtok', 'cache_write_mtok', 'input_mtok', 'output_mtok'])
+    const usage = { cache_read_tokens: 200, cache_write_tokens: 100, input_tokens: 1000, output_tokens: 500 }
+    expect(computeLeafValues(priced, usage, family)).toEqual({
+      cache_read_mtok: 200,
+      cache_write_mtok: 100,
+      input_mtok: 700,
+      output_mtok: 500,
+    })
+  })
+
+  it('with audio', () => {
+    const priced = new Set(['input_audio_mtok', 'input_mtok', 'output_mtok'])
+    const usage = { input_audio_tokens: 300, input_tokens: 1000, output_tokens: 500 }
+    expect(computeLeafValues(priced, usage, family)).toEqual({
+      input_audio_mtok: 300,
+      input_mtok: 700,
+      output_mtok: 500,
+    })
+  })
+
+  it('lattice: cache_read_audio carved from both', () => {
+    const priced = new Set(['cache_audio_read_mtok', 'cache_read_mtok', 'input_audio_mtok', 'input_mtok'])
+    const usage = {
+      cache_audio_read_tokens: 50,
+      cache_read_tokens: 200,
+      input_audio_tokens: 300,
+      input_tokens: 1000,
+    }
+    expect(computeLeafValues(priced, usage, family)).toEqual({
+      cache_audio_read_mtok: 50,
+      cache_read_mtok: 150,
+      input_audio_mtok: 250,
+      input_mtok: 550,
+    })
+  })
+
+  it('unpriced audio stays in catch-all', () => {
+    const priced = new Set(['cache_read_mtok', 'input_mtok', 'output_mtok'])
+    const usage = {
+      cache_read_tokens: 200,
+      input_audio_tokens: 300,
+      input_tokens: 1000,
+      output_tokens: 500,
+    }
+    expect(computeLeafValues(priced, usage, family)).toEqual({
+      cache_read_mtok: 200,
+      input_mtok: 800,
+      output_mtok: 500,
+    })
+  })
+
+  it('unpriced cache stays in catch-all', () => {
+    const priced = new Set(['input_mtok', 'output_mtok'])
+    const usage = { cache_read_tokens: 200, input_tokens: 1000, output_tokens: 500 }
+    expect(computeLeafValues(priced, usage, family)).toEqual({
+      input_mtok: 1000,
+      output_mtok: 500,
+    })
+  })
+
+  it('negative leaf raises error', () => {
+    const priced = new Set(['cache_read_mtok', 'input_mtok'])
+    const usage = { cache_read_tokens: 200, input_tokens: 100 }
+    expect(() => computeLeafValues(priced, usage, family)).toThrow(/Negative leaf value.*input_mtok/)
+  })
+
+  it('full 7-unit model', () => {
+    const priced = new Set([
+      'cache_audio_read_mtok',
+      'cache_read_mtok',
+      'cache_write_mtok',
+      'input_audio_mtok',
+      'input_mtok',
+      'output_audio_mtok',
+      'output_mtok',
+    ])
+    const usage = {
+      cache_audio_read_tokens: 50,
+      cache_read_tokens: 200,
+      cache_write_tokens: 100,
+      input_audio_tokens: 300,
+      input_tokens: 1000,
+      output_audio_tokens: 150,
+      output_tokens: 800,
+    }
+    expect(computeLeafValues(priced, usage, family)).toEqual({
+      cache_audio_read_mtok: 50,
+      cache_read_mtok: 150,
+      cache_write_mtok: 100,
+      input_audio_mtok: 250,
+      input_mtok: 450,
+      output_audio_mtok: 150,
+      output_mtok: 650,
+    })
+  })
+
+  it('accepts Usage object (optional fields)', () => {
+    const priced = new Set(['cache_read_mtok', 'input_mtok', 'output_mtok'])
+    const usage: Usage = { cache_read_tokens: 200, input_tokens: 1000, output_tokens: 500 }
+    expect(computeLeafValues(priced, usage as Record<string, unknown>, family)).toEqual({
+      cache_read_mtok: 200,
+      input_mtok: 800,
+      output_mtok: 500,
+    })
+  })
+})
+
+describe('Ancestor Coverage', () => {
+  it('valid: input + output', () => {
+    expect(() => {
+      validateAncestorCoverage(new Set(['input_mtok', 'output_mtok']), TOKENS_FAMILY)
+    }).not.toThrow()
+  })
+
+  it('valid: with cache', () => {
+    expect(() => {
+      validateAncestorCoverage(new Set(['cache_read_mtok', 'input_mtok', 'output_mtok']), TOKENS_FAMILY)
+    }).not.toThrow()
+  })
+
+  it('missing ancestor: cache_read without input', () => {
+    expect(() => {
+      validateAncestorCoverage(new Set(['cache_read_mtok', 'output_mtok']), TOKENS_FAMILY)
+    }).toThrow(/ancestor.*input_mtok/)
+  })
+
+  it('missing intermediate ancestor', () => {
+    expect(() => {
+      validateAncestorCoverage(new Set(['cache_audio_read_mtok', 'input_mtok', 'output_mtok']), TOKENS_FAMILY)
+    }).toThrow(/ancestor/)
+  })
+})

packages/js/src/__tests__/units.test.ts

@@ -0,0 +1,60 @@
+import { describe, expect, it } from 'vitest'
+
+import { getFamily, getUnit, TOKENS_FAMILY } from '../units'
+
+describe('Unit Registry', () => {
+  it('should load the tokens family', () => {
+    const family = getFamily('tokens')
+    expect(family.id).toBe('tokens')
+    expect(family.per).toBe(1_000_000)
+  })
+
+  it('should have 20 token units', () => {
+    const family = getFamily('tokens')
+    expect(Object.keys(family.units)).toHaveLength(20)
+  })
+
+  it('should look up a unit by ID', () => {
+    const unit = getUnit('input_mtok')
+    expect(unit.familyId).toBe('tokens')
+    expect(unit.usageKey).toBe('input_tokens')
+    expect(unit.dimensions).toEqual({ direction: 'input' })
+  })
+
+  it('should throw on unknown unit', () => {
+    expect(() => getUnit('nonexistent')).toThrow()
+  })
+
+  it('should throw on unknown family', () => {
+    expect(() => getFamily('nonexistent')).toThrow()
+  })
+
+  it('should have all 7 currently-used units', () => {
+    for (const unitId of [
+      'input_mtok',
+      'output_mtok',
+      'cache_read_mtok',
+      'cache_write_mtok',
+      'input_audio_mtok',
+      'cache_audio_read_mtok',
+      'output_audio_mtok',
+    ]) {
+      expect(TOKENS_FAMILY.units[unitId]).toBeDefined()
+    }
+  })
+
+  it('should have correct usage_key for cache_audio_read_mtok', () => {
+    const unit = getUnit('cache_audio_read_mtok')
+    expect(unit.usageKey).toBe('cache_audio_read_tokens')
+  })
+
+  it('should validate all unit dimensions against family dimensions', () => {
+    const family = getFamily('tokens')
+    for (const unit of Object.values(family.units)) {
+      for (const [dimKey, dimVal] of Object.entries(unit.dimensions)) {
+        expect(family.dimensions[dimKey]).toBeDefined()
+        expect(family.dimensions[dimKey]).toContain(dimVal)
+      }
+    }
+  })
+})

packages/js/src/data.ts

@@ -9432,7 +9432,9 @@ export const data: Provider[] = [
           starts_with: 'gpt-4o-audio-preview',
         },
         context_window: 128000,
+        price_comments: 'input_mtok set equal to input_audio_mtok (audio-only input model, catch-all required by ancestor coverage rule)',
         prices: {
+          input_mtok: 2.5,
           output_mtok: 10,
           input_audio_mtok: 2.5,
         },
@@ -9483,7 +9485,9 @@ export const data: Provider[] = [
         match: {
           starts_with: 'gpt-4o-mini-audio',
         },
+        price_comments: 'input_mtok set equal to input_audio_mtok (audio-only input model, catch-all required by ancestor coverage rule)',
         prices: {
+          input_mtok: 0.15,
           output_mtok: 0.6,
           input_audio_mtok: 0.15,
         },
@@ -9518,8 +9522,11 @@ export const data: Provider[] = [
         match: {
           equals: 'gpt-4o-mini-tts',
         },
+        price_comments:
+          'output_mtok set equal to output_audio_mtok (audio-only output model, catch-all required by ancestor coverage rule)',
         prices: {
           input_mtok: 0.6,
+          output_mtok: 12,
           output_audio_mtok: 12,
         },
       },

packages/js/src/decompose.ts

@@ -0,0 +1,83 @@
+import type { UnitDef, UnitFamily } from './units'
+
+/**
+ * True if candidate's dimensions are a (non-strict) superset of ancestor's.
+ */
+export function isDescendantOrSelf(ancestor: UnitDef, candidate: UnitDef): boolean {
+  if (ancestor.familyId !== candidate.familyId) return false
+  return Object.entries(ancestor.dimensions).every(([k, v]) => candidate.dimensions[k] === v)
+}
+
+/**
+ * Get a usage value by key from a plain object. Returns 0 for missing/undefined/null.
+ */
+function getUsageValue(usage: Record<string, unknown>, key: string): number {
+  const val = usage[key]
+  return typeof val === 'number' ? val : 0
+}
+
+/**
+ * Validate that every priced unit has all its ancestors also priced.
+ * Throws if any ancestor is missing.
+ */
+export function validateAncestorCoverage(pricedUnitIds: Set<string>, family: UnitFamily): void {
+  for (const unitId of pricedUnitIds) {
+    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+    const unit = family.units[unitId]!
+    for (const [otherId, other] of Object.entries(family.units)) {
+      if (otherId !== unitId && isDescendantOrSelf(other, unit) && !pricedUnitIds.has(otherId)) {
+        throw new Error(
+          `Unit '${unitId}' is priced but its ancestor '${otherId}' is not. ` + `All ancestors of a priced unit must also be priced.`
+        )
+      }
+    }
+  }
+}
+
+/**
+ * Compute leaf values for each priced unit via Mobius inversion on the containment poset.
+ *
+ * Only priced units participate. Unpriced units' tokens stay in the nearest
+ * priced ancestor's catch-all. Throws on negative leaf values.
+ *
+ * Precondition: ancestor coverage — if a unit is priced, all its ancestors must
+ * also be priced. Validated by validateAncestorCoverage() called from calcPrice.
+ */
+export function computeLeafValues(pricedUnitIds: Set<string>, usage: Record<string, unknown>, family: UnitFamily): Record<string, number> {
+  const result: Record<string, number> = {}
+
+  for (const unitId of pricedUnitIds) {
+    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+    const unit = family.units[unitId]!
+    const targetDepth = Object.keys(unit.dimensions).length
+
+    let leafValue = 0
+    for (const otherId of pricedUnitIds) {
+      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+      const other = family.units[otherId]!
+      if (!isDescendantOrSelf(unit, other)) continue
+      const depthDiff = Object.keys(other.dimensions).length - targetDepth
+      const coefficient = (-1) ** depthDiff
+      leafValue += coefficient * getUsageValue(usage, other.usageKey)
+    }
+
+    if (leafValue < 0) {
+      const involved = [...pricedUnitIds]
+        .filter((oid) => {
+          // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+          const o = family.units[oid]!
+          return isDescendantOrSelf(unit, o) && getUsageValue(usage, o.usageKey) !== 0
+        })
+        .map((oid) => {
+          // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+          const o = family.units[oid]!
+          return `${o.usageKey}=${String(getUsageValue(usage, o.usageKey))}`
+        })
+      throw new Error(`Negative leaf value (${String(leafValue)}) for ${unitId}: inconsistent usage values: ${involved.join(', ')}`)
+    }
+
+    result[unitId] = leafValue
+  }
+
+  return result
+}