【新機能】Gemini API や Vertex AI で利用費の上限を設定できる機能が公開されました

import { GoogleGenAI } from "@google/genai";
import { appendFileSync, writeFileSync } from "node:fs";
import { parseArgs } from "node:util";

const LONG_PROMPT =
  "Please write a detailed 2000-word essay about the history of " +
  "artificial intelligence, covering key milestones from the 1950s " +
  "to the present day. Include discussion of neural networks, " +
  "expert systems, machine learning, and large language models.";

const { values: args } = parseArgs({
  options: {
    model: { type: "string", default: "gemini-3.1-pro-preview" },
    interval: { type: "string", default: "10" },
    "max-requests": { type: "string", default: "500" },
    "max-errors": { type: "string", default: "10" },
    output: { type: "string", default: "spend_cap_results.jsonl" },
  },
});

const model = args.model;
const interval = Number(args.interval);
const maxRequests = Number(args["max-requests"]);
const maxErrors = Number(args["max-errors"]);
const outputPath = args.output;

const ai = new GoogleGenAI({ apiKey: process.env.GEMINI_API_KEY });

const sleep = (sec) => new Promise((r) => setTimeout(r, sec * 1000));

async function callGemini() {
  const timestamp = new Date().toISOString();
  try {
    const response = await ai.models.generateContent({
      model,
      contents: LONG_PROMPT,
    });
    const usage = response.usageMetadata;
    const inputTokens = usage.promptTokenCount;
    const outputTokens = usage.candidatesTokenCount;
    const thinkingTokens = usage.thoughtsTokenCount ?? 0;
    return {
      timestamp,
      status: "success",
      input_tokens: inputTokens,
      output_tokens: outputTokens,
      thinking_tokens: thinkingTokens,
      total_tokens: inputTokens + outputTokens + thinkingTokens,
    };
  } catch (e) {
    return {
      timestamp,
      status: "error",
      error_type: e.constructor.name,
      error_message: String(e.message ?? e),
    };
  }
}

console.log(`Model: ${model}`);
console.log(`Interval: ${interval}s`);
console.log(`Output: ${outputPath}`);
console.log(`Max requests: ${maxRequests}`);
console.log("---");

let consecutiveErrors = 0;
let totalTokens = 0;
let successCount = 0;
let errorCount = 0;

for (let i = 1; i <= maxRequests; i++) {
  const result = await callGemini();
  result.request_number = i;
  appendFileSync(outputPath, JSON.stringify(result) + "\n");

  if (result.status === "success") {
    consecutiveErrors = 0;
    successCount++;
    totalTokens += result.total_tokens;
    console.log(
      `[${i}] OK | tokens: ${result.total_tokens.toLocaleString()} | cumulative: ${totalTokens.toLocaleString()}`
    );
  } else {
    consecutiveErrors++;
    errorCount++;
    console.log(
      `[${i}] ERROR | ${result.error_type}: ${result.error_message.slice(0, 120)}`
    );
  }

  if (consecutiveErrors >= maxErrors) {
    console.log(`\n${maxErrors} consecutive errors. Stopping.`);
    break;
  }

  if (i < maxRequests) {
    await sleep(interval);
  }
}

console.log("\n--- Summary ---");
console.log(`Total requests: ${successCount + errorCount}`);
console.log(`Success: ${successCount}`);
console.log(`Errors: ${errorCount}`);
console.log(`Total tokens: ${totalTokens.toLocaleString()}`);