feat: Anthropic API budget checker, fails open on error
This commit is contained in:
@@ -0,0 +1,36 @@
|
|||||||
|
"""Check Anthropic API budget usage. Fails open (returns True) on errors."""
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
import urllib.request
|
||||||
|
import urllib.error
|
||||||
|
import json
|
||||||
|
|
||||||
|
USAGE_API = "https://api.anthropic.com/v1/usage"
|
||||||
|
|
||||||
|
def get_used_usd() -> float:
|
||||||
|
api_key = os.environ.get('ANTHROPIC_API_KEY', '')
|
||||||
|
if not api_key:
|
||||||
|
raise EnvironmentError("ANTHROPIC_API_KEY not set")
|
||||||
|
req = urllib.request.Request(
|
||||||
|
USAGE_API,
|
||||||
|
headers={
|
||||||
|
"x-api-key": api_key,
|
||||||
|
"anthropic-version": "2023-06-01",
|
||||||
|
}
|
||||||
|
)
|
||||||
|
with urllib.request.urlopen(req, timeout=5) as resp:
|
||||||
|
data = json.loads(resp.read())
|
||||||
|
return sum(item.get('cost_usd', 0.0) for item in data.get('data', []))
|
||||||
|
|
||||||
|
def is_budget_ok(limit_usd: float) -> bool:
|
||||||
|
try:
|
||||||
|
used = get_used_usd()
|
||||||
|
return used < limit_usd
|
||||||
|
except Exception:
|
||||||
|
return True # fail open
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
limit = float(sys.argv[1]) if len(sys.argv) > 1 else 5.0
|
||||||
|
ok = is_budget_ok(limit)
|
||||||
|
print("ok" if ok else "low")
|
||||||
|
sys.exit(0 if ok else 1)
|
||||||
@@ -0,0 +1,22 @@
|
|||||||
|
import sys
|
||||||
|
sys.path.insert(0, 'bin')
|
||||||
|
from unittest.mock import patch
|
||||||
|
import pytest
|
||||||
|
from check_budget import is_budget_ok, get_used_usd
|
||||||
|
|
||||||
|
def test_budget_ok_when_used_less_than_limit():
|
||||||
|
with patch('check_budget.get_used_usd', return_value=2.0):
|
||||||
|
assert is_budget_ok(limit_usd=5.0) is True
|
||||||
|
|
||||||
|
def test_budget_not_ok_when_used_exceeds_limit():
|
||||||
|
with patch('check_budget.get_used_usd', return_value=5.5):
|
||||||
|
assert is_budget_ok(limit_usd=5.0) is False
|
||||||
|
|
||||||
|
def test_budget_ok_on_api_error():
|
||||||
|
"""Fail open: if we can't check the budget, assume OK to avoid blocking research."""
|
||||||
|
with patch('check_budget.get_used_usd', side_effect=Exception("network error")):
|
||||||
|
assert is_budget_ok(limit_usd=5.0) is True
|
||||||
|
|
||||||
|
def test_budget_exactly_at_limit_is_not_ok():
|
||||||
|
with patch('check_budget.get_used_usd', return_value=5.0):
|
||||||
|
assert is_budget_ok(limit_usd=5.0) is False
|
||||||
Reference in New Issue
Block a user