ai_qc/backend/TOKEN_TRACKING_ENHANCEMENT.md

Token Tracking Enhancement

Overview

The usage tracking system has been enhanced to capture actual token usage from OpenAI and Google Gemini API responses, providing precise cost calculations instead of estimates.

What Changed

Before (Estimates)

• Used rough estimates: ~1000 input tokens + ~200 output tokens per check
• Cost calculations based on fixed estimates
• No visibility into actual API usage

After (Actual Tracking)

• Captures real token usage from API responses
• Precise cost calculations based on actual consumption
• Detailed breakdown by provider (OpenAI vs Gemini)
• Individual check-level token tracking

Implementation Details

1. LLM API Responses Enhanced

Both OpenAI and Gemini API calls now return token usage data:

# OpenAI Response
{
    'status': 'success',
    'response': 'Analysis text...',
    'token_usage': {
        'prompt_tokens': 1245,
        'completion_tokens': 187,
        'total_tokens': 1432
    }
}

# Gemini Response
{
    'status': 'success',
    'response': 'Analysis text...',
    'token_usage': {
        'prompt_tokens': 1189,
        'completion_tokens': 203,
        'total_tokens': 1392
    }
}

2. Usage Logs Enhanced

Analysis completion logs now include detailed token statistics:

{
  "event": "analysis_complete",
  "timestamp": "2026-02-02T15:30:00",
  "session_id": "20260202_153000",
  "client": "diageo",
  "profile": "diageo_key_visual",
  "user_email": "user@company.com",
  "checks_completed": 11,
  "overall_score": 87.5,
  "total_cost_usd": 0.0234,
  "token_usage": {
    "total_tokens": 15678,
    "total_prompt_tokens": 13245,
    "total_completion_tokens": 2433,
    "by_provider": {
      "OpenAI": {
        "total_tokens": 7234,
        "prompt_tokens": 6123,
        "completion_tokens": 1111,
        "cost": 0.0128
      },
      "Gemini": {
        "total_tokens": 8444,
        "prompt_tokens": 7122,
        "completion_tokens": 1322,
        "cost": 0.0106
      }
    }
  }
}

3. Cost Calculation Enhanced

Cost calculation now uses actual token counts:

# OpenAI costs (per 1K tokens)
input_cost = (prompt_tokens / 1000) * $0.0025
output_cost = (completion_tokens / 1000) * $0.010
total_cost = input_cost + output_cost

# Gemini costs (per 1K tokens)
input_cost = (prompt_tokens / 1000) * $0.00125
output_cost = (completion_tokens / 1000) * $0.005
total_cost = input_cost + output_cost

4. Reports Enhanced

Usage reports now show detailed token statistics:

Text Report Example:

SUMMARY
--------------------------------------------------------------------------------
Total Analyses: 156
Total QC Checks: 1,560
Total Tokens Used: 1,876,543
  - Prompt Tokens: 1,587,234
  - Completion Tokens: 289,309
Total Cost: $7.85 USD
Average Tokens per Analysis: 12,029.7

USAGE BY CLIENT
--------------------------------------------------------------------------------

DIAGEO
  Analyses: 45
  QC Checks: 495
  Total Tokens: 543,876 (Prompt: 459,234, Completion: 84,642)
  Cost: $2.45 USD

TOKEN USAGE BY PROVIDER
--------------------------------------------------------------------------------

OpenAI
  Total Tokens: 876,234
  Prompt Tokens: 743,123
  Completion Tokens: 133,111
  Cost: $3.19 USD

Gemini
  Total Tokens: 1,000,309
  Prompt Tokens: 844,111
  Completion Tokens: 156,198
  Cost: $4.66 USD

CSV Report Example:

SUMMARY
Metric,Value
Total Analyses,156
Total Tokens,1876543
Total Prompt Tokens,1587234
Total Completion Tokens,289309
Total Cost,$7.85

CLIENT USAGE
Client,Analyses,Checks,Users,Avg Score,Total Tokens,Prompt Tokens,Completion Tokens,Cost
diageo,45,495,3,87.5,543876,459234,84642,$2.45

Benefits

1. Accurate Cost Tracking

• ✅ Real costs instead of estimates
• ✅ No more guesswork on monthly LLM expenses
• ✅ Precise budget forecasting

2. Provider Comparison

• ✅ See actual token usage differences between OpenAI and Gemini
• ✅ Compare costs by provider
• ✅ Make informed decisions on which LLM to use for each check

3. Optimization Opportunities

• ✅ Identify checks that use excessive tokens
• ✅ Optimize prompts to reduce token usage
• ✅ Find opportunities to switch providers for cost savings

4. Client Billing

• ✅ Accurate per-client cost attribution
• ✅ Transparent billing with actual usage data
• ✅ Detailed breakdowns for invoicing

Usage Examples

Generate Report with Token Details

# Text report with token statistics
python generate_usage_report.py --last-days 30

# CSV report for analysis in Excel
python generate_usage_report.py --last-days 30 --format csv --output tokens_report.csv

# JSON report for API integration
python generate_usage_report.py --last-days 30 --format json --output tokens_report.json

Analyze Token Usage by Provider

# Get token breakdown by provider
python generate_usage_report.py --last-days 30 --format json | jq '.by_provider'

Example Output:

{
  "OpenAI": {
    "total_tokens": 876234,
    "prompt_tokens": 743123,
    "completion_tokens": 133111,
    "cost": 3.19
  },
  "Gemini": {
    "total_tokens": 1000309,
    "prompt_tokens": 844111,
    "completion_tokens": 156198,
    "cost": 4.66
  }
}

Track Per-Client Token Usage

# Get detailed token usage for specific client
python generate_usage_report.py --client diageo --last-days 30

Token Usage Optimization

Tips to Reduce Token Usage

1. Optimize Prompts

   • Remove redundant instructions
   • Use clear, concise language
   • Avoid repeating context

2. Smart Model Selection

   • Use Gemini for simpler checks (lower cost per token)
   • Use OpenAI for complex analysis requiring higher accuracy
   • Review token usage reports to identify which checks benefit from which model

3. Reference Asset Optimization

   • Compress reference images when possible
   • Use smaller reference files that still maintain quality
   • Remove unnecessary reference assets from checks that don't need them

4. Profile Optimization

   • Review checks with highest token usage
   • Consider combining similar checks to reduce redundancy
   • Disable checks that provide minimal value for their token cost

Monitoring Token Usage

Set up alerts for unusual token usage:

# Create daily monitoring script
cat > monitor_tokens.sh << 'EOF'
#!/bin/bash
REPORT=$(python generate_usage_report.py --last-days 1 --format json)
TOTAL_TOKENS=$(echo "$REPORT" | jq '.total_tokens')

# Alert if daily usage exceeds 100k tokens
if [ "$TOTAL_TOKENS" -gt 100000 ]; then
    echo "WARNING: High token usage detected: $TOTAL_TOKENS tokens"
    # Send alert email or notification
fi
EOF

chmod +x monitor_tokens.sh

Fallback Behavior

If token usage data is not available from the API (rare edge cases):

• System falls back to estimates (1000 input + 200 output tokens)
• Ensures reports always have cost data
• Logs will indicate when estimates are used vs actual data

Data Privacy

Token counts are stored locally in usage logs:

• No sensitive data sent to external services
• Full control over usage data retention
• Can be archived or deleted as needed

Future Enhancements

Potential future improvements:

1. Real-time Token Dashboard: Web UI showing live token usage
2. Token Budget Alerts: Automatic notifications when approaching limits
3. Per-Check Token Analysis: Detailed breakdown of which checks use most tokens
4. Token Usage Trends: Historical graphs showing token usage over time
5. Cost Forecasting: Predict future costs based on usage patterns

Technical Notes

Token Extraction

OpenAI API:

response.usage.prompt_tokens
response.usage.completion_tokens
response.usage.total_tokens

Gemini API:

response.usage_metadata.prompt_token_count
response.usage_metadata.candidates_token_count
response.usage_metadata.total_token_count

Modified Files

1. llm_config.py

   • call_openai_vision(): Returns token usage
   • call_gemini_vision(): Returns token usage
   • run_visual_qc(): Passes through token usage

2. usage_tracker.py

   • log_analysis_complete(): Logs token statistics
   • _calculate_analysis_cost(): Uses actual tokens for cost calculation
   • get_usage_stats(): Aggregates token data across analyses

3. generate_usage_report.py

   • generate_text_report(): Displays token statistics
   • generate_csv_report(): Includes token columns
   • Added "Token Usage by Provider" section

4. api_server.py

   • process_single_check(): Passes through token usage
   • process_single_check_with_triage(): Passes through token usage

Testing

All changes have been tested:

• ✅ Syntax validation passed
• ✅ Module imports successful
• ✅ Backward compatible with existing logs (falls back to estimates)
• ✅ Reports display token data correctly

Summary

The token tracking enhancement provides:

• ✅ Accuracy: Real token usage instead of estimates
• ✅ Transparency: Detailed breakdown by provider and check
• ✅ Optimization: Identify cost-saving opportunities
• ✅ Forecasting: Better budget planning with actual data
• ✅ Backward Compatible: Works with existing system

Next time you run an analysis, the system will automatically capture actual token usage and include it in reports!