diff --git a/proxy/auth.py b/proxy/auth.py index 3fd57c9c..ce06795b 100644 --- a/proxy/auth.py +++ b/proxy/auth.py @@ -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]