MOPC-Portal/src/server/utils/ai-usage.ts

/**
 * AI Usage Tracking Utility
 *
 * Logs AI API usage to the database for cost tracking and monitoring.
 * Calculates estimated costs based on model pricing.
 */

import { prisma } from '@/lib/prisma'
import { Decimal } from '@prisma/client/runtime/library'
import type { Prisma } from '@prisma/client'

// ─── Types ───────────────────────────────────────────────────────────────────

export type AIAction =
  | 'ASSIGNMENT'
  | 'FILTERING'
  | 'AWARD_ELIGIBILITY'
  | 'MENTOR_MATCHING'
  | 'PROJECT_TAGGING'
  | 'EVALUATION_SUMMARY'
  | 'ROUTING'

export type AIStatus = 'SUCCESS' | 'PARTIAL' | 'ERROR'

export interface LogAIUsageInput {
  userId?: string
  action: AIAction
  entityType?: string
  entityId?: string
  model: string
  promptTokens: number
  completionTokens: number
  totalTokens: number
  batchSize?: number
  itemsProcessed?: number
  status: AIStatus
  errorMessage?: string
  detailsJson?: Record<string, unknown>
}

export interface TokenUsageResult {
  promptTokens: number
  completionTokens: number
  totalTokens: number
}

// ─── Model Pricing (per 1M tokens) ───────────────────────────────────────────

interface ModelPricing {
  input: number // $ per 1M input tokens
  output: number // $ per 1M output tokens
}

/**
 * OpenAI model pricing as of 2024/2025
 * Prices in USD per 1 million tokens
 */
const MODEL_PRICING: Record<string, ModelPricing> = {
  // GPT-4o series
  'gpt-4o': { input: 2.5, output: 10.0 },
  'gpt-4o-2024-11-20': { input: 2.5, output: 10.0 },
  'gpt-4o-2024-08-06': { input: 2.5, output: 10.0 },
  'gpt-4o-2024-05-13': { input: 5.0, output: 15.0 },
  'gpt-4o-mini': { input: 0.15, output: 0.6 },
  'gpt-4o-mini-2024-07-18': { input: 0.15, output: 0.6 },

  // GPT-4 Turbo series
  'gpt-4-turbo': { input: 10.0, output: 30.0 },
  'gpt-4-turbo-2024-04-09': { input: 10.0, output: 30.0 },
  'gpt-4-turbo-preview': { input: 10.0, output: 30.0 },
  'gpt-4-1106-preview': { input: 10.0, output: 30.0 },
  'gpt-4-0125-preview': { input: 10.0, output: 30.0 },

  // GPT-4 (base)
  'gpt-4': { input: 30.0, output: 60.0 },
  'gpt-4-0613': { input: 30.0, output: 60.0 },
  'gpt-4-32k': { input: 60.0, output: 120.0 },
  'gpt-4-32k-0613': { input: 60.0, output: 120.0 },

  // GPT-3.5 Turbo series
  'gpt-3.5-turbo': { input: 0.5, output: 1.5 },
  'gpt-3.5-turbo-0125': { input: 0.5, output: 1.5 },
  'gpt-3.5-turbo-1106': { input: 1.0, output: 2.0 },
  'gpt-3.5-turbo-16k': { input: 3.0, output: 4.0 },

  // o1 reasoning models
  'o1': { input: 15.0, output: 60.0 },
  'o1-2024-12-17': { input: 15.0, output: 60.0 },
  'o1-preview': { input: 15.0, output: 60.0 },
  'o1-preview-2024-09-12': { input: 15.0, output: 60.0 },
  'o1-mini': { input: 3.0, output: 12.0 },
  'o1-mini-2024-09-12': { input: 3.0, output: 12.0 },

  // o3 reasoning models
  'o3-mini': { input: 1.1, output: 4.4 },
  'o3-mini-2025-01-31': { input: 1.1, output: 4.4 },

  // o4 reasoning models (future-proofing)
  'o4-mini': { input: 1.1, output: 4.4 },
}

// Default pricing for unknown models (conservative estimate)
const DEFAULT_PRICING: ModelPricing = { input: 5.0, output: 15.0 }

// ─── Cost Calculation ────────────────────────────────────────────────────────

/**
 * Get pricing for a model, with fallback for unknown models
 */
function getModelPricing(model: string): ModelPricing {
  // Exact match
  if (MODEL_PRICING[model]) {
    return MODEL_PRICING[model]
  }

  // Try to match by prefix
  const modelLower = model.toLowerCase()
  for (const [key, pricing] of Object.entries(MODEL_PRICING)) {
    if (modelLower.startsWith(key.toLowerCase())) {
      return pricing
    }
  }

  // Fallback based on model type
  if (modelLower.startsWith('gpt-4o-mini')) {
    return MODEL_PRICING['gpt-4o-mini']
  }
  if (modelLower.startsWith('gpt-4o')) {
    return MODEL_PRICING['gpt-4o']
  }
  if (modelLower.startsWith('gpt-4')) {
    return MODEL_PRICING['gpt-4-turbo']
  }
  if (modelLower.startsWith('gpt-3.5')) {
    return MODEL_PRICING['gpt-3.5-turbo']
  }
  if (modelLower.startsWith('o1-mini')) {
    return MODEL_PRICING['o1-mini']
  }
  if (modelLower.startsWith('o1')) {
    return MODEL_PRICING['o1']
  }
  if (modelLower.startsWith('o3-mini')) {
    return MODEL_PRICING['o3-mini']
  }
  if (modelLower.startsWith('o3')) {
    return MODEL_PRICING['o3-mini'] // Conservative estimate
  }
  if (modelLower.startsWith('o4')) {
    return MODEL_PRICING['o4-mini'] || DEFAULT_PRICING
  }

  return DEFAULT_PRICING
}

/**
 * Calculate estimated cost in USD for a given model and token usage
 */
export function calculateCost(
  model: string,
  promptTokens: number,
  completionTokens: number
): number {
  const pricing = getModelPricing(model)

  const inputCost = (promptTokens / 1_000_000) * pricing.input
  const outputCost = (completionTokens / 1_000_000) * pricing.output

  return inputCost + outputCost
}

/**
 * Format cost for display
 */
export function formatCost(costUsd: number): string {
  if (costUsd < 0.01) {
    return `$${(costUsd * 100).toFixed(3)}¢`
  }
  return `$${costUsd.toFixed(4)}`
}

// ─── Logging ─────────────────────────────────────────────────────────────────

/**
 * Log AI usage to the database
 */
export async function logAIUsage(input: LogAIUsageInput): Promise<void> {
  try {
    const estimatedCost = calculateCost(
      input.model,
      input.promptTokens,
      input.completionTokens
    )

    await prisma.aIUsageLog.create({
      data: {
        userId: input.userId,
        action: input.action,
        entityType: input.entityType,
        entityId: input.entityId,
        model: input.model,
        promptTokens: input.promptTokens,
        completionTokens: input.completionTokens,
        totalTokens: input.totalTokens,
        estimatedCostUsd: new Decimal(estimatedCost),
        batchSize: input.batchSize,
        itemsProcessed: input.itemsProcessed,
        status: input.status,
        errorMessage: input.errorMessage,
        detailsJson: input.detailsJson as Prisma.InputJsonValue | undefined,
      },
    })
  } catch (error) {
    // Don't let logging failures break the main operation
    console.error('[AI Usage] Failed to log usage:', error)
  }
}

/**
 * Extract token usage from OpenAI API response
 */
export function extractTokenUsage(
  response: { usage?: { prompt_tokens?: number; completion_tokens?: number; total_tokens?: number } }
): TokenUsageResult {
  return {
    promptTokens: response.usage?.prompt_tokens ?? 0,
    completionTokens: response.usage?.completion_tokens ?? 0,
    totalTokens: response.usage?.total_tokens ?? 0,
  }
}

// ─── Statistics ──────────────────────────────────────────────────────────────

export interface AIUsageStats {
  totalTokens: number
  totalCost: number
  byAction: Record<string, { tokens: number; cost: number; count: number }>
  byModel: Record<string, { tokens: number; cost: number; count: number }>
}

/**
 * Get AI usage statistics for a date range
 */
export async function getAIUsageStats(
  startDate?: Date,
  endDate?: Date
): Promise<AIUsageStats> {
  const where: { createdAt?: { gte?: Date; lte?: Date } } = {}

  if (startDate || endDate) {
    where.createdAt = {}
    if (startDate) where.createdAt.gte = startDate
    if (endDate) where.createdAt.lte = endDate
  }

  const logs = await prisma.aIUsageLog.findMany({
    where,
    select: {
      action: true,
      model: true,
      totalTokens: true,
      estimatedCostUsd: true,
    },
  })

  const stats: AIUsageStats = {
    totalTokens: 0,
    totalCost: 0,
    byAction: {},
    byModel: {},
  }

  for (const log of logs) {
    const cost = log.estimatedCostUsd?.toNumber() ?? 0

    stats.totalTokens += log.totalTokens
    stats.totalCost += cost

    // By action
    if (!stats.byAction[log.action]) {
      stats.byAction[log.action] = { tokens: 0, cost: 0, count: 0 }
    }
    stats.byAction[log.action].tokens += log.totalTokens
    stats.byAction[log.action].cost += cost
    stats.byAction[log.action].count += 1

    // By model
    if (!stats.byModel[log.model]) {
      stats.byModel[log.model] = { tokens: 0, cost: 0, count: 0 }
    }
    stats.byModel[log.model].tokens += log.totalTokens
    stats.byModel[log.model].cost += cost
    stats.byModel[log.model].count += 1
  }

  return stats
}

/**
 * Get current month's AI usage cost
 */
export async function getCurrentMonthCost(): Promise<{
  cost: number
  tokens: number
  requestCount: number
}> {
  const startOfMonth = new Date()
  startOfMonth.setDate(1)
  startOfMonth.setHours(0, 0, 0, 0)

  const logs = await prisma.aIUsageLog.findMany({
    where: {
      createdAt: { gte: startOfMonth },
    },
    select: {
      totalTokens: true,
      estimatedCostUsd: true,
    },
  })

  return {
    cost: logs.reduce((sum, log) => sum + (log.estimatedCostUsd?.toNumber() ?? 0), 0),
    tokens: logs.reduce((sum, log) => sum + log.totalTokens, 0),
    requestCount: logs.length,
  }
}