マルチテナント SaaS で Gemini API を安全に提供する設計 — テナント分離・課金計測・暴走防止の本番アーキテクチャ

# tenant_context.py
from dataclasses import dataclass
from typing import Literal
 
PlanTier = Literal["free", "starter", "pro", "enterprise"]
 
@dataclass(frozen=True)
class TenantContext:
    """1リクエストのライフサイクル中だけ生きるテナント情報。
    認証ミドルウェアで生成し、Gemini 呼び出し関数まで引き回す。"""
 
    tenant_id: str            # 例: "tnt_01HX..."
    user_id: str              # 例: "usr_01HX..."
    plan: PlanTier            # 課金プラン
    monthly_token_budget: int # この月の残り使用可能トークン数
    request_id: str           # トレース用の一意ID
 
    @property
    def has_budget(self) -> bool:
        return self.monthly_token_budget > 0

# tenant_rate_limiter.py
import time
import redis.asyncio as redis
 
# Lua: 1呼び出しの中で「補充→残量取得→残ってれば1減らす」を原子的に行う
# return: 残りトークン（負なら拒否すべき）
LUA_TOKEN_BUCKET = """
local key = KEYS[1]
local rate_per_sec = tonumber(ARGV[1])
local capacity = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local cost = tonumber(ARGV[4])
 
local data = redis.call('HMGET', key, 'tokens', 'last_refill')
local tokens = tonumber(data[1]) or capacity
local last_refill = tonumber(data[2]) or now
 
-- 経過時間ぶん補充
local elapsed = math.max(0, now - last_refill)
tokens = math.min(capacity, tokens + elapsed * rate_per_sec)
 
local allowed = tokens >= cost
if allowed then
  tokens = tokens - cost
end
 
redis.call('HMSET', key, 'tokens', tokens, 'last_refill', now)
redis.call('EXPIRE', key, 3600)
return allowed and tokens or -1
"""
 
class TenantRateLimiter:
    # プランごとに rate(tokens/sec) と burst(capacity) を分ける
    PLAN_LIMITS = {
        "free":       (1.0,   10),   # 1 req/sec、burst 10
        "starter":    (5.0,   50),
        "pro":        (20.0,  200),
        "enterprise": (100.0, 1000),
    }
 
    def __init__(self, redis_client: redis.Redis):
        self.redis = redis_client
        self._sha: str | None = None
 
    async def _ensure_loaded(self) -> str:
        if self._sha is None:
            self._sha = await self.redis.script_load(LUA_TOKEN_BUCKET)
        return self._sha
 
    async def try_acquire(self, ctx, cost: int = 1) -> bool:
        """このテナントが今すぐ1リクエスト送ってよいかを判定する。
        cost を可変にすると「重いリクエストは2消費」のような重み付けも可能。"""
        rate, capacity = self.PLAN_LIMITS[ctx.plan]
        sha = await self._ensure_loaded()
        key = f"ratelimit:tenant:{ctx.tenant_id}"
        result = await self.redis.evalsha(
            sha, 1, key, rate, capacity, time.time(), cost
        )
        return result != -1

# usage_recorder.py
import asyncio
from dataclasses import dataclass
from datetime import datetime, timezone
 
from google import genai
from google.genai import types
 
# モデル別の単価（2026年4月時点・USD / 1M tokens、サンプル値）
MODEL_PRICING = {
    "gemini-2.5-pro":    {"input": 1.25, "output": 10.00},
    "gemini-2.5-flash":  {"input": 0.10, "output": 0.40},
    "gemini-2.5-flash-lite": {"input": 0.05, "output": 0.20},
}
 
@dataclass
class UsageRecord:
    tenant_id: str
    user_id: str
    request_id: str
    model: str
    input_tokens: int
    output_tokens: int
    cost_usd: float
    created_at: datetime
    success: bool
    error_code: str | None = None
 
def _calc_cost(model: str, input_tokens: int, output_tokens: int) -> float:
    p = MODEL_PRICING.get(model, {"input": 1.0, "output": 5.0})
    return (input_tokens * p["input"] + output_tokens * p["output"]) / 1_000_000
 
class GeminiClientWithMetering:
    """Gemini 呼び出しと課金記録を同じトランザクションにまとめるラッパー。
    呼び出し側は ctx と prompt を渡すだけでよい。"""
 
    def __init__(self, client: genai.Client, recorder, budget_guard):
        self.client = client
        self.recorder = recorder           # 後述: DB 書き込み担当
        self.budget_guard = budget_guard   # 後述: 月予算チェック担当
 
    async def generate(self, ctx, model: str, prompt: str) -> str:
        # 月予算チェック（残り0なら即拒否）
        if not await self.budget_guard.allow(ctx):
            raise PermissionError(f"tenant {ctx.tenant_id} over monthly budget")
 
        record = UsageRecord(
            tenant_id=ctx.tenant_id,
            user_id=ctx.user_id,
            request_id=ctx.request_id,
            model=model,
            input_tokens=0,
            output_tokens=0,
            cost_usd=0.0,
            created_at=datetime.now(timezone.utc),
            success=False,
        )
 
        try:
            resp = await self.client.aio.models.generate_content(
                model=model,
                contents=prompt,
            )
            usage = resp.usage_metadata
            record.input_tokens = usage.prompt_token_count or 0
            # 思考トークンは出力と同じ単価で課金されるが、
            # candidates_token_count には含まれない。別枠の値を足す
            record.output_tokens = (
                (usage.candidates_token_count or 0)
                + (usage.thoughts_token_count or 0)
            )
            record.cost_usd = _calc_cost(
                model, record.input_tokens, record.output_tokens
            )
            record.success = True
            return resp.text
 
        except Exception as e:
            record.error_code = type(e).__name__
            raise
        finally:
            # 成功・失敗を問わず必ず記録する。失敗した呼び出しでも、
            # Gemini 側で課金が発生する場合があるため。
            await self.recorder.persist(record)

-- migrations/001_usage_records.sql
CREATE TABLE usage_records (
    id              BIGSERIAL PRIMARY KEY,
    tenant_id       TEXT NOT NULL,
    user_id         TEXT NOT NULL,
    request_id      TEXT NOT NULL UNIQUE,  -- 二重課金防止
    model           TEXT NOT NULL,
    input_tokens    INTEGER NOT NULL,
    output_tokens   INTEGER NOT NULL,
    cost_usd        NUMERIC(10, 6) NOT NULL,
    success         BOOLEAN NOT NULL,
    error_code      TEXT,
    created_at      TIMESTAMPTZ NOT NULL DEFAULT NOW()
);
 
CREATE INDEX idx_usage_tenant_created
    ON usage_records (tenant_id, created_at DESC);
 
-- 月次集計を高速化したいなら以下のような Materialized View を併用
CREATE MATERIALIZED VIEW usage_monthly AS
SELECT
    tenant_id,
    date_trunc('month', created_at) AS month,
    SUM(input_tokens) AS input_tokens,
    SUM(output_tokens) AS output_tokens,
    SUM(cost_usd) AS cost_usd,
    COUNT(*) FILTER (WHERE success) AS ok_count,
    COUNT(*) FILTER (WHERE NOT success) AS error_count
FROM usage_records
GROUP BY tenant_id, date_trunc('month', created_at);

# budget_guard.py
import asyncio
from datetime import datetime, timezone
 
import asyncpg
 
class TenantBudgetGuard:
    """テナント単位で『今月使ったコスト』を短時間キャッシュし、
    月予算の80%で警告、105%で拒否する。"""
 
    # プランごとの月間予算（USD）
    PLAN_BUDGETS = {
        "free":       0.50,
        "starter":    20.0,
        "pro":        200.0,
        "enterprise": 2000.0,  # ハードキャップ。事前合意が必要
    }
 
    GRACE_RATIO = 1.05  # 5%までは許容（集計の遅れで境界の呼び出しを不当に弾かないため）
 
    def __init__(self, pool: asyncpg.Pool, alerter):
        self.pool = pool
        self.alerter = alerter
        self._cache: dict[str, tuple[float, datetime]] = {}
        self._locks: dict[str, asyncio.Lock] = {}              # ロックはテナントごと
        self._alerted: set[tuple[str, str, str]] = set()       # (テナント, 月, 種別)
 
    async def _current_spend(self, tenant_id: str) -> float:
        # 60秒キャッシュ。完全な厳密性より、頻繁な DB 叩きの抑制を優先。
        lock = self._locks.setdefault(tenant_id, asyncio.Lock())
        async with lock:
            cached = self._cache.get(tenant_id)
            now = datetime.now(timezone.utc)
            if cached and (now - cached[1]).total_seconds() < 60:
                return cached[0]
 
            month_start = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
            row = await self.pool.fetchrow(
                "SELECT COALESCE(SUM(cost_usd), 0) AS spend "
                "FROM usage_records "
                "WHERE tenant_id = $1 AND created_at >= $2",
                tenant_id, month_start,
            )
            spend = float(row["spend"])
            self._cache[tenant_id] = (spend, now)
            return spend
 
    async def _alert_once(self, kind: str, tenant_id: str, spend: float, budget: float):
        # 同じ月・同じ種別の通知は1回だけ。kind は "warning" か "blocked"
        key = (tenant_id, datetime.now(timezone.utc).strftime("%Y-%m"), kind)
        if key in self._alerted:
            return
        self._alerted.add(key)
        await getattr(self.alerter, f"notify_{kind}")(tenant_id, spend, budget)
 
    async def allow(self, ctx) -> bool:
        budget = self.PLAN_BUDGETS[ctx.plan]
        spend = await self._current_spend(ctx.tenant_id)
 
        if spend >= budget * self.GRACE_RATIO:
            await self._alert_once("blocked", ctx.tenant_id, spend, budget)
            return False
 
        if spend >= budget * 0.80:
            await self._alert_once("warning", ctx.tenant_id, spend, budget)
 
        return True

# spend_ledger.py
from datetime import datetime, timezone
 
class SpendLedger:
    """Redis 上の『今月の概算支出』。呼び出し前に最悪ケースを予約し、
    呼び出し後に実績で精算する。"""
 
    def __init__(self, redis_client, db_spend, budgets: dict, grace: float = 1.05):
        self.redis = redis_client
        self.db_spend = db_spend      # async (tenant_id) -> 今月の確定支出（DB 集計）
        self.budgets = budgets
        self.grace = grace
 
    async def reserve(self, ctx, est_usd: float) -> str | None:
        """予約できれば Redis キーを返し、予算超過なら None を返す。"""
        key = f"spend:{ctx.tenant_id}:{datetime.now(timezone.utc):%Y%m}"
        if not await self.redis.exists(key):
            # 初回だけ DB の集計値で初期化。NX なので並行しても1回だけ入る
            await self.redis.set(key, await self.db_spend(ctx.tenant_id),
                                 nx=True, ex=40 * 86400)
        after = float(await self.redis.incrbyfloat(key, est_usd))
        if after > self.budgets[ctx.plan] * self.grace:
            await self.redis.incrbyfloat(key, -est_usd)   # 予約を戻す
            return None
        return key
 
    async def settle(self, key: str, est_usd: float, actual_usd: float) -> None:
        """見積りと実績の差を反映する。失敗した呼び出しは actual_usd=0 で予約を解放。"""
        await self.redis.incrbyfloat(key, actual_usd - est_usd)

MAX_OUT = 4096   # max_output_tokens。思考トークンを含む上限
 
p = MODEL_PRICING[model]
est = (len(prompt) * p["input"] + MAX_OUT * p["output"]) / 1_000_000
key = await ledger.reserve(ctx, est)
if key is None:
    raise PermissionError(f"tenant {ctx.tenant_id} over monthly budget")
try:
    resp = await client.aio.models.generate_content(
        model=model, contents=prompt,
        config=types.GenerateContentConfig(max_output_tokens=MAX_OUT),
    )
    ...  # 実績コストを record.cost_usd に入れる
finally:
    await ledger.settle(key, est, record.cost_usd)   # 失敗時は 0 で解放

# 悪い例
async for chunk in client.aio.models.generate_content_stream(...):
    yield chunk.text
    if some_condition:
        break  # ここで抜けると usage_metadata を取り損ねる
 
# 良い例
last_chunk = None
async for chunk in client.aio.models.generate_content_stream(...):
    last_chunk = chunk
    yield chunk.text
if last_chunk and last_chunk.usage_metadata:
    record.input_tokens = last_chunk.usage_metadata.prompt_token_count
    record.output_tokens = last_chunk.usage_metadata.candidates_token_count

# app.py
from fastapi import FastAPI, Depends, HTTPException, Request
from contextvars import ContextVar
import uuid
 
from tenant_context import TenantContext
from tenant_rate_limiter import TenantRateLimiter
from usage_recorder import GeminiClientWithMetering
from budget_guard import TenantBudgetGuard
 
# 実運用では起動時に1度だけ構築します
limiter: TenantRateLimiter
gemini: GeminiClientWithMetering
guard: TenantBudgetGuard
 
# ログ付与のため、リクエストごとにテナントコンテキストを束ねます
current_ctx: ContextVar[TenantContext | None] = ContextVar("current_ctx", default=None)
 
async def get_tenant_context(request: Request) -> TenantContext:
    api_key = request.headers.get("Authorization", "").removeprefix("Bearer ").strip()
    if not api_key:
        raise HTTPException(401, "missing bearer token")
 
    # 自前のテナントテーブルを引きます。この前に Gemini を呼んではいけません
    row = await app.state.db.fetchrow(
        "SELECT tenant_id, user_id, plan FROM api_keys WHERE token_hash = $1",
        sha256_hex(api_key),
    )
    if row is None:
        raise HTTPException(401, "invalid token")
 
    ctx = TenantContext(
        tenant_id=row["tenant_id"],
        user_id=row["user_id"],
        plan=row["plan"],
        monthly_token_budget=10**9,  # 概算。実際の上限は予算ガードが強制します
        request_id=str(uuid.uuid4()),
    )
    current_ctx.set(ctx)
    return ctx
 
app = FastAPI()
 
@app.post("/v1/chat")
async def chat(payload: dict, ctx: TenantContext = Depends(get_tenant_context)):
    if not await limiter.try_acquire(ctx, cost=1):
        # 行儀の良いクライアントが正しくバックオフできるようヒントを返します
        raise HTTPException(429, "tenant rate limit exceeded", headers={"Retry-After": "1"})
 
    try:
        text = await gemini.generate(ctx, model=payload["model"], prompt=payload["prompt"])
    except PermissionError as e:
        raise HTTPException(402, str(e))  # 予算超過は 402 Payment Required で返す
 
    return {"text": text, "request_id": ctx.request_id}

# tests/fakes.py
from dataclasses import dataclass
from types import SimpleNamespace
 
@dataclass
class FakeUsage:
    prompt_token_count: int
    candidates_token_count: int
    total_token_count: int
 
class FakeGeminiAio:
    def __init__(self, prompt_tokens=100, output_tokens=200):
        self.prompt_tokens = prompt_tokens
        self.output_tokens = output_tokens
        self.calls = 0
 
    @property
    def models(self):
        outer = self
 
        class _Models:
            async def generate_content(self, model, contents):
                outer.calls += 1
                return SimpleNamespace(
                    text="ok",
                    usage_metadata=FakeUsage(
                        outer.prompt_tokens,
                        outer.output_tokens,
                        outer.prompt_tokens + outer.output_tokens,
                    ),
                )
 
        return _Models()
 
class FakeGeminiClient:
    def __init__(self):
        self.aio = FakeGeminiAio()