ed
80eebfb83b
Both sites 9 (gemini) and 10 (gemini_cli) in get_token_stats had:
try: _ensure_gemini_client()
if _gemini_client:
resp = _gemini_client.models.count_tokens(model=_model, contents=md_content)
total_tokens = cast(int, resp.total_tokens)
except Exception: pass
Body: pass = SS violation.
New helper _count_gemini_tokens_for_stats_result(md_content) -> Result[int]:
- Returns Result(data=token_count) on success
- Returns Result(data=0, errors=[ErrorInfo]) on SDK failure or warmup failure
- Caller treats 0 as 'token count unavailable' and falls back to
character-based estimation
Legacy get_token_stats now uses:
if p in ('gemini', 'gemini_cli'):
total_tokens = _count_gemini_tokens_for_stats_result(md_content).data
(combined both branches into one since the logic was identical)
Audit: ai_client SS 5 -> 3. COMPLIANT 31 -> 32.
35 lines
1.1 KiB
Python
35 lines
1.1 KiB
Python
"""Phase 11 sites 9+10: get_token_stats count_tokens (gemini + gemini_cli).
|
|
|
|
Both have:
|
|
try:
|
|
_ensure_gemini_client()
|
|
if _gemini_client:
|
|
resp = _gemini_client.models.count_tokens(model=_model, contents=md_content)
|
|
total_tokens = cast(int, resp.total_tokens)
|
|
except Exception:
|
|
pass
|
|
|
|
Body: pass = SS violation. Migrate via Result[int] helper.
|
|
"""
|
|
import sys
|
|
sys.path.insert(0, ".")
|
|
|
|
|
|
def test_phase11_sites910_count_gemini_tokens_for_stats_result_exists():
|
|
import src.ai_client
|
|
assert hasattr(src.ai_client, "_count_gemini_tokens_for_stats_result"), \
|
|
"_count_gemini_tokens_for_stats_result helper missing"
|
|
|
|
|
|
def test_phase11_sites910_helper_returns_result():
|
|
import src.ai_client
|
|
import inspect
|
|
fn = src.ai_client._count_gemini_tokens_for_stats_result
|
|
sig = inspect.signature(fn)
|
|
assert "Result" in str(sig.return_annotation), \
|
|
f"_count_gemini_tokens_for_stats_result return must be Result, got {sig.return_annotation}"
|
|
|
|
|
|
def test_phase11_sites910_get_token_stats_legacy_preserved():
|
|
import src.ai_client
|
|
assert callable(getattr(src.ai_client, "get_token_stats", None)) |