Enrichment Orchestrator

The per-run token + cost tracker the enrichment pipeline records usage with — self-contained and reusable on any Claude pipeline.

// ---------------------------------------------------------------------------
// Usage Tracker — track Claude API token usage and estimated cost per run
// ---------------------------------------------------------------------------

// Pricing per million tokens (USD) as of 2025-05.
const MODEL_PRICING = {
  'claude-sonnet-4-20250514': { input: 3, output: 15 },
  // Fallback for unknown models
  default: { input: 3, output: 15 },
};

export class UsageTracker {
  constructor() {
    this.calls = [];
  }

  /**
   * Record a single API call's token usage.
   * @param {{ agent: string, inputTokens: number, outputTokens: number, model: string }} usage
   */
  record({ agent, inputTokens, outputTokens, model }) {
    this.calls.push({
      agent,
      inputTokens,
      outputTokens,
      model,
      timestamp: new Date().toISOString(),
    });
  }

  /**
   * Get aggregated usage summary.
   */
  getSummary() {
    const byAgent = {};
    let totalInput = 0;
    let totalOutput = 0;
    let totalCalls = 0;

    for (const call of this.calls) {
      if (!byAgent[call.agent]) {
        byAgent[call.agent] = { calls: 0, inputTokens: 0, outputTokens: 0 };
      }
      byAgent[call.agent].calls++;
      byAgent[call.agent].inputTokens += call.inputTokens;
      byAgent[call.agent].outputTokens += call.outputTokens;
      totalInput += call.inputTokens;
      totalOutput += call.outputTokens;
      totalCalls++;
    }

    const estimatedCost = this._estimateCost(totalInput, totalOutput);

    return {
      byAgent,
      totals: { calls: totalCalls, inputTokens: totalInput, outputTokens: totalOutput },
      estimatedCost,
    };
  }

  /**
   * Persist usage data to usage.json in the enrichments directory, appending to a runs array.
   * @param {import('./store.mjs').EnrichmentStore|{readEnrichment: Function, writeEnrichment: Function}} store
   */
  persist(store) {
    const existing = store.readEnrichment('usage') || { runs: [] };
    const summary = this.getSummary();
    existing.runs.push({
      ...summary,
      timestamp: new Date().toISOString(),
    });
    store.writeEnrichment('usage', existing);
  }

  /**
   * Create a usageCallback function bound to this tracker for a specific agent.
   * Pass this to callClaude/callClaudeJson.
   */
  callbackFor(agentName) {
    return ({ inputTokens, outputTokens, model }) => {
      this.record({ agent: agentName, inputTokens, outputTokens, model });
    };
  }

  _estimateCost(inputTokens, outputTokens) {
    // Use default pricing — all calls currently use the same model
    const pricing = MODEL_PRICING.default;
    const inputCost = (inputTokens / 1_000_000) * pricing.input;
    const outputCost = (outputTokens / 1_000_000) * pricing.output;
    return Math.round((inputCost + outputCost) * 10000) / 10000; // 4 decimal places
  }
}