content-find-agent-new/content_agent/integrations/gemini_quota.py

"""单 run Gemini 调用配额闸(V3-M5C)。

实现 GeminiVideoClient Protocol,包装真/假 client 在 run_service 单点注入——
对 recall/walk_engine/graph 全部 analyze 透明,零签名改。每 run new 一个实例,
计数跨"初始 recall + walk 内两次 recall"累计。截断的正常路径是 recall 提交前
按 offset 预判(remaining_quota/consume);analyze 内超额返回 _fail 仅作 backstop。
"""

from __future__ import annotations

import threading
from typing import Any

from content_agent.integrations.gemini_video import _fail
from content_agent.interfaces import GeminiVideoClient

_UNLIMITED = 1_000_000_000


class QuotaCappedGeminiVideoClient:
    def __init__(self, inner: GeminiVideoClient, cap: int | None) -> None:
        self.inner = inner
        self.cap = cap
        self.used = 0
        self._lock = threading.Lock()

    def remaining_quota(self) -> int:
        with self._lock:
            return (self.cap - self.used) if self.cap is not None else _UNLIMITED

    def consume(self, count: int) -> None:
        with self._lock:
            self.used += count

    def analyze(
        self,
        content: dict[str, Any],
        media: dict[str, Any],
        source_context: dict[str, Any],
    ) -> dict[str, Any]:
        if self.cap is not None and self.remaining_quota() < 0:
            return _fail("gemini_quota_exhausted")
        return self.inner.analyze(content, media, source_context)