stack/apps/api/src/llm-usage/llm-usage.service.ts

import { Injectable, Logger } from "@nestjs/common";
import { PrismaService } from "../prisma/prisma.service";
import type {
  TrackUsageDto,
  UsageAnalyticsQueryDto,
  UsageAnalyticsResponseDto,
  ProviderUsageDto,
  ModelUsageDto,
  TaskTypeUsageDto,
} from "./dto";

/**
 * LLM Usage Service
 *
 * Tracks and analyzes LLM usage across workspaces, providers, and models.
 * Provides analytics for cost tracking, token usage, and performance metrics.
 */
@Injectable()
export class LlmUsageService {
  private readonly logger = new Logger(LlmUsageService.name);

  constructor(private readonly prisma: PrismaService) {}

  /**
   * Track a single LLM usage event.
   * Records token counts, cost, duration, and metadata.
   *
   * @param dto - Usage tracking data
   * @returns The created usage log entry
   */
  async trackUsage(dto: TrackUsageDto) {
    this.logger.debug(
      `Tracking usage: ${dto.provider}/${dto.model} - ${String(dto.totalTokens)} tokens`
    );

    return this.prisma.llmUsageLog.create({
      data: dto,
    });
  }

  /**
   * Get aggregated usage analytics based on query filters.
   * Supports filtering by workspace, provider, model, user, and date range.
   *
   * @param query - Analytics query filters
   * @returns Aggregated usage analytics
   */
  async getUsageAnalytics(query: UsageAnalyticsQueryDto): Promise<UsageAnalyticsResponseDto> {
    const where: Record<string, unknown> = {};

    if (query.workspaceId) {
      where.workspaceId = query.workspaceId;
    }
    if (query.provider) {
      where.provider = query.provider;
    }
    if (query.model) {
      where.model = query.model;
    }
    if (query.userId) {
      where.userId = query.userId;
    }
    if (query.startDate || query.endDate) {
      where.createdAt = {};
      if (query.startDate) {
        (where.createdAt as Record<string, Date>).gte = new Date(query.startDate);
      }
      if (query.endDate) {
        (where.createdAt as Record<string, Date>).lte = new Date(query.endDate);
      }
    }

    const usageLogs = await this.prisma.llmUsageLog.findMany({ where });

    // Aggregate totals
    const totalCalls = usageLogs.length;
    const totalPromptTokens = usageLogs.reduce((sum, log) => sum + log.promptTokens, 0);
    const totalCompletionTokens = usageLogs.reduce((sum, log) => sum + log.completionTokens, 0);
    const totalTokens = usageLogs.reduce((sum, log) => sum + log.totalTokens, 0);
    const totalCostCents = usageLogs.reduce((sum, log) => sum + (log.costCents ?? 0), 0);

    const durations = usageLogs.map((log) => log.durationMs).filter((d): d is number => d !== null);
    const averageDurationMs =
      durations.length > 0 ? durations.reduce((sum, d) => sum + d, 0) / durations.length : 0;

    // Group by provider
    const byProviderMap = new Map<string, ProviderUsageDto>();
    for (const log of usageLogs) {
      const existing = byProviderMap.get(log.provider);
      if (existing) {
        existing.calls += 1;
        existing.promptTokens += log.promptTokens;
        existing.completionTokens += log.completionTokens;
        existing.totalTokens += log.totalTokens;
        existing.costCents += log.costCents ?? 0;
        if (log.durationMs) {
          const count = existing.calls === 1 ? 1 : existing.calls - 1;
          existing.averageDurationMs =
            (existing.averageDurationMs * (count - 1) + log.durationMs) / count;
        }
      } else {
        byProviderMap.set(log.provider, {
          provider: log.provider,
          calls: 1,
          promptTokens: log.promptTokens,
          completionTokens: log.completionTokens,
          totalTokens: log.totalTokens,
          costCents: log.costCents ?? 0,
          averageDurationMs: log.durationMs ?? 0,
        });
      }
    }

    // Group by model
    const byModelMap = new Map<string, ModelUsageDto>();
    for (const log of usageLogs) {
      const existing = byModelMap.get(log.model);
      if (existing) {
        existing.calls += 1;
        existing.promptTokens += log.promptTokens;
        existing.completionTokens += log.completionTokens;
        existing.totalTokens += log.totalTokens;
        existing.costCents += log.costCents ?? 0;
        if (log.durationMs) {
          const count = existing.calls === 1 ? 1 : existing.calls - 1;
          existing.averageDurationMs =
            (existing.averageDurationMs * (count - 1) + log.durationMs) / count;
        }
      } else {
        byModelMap.set(log.model, {
          model: log.model,
          calls: 1,
          promptTokens: log.promptTokens,
          completionTokens: log.completionTokens,
          totalTokens: log.totalTokens,
          costCents: log.costCents ?? 0,
          averageDurationMs: log.durationMs ?? 0,
        });
      }
    }

    // Group by task type
    const byTaskTypeMap = new Map<string, TaskTypeUsageDto>();
    for (const log of usageLogs) {
      const taskType = log.taskType ?? "unknown";
      const existing = byTaskTypeMap.get(taskType);
      if (existing) {
        existing.calls += 1;
        existing.promptTokens += log.promptTokens;
        existing.completionTokens += log.completionTokens;
        existing.totalTokens += log.totalTokens;
        existing.costCents += log.costCents ?? 0;
        if (log.durationMs) {
          const count = existing.calls === 1 ? 1 : existing.calls - 1;
          existing.averageDurationMs =
            (existing.averageDurationMs * (count - 1) + log.durationMs) / count;
        }
      } else {
        byTaskTypeMap.set(taskType, {
          taskType,
          calls: 1,
          promptTokens: log.promptTokens,
          completionTokens: log.completionTokens,
          totalTokens: log.totalTokens,
          costCents: log.costCents ?? 0,
          averageDurationMs: log.durationMs ?? 0,
        });
      }
    }

    return {
      totalCalls,
      totalPromptTokens,
      totalCompletionTokens,
      totalTokens,
      totalCostCents,
      averageDurationMs,
      byProvider: Array.from(byProviderMap.values()),
      byModel: Array.from(byModelMap.values()),
      byTaskType: Array.from(byTaskTypeMap.values()),
    };
  }

  /**
   * Get all usage logs for a specific workspace.
   *
   * @param workspaceId - Workspace UUID
   * @returns Array of usage logs
   */
  async getUsageByWorkspace(workspaceId: string) {
    return this.prisma.llmUsageLog.findMany({
      where: { workspaceId },
      orderBy: { createdAt: "desc" },
    });
  }

  /**
   * Get usage logs for a specific provider within a workspace.
   *
   * @param workspaceId - Workspace UUID
   * @param provider - Provider name
   * @returns Array of usage logs
   */
  async getUsageByProvider(workspaceId: string, provider: string) {
    return this.prisma.llmUsageLog.findMany({
      where: { workspaceId, provider },
      orderBy: { createdAt: "desc" },
    });
  }

  /**
   * Get usage logs for a specific model within a workspace.
   *
   * @param workspaceId - Workspace UUID
   * @param model - Model name
   * @returns Array of usage logs
   */
  async getUsageByModel(workspaceId: string, model: string) {
    return this.prisma.llmUsageLog.findMany({
      where: { workspaceId, model },
      orderBy: { createdAt: "desc" },
    });
  }
}