61 lines
2.0 KiB
Python
61 lines
2.0 KiB
Python
"""Append the co-op credit-pack hint to LiteLLM's hardcoded budget-exceeded messages.
|
|
|
|
Run at BUILD time (the runtime rootfs is read-only under Cloudron). Non-fatal:
|
|
if the target file or message strings are missing (e.g. after an upstream
|
|
version bump moves or rewrites them), it prints a warning and exits 0 so the
|
|
build — and the gateway — still succeeds. The hint is a nice-to-have, never a
|
|
reason to take the gateway down.
|
|
"""
|
|
import sys
|
|
import glob
|
|
|
|
HINT = '. You can add credit packs at https://opencollective.com/inference-cooperative'
|
|
|
|
TEAM_MSG = (
|
|
'message=f"Budget has been exceeded! Team={team_object.team_id} '
|
|
'Current cost: {spend}, Max budget: {team_object.max_budget}"'
|
|
)
|
|
KEY_MSG = (
|
|
'message=f"Budget has been exceeded! Key={key_descriptor} '
|
|
'Current cost: {spend}, Max budget: {valid_token.max_budget}"'
|
|
)
|
|
|
|
# Find the proxy auth_checks.py wherever it lives in this image (venv path
|
|
# varies by Python minor version; the running CLI imports from the image's
|
|
# venv). Search the site-packages copies, not the type-stub one under
|
|
# litellm/types/.
|
|
candidates = glob.glob(
|
|
"/app/.venv/lib/python*/site-packages/litellm/proxy/auth/auth_checks.py"
|
|
)
|
|
|
|
patched_any = False
|
|
for p in candidates:
|
|
with open(p) as f:
|
|
src = f.read()
|
|
|
|
if "opencollective.com/inference-cooperative" in src:
|
|
print(f"INFO: credit-pack hint already present in {p}")
|
|
continue
|
|
|
|
new = src.replace(TEAM_MSG, TEAM_MSG + ' + "' + HINT + '"')
|
|
new = new.replace(KEY_MSG, KEY_MSG + ' + "' + HINT + '"')
|
|
|
|
if new == src:
|
|
print(f"WARN: budget-exceeded message format changed in {p}; "
|
|
"hint not applied (non-fatal)", file=sys.stderr)
|
|
continue
|
|
|
|
with open(p, "w") as f:
|
|
f.write(new)
|
|
print(f"INFO: patched budget-exceeded messages in {p}")
|
|
patched_any = True
|
|
|
|
if not candidates:
|
|
print("WARN: auth_checks.py not found; hint not applied (non-fatal)",
|
|
file=sys.stderr)
|
|
elif not patched_any:
|
|
print("WARN: no budget-exceeded messages patched (non-fatal)",
|
|
file=sys.stderr)
|
|
|
|
sys.exit(0)
|