mirror of
https://github.com/Routstr/routstr-core.git
synced 2026-08-10 11:13:08 +00:00
137 lines
4.3 KiB
Python
137 lines
4.3 KiB
Python
import base64
|
||
import json
|
||
import os
|
||
|
||
import cbor2
|
||
from fastapi import HTTPException, Response
|
||
from sixty_nuts.types import CurrencyUnit
|
||
|
||
from router.models import MODELS
|
||
from router.payment.cost_caculation import COST_PER_REQUEST, MODEL_BASED_PRICING
|
||
|
||
UPSTREAM_BASE_URL = os.environ["UPSTREAM_BASE_URL"]
|
||
UPSTREAM_API_KEY = os.environ.get("UPSTREAM_API_KEY", "")
|
||
|
||
|
||
def get_cost_per_request(model: str | None = None) -> int:
|
||
if MODEL_BASED_PRICING and MODELS and model:
|
||
return get_max_cost_for_model(model=model)
|
||
return COST_PER_REQUEST
|
||
|
||
|
||
def check_token_balance(headers: dict, body: dict) -> CurrencyUnit:
|
||
if x_cashu := headers.get("x-cashu", None):
|
||
cashu_token = x_cashu
|
||
elif auth := headers.get("authorization", None):
|
||
cashu_token = auth.split(" ")[1] if len(auth.split(" ")) > 1 else ""
|
||
else:
|
||
raise HTTPException(status_code=401, detail="Unauthorized")
|
||
|
||
# Handle empty token
|
||
if not cashu_token:
|
||
raise HTTPException(
|
||
status_code=401,
|
||
detail={
|
||
"error": {
|
||
"message": "API key or Cashu token required",
|
||
"type": "invalid_request_error",
|
||
"code": "missing_api_key",
|
||
}
|
||
},
|
||
)
|
||
|
||
# Handle regular API keys (sk-*)
|
||
if cashu_token.startswith("sk-"):
|
||
# For regular API keys, return default unit
|
||
return "sat"
|
||
|
||
cost = get_cost_per_request(model=body.get("model", None))
|
||
if cashu_token.startswith("cashuA"):
|
||
_token = base64_token_json(cashu_token)
|
||
amount = sum(p["amount"] for t in _token["token"] for p in t["proofs"])
|
||
unit: CurrencyUnit = _token["unit"]
|
||
if unit == "sat":
|
||
amount *= 1000
|
||
if amount < cost:
|
||
raise HTTPException(status_code=413, detail="Insufficient balance")
|
||
elif cashu_token.startswith("cashuB"):
|
||
_token = base64_token_cbor(cashu_token)
|
||
amount = sum(p["a"] for t in _token["t"] for p in t["p"])
|
||
unit = _token["u"]
|
||
if unit == "sat":
|
||
amount *= 1000
|
||
if amount < cost:
|
||
raise HTTPException(status_code=413, detail="Insufficient balance")
|
||
else:
|
||
raise HTTPException(status_code=401, detail="Unauthorized")
|
||
return unit
|
||
|
||
|
||
def base64_token_json(cashu_token: str) -> dict:
|
||
# Version 3 - JSON format
|
||
encoded = cashu_token[6:] # Remove "cashuA"
|
||
# Add correct padding – (-len) % 4 equals 0,1,2,3
|
||
encoded += "=" * ((-len(encoded)) % 4)
|
||
|
||
decoded = base64.urlsafe_b64decode(encoded).decode()
|
||
token_data = json.loads(decoded)
|
||
|
||
return token_data
|
||
|
||
|
||
def base64_token_cbor(cashu_token: str) -> dict:
|
||
encoded = cashu_token[6:] # Remove "cashuB"
|
||
encoded += "=" * ((-len(encoded)) % 4)
|
||
decoded_bytes = base64.urlsafe_b64decode(encoded)
|
||
token_data = cbor2.loads(decoded_bytes)
|
||
return token_data
|
||
|
||
|
||
def get_max_cost_for_model(model: str) -> int:
|
||
if not MODEL_BASED_PRICING or not MODELS:
|
||
return COST_PER_REQUEST
|
||
if model not in [model.id for model in MODELS]:
|
||
return COST_PER_REQUEST
|
||
for m in MODELS:
|
||
if m.id == model:
|
||
return m.sats_pricing.max_cost * 1000 # type: ignore
|
||
return COST_PER_REQUEST
|
||
|
||
|
||
def create_error_response(error_type: str, message: str, status_code: int) -> Response:
|
||
"""Create a standardized error response."""
|
||
return Response(
|
||
content=json.dumps(
|
||
{
|
||
"error": {
|
||
"message": message,
|
||
"type": error_type,
|
||
"code": status_code,
|
||
}
|
||
}
|
||
),
|
||
status_code=status_code,
|
||
media_type="application/json",
|
||
)
|
||
|
||
|
||
def prepare_upstream_headers(request_headers: dict) -> dict:
|
||
"""Prepare headers for upstream request, removing sensitive/problematic ones."""
|
||
headers = dict(request_headers)
|
||
# Remove headers that shouldn't be forwarded
|
||
headers.pop("host", None)
|
||
headers.pop("content-length", None)
|
||
headers.pop("refund-lnurl", None)
|
||
headers.pop("key-expiry-time", None)
|
||
headers.pop("x-cashu", None)
|
||
|
||
# Handle authorization
|
||
if UPSTREAM_API_KEY:
|
||
headers["Authorization"] = f"Bearer {UPSTREAM_API_KEY}"
|
||
headers.pop("authorization", None)
|
||
else:
|
||
headers.pop("Authorization", None)
|
||
headers.pop("authorization", None)
|
||
|
||
return headers
|