Prevent runaway AI costs with pre-flight checks and token budgeting

import type { ModelMessage } from "ai";
 
/**
 * Estimate tokens for messages array using rough character-to-token ratio
 * Use this for pre-flight cost estimation
 *
 * Note: This is an approximation. Actual tokens depend on tokenizer.
 * Rule of thumb: 1 token ≈ 4 characters for English text.
 * In production, use a real tokenizer like OpenAI's tiktoken for more accurate estimates:
 * https://github.com/openai/tiktoken
 */
export function estimateTokens(messages: ModelMessage[]): number {
  return messages.reduce((sum, msg) => {
    const content = typeof msg.content === 'string'
      ? msg.content
      : JSON.stringify(msg.content);
    // Rough estimate: 1 token per 4 characters
    return sum + Math.ceil(content.length / 4);
  }, 0);
}
 
/**
 * Calculate request cost before making it
 * Use for budget decisions and request rejection
 *
 * January 2025 pricing (per 1K tokens):
 * - gpt-4o-mini: $0.00015 input, $0.0006 output
 * - gpt-3.5-turbo: $0.0005 input, $0.0015 output
 * - gpt-4o: $0.0025 input, $0.01 output
 */
export function estimateRequestCost(
  inputTokens: number,
  outputTokens: number,
  model: string
): number {
  const costs: Record<string, { input: number; output: number }> = {
    'gpt-4o-mini': { input: 0.00015, output: 0.0006 },
    'gpt-3.5-turbo': { input: 0.0005, output: 0.0015 },
    'gpt-4o': { input: 0.0025, output: 0.01 },
  };
 
  const modelCost = costs[model] || costs['gpt-3.5-turbo'];
  return (
    (inputTokens / 1000) * modelCost.input +
    (outputTokens / 1000) * modelCost.output
  );
}
 
/**
 * Calculate actual cost from AI SDK usage response
 * Use to compare estimated vs actual and tune your estimates
 */
export function calculateActualCost(
  promptTokens: number,
  completionTokens: number,
  model: string
): number {
  return estimateRequestCost(promptTokens, completionTokens, model);
}
 
/**
 * Pre-flight check: Should we reject this request due to cost?
 * Reject obviously expensive requests BEFORE they burn money
 */
export function shouldRejectRequest(estimatedCost: number): {
  reject: boolean;
  reason?: string;
} {
  const MAX_REQUEST_COST = 0.10; // $0.10 per request cap
 
  if (estimatedCost > MAX_REQUEST_COST) {
    return {
      reject: true,
      reason: `Request would cost $${estimatedCost.toFixed(4)}, exceeds $${MAX_REQUEST_COST} limit`
    };
  }
 
  return { reject: false };
}
 
/**
 * Check if we should use cheap mode based on environment variable
 * In production, set this based on your own budget tracking
 */
export function shouldUseCheapModel(): boolean {
  return process.env.FORCE_CHEAP_MODEL === 'true';
}

import type { KnownEventFromType } from "@slack/bolt";
import { generateText, type ModelMessage, stepCountIs } from "ai";
import { app } from "~/app";
import {
  calculateActualCost,
  estimateRequestCost,
  estimateTokens,
  shouldRejectRequest,
  shouldUseCheapModel,
} from "./cost-control";
import { simulateFailure, withRetry } from "./retry-wrapper";
// ... rest of imports ...
 
export const respondToMessage = async ({
  messages,
  event,
  channel,
  thread_ts,
  botId,
  correlation,
}: RespondToMessageOptions) => {
  // Pre-flight cost estimation - catch expensive requests BEFORE they cost money
  const inputTokens = estimateTokens(messages);
  const estimatedOutputTokens = 500; // Conservative estimate for response
  const primaryModel = "gpt-4o-mini";
  const estimatedCost = estimateRequestCost(
    inputTokens,
    estimatedOutputTokens,
    primaryModel
  );
 
  app.logger.info("Pre-flight cost check", {
    ...correlation,
    inputTokens,
    estimatedOutputTokens,
    estimatedCost,
    model: primaryModel,
    status: estimatedCost > 0.10 ? "rejected" : "approved",
  });
 
  // Reject expensive requests BEFORE making them
  const rejectCheck = shouldRejectRequest(estimatedCost);
  if (rejectCheck.reject) {
    app.logger.warn("Request rejected - cost limit exceeded", {
      ...correlation,
      estimatedCost,
      limit: 0.1,
      inputTokens,
      estimatedOutputTokens,
      model: primaryModel,
    });
 
    // Return user-friendly error instead of making expensive request
    return `Your request is too large. Please break it into smaller questions or provide less context. (Estimated cost: $${estimatedCost.toFixed(4)})`;
  }
 
  // Model selection with budget awareness
  // In production, set FORCE_CHEAP_MODEL based on your own budget tracking
  const useCheapMode = shouldUseCheapModel();
  const models = useCheapMode
    ? ["gpt-3.5-turbo"] // Force cheap model only
    : ["gpt-4o-mini", "gpt-3.5-turbo"];
 
  if (useCheapMode) {
    app.logger.warn("Budget threshold triggered - using cheap model", {
      ...correlation,
      reason: "FORCE_CHEAP_MODEL environment variable set",
      selectedModel: "gpt-3.5-turbo",
      normalModel: "gpt-4o-mini",
    });
  }
 
  let lastError: unknown;
 
  for (const currentModel of models) {
    try {
      const result = await withRetry(
        async () => {
          app.logger.info("Attempting AI request", {
            ...correlation,
            model: currentModel,
            inputTokens,
            estimatedCost,
          });
 
          return await generateText({
            model: currentModel,
            system: `You are Slack Agent, a helpful assistant in Slack.
            // ... rest of system prompt ...
            `,
            messages,
            // ... rest of config ...
          });
        },
        {
          maxRetries: 3,
          initialDelayMs: 1000,
        }
      );
 
      // Log actual cost after request completes
      const actualCost = calculateActualCost(
        result.usage.promptTokens,
        result.usage.completionTokens,
        currentModel
      );
 
      const estimationAccuracy = ((actualCost / estimatedCost) * 100).toFixed(0);
 
      app.logger.info("AI request completed", {
        ...correlation,
        model: currentModel,
        usage: {
          promptTokens: result.usage.promptTokens,
          completionTokens: result.usage.completionTokens,
          totalTokens: result.usage.totalTokens,
        },
        actualCost,
        estimatedCost,
        estimationAccuracy: `${estimationAccuracy}%`,
      });
 
      return result.text;
    } catch (error) {
      // ... existing error handling ...
      lastError = error;
    }
  }
 
  throw lastError;
};