nxtgauge-backend-rust/crates/db/migrations/20260706500000_ai_token_cost_engine.up.sql

-- Token cost engine and provider cost tracking (Task 9)
-- Configures model costs and tracks actual usage costs for margin analysis

BEGIN;

-- Model cost configuration table
CREATE TABLE ai_model_cost_config (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    model_alias VARCHAR(100) NOT NULL,
    
    -- Cost per 1K tokens
    cost_per_1k_input_tokens NUMERIC(12, 8) NOT NULL,
    cost_per_1k_output_tokens NUMERIC(12, 8) NOT NULL,
    
    -- Cost basis: compute_amortized (self-hosted) or provider_metered (paid APIs)
    cost_basis VARCHAR(30) NOT NULL CHECK (cost_basis IN ('compute_amortized', 'provider_metered')),
    
    -- For compute_amortized: optional metadata about calculation
    cost_calculation_notes TEXT,
    
    -- Effective date range
    effective_from TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    effective_until TIMESTAMPTZ,
    
    -- Status
    is_active BOOLEAN NOT NULL DEFAULT TRUE,
    
    -- Audit
    created_by UUID REFERENCES users(id),
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    
    UNIQUE(model_alias, effective_from)
);

-- Add cost tracking columns to ai_usage_logs if not exist
DO $$
BEGIN
    -- These columns already exist per the initial migration, but verify they're populated
    IF NOT EXISTS (SELECT 1 FROM information_schema.columns 
                   WHERE table_name = 'ai_usage_logs' 
                   AND column_name = 'input_tokens') THEN
        ALTER TABLE ai_usage_logs ADD COLUMN input_tokens INT;
    END IF;
    
    IF NOT EXISTS (SELECT 1 FROM information_schema.columns 
                   WHERE table_name = 'ai_usage_logs' 
                   AND column_name = 'output_tokens') THEN
        ALTER TABLE ai_usage_logs ADD COLUMN output_tokens INT;
    END IF;
    
    IF NOT EXISTS (SELECT 1 FROM information_schema.columns 
                   WHERE table_name = 'ai_usage_logs' 
                   AND column_name = 'total_tokens') THEN
        ALTER TABLE ai_usage_logs ADD COLUMN total_tokens INT;
    END IF;
END $$;

-- Create or replace view for margin calculation
-- This provides real-time margin analysis without storing computed values
CREATE OR REPLACE VIEW ai_usage_margin_view AS
SELECT 
    l.id,
    l.user_id,
    l.feature_code,
    l.model_alias,
    l.credits_charged,
    l.input_tokens,
    l.output_tokens,
    l.total_tokens,
    l.status,
    l.created_at,
    
    -- Calculate internal cost using current cost config
    CASE 
        WHEN l.input_tokens IS NOT NULL AND l.output_tokens IS NOT NULL THEN
            (l.input_tokens::NUMERIC / 1000.0 * COALESCE(c.cost_per_1k_input_tokens, 0)) +
            (l.output_tokens::NUMERIC / 1000.0 * COALESCE(c.cost_per_1k_output_tokens, 0))
        ELSE NULL
    END as calculated_internal_cost,
    
    -- Cost basis
    c.cost_basis,
    
    -- Margin (credits charged are in some credit unit, convert to currency for margin calc)
    -- Note: This assumes credits have a monetary value; adjust conversion rate as needed
    CASE 
        WHEN l.input_tokens IS NOT NULL AND l.output_tokens IS NOT NULL THEN
            l.credits_charged - (
                (l.input_tokens::NUMERIC / 1000.0 * COALESCE(c.cost_per_1k_input_tokens, 0)) +
                (l.output_tokens::NUMERIC / 1000.0 * COALESCE(c.cost_per_1k_output_tokens, 0))
            )
        ELSE NULL
    END as margin_credits
    
FROM ai_usage_logs l
LEFT JOIN ai_model_cost_config c ON l.model_alias = c.model_alias 
    AND c.is_active = TRUE 
    AND c.effective_from <= l.created_at
    AND (c.effective_until IS NULL OR c.effective_until > l.created_at)
WHERE l.status = 'success';

-- Indexes
CREATE INDEX idx_ai_model_cost_config_alias ON ai_model_cost_config(model_alias) WHERE is_active = TRUE;
CREATE INDEX idx_ai_usage_logs_tokens ON ai_usage_logs(user_id, created_at DESC) 
    WHERE input_tokens IS NOT NULL OR output_tokens IS NOT NULL;

-- Seed initial cost config for self-hosted Ollama models (compute amortized)
-- These are example values - adjust based on actual GPU costs and throughput
INSERT INTO ai_model_cost_config (model_alias, cost_per_1k_input_tokens, cost_per_1k_output_tokens, cost_basis, cost_calculation_notes)
VALUES 
    ('askash-fast', 0.0001, 0.0002, 'compute_amortized', 'qwen3:4b on shared GPU - cost from GPU node amortization'),
    ('askash-main', 0.0002, 0.0004, 'compute_amortized', 'qwen3:8b on shared GPU - cost from GPU node amortization'),
    ('ultra-fast', 0.00005, 0.0001, 'compute_amortized', 'gemma3:270m on shared GPU - cost from GPU node amortization'),
    ('jd-generator', 0.0002, 0.0004, 'compute_amortized', 'qwen3:8b - specialized for job descriptions'),
    ('profile-writer', 0.0002, 0.0004, 'compute_amortized', 'qwen3:8b - specialized for profile writing')
ON CONFLICT (model_alias, effective_from) DO NOTHING;

COMMIT;