from __future__ import annotations """일일 토큰 상한을 강제한다(아키텍처 §4.1 enrich 단계, agy SSOT §14). agy 의 크레딧·쿼터 정책은 공식 문서에 수치가 없다(SSOT §14). 그래서 **한도를 우리가 직접 센다.** ``agy_calls`` 에 남은 그날의 ``total_tokens`` 합계가 ``agy.daily_token_cap`` 이상이면 호출 자체를 막는다. 쿼터가 실제로 소진되기 전에 우리 쪽에서 먼저 멈추는 것이 목적이다 — 소진된 뒤에는 다음 배치까지 AI 가 통째로 죽기 때문이다. 이 모듈은 **읽기 전용**이다. ``agy_calls`` 에 쓰는 것은 저장 계층(``storage/repo.py``)의 몫이며, 여기서는 집계만 한다. DB 가 없거나 테이블이 아직 없으면 "오늘 0 토큰" 으로 본다 — 예산 조회 실패가 AI 를 막는 이유가 되면 안 된다. """ import logging import sqlite3 from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import Any, Optional logger = logging.getLogger(__name__) __all__ = [ "BudgetVerdict", "DEFAULT_DAILY_TOKEN_CAP", "ESTIMATED_CALL_TOKENS", "calls_today", "check", "today", "tokens_used_today", ] #: 한국 표준시. 하루 경계는 KST 자정이다(배치가 06:00 KST 에 돈다). _KST = timezone(timedelta(hours=9)) #: 설정을 못 읽었을 때 쓰는 기본 상한(``config.AgyCfg.daily_token_cap`` 과 같은 값). DEFAULT_DAILY_TOKEN_CAP = 300_000 #: 한 번 호출하면 최소 이만큼은 쓴다는 보수적 추정치. #: SSOT §4.3 실측: "OK 한 단어" 프롬프트에도 input 28,317 토큰이 들었다. #: 잔여 예산이 이보다 적으면 호출해 봐야 도중에 끊길 뿐이므로 미리 막는다. ESTIMATED_CALL_TOKENS = 30_000 #: ``agy_calls`` 에서 날짜로 쓸 수 있는 컬럼 후보. 스키마(``started_at``)와 #: ``models.AgyCallRow``(``called_at``) 의 이름이 달라 양쪽을 모두 받는다. _DATE_COLUMNS = ("started_at", "called_at") #: 토큰 합계 컬럼 후보. _TOKEN_COLUMNS = ("total_tokens",) _TABLE = "agy_calls" @dataclass(frozen=True, slots=True) class BudgetVerdict: """예산 판정 결과. ``allowed`` 가 False 면 enrich 단계를 건너뛴다.""" allowed: bool # 오늘 호출해도 되는가 used_today: int # 오늘 이미 쓴 토큰 합계 cap: int # 일일 상한 remaining: int # 남은 토큰(음수가 되지 않게 0 으로 바닥을 친다) calls_today: int = 0 # 오늘 호출 횟수(재시도 포함) reason: str = "" # 차단 사유. 허용이면 빈 문자열 day: str = "" # 판정 기준 날짜(YYYY-MM-DD, KST) @property def skip_reason(self) -> str: """``enrichment_run.skip_reason`` 에 넣을 코드. 허용이면 빈 문자열.""" return "token_cap" if not self.allowed else "" @property def usage_ratio(self) -> float: """상한 대비 사용률(0.0~). 상한이 0 이하면 0.0.""" return (self.used_today / self.cap) if self.cap > 0 else 0.0 def today(now: Optional[datetime] = None) -> str: """KST 기준 오늘 날짜 문자열을 돌려준다. Args: now: 기준 시각. 생략하면 현재 시각. Returns: str: ``YYYY-MM-DD``. """ moment = now or datetime.now(_KST) if moment.tzinfo is None: moment = moment.replace(tzinfo=_KST) return moment.astimezone(_KST).strftime("%Y-%m-%d") def _columns(conn: sqlite3.Connection) -> tuple[str, ...]: """``agy_calls`` 의 컬럼 이름 목록을 읽는다. Args: conn: 열린 SQLite 연결. Returns: tuple[str, ...]: 컬럼 이름. 테이블이 없거나 조회 실패면 빈 튜플. """ try: rows = conn.execute(f"PRAGMA table_info({_TABLE})").fetchall() except sqlite3.Error as exc: logger.debug("%s 컬럼 정보를 읽지 못했다: %s", _TABLE, exc) return () out: list[str] = [] for row in rows: try: out.append(str(row[1])) except (IndexError, TypeError): continue return tuple(out) def _pick(available: tuple[str, ...], candidates: tuple[str, ...]) -> Optional[str]: """후보 중 실제로 존재하는 첫 컬럼 이름을 고른다. Args: available: 테이블의 실제 컬럼 목록. candidates: 우선순위대로 나열한 후보. Returns: Optional[str]: 고른 컬럼 이름. 없으면 None. """ for name in candidates: if name in available: return name return None def _aggregate(conn: Any, day: str) -> tuple[int, int]: """그날의 (토큰 합계, 호출 횟수) 를 구한다. Args: conn: 열린 SQLite 연결. None 이면 (0, 0). day: ``YYYY-MM-DD``. Returns: tuple[int, int]: (토큰 합계, 호출 횟수). 조회 불가면 (0, 0). """ if conn is None: return 0, 0 available = _columns(conn) if not available: logger.debug("%s 테이블이 아직 없다. 오늘 사용량을 0 으로 본다.", _TABLE) return 0, 0 date_col = _pick(available, _DATE_COLUMNS) token_col = _pick(available, _TOKEN_COLUMNS) if date_col is None or token_col is None: logger.warning( "%s 에서 날짜/토큰 컬럼을 찾지 못했다(컬럼: %s). 사용량을 0 으로 본다.", _TABLE, ", ".join(available), ) return 0, 0 sql = ( f"SELECT COALESCE(SUM({token_col}), 0), COUNT(*) " f"FROM {_TABLE} WHERE substr({date_col}, 1, 10) = ?" ) try: row = conn.execute(sql, (day,)).fetchone() except sqlite3.Error as exc: logger.warning("일일 토큰 사용량 집계에 실패했다: %s", exc) return 0, 0 if not row: return 0, 0 try: return int(row[0] or 0), int(row[1] or 0) except (TypeError, ValueError): return 0, 0 def tokens_used_today(conn: Any, day: Optional[str] = None) -> int: """그날 이미 쓴 토큰 합계를 돌려준다(아키텍처 §3.8 ``repo.tokens_used_today`` 대응). Args: conn: 열린 SQLite 연결. day: ``YYYY-MM-DD``. 생략하면 KST 오늘. Returns: int: 토큰 합계. 조회할 수 없으면 0. """ return _aggregate(conn, day or today())[0] def calls_today(conn: Any, day: Optional[str] = None) -> int: """그날의 agy 호출 횟수를 돌려준다(재시도 포함). Args: conn: 열린 SQLite 연결. day: ``YYYY-MM-DD``. 생략하면 KST 오늘. Returns: int: 호출 횟수. 조회할 수 없으면 0. """ return _aggregate(conn, day or today())[1] def check( conn: Any, cfg: Any, *, day: Optional[str] = None, estimated_tokens: int = ESTIMATED_CALL_TOKENS, ) -> BudgetVerdict: """오늘 agy 를 호출해도 되는지 판정한다. 남은 예산이 ``estimated_tokens`` 보다 적으면 **호출하기 전에** 막는다. 한 번 부르면 최소 3만 토큰이 나가므로(SSOT §4.3 실측), 남은 게 그보다 적을 때 부르는 것은 상한을 넘기겠다고 예고하는 것과 같다. Args: conn: 열린 SQLite 연결. None 이어도 동작한다(사용량 0 으로 간주). cfg: ``config.AgyCfg``. ``daily_token_cap`` 을 읽는다. day: ``YYYY-MM-DD``. 생략하면 KST 오늘. estimated_tokens: 이번 호출의 보수적 예상 소모량. 0 이면 잔여 검사만 한다. Returns: BudgetVerdict: ``allowed=False`` 면 ``reason`` 에 사용자에게 보여줄 사유가 담긴다. """ target_day = day or today() try: cap = int(getattr(cfg, "daily_token_cap", DEFAULT_DAILY_TOKEN_CAP)) except (TypeError, ValueError): cap = DEFAULT_DAILY_TOKEN_CAP used, count = _aggregate(conn, target_day) remaining = max(0, cap - used) if cap <= 0: return BudgetVerdict( allowed=False, used_today=used, cap=cap, remaining=0, calls_today=count, reason="일일 토큰 상한이 0 이하로 설정돼 AI 계층이 꺼져 있습니다.", day=target_day, ) if used >= cap: return BudgetVerdict( allowed=False, used_today=used, cap=cap, remaining=0, calls_today=count, reason=( f"오늘 AI 토큰 사용량이 상한에 도달했습니다({used:,} / {cap:,}). " "내일 자동으로 다시 시도합니다." ), day=target_day, ) if estimated_tokens > 0 and remaining < estimated_tokens: return BudgetVerdict( allowed=False, used_today=used, cap=cap, remaining=remaining, calls_today=count, reason=( f"남은 AI 토큰이 {remaining:,} 개로 1회 호출 예상치({estimated_tokens:,})보다 " f"적어 건너뜁니다(상한 {cap:,})." ), day=target_day, ) return BudgetVerdict( allowed=True, used_today=used, cap=cap, remaining=remaining, calls_today=count, reason="", day=target_day, )