laantungir/routstr-core
1
0
Fork 0
mirror of https://github.com/Routstr/routstr-core.git synced 2026-10-05 20:28:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

update cost calculation and auth

Browse Source
This commit is contained in:
Shroominic
2025-04-25 13:27:20 +08:00
parent 63a4bd6fcb
commit 139f246b9e
1 changed files with 60 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files
proxy/auth.py
+60 -95
View File
@@ -4,18 +4,21 @@ import os
from typing import Literal
from fastapi import HTTPException
from pydantic import BaseModel
from pydantic import BaseModel, Field
from .redeem import redeem
from .db import (
ApiKey,
create_session,
RECIEIVE_LN_ADDRESS,
COST_PER_REQUEST,
COST_PER_1K_PROMPT_TOKENS,
COST_PER_1K_COMPLETION_TOKENS,
MODEL_BASED_PRICING,
)
from .db import ApiKey, create_session
from .price import btc_usd_ask_price
RECIEIVE_LN_ADDRESS = os.environ["RECIEIVE_LN_ADDRESS"]
COST_PER_REQUEST = int(os.environ["COST_PER_REQUEST"]) * 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"
def _hash_api_key(api_key: str) -> str:
@@ -23,7 +26,7 @@ def _hash_api_key(api_key: str) -> str:
return hashlib.sha256(api_key.encode()).hexdigest()
async def validate_api_key(api_key: str) -> None:
async def validate_api_key(api_key: str) -> ApiKey:
"""
Validates the provided API key using SQLModel.
If it's a cashu key, it redeems it and stores its hash and balance.
@@ -35,36 +38,29 @@ async def validate_api_key(api_key: str) -> None:
hashed_key = _hash_api_key(api_key)
async with create_session() as session:
# check if key exists
if await session.get(ApiKey, hashed_key):
return
if key := await session.get(ApiKey, hashed_key):
return key
# If hash not found, check if it's a potentially new cashu key
if api_key.startswith("cashu"):
try:
print(
f"Attempting to redeem cashu key: {api_key[:15]}...{api_key[-15:]}"
)
# Redeem the original cashu key
amount = await redeem(api_key, RECIEIVE_LN_ADDRESS)
amount_msats = amount * 1000 # Convert sats to msats
print(
f"Redeemed successfully. Amount: {amount} sats ({amount_msats} msats)"
)
# Store the hash and the redeemed amount using SQLModel
new_key = ApiKey(hashed_key=hashed_key, balance=amount_msats)
session.add(new_key)
await session.commit()
await session.refresh(new_key)
return
return new_key
except Exception as e:
print(f"Redemption failed: {e}")
# Include the redemption error message for better debugging
raise HTTPException(
status_code=401, detail=f"Invalid or expired cashu key: {e}"
)
if api_key.startswith("sk-"):
if exsisting_key := await session.get(ApiKey, api_key[3:]):
return exsisting_key
# If it's not a known hash and not a valid new cashu key
raise HTTPException(status_code=401, detail="Invalid API key")
@@ -72,13 +68,10 @@ async def pay_for_request(api_key: str) -> None:
"""Deducts the cost of a request from the balance associated with the API key hash using SQLModel."""
hashed_key = _hash_api_key(api_key)
# Get the key record using SQLModel
async with create_session() as session:
key_record = await session.get(ApiKey, hashed_key)
if key_record is None:
# This should theoretically not happen if validate_api_key was called first
# Consider adding a check or relying on validate_api_key structure
if not key_record: # This should not happen
raise HTTPException(status_code=401, detail="API key not validated")
if key_record.balance < COST_PER_REQUEST:
@@ -90,14 +83,10 @@ async def pay_for_request(api_key: str) -> None:
key_record.balance -= COST_PER_REQUEST
key_record.total_spent += COST_PER_REQUEST
key_record.total_requests += 1
session.add(key_record) # Mark the object as changed
session.add(key_record)
await session.commit()
await session.refresh(key_record)
print(
f"Charged {COST_PER_REQUEST} msats. New balance for key hash {hashed_key[:10]}...: {key_record.balance} msats"
)
async def adjust_payment_for_tokens(api_key: str, response_data: dict) -> dict:
"""
@@ -106,51 +95,34 @@ async def adjust_payment_for_tokens(api_key: str, response_data: dict) -> dict:
Returns cost data to be included in the response.
"""
cost_data = {
"base_cost_msats": COST_PER_REQUEST,
"prompt_cost_msats": 0,
"completion_cost_msats": 0,
"total_cost_msats": COST_PER_REQUEST,
"base_msats": COST_PER_REQUEST,
"input_msats": 0,
"output_msats": 0,
"total_msats": COST_PER_REQUEST,
}
if MODEL_BASED_PRICING and os.path.exists("models.json"):
offering = Offering.validate(json.load(open("models.json")))
models = offering.models
models = read_models()
response_model = response_data.get("model", "")
if response_model not in [model.name for model in models]:
raise HTTPException(status_code=400, detail="Invalid model")
model = next(model for model in models if model.name == response_model)
prompt_tokens = response_data.get("usage", {}).get("prompt_tokens", 0)
completion_tokens = response_data.get("usage", {}).get("completion_tokens", 0)
cost_data["base_cost_msats"] = 0
cost_data["prompt_cost_msats"] = int(
prompt_tokens / 1000 * model.msats_per_1k_prompt_tokens + 0.999
)
cost_data["completion_cost_msats"] = int(
completion_tokens / 1000 * model.msats_per_1k_completion_tokens + 0.999
)
cost_data["total_cost_msats"] = int(
cost_data["prompt_cost_msats"] + cost_data["completion_cost_msats"]
)
print(cost_data)
if not (COST_PER_1K_PROMPT_TOKENS or COST_PER_1K_COMPLETION_TOKENS):
return cost_data # Skip if token-based pricing is not enabled
MSATS_PER_1K_INPUT_TOKENS = await model.msats_per_1k_input_tokens()
MSATS_PER_1K_OUTPUT_TOKENS = await model.msats_per_1k_output_tokens()
# Extract token usage from response
usage = response_data.get("usage", {})
if not usage:
return cost_data # No token usage information available
if not (MSATS_PER_1K_OUTPUT_TOKENS and MSATS_PER_1K_INPUT_TOKENS):
raise HTTPException(status_code=400, detail="Model pricing not defined")
prompt_tokens = usage.get("prompt_tokens", 0)
completion_tokens = usage.get("completion_tokens", 0)
input_tokens = response_data.get("usage", {}).get("prompt_tokens", 0)
output_tokens = response_data.get("usage", {}).get("completion_tokens", 0)
# Calculate token-based cost
prompt_cost = (prompt_tokens / 1000) * COST_PER_1K_PROMPT_TOKENS
completion_cost = (completion_tokens / 1000) * COST_PER_1K_COMPLETION_TOKENS
token_based_cost = int(prompt_cost + completion_cost)
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_cost_msats"] = 0
cost_data["prompt_cost_msats"] = int(prompt_cost)
cost_data["completion_cost_msats"] = int(completion_cost)
cost_data["total_cost_msats"] = token_based_cost
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
@@ -180,51 +152,44 @@ async def adjust_payment_for_tokens(api_key: str, response_data: dict) -> dict:
else:
key_record.balance -= cost_difference
key_record.total_spent += cost_difference
print(
f"Additional charge for tokens: {cost_difference} msats. New balance: {key_record.balance} msats"
)
cost_data["total_cost_msats"] = COST_PER_REQUEST + cost_difference
cost_data["total_msats"] = COST_PER_REQUEST + cost_difference
else:
# Refund some of the base cost
refund = abs(cost_difference)
key_record.balance += refund
key_record.total_spent -= refund
print(
f"Refund for tokens: {refund} msats. New balance: {key_record.balance} msats"
)
cost_data["total_cost_msats"] = COST_PER_REQUEST - refund
cost_data["total_msats"] = COST_PER_REQUEST - refund
session.add(key_record)
await session.commit()
print("cost_data:", cost_data)
return cost_data
def convert_usd_to_btc(usd: float) -> float:
EXCHANGE_FEE = 0.005
BTC_PRICE = 93000
return usd / BTC_PRICE * (1 - EXCHANGE_FEE)
class LLModel(BaseModel):
name: str
cost_per_1m_prompt_tokens: float
cost_per_1m_completion_tokens: float
cost_per_1m_input_tokens: float = Field(alias="cost_per_1m_prompt_tokens")
cost_per_1m_output_tokens: float = Field(alias="cost_per_1m_completion_tokens")
currency: Literal["btc", "usd"]
@property
def msats_per_1k_prompt_tokens(self) -> float:
async def msats_per_1k_input_tokens(self) -> float:
if self.currency == "btc":
return self.cost_per_1m_prompt_tokens * 100_000_000
return convert_usd_to_btc(self.cost_per_1m_prompt_tokens) * 100_000_000
return self.cost_per_1m_input_tokens * 100_000
btc_price = await btc_usd_ask_price()
return (self.cost_per_1m_input_tokens / 1000) * (100_000_000_000 / btc_price)
@property
def msats_per_1k_completion_tokens(self) -> float:
async def msats_per_1k_output_tokens(self) -> float:
if self.currency == "btc":
return self.cost_per_1m_completion_tokens * 100_000_000
return convert_usd_to_btc(self.cost_per_1m_completion_tokens) * 100_000_000
return self.cost_per_1m_output_tokens * 100_000
btc_price = await btc_usd_ask_price()
return (self.cost_per_1m_output_tokens / 1000) * (100_000_000_000 / btc_price)
class Offering(BaseModel):
# npub: str
models: list[LLModel]
def read_models() -> list[LLModel]:
if not os.path.exists("models.json"):
raise HTTPException(status_code=400, detail="Models not defined")
with open("models.json", "r") as f:
models = json.load(f)["models"]
return [LLModel(**model) for model in models]