mirror of
https://github.com/Routstr/routstr-core.git
synced 2026-08-11 19:57:32 +00:00
refactor for x-cashu
This commit is contained in:
9qeklajc
@@ -11,6 +11,7 @@ dependencies = [
|
||||
"sqlmodel>=0.0.24",
|
||||
"httpx[socks]>=0.25.2",
|
||||
"greenlet>=3.2.1",
|
||||
"pydantic>=1.10.22",
|
||||
]
|
||||
|
||||
[dependency-groups]
|
||||
|
||||
@@ -4,5 +4,4 @@ dotenv.load_dotenv()
|
||||
|
||||
from .main import app as fastapi_app # noqa
|
||||
|
||||
|
||||
__all__ = ["fastapi_app"]
|
||||
|
||||
+65
-180
@@ -1,26 +1,19 @@
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import HTTPException
|
||||
from sqlmodel import col, update
|
||||
|
||||
from router.payment.cost_caculation import (COST_PER_REQUEST,
|
||||
MODEL_BASED_PRICING, CostData,
|
||||
CostDataError, MaxCostData,
|
||||
calculate_cost)
|
||||
from router.payment.helpers import get_max_cost_for_model
|
||||
|
||||
from .cashu import credit_balance
|
||||
from .db import ApiKey, AsyncSession
|
||||
from .models import MODELS
|
||||
|
||||
COST_PER_REQUEST = (
|
||||
int(os.environ.get("COST_PER_REQUEST", "1")) * 1000
|
||||
) # Convert to msats
|
||||
COST_PER_1K_INPUT_TOKENS = (
|
||||
int(os.environ.get("COST_PER_1K_INPUT_TOKENS", "0")) * 1000
|
||||
) # Convert to msats
|
||||
COST_PER_1K_OUTPUT_TOKENS = (
|
||||
int(os.environ.get("COST_PER_1K_OUTPUT_TOKENS", "0")) * 1000
|
||||
) # Convert to msats
|
||||
MODEL_BASED_PRICING = os.environ.get("MODEL_BASED_PRICING", "false").lower() == "true"
|
||||
|
||||
# TODO: implement prepaid api key (not like it was before)
|
||||
# PREPAID_API_KEY = os.environ.get("PREPAID_API_KEY", None)
|
||||
# PREPAID_BALANCE = int(os.environ.get("PREPAID_BALANCE", "0")) * 1000 # Convert to msats
|
||||
@@ -105,68 +98,6 @@ async def validate_bearer_key(
|
||||
)
|
||||
|
||||
|
||||
def base64_token_json(cashu_token: str) -> dict:
|
||||
import base64
|
||||
|
||||
# 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:
|
||||
import base64
|
||||
|
||||
import cbor2
|
||||
|
||||
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 check_token_balance(headers: dict, body: dict) -> None:
|
||||
if x_cashu := headers.get("x-cashu", None):
|
||||
cashu_token = x_cashu
|
||||
elif auth := headers.get("authorization", None):
|
||||
cashu_token = auth.split(" ")[1]
|
||||
else:
|
||||
raise HTTPException(status_code=401, detail="Unauthorized")
|
||||
COST_PER_REQUEST = get_max_cost_for_model(model=body["model"])
|
||||
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 = _token["unit"]
|
||||
if unit == "sat":
|
||||
amount *= 1000
|
||||
if amount < COST_PER_REQUEST:
|
||||
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_PER_REQUEST:
|
||||
raise HTTPException(status_code=413, detail="Insufficient balance")
|
||||
else:
|
||||
raise HTTPException(status_code=401, detail="Unauthorized")
|
||||
|
||||
|
||||
def get_max_cost_for_model(model: str) -> int:
|
||||
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
|
||||
|
||||
|
||||
async def pay_for_request(
|
||||
key: ApiKey,
|
||||
@@ -220,120 +151,74 @@ async def pay_for_request(
|
||||
|
||||
async def adjust_payment_for_tokens(
|
||||
key: ApiKey, response_data: dict, session: AsyncSession
|
||||
) -> dict:
|
||||
) -> dict | None:
|
||||
"""
|
||||
Adjusts the payment based on token usage in the response.
|
||||
This is called after the initial payment and the upstream request is complete.
|
||||
Returns cost data to be included in the response.
|
||||
"""
|
||||
max_cost = get_max_cost_for_model(model=response_data["model"])
|
||||
cost_data: dict = {
|
||||
"base_msats": max_cost,
|
||||
"input_msats": 0,
|
||||
"output_msats": 0,
|
||||
"total_msats": max_cost,
|
||||
}
|
||||
cost_data = calculate_cost(response_data, max_cost)
|
||||
|
||||
# Check if we have usage data
|
||||
if "usage" not in response_data or response_data["usage"] is None:
|
||||
print("No usage data in response, using base cost only")
|
||||
return cost_data
|
||||
match calculate_cost(response_data, max_cost):
|
||||
case MaxCostData() as cost:
|
||||
return cost.dict()
|
||||
case CostData() as cost:
|
||||
# If token-based pricing is enabled and base cost is 0, use token-based cost
|
||||
# Otherwise, token cost is additional to the base cost
|
||||
cost_difference = cost.total_msats - max_cost
|
||||
|
||||
# Default to configured pricing
|
||||
MSATS_PER_1K_INPUT_TOKENS = COST_PER_1K_INPUT_TOKENS
|
||||
MSATS_PER_1K_OUTPUT_TOKENS = COST_PER_1K_OUTPUT_TOKENS
|
||||
if cost_difference == 0:
|
||||
await session.commit()
|
||||
return cost.dict()
|
||||
|
||||
if MODEL_BASED_PRICING and MODELS:
|
||||
response_model = response_data.get("model", "")
|
||||
if response_model not in [model.id for model in MODELS]:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": {
|
||||
"message": f"Invalid model in response: {response_model}",
|
||||
"type": "invalid_request_error",
|
||||
"code": "model_not_found",
|
||||
}
|
||||
},
|
||||
)
|
||||
model = next(model for model in MODELS if model.id == response_model)
|
||||
if model.sats_pricing is None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": {
|
||||
"message": "Model pricing not defined",
|
||||
"type": "invalid_request_error",
|
||||
"code": "pricing_not_found",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
MSATS_PER_1K_INPUT_TOKENS = model.sats_pricing.prompt * 1_000_000 # type: ignore
|
||||
MSATS_PER_1K_OUTPUT_TOKENS = model.sats_pricing.completion * 1_000_000 # type: ignore
|
||||
|
||||
if not (MSATS_PER_1K_OUTPUT_TOKENS and MSATS_PER_1K_INPUT_TOKENS):
|
||||
# If no token pricing is configured, just return base cost
|
||||
return cost_data
|
||||
|
||||
input_tokens = response_data.get("usage", {}).get("prompt_tokens", 0)
|
||||
output_tokens = response_data.get("usage", {}).get("completion_tokens", 0)
|
||||
|
||||
input_msats = int(round(input_tokens / 1000 * MSATS_PER_1K_INPUT_TOKENS, 0))
|
||||
output_msats = int(round(output_tokens / 1000 * MSATS_PER_1K_OUTPUT_TOKENS, 0))
|
||||
token_based_cost = int(round(input_msats + output_msats, 0))
|
||||
|
||||
cost_data["base_msats"] = 0
|
||||
cost_data["input_msats"] = input_msats
|
||||
cost_data["output_msats"] = output_msats
|
||||
cost_data["total_msats"] = token_based_cost
|
||||
|
||||
# If token-based pricing is enabled and base cost is 0, use token-based cost
|
||||
# Otherwise, token cost is additional to the base cost
|
||||
cost_difference = token_based_cost - max_cost
|
||||
|
||||
if cost_difference == 0:
|
||||
await session.commit()
|
||||
return cost_data # No adjustment needed
|
||||
|
||||
if cost_difference > 0:
|
||||
# Need to charge more
|
||||
if key.balance < cost_difference:
|
||||
print(
|
||||
f"Warning: Insufficient balance for token-based pricing adjustment: {key.hashed_key[:10]}..."
|
||||
)
|
||||
cost_data["warning"] = "Insufficient balance for full token-based pricing"
|
||||
cost_data["balance_shortage_msats"] = cost_difference - key.balance
|
||||
await session.commit()
|
||||
else:
|
||||
charge_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.where(col(ApiKey.balance) >= cost_difference)
|
||||
.values(
|
||||
balance=col(ApiKey.balance) - cost_difference,
|
||||
total_spent=col(ApiKey.total_spent) + cost_difference,
|
||||
if cost_difference > 0:
|
||||
# Need to charge more
|
||||
if key.balance < cost_difference:
|
||||
print(
|
||||
f"Warning: Insufficient balance for token-based pricing adjustment: {key.hashed_key[:10]}..."
|
||||
)
|
||||
await session.commit()
|
||||
else:
|
||||
charge_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.where(col(ApiKey.balance) >= cost_difference)
|
||||
.values(
|
||||
balance=col(ApiKey.balance) - cost_difference,
|
||||
total_spent=col(ApiKey.total_spent) + cost_difference,
|
||||
)
|
||||
)
|
||||
result = await session.exec(charge_stmt) # type: ignore[call-overload]
|
||||
await session.commit()
|
||||
if result.rowcount:
|
||||
cost.total_msats = max_cost + cost_difference
|
||||
await session.refresh(key)
|
||||
else:
|
||||
# Refund some of the base cost
|
||||
refund = abs(cost_difference)
|
||||
refund_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(
|
||||
balance=col(ApiKey.balance) + refund,
|
||||
total_spent=col(ApiKey.total_spent) - refund,
|
||||
)
|
||||
)
|
||||
)
|
||||
result = await session.exec(charge_stmt) # type: ignore[call-overload]
|
||||
await session.commit()
|
||||
if result.rowcount:
|
||||
cost_data["total_msats"] = max_cost + cost_difference
|
||||
await session.exec(refund_stmt) # type: ignore[call-overload]
|
||||
await session.commit()
|
||||
cost_data.total_msats = max_cost - refund
|
||||
await session.refresh(key)
|
||||
else:
|
||||
# Refund some of the base cost
|
||||
refund = abs(cost_difference)
|
||||
refund_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(
|
||||
balance=col(ApiKey.balance) + refund,
|
||||
total_spent=col(ApiKey.total_spent) - refund,
|
||||
)
|
||||
)
|
||||
await session.exec(refund_stmt) # type: ignore[call-overload]
|
||||
await session.commit()
|
||||
cost_data["total_msats"] = max_cost - refund
|
||||
await session.refresh(key)
|
||||
|
||||
return cost_data
|
||||
return cost.dict()
|
||||
case CostDataError() as error:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": {
|
||||
"message": error.message,
|
||||
"type": "invalid_request_error",
|
||||
"code": error.code,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
+1
-1
@@ -21,7 +21,7 @@ WALLET: Wallet | None = None
|
||||
|
||||
async def init_wallet() -> None:
|
||||
global WALLET
|
||||
WALLET = await Wallet.create(nsec=NSEC)
|
||||
WALLET = await Wallet.create(nsec=NSEC, mint_urls=[MINT], currency="msat")
|
||||
|
||||
|
||||
def wallet() -> Wallet:
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
from . import cost_caculation, x_cashu
|
||||
@@ -0,0 +1,85 @@
|
||||
import os
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from router.models import MODELS
|
||||
|
||||
COST_PER_REQUEST = (
|
||||
int(os.environ.get("COST_PER_REQUEST", "1")) * 500
|
||||
) # Convert to msats
|
||||
COST_PER_1K_INPUT_TOKENS = (
|
||||
int(os.environ.get("COST_PER_1K_INPUT_TOKENS", "0")) * 1000
|
||||
) # Convert to msats
|
||||
COST_PER_1K_OUTPUT_TOKENS = (
|
||||
int(os.environ.get("COST_PER_1K_OUTPUT_TOKENS", "0")) * 1000
|
||||
) # Convert to msats
|
||||
MODEL_BASED_PRICING = os.environ.get("MODEL_BASED_PRICING", "false").lower() == "true"
|
||||
|
||||
|
||||
class CostData(BaseModel):
|
||||
base_msats: int
|
||||
input_msats: int
|
||||
output_msats: int
|
||||
total_msats: int
|
||||
|
||||
|
||||
class MaxCostData(CostData):
|
||||
pass
|
||||
|
||||
|
||||
class CostDataError(BaseModel):
|
||||
message: str
|
||||
code: str
|
||||
|
||||
|
||||
def calculate_cost(
|
||||
response_data: dict, max_cost: int
|
||||
) -> CostData | MaxCostData | CostDataError:
|
||||
cost_data = MaxCostData(
|
||||
base_msats=max_cost,
|
||||
input_msats=0,
|
||||
output_msats=0,
|
||||
total_msats=max_cost,
|
||||
)
|
||||
|
||||
if "usage" not in response_data or response_data["usage"] is None:
|
||||
print("No usage data in response, using base cost only")
|
||||
return cost_data
|
||||
|
||||
MSATS_PER_1K_INPUT_TOKENS = COST_PER_1K_INPUT_TOKENS
|
||||
MSATS_PER_1K_OUTPUT_TOKENS = COST_PER_1K_OUTPUT_TOKENS
|
||||
|
||||
if MODEL_BASED_PRICING and MODELS:
|
||||
response_model = response_data.get("model", "")
|
||||
if response_model not in [model.id for model in MODELS]:
|
||||
return CostDataError(
|
||||
message=f"Invalid model in response: {response_model}",
|
||||
code="model_not_found",
|
||||
)
|
||||
|
||||
model = next(model for model in MODELS if model.id == response_model)
|
||||
if model.sats_pricing is None:
|
||||
return CostDataError(
|
||||
message="Model pricing not defined", code="pricing_not_found"
|
||||
)
|
||||
|
||||
MSATS_PER_1K_INPUT_TOKENS = model.sats_pricing.prompt * 1_000_000 # type: ignore
|
||||
MSATS_PER_1K_OUTPUT_TOKENS = model.sats_pricing.completion * 1_000_000 # type: ignore
|
||||
|
||||
if not (MSATS_PER_1K_OUTPUT_TOKENS and MSATS_PER_1K_INPUT_TOKENS):
|
||||
# If no token pricing is configured, just return base cost
|
||||
return cost_data
|
||||
|
||||
input_tokens = response_data.get("usage", {}).get("prompt_tokens", 0)
|
||||
output_tokens = response_data.get("usage", {}).get("completion_tokens", 0)
|
||||
|
||||
input_msats = int(round(input_tokens / 1000 * MSATS_PER_1K_INPUT_TOKENS, 0))
|
||||
output_msats = int(round(output_tokens / 1000 * MSATS_PER_1K_OUTPUT_TOKENS, 0))
|
||||
token_based_cost = int(round(input_msats + output_msats, 0))
|
||||
|
||||
return CostData(
|
||||
base_msats=0,
|
||||
input_msats=input_msats,
|
||||
output_msats=output_msats,
|
||||
total_msats=token_based_cost,
|
||||
)
|
||||
@@ -0,0 +1,112 @@
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
from typing import Literal
|
||||
|
||||
import cbor2
|
||||
from fastapi import HTTPException, Response
|
||||
|
||||
from router.models import MODELS
|
||||
from router.payment.cost_caculation import COST_PER_REQUEST
|
||||
|
||||
UPSTREAM_BASE_URL = os.environ["UPSTREAM_BASE_URL"]
|
||||
UPSTREAM_API_KEY = os.environ.get("UPSTREAM_API_KEY", "")
|
||||
|
||||
|
||||
def check_token_balance(headers: dict, body: dict, unit: Literal['sat', 'msat']) -> None:
|
||||
if x_cashu := headers.get("x-cashu", None):
|
||||
cashu_token = x_cashu
|
||||
elif auth := headers.get("authorization", None):
|
||||
cashu_token = auth.split(" ")[1]
|
||||
else:
|
||||
raise HTTPException(status_code=401, detail="Unauthorized")
|
||||
COST_PER_REQUEST = get_max_cost_for_model(model=body["model"])
|
||||
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 = _token["unit"]
|
||||
if unit == "msat":
|
||||
pass
|
||||
elif unit == "sat":
|
||||
amount *= 1000
|
||||
if amount < COST_PER_REQUEST:
|
||||
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_PER_REQUEST:
|
||||
raise HTTPException(status_code=413, detail="Insufficient balance")
|
||||
else:
|
||||
raise HTTPException(status_code=401, detail="Unauthorized")
|
||||
|
||||
|
||||
|
||||
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 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
|
||||
|
||||
|
||||
@@ -0,0 +1,260 @@
|
||||
import json
|
||||
import re
|
||||
import traceback
|
||||
from typing import AsyncGenerator, Literal, cast
|
||||
|
||||
import httpx
|
||||
from fastapi import BackgroundTasks, HTTPException, Request
|
||||
from fastapi.responses import Response, StreamingResponse
|
||||
|
||||
from router.cashu import wallet
|
||||
from router.payment.cost_caculation import (CostData, CostDataError,
|
||||
MaxCostData, calculate_cost)
|
||||
from router.payment.helpers import (UPSTREAM_BASE_URL, check_token_balance,
|
||||
create_error_response,
|
||||
get_max_cost_for_model,
|
||||
prepare_upstream_headers)
|
||||
|
||||
type Currency = Literal["sat", "msat"]
|
||||
|
||||
async def x_cashu_handler(
|
||||
request: Request, x_cashu_token: str, path: str
|
||||
) -> Response | StreamingResponse:
|
||||
print(x_cashu_token)
|
||||
headers = dict(request.headers)
|
||||
|
||||
# amount, _ = await redeem_token(x_cashu_token)
|
||||
# print(amount)
|
||||
|
||||
headers = prepare_upstream_headers(dict(request.headers))
|
||||
return await forward_to_upstream(request, path, headers, 1000)
|
||||
|
||||
|
||||
async def forward_to_upstream(
|
||||
request: Request, path: str, headers: dict, amount: int
|
||||
) -> Response | StreamingResponse:
|
||||
print(path, amount)
|
||||
"""Forward request to upstream and handle the response."""
|
||||
if path.startswith("v1/"):
|
||||
path = path.replace("v1/", "")
|
||||
|
||||
url = f"{UPSTREAM_BASE_URL}/{path}"
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.AsyncHTTPTransport(retries=1),
|
||||
timeout=None,
|
||||
) as client:
|
||||
print(url)
|
||||
try:
|
||||
response = await client.send(
|
||||
client.build_request(
|
||||
request.method,
|
||||
url,
|
||||
headers=headers,
|
||||
content=request.stream(),
|
||||
params=request.query_params,
|
||||
),
|
||||
stream=True,
|
||||
)
|
||||
|
||||
if path.endswith("chat/completions"):
|
||||
result = await handle_streaming_chat_completion(response, amount)
|
||||
background_tasks = BackgroundTasks()
|
||||
background_tasks.add_task(response.aclose)
|
||||
result.background = background_tasks
|
||||
return result
|
||||
|
||||
background_tasks = BackgroundTasks()
|
||||
background_tasks.add_task(response.aclose)
|
||||
background_tasks.add_task(client.aclose)
|
||||
|
||||
return StreamingResponse(
|
||||
response.aiter_bytes(),
|
||||
status_code=response.status_code,
|
||||
headers=dict(response.headers),
|
||||
background=background_tasks,
|
||||
)
|
||||
except Exception as exc:
|
||||
tb = traceback.format_exc()
|
||||
print(
|
||||
f"Unexpected error: {exc}\n"
|
||||
f"Request details: method={request.method}, url={url}, headers={headers}, "
|
||||
f"path={path}, query_params={dict(request.query_params)}\n"
|
||||
f"Traceback:\n{tb}"
|
||||
)
|
||||
return create_error_response(
|
||||
"internal_error", "An unexpected server error occurred", 500
|
||||
)
|
||||
|
||||
|
||||
async def handle_x_cashu_chat_completion(
|
||||
response: httpx.Response, amount: int
|
||||
) -> StreamingResponse | Response:
|
||||
"""Handle non-streaming chat completion responses with token-based pricing."""
|
||||
try:
|
||||
content = await response.aread()
|
||||
print(content)
|
||||
response_json = json.loads(content)
|
||||
print(response_json, amount)
|
||||
cost_data = await get_cost(response_json)
|
||||
|
||||
if not cost_data:
|
||||
# response.headers["X-Cashu"] = await send_refund(amount)
|
||||
return Response(
|
||||
content=json.dumps(
|
||||
{
|
||||
"error": {
|
||||
"message": "Error forwarding request to upstream",
|
||||
"type": "upstream_error",
|
||||
"code": response.status_code,
|
||||
}
|
||||
}
|
||||
),
|
||||
status_code=response.status_code,
|
||||
media_type="application/json",
|
||||
)
|
||||
|
||||
response_headers = dict(response.headers)
|
||||
|
||||
if "transfer-encoding" in response_headers:
|
||||
del response_headers["transfer-encoding"]
|
||||
|
||||
if "content-encoding" in response_headers:
|
||||
del response_headers["content-encoding"]
|
||||
|
||||
refund_token = await send_refund(amount - cost_data.total_msats)
|
||||
# response.headers["X-Cashu"] = refund_token
|
||||
|
||||
return StreamingResponse(
|
||||
content=response.aiter_bytes(),
|
||||
status_code=response.status_code,
|
||||
headers=response_headers,
|
||||
media_type="application/json",
|
||||
)
|
||||
except json.JSONDecodeError as e:
|
||||
print(f"Failed to parse JSON from upstream response: {e}")
|
||||
raise
|
||||
except Exception as e:
|
||||
print(f"Error adjusting payment for tokens: {e}")
|
||||
raise
|
||||
|
||||
|
||||
async def get_cost(response_data: dict) -> MaxCostData | CostData | None:
|
||||
"""
|
||||
Adjusts the payment based on token usage in the response.
|
||||
This is called after the initial payment and the upstream request is complete.
|
||||
Returns cost data to be included in the response.
|
||||
"""
|
||||
max_cost = get_max_cost_for_model(model=response_data["model"])
|
||||
|
||||
match calculate_cost(response_data, max_cost):
|
||||
case MaxCostData() as cost:
|
||||
return cost
|
||||
case CostData() as cost:
|
||||
return cost
|
||||
case CostDataError() as error:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": {
|
||||
"message": error.message,
|
||||
"type": "invalid_request_error",
|
||||
"code": error.code,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def redeem_token(x_cashu_token) -> tuple[int, Currency]:
|
||||
try:
|
||||
result = await wallet().redeem(x_cashu_token)
|
||||
return cast(tuple[int, Currency], result)
|
||||
except Exception as e:
|
||||
print(f"Redemption failed: {e}")
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail={
|
||||
"error": {
|
||||
"message": f"Invalid or expired Cashu key: {str(e)}",
|
||||
"type": "invalid_request_error",
|
||||
"code": "invalid_api_key",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def send_refund(amount) -> str:
|
||||
try:
|
||||
return await wallet().send(amount)
|
||||
except Exception as e:
|
||||
print(f"send failed: {e}")
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail={
|
||||
"error": {
|
||||
"message": f"failed to create refund: {str(e)}",
|
||||
"type": "invalid_request_error",
|
||||
"code": "send_token_failed",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def handle_streaming_chat_completion(
|
||||
response: httpx.Response,
|
||||
amount: int
|
||||
) -> StreamingResponse:
|
||||
"""Handle streaming chat completion responses with token-based pricing."""
|
||||
|
||||
async def stream_with_cost() -> AsyncGenerator[bytes, None]:
|
||||
# Store all chunks to analyze
|
||||
stored_chunks = []
|
||||
|
||||
async for chunk in response.aiter_bytes():
|
||||
# Store chunk for later analysis
|
||||
stored_chunks.append(chunk)
|
||||
|
||||
# Pass through each chunk to client
|
||||
yield chunk
|
||||
|
||||
# Process stored chunks to find usage data
|
||||
# Start from the end and work backwards
|
||||
for i in range(len(stored_chunks) - 1, -1, -1):
|
||||
chunk = stored_chunks[i]
|
||||
if not chunk or chunk == b"":
|
||||
continue
|
||||
|
||||
try:
|
||||
# Split by "data: " to get individual SSE events
|
||||
events = re.split(b"data: ", chunk)
|
||||
for event_data in events:
|
||||
if (
|
||||
not event_data
|
||||
or event_data.strip() == b"[DONE]"
|
||||
or event_data.strip() == b""
|
||||
):
|
||||
continue
|
||||
|
||||
try:
|
||||
data = json.loads(event_data)
|
||||
if (
|
||||
"usage" in data
|
||||
and data["usage"] is not None
|
||||
and isinstance(data["usage"], dict)
|
||||
):
|
||||
|
||||
cost_data = await get_cost(data)
|
||||
cost_json = json.dumps({"cost": cost_data})
|
||||
yield f"data: {cost_json}\n\n".encode()
|
||||
break
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing streaming response for cost: {e}")
|
||||
|
||||
return StreamingResponse(
|
||||
stream_with_cost(),
|
||||
status_code=response.status_code,
|
||||
headers=dict(response.headers),
|
||||
)
|
||||
|
||||
+9
-49
@@ -8,59 +8,19 @@ import httpx
|
||||
from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Request
|
||||
from fastapi.responses import Response, StreamingResponse
|
||||
|
||||
from .auth import (
|
||||
adjust_payment_for_tokens,
|
||||
check_token_balance,
|
||||
pay_for_request,
|
||||
validate_bearer_key,
|
||||
)
|
||||
from router.payment.helpers import (UPSTREAM_API_KEY, UPSTREAM_BASE_URL,
|
||||
check_token_balance, create_error_response,
|
||||
prepare_upstream_headers)
|
||||
from router.payment.x_cashu import x_cashu_handler
|
||||
|
||||
from .auth import (adjust_payment_for_tokens, pay_for_request,
|
||||
validate_bearer_key)
|
||||
from .cashu import x_cashu_refund
|
||||
from .db import ApiKey, AsyncSession, create_session, get_session
|
||||
|
||||
UPSTREAM_BASE_URL = os.environ["UPSTREAM_BASE_URL"]
|
||||
UPSTREAM_API_KEY = os.environ.get("UPSTREAM_API_KEY", "")
|
||||
|
||||
proxy_router = APIRouter()
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
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",
|
||||
)
|
||||
|
||||
|
||||
async def handle_streaming_chat_completion(
|
||||
response: httpx.Response, key: ApiKey, session: AsyncSession
|
||||
) -> StreamingResponse:
|
||||
@@ -307,8 +267,8 @@ async def proxy(
|
||||
if x_cashu := headers.get("x-cashu", None):
|
||||
# Check token balance before authentication for cashu tokens
|
||||
if request_body_dict:
|
||||
check_token_balance(headers, request_body_dict)
|
||||
key = await validate_bearer_key(x_cashu, session, "X-CASHU")
|
||||
check_token_balance(headers, request_body_dict, "msat")
|
||||
return await x_cashu_handler(request, x_cashu, path)
|
||||
|
||||
elif auth := headers.get("authorization", None):
|
||||
key = await get_bearer_token_key(headers, path, session, auth)
|
||||
|
||||
Reference in New Issue
Block a user