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

import { Injectable, Logger } from "@nestjs/common";
import type { LlmUsageLog, Prisma } from "@prisma/client";
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): Promise<LlmUsageLog> {
    this.logger.debug(
      `Tracking usage: ${dto.provider}/${dto.model} - ${String(dto.totalTokens)} tokens`
    );

    return await 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: Prisma.LlmUsageLogWhereInput = {};

    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.gte = new Date(query.startDate);
      }
      if (query.endDate) {
        where.createdAt.lte = new Date(query.endDate);
      }
    }

    const usageLogs: LlmUsageLog[] = await this.prisma.llmUsageLog.findMany({ where });

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

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

    // Group by provider
    const byProviderMap = new Map<string, ProviderUsageDto>();
    for (const log of usageLogs) {
      const existing: ProviderUsageDto | undefined = 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 !== null) {
          const count: number = existing.calls === 1 ? 1 : existing.calls - 1;
          existing.averageDurationMs =
            (existing.averageDurationMs * (count - 1) + log.durationMs) / count;
        }
      } else {
        const newProvider: ProviderUsageDto = {
          provider: log.provider,
          calls: 1,
          promptTokens: log.promptTokens,
          completionTokens: log.completionTokens,
          totalTokens: log.totalTokens,
          costCents: log.costCents ?? 0,
          averageDurationMs: log.durationMs ?? 0,
        };
        byProviderMap.set(log.provider, newProvider);
      }
    }

    // Group by model
    const byModelMap = new Map<string, ModelUsageDto>();
    for (const log of usageLogs) {
      const existing: ModelUsageDto | undefined = 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 !== null) {
          const count: number = existing.calls === 1 ? 1 : existing.calls - 1;
          existing.averageDurationMs =
            (existing.averageDurationMs * (count - 1) + log.durationMs) / count;
        }
      } else {
        const newModel: ModelUsageDto = {
          model: log.model,
          calls: 1,
          promptTokens: log.promptTokens,
          completionTokens: log.completionTokens,
          totalTokens: log.totalTokens,
          costCents: log.costCents ?? 0,
          averageDurationMs: log.durationMs ?? 0,
        };
        byModelMap.set(log.model, newModel);
      }
    }

    // Group by task type
    const byTaskTypeMap = new Map<string, TaskTypeUsageDto>();
    for (const log of usageLogs) {
      const taskType: string = log.taskType ?? "unknown";
      const existing: TaskTypeUsageDto | undefined = 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 !== null) {
          const count: number = existing.calls === 1 ? 1 : existing.calls - 1;
          existing.averageDurationMs =
            (existing.averageDurationMs * (count - 1) + log.durationMs) / count;
        }
      } else {
        const newTaskType: TaskTypeUsageDto = {
          taskType,
          calls: 1,
          promptTokens: log.promptTokens,
          completionTokens: log.completionTokens,
          totalTokens: log.totalTokens,
          costCents: log.costCents ?? 0,
          averageDurationMs: log.durationMs ?? 0,
        };
        byTaskTypeMap.set(taskType, newTaskType);
      }
    }

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

    return response;
  }

  /**
   * Get all usage logs for a specific workspace.
   *
   * @param workspaceId - Workspace UUID
   * @returns Array of usage logs
   */
  async getUsageByWorkspace(workspaceId: string): Promise<LlmUsageLog[]> {
    return await 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): Promise<LlmUsageLog[]> {
    return await 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): Promise<LlmUsageLog[]> {
    return await this.prisma.llmUsageLog.findMany({
      where: { workspaceId, model },
      orderBy: { createdAt: "desc" },
    });
  }
}