mirror of
https://github.com/Routstr/routstr-core.git
synced 2026-07-31 07:46:15 +00:00
Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Shroominic
|
35c9ce5699 | ||
|
Shroominic
|
24015ebec1 |
@@ -1,45 +0,0 @@
|
||||
import json
|
||||
import sys
|
||||
|
||||
import httpx
|
||||
|
||||
|
||||
def create_child_keys(base_url: str, api_key: str, count: int = 3) -> list[str]:
|
||||
headers = {"Authorization": f"Bearer {api_key}"}
|
||||
|
||||
print(f"Requesting {count} child keys from {base_url}...")
|
||||
|
||||
child_keys = []
|
||||
|
||||
for i in range(count):
|
||||
try:
|
||||
response = httpx.post(f"{base_url}/v1/balance/child-key", headers=headers)
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
child_keys.append(data["api_key"])
|
||||
print(
|
||||
f" [{i + 1}] Created: {data['api_key']} (Cost: {data['cost_msats']} msats)"
|
||||
)
|
||||
else:
|
||||
print(f" [{i + 1}] Failed: {response.status_code} - {response.text}")
|
||||
except Exception as e:
|
||||
print(f" [{i + 1}] Error: {str(e)}")
|
||||
|
||||
return child_keys
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if len(sys.argv) < 2:
|
||||
print("Usage: python create_child_keys.py <api_key_or_cashu_token> [base_url]")
|
||||
sys.exit(1)
|
||||
|
||||
auth_key = sys.argv[1]
|
||||
base_url = sys.argv[2] if len(sys.argv) > 2 else "http://localhost:8000"
|
||||
|
||||
keys = create_child_keys(base_url, auth_key)
|
||||
|
||||
if keys:
|
||||
print("\nSuccessfully created child keys:")
|
||||
print(json.dumps(keys, indent=2))
|
||||
else:
|
||||
print("\nNo child keys were created.")
|
||||
@@ -1,42 +0,0 @@
|
||||
"""
|
||||
|
||||
Revision ID: a86e5348850b
|
||||
Revises: b9667ffc5701
|
||||
Create Date: 2026-01-10 18:57:48.475781
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
import sqlmodel
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "a86e5348850b"
|
||||
down_revision = "b9667ffc5701"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# Use batch_alter_table for SQLite compatibility
|
||||
with op.batch_alter_table("api_keys", schema=None) as batch_op:
|
||||
batch_op.add_column(
|
||||
sa.Column(
|
||||
"parent_key_hash", sqlmodel.sql.sqltypes.AutoString(), nullable=True
|
||||
)
|
||||
)
|
||||
batch_op.create_index(
|
||||
batch_op.f("ix_api_keys_parent_key_hash"), ["parent_key_hash"], unique=False
|
||||
)
|
||||
batch_op.create_foreign_key(
|
||||
"fk_api_keys_parent_key_hash",
|
||||
"api_keys",
|
||||
["parent_key_hash"],
|
||||
["hashed_key"],
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
with op.batch_alter_table("api_keys", schema=None) as batch_op:
|
||||
batch_op.drop_constraint("fk_api_keys_parent_key_hash", type_="foreignkey")
|
||||
batch_op.drop_index(batch_op.f("ix_api_keys_parent_key_hash"))
|
||||
batch_op.drop_column("parent_key_hash")
|
||||
+40
-166
@@ -286,55 +286,30 @@ async def validate_bearer_key(
|
||||
)
|
||||
|
||||
|
||||
async def get_billing_key(key: ApiKey, session: AsyncSession) -> ApiKey:
|
||||
"""Returns the key that should be charged for the request."""
|
||||
if key.parent_key_hash:
|
||||
parent = await session.get(ApiKey, key.parent_key_hash)
|
||||
if parent:
|
||||
# We want to keep the total_requests and total_spent on the child key
|
||||
# but use the balance and reserved_balance of the parent.
|
||||
# However, pay_for_request updates reserved_balance and total_requests.
|
||||
# To stay simple, we charge the parent's balance and update parent's total_requests.
|
||||
return parent
|
||||
else:
|
||||
logger.error(
|
||||
"Parent key not found for child key",
|
||||
extra={
|
||||
"child_key_hash": key.hashed_key[:8] + "...",
|
||||
"parent_key_hash": key.parent_key_hash[:8] + "...",
|
||||
},
|
||||
)
|
||||
return key
|
||||
|
||||
|
||||
async def pay_for_request(
|
||||
key: ApiKey, cost_per_request: int, session: AsyncSession
|
||||
) -> int:
|
||||
"""Process payment for a request."""
|
||||
|
||||
billing_key = await get_billing_key(key, session)
|
||||
|
||||
logger.info(
|
||||
"Processing payment for request",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"current_balance": billing_key.balance,
|
||||
"current_balance": key.balance,
|
||||
"required_cost": cost_per_request,
|
||||
"sufficient_balance": billing_key.balance >= cost_per_request,
|
||||
"sufficient_balance": key.balance >= cost_per_request,
|
||||
},
|
||||
)
|
||||
|
||||
if billing_key.total_balance < cost_per_request:
|
||||
if key.total_balance < cost_per_request:
|
||||
logger.warning(
|
||||
"Insufficient balance for request",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"balance": billing_key.balance,
|
||||
"reserved_balance": billing_key.reserved_balance,
|
||||
"balance": key.balance,
|
||||
"reserved_balance": key.reserved_balance,
|
||||
"required": cost_per_request,
|
||||
"shortfall": cost_per_request - billing_key.total_balance,
|
||||
"shortfall": cost_per_request - key.total_balance,
|
||||
},
|
||||
)
|
||||
|
||||
@@ -342,7 +317,7 @@ async def pay_for_request(
|
||||
status_code=402,
|
||||
detail={
|
||||
"error": {
|
||||
"message": f"Insufficient balance: {cost_per_request} mSats required. {billing_key.total_balance} available. (reserved: {billing_key.reserved_balance})",
|
||||
"message": f"Insufficient balance: {cost_per_request} mSats required. {key.total_balance} available. (reserved: {key.reserved_balance})",
|
||||
"type": "insufficient_quota",
|
||||
"code": "insufficient_balance",
|
||||
}
|
||||
@@ -353,16 +328,15 @@ async def pay_for_request(
|
||||
"Charging base cost for request",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"cost": cost_per_request,
|
||||
"balance_before": billing_key.balance,
|
||||
"balance_before": key.balance,
|
||||
},
|
||||
)
|
||||
|
||||
# Charge the base cost for the request atomically to avoid race conditions
|
||||
stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.where(col(ApiKey.balance) - col(ApiKey.reserved_balance) >= cost_per_request)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance) + cost_per_request,
|
||||
@@ -370,16 +344,6 @@ async def pay_for_request(
|
||||
)
|
||||
)
|
||||
result = await session.exec(stmt) # type: ignore[call-overload]
|
||||
|
||||
# Also increment total_requests on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(total_requests=col(ApiKey.total_requests) + 1)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
|
||||
await session.commit()
|
||||
|
||||
if result.rowcount == 0:
|
||||
@@ -387,9 +351,8 @@ async def pay_for_request(
|
||||
"Concurrent request depleted balance",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"required_cost": cost_per_request,
|
||||
"current_balance": billing_key.balance,
|
||||
"current_balance": key.balance,
|
||||
},
|
||||
)
|
||||
|
||||
@@ -398,26 +361,23 @@ async def pay_for_request(
|
||||
status_code=402,
|
||||
detail={
|
||||
"error": {
|
||||
"message": f"Insufficient balance: {cost_per_request} mSats required. {billing_key.balance} available.",
|
||||
"message": f"Insufficient balance: {cost_per_request} mSats required. {key.balance} available.",
|
||||
"type": "insufficient_quota",
|
||||
"code": "insufficient_balance",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
await session.refresh(billing_key)
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
await session.refresh(key)
|
||||
await session.refresh(key)
|
||||
|
||||
logger.info(
|
||||
"Payment processed successfully",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"charged_amount": cost_per_request,
|
||||
"new_balance": billing_key.balance,
|
||||
"total_spent": billing_key.total_spent,
|
||||
"total_requests": billing_key.total_requests,
|
||||
"new_balance": key.balance,
|
||||
"total_spent": key.total_spent,
|
||||
"total_requests": key.total_requests,
|
||||
},
|
||||
)
|
||||
|
||||
@@ -427,11 +387,9 @@ async def pay_for_request(
|
||||
async def revert_pay_for_request(
|
||||
key: ApiKey, session: AsyncSession, cost_per_request: int
|
||||
) -> None:
|
||||
billing_key = await get_billing_key(key, session)
|
||||
|
||||
stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance) - cost_per_request,
|
||||
total_requests=col(ApiKey.total_requests) - 1,
|
||||
@@ -439,40 +397,27 @@ async def revert_pay_for_request(
|
||||
)
|
||||
|
||||
result = await session.exec(stmt) # type: ignore[call-overload]
|
||||
|
||||
# Also decrement total_requests on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(total_requests=col(ApiKey.total_requests) - 1)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
|
||||
await session.commit()
|
||||
if result.rowcount == 0:
|
||||
logger.error(
|
||||
"Failed to revert payment - insufficient reserved balance",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"cost_to_revert": cost_per_request,
|
||||
"current_reserved_balance": billing_key.reserved_balance,
|
||||
"current_reserved_balance": key.reserved_balance,
|
||||
},
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=402,
|
||||
detail={
|
||||
"error": {
|
||||
"message": f"failed to revert request payment: {cost_per_request} mSats required. {billing_key.balance} available.",
|
||||
"message": f"failed to revert request payment: {cost_per_request} mSats required. {key.balance} available.",
|
||||
"type": "payment_error",
|
||||
"code": "payment_error",
|
||||
}
|
||||
},
|
||||
)
|
||||
await session.refresh(billing_key)
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
await session.refresh(key)
|
||||
await session.refresh(key)
|
||||
|
||||
|
||||
async def adjust_payment_for_tokens(
|
||||
@@ -483,17 +428,15 @@ async def adjust_payment_for_tokens(
|
||||
This is called after the initial payment and the upstream request is complete.
|
||||
Returns cost data to be included in the response.
|
||||
"""
|
||||
billing_key = await get_billing_key(key, session)
|
||||
model = response_data.get("model", "unknown")
|
||||
|
||||
logger.debug(
|
||||
"Starting payment adjustment for tokens",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"model": model,
|
||||
"deducted_max_cost": deducted_max_cost,
|
||||
"current_balance": billing_key.balance,
|
||||
"current_balance": key.balance,
|
||||
"has_usage": "usage" in response_data,
|
||||
},
|
||||
)
|
||||
@@ -503,10 +446,8 @@ async def adjust_payment_for_tokens(
|
||||
try:
|
||||
release_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost
|
||||
)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost)
|
||||
)
|
||||
await session.exec(release_stmt) # type: ignore[call-overload]
|
||||
await session.commit()
|
||||
@@ -514,18 +455,13 @@ async def adjust_payment_for_tokens(
|
||||
"Released reservation without charging (fallback)",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"deducted_max_cost": deducted_max_cost,
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to release reservation in fallback",
|
||||
extra={
|
||||
"error": str(e),
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
},
|
||||
extra={"error": str(e), "key_hash": key.hashed_key[:8] + "..."},
|
||||
)
|
||||
|
||||
match await calculate_cost(response_data, deducted_max_cost, session):
|
||||
@@ -534,7 +470,6 @@ async def adjust_payment_for_tokens(
|
||||
"Using max cost data (no token adjustment)",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"model": model,
|
||||
"max_cost": cost.total_msats,
|
||||
},
|
||||
@@ -542,7 +477,7 @@ async def adjust_payment_for_tokens(
|
||||
# Finalize by releasing reservation and charging max cost
|
||||
finalize_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost,
|
||||
balance=col(ApiKey.balance) - cost.total_msats,
|
||||
@@ -550,41 +485,27 @@ async def adjust_payment_for_tokens(
|
||||
)
|
||||
)
|
||||
result = await session.exec(finalize_stmt) # type: ignore[call-overload]
|
||||
|
||||
# Also update total_spent on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(total_spent=col(ApiKey.total_spent) + cost.total_msats)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
|
||||
await session.commit()
|
||||
if result.rowcount == 0:
|
||||
logger.error(
|
||||
"Failed to finalize max-cost payment - retrying reservation release",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"deducted_max_cost": deducted_max_cost,
|
||||
"current_reserved_balance": billing_key.reserved_balance,
|
||||
"current_reserved_balance": key.reserved_balance,
|
||||
"total_cost": cost.total_msats,
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
await release_reservation_only()
|
||||
else:
|
||||
await session.refresh(billing_key)
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
await session.refresh(key)
|
||||
await session.refresh(key)
|
||||
logger.info(
|
||||
"Max cost payment finalized",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"charged_amount": cost.total_msats,
|
||||
"new_balance": billing_key.balance,
|
||||
"new_balance": key.balance,
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
@@ -600,7 +521,6 @@ async def adjust_payment_for_tokens(
|
||||
"Calculated token-based cost",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"model": model,
|
||||
"token_cost": cost.total_msats,
|
||||
"deducted_max_cost": deducted_max_cost,
|
||||
@@ -613,15 +533,11 @@ async def adjust_payment_for_tokens(
|
||||
if cost_difference == 0:
|
||||
logger.debug(
|
||||
"Finalizing with exact reserved cost",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"model": model,
|
||||
},
|
||||
extra={"key_hash": key.hashed_key[:8] + "...", "model": model},
|
||||
)
|
||||
finalize_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance)
|
||||
- deducted_max_cost,
|
||||
@@ -630,20 +546,8 @@ async def adjust_payment_for_tokens(
|
||||
)
|
||||
)
|
||||
await session.exec(finalize_stmt) # type: ignore[call-overload]
|
||||
|
||||
# Also update total_spent on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(total_spent=col(ApiKey.total_spent) + total_cost_msats)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
|
||||
await session.commit()
|
||||
await session.refresh(billing_key)
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
await session.refresh(key)
|
||||
await session.refresh(key)
|
||||
return cost.dict()
|
||||
|
||||
# this should never happen why do we handle this???
|
||||
@@ -653,17 +557,16 @@ async def adjust_payment_for_tokens(
|
||||
"Additional charge required for token usage",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"additional_charge": cost_difference,
|
||||
"current_balance": billing_key.balance,
|
||||
"sufficient_balance": billing_key.balance >= cost_difference,
|
||||
"current_balance": key.balance,
|
||||
"sufficient_balance": key.balance >= cost_difference,
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
|
||||
finalize_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance)
|
||||
- deducted_max_cost,
|
||||
@@ -672,31 +575,18 @@ async def adjust_payment_for_tokens(
|
||||
)
|
||||
)
|
||||
result = await session.exec(finalize_stmt) # type: ignore[call-overload]
|
||||
|
||||
# Also update total_spent on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(total_spent=col(ApiKey.total_spent) + total_cost_msats)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
|
||||
await session.commit()
|
||||
|
||||
if result.rowcount:
|
||||
cost.total_msats = total_cost_msats
|
||||
await session.refresh(billing_key)
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
await session.refresh(key)
|
||||
await session.refresh(key)
|
||||
|
||||
logger.info(
|
||||
"Finalized payment with additional charge",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"charged_amount": total_cost_msats,
|
||||
"new_balance": billing_key.balance,
|
||||
"new_balance": key.balance,
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
@@ -705,7 +595,6 @@ async def adjust_payment_for_tokens(
|
||||
"Failed to finalize additional charge - releasing reservation",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"attempted_charge": total_cost_msats,
|
||||
"model": model,
|
||||
},
|
||||
@@ -718,16 +607,15 @@ async def adjust_payment_for_tokens(
|
||||
"Refunding excess payment",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"refund_amount": refund,
|
||||
"current_balance": billing_key.balance,
|
||||
"current_balance": key.balance,
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
|
||||
refund_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance)
|
||||
- deducted_max_cost,
|
||||
@@ -736,16 +624,6 @@ async def adjust_payment_for_tokens(
|
||||
)
|
||||
)
|
||||
result = await session.exec(refund_stmt) # type: ignore[call-overload]
|
||||
|
||||
# Also update total_spent on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(total_spent=col(ApiKey.total_spent) + total_cost_msats)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
|
||||
await session.commit()
|
||||
|
||||
if result.rowcount == 0:
|
||||
@@ -753,9 +631,8 @@ async def adjust_payment_for_tokens(
|
||||
"Failed to finalize payment - releasing reservation",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"deducted_max_cost": deducted_max_cost,
|
||||
"current_reserved_balance": billing_key.reserved_balance,
|
||||
"current_reserved_balance": key.reserved_balance,
|
||||
"total_cost": total_cost_msats,
|
||||
"model": model,
|
||||
},
|
||||
@@ -763,17 +640,14 @@ async def adjust_payment_for_tokens(
|
||||
await release_reservation_only()
|
||||
else:
|
||||
cost.total_msats = total_cost_msats
|
||||
await session.refresh(billing_key)
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
await session.refresh(key)
|
||||
await session.refresh(key)
|
||||
|
||||
logger.info(
|
||||
"Refund processed successfully",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"refunded_amount": refund,
|
||||
"new_balance": billing_key.balance,
|
||||
"new_balance": key.balance,
|
||||
"final_cost": cost.total_msats,
|
||||
"model": model,
|
||||
},
|
||||
|
||||
+37
-85
@@ -32,28 +32,14 @@ async def get_key_from_header(
|
||||
)
|
||||
|
||||
|
||||
async def get_balance_info(key: ApiKey, session: AsyncSession) -> dict:
|
||||
from .auth import get_billing_key
|
||||
|
||||
billing_key = await get_billing_key(key, session)
|
||||
return {
|
||||
"api_key": "sk-" + key.hashed_key,
|
||||
"balance": billing_key.balance,
|
||||
"reserved": billing_key.reserved_balance,
|
||||
"is_child": key.parent_key_hash is not None,
|
||||
"parent_key": "sk-" + key.parent_key_hash if key.parent_key_hash else None,
|
||||
"total_requests": key.total_requests,
|
||||
"total_spent": key.total_spent,
|
||||
}
|
||||
|
||||
|
||||
# TODO: remove this endpoint when frontend is updated
|
||||
@router.get("/", include_in_schema=False)
|
||||
async def account_info(
|
||||
key: ApiKey = Depends(get_key_from_header),
|
||||
session: AsyncSession = Depends(get_session),
|
||||
) -> dict:
|
||||
return await get_balance_info(key, session)
|
||||
async def account_info(key: ApiKey = Depends(get_key_from_header)) -> dict:
|
||||
return {
|
||||
"api_key": "sk-" + key.hashed_key,
|
||||
"balance": key.balance,
|
||||
"reserved": key.reserved_balance,
|
||||
}
|
||||
|
||||
|
||||
# TODO: Implement POST /v1/wallet/create endpoint
|
||||
@@ -80,11 +66,12 @@ async def create_balance(
|
||||
|
||||
|
||||
@router.get("/info")
|
||||
async def wallet_info(
|
||||
key: ApiKey = Depends(get_key_from_header),
|
||||
session: AsyncSession = Depends(get_session),
|
||||
) -> dict:
|
||||
return await get_balance_info(key, session)
|
||||
async def wallet_info(key: ApiKey = Depends(get_key_from_header)) -> dict:
|
||||
return {
|
||||
"api_key": "sk-" + key.hashed_key,
|
||||
"balance": key.balance,
|
||||
"reserved": key.reserved_balance,
|
||||
}
|
||||
|
||||
|
||||
class TopupRequest(BaseModel):
|
||||
@@ -98,10 +85,6 @@ async def topup_wallet_endpoint(
|
||||
key: ApiKey = Depends(get_key_from_header),
|
||||
session: AsyncSession = Depends(get_session),
|
||||
) -> dict[str, int]:
|
||||
from .auth import get_billing_key
|
||||
|
||||
billing_key = await get_billing_key(key, session)
|
||||
|
||||
if topup_request is not None:
|
||||
cashu_token = topup_request.cashu_token
|
||||
if cashu_token is None:
|
||||
@@ -111,7 +94,7 @@ async def topup_wallet_endpoint(
|
||||
if len(cashu_token) < 10 or "cashu" not in cashu_token:
|
||||
raise HTTPException(status_code=400, detail="Invalid token format")
|
||||
try:
|
||||
amount_msats = await credit_balance(cashu_token, billing_key, session)
|
||||
amount_msats = await credit_balance(cashu_token, key, session)
|
||||
except ValueError as e:
|
||||
error_msg = str(e)
|
||||
if "already spent" in error_msg.lower():
|
||||
@@ -159,6 +142,7 @@ async def refund_wallet_endpoint(
|
||||
authorization: Annotated[str, Header(...)],
|
||||
session: AsyncSession = Depends(get_session),
|
||||
) -> dict[str, str]:
|
||||
logger.info("Refund request received", extra={"authorization": authorization})
|
||||
if not authorization.startswith("Bearer "):
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
@@ -172,14 +156,9 @@ async def refund_wallet_endpoint(
|
||||
|
||||
key: ApiKey = await validate_bearer_key(bearer_value, session)
|
||||
|
||||
if key.parent_key_hash:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Cannot refund child key. Please refund the parent key instead.",
|
||||
)
|
||||
|
||||
remaining_balance_msats: int = key.total_balance
|
||||
|
||||
logger.info("Refunding key: %s", key.dict())
|
||||
if key.refund_currency == "sat":
|
||||
remaining_balance = remaining_balance_msats // 1000
|
||||
else:
|
||||
@@ -195,6 +174,15 @@ async def refund_wallet_endpoint(
|
||||
if key.refund_address:
|
||||
from .core.settings import settings as global_settings
|
||||
|
||||
logger.info(
|
||||
"Sending refund to lnurl",
|
||||
extra={
|
||||
"remaining_balance": remaining_balance,
|
||||
"refund_currency": key.refund_currency,
|
||||
"refund_mint_url": key.refund_mint_url,
|
||||
"refund_address": key.refund_address,
|
||||
},
|
||||
)
|
||||
await send_to_lnurl(
|
||||
remaining_balance,
|
||||
key.refund_currency or "sat",
|
||||
@@ -203,7 +191,16 @@ async def refund_wallet_endpoint(
|
||||
)
|
||||
result = {"recipient": key.refund_address}
|
||||
else:
|
||||
logger.info(
|
||||
"Sending refund as token",
|
||||
extra={
|
||||
"remaining_balance": remaining_balance,
|
||||
"refund_currency": key.refund_currency,
|
||||
"refund_mint_url": key.refund_mint_url,
|
||||
},
|
||||
)
|
||||
refund_currency = key.refund_currency or "sat"
|
||||
|
||||
token = await send_token(
|
||||
remaining_balance, refund_currency, key.refund_mint_url
|
||||
)
|
||||
@@ -214,10 +211,12 @@ async def refund_wallet_endpoint(
|
||||
else:
|
||||
result["msats"] = str(remaining_balance_msats)
|
||||
|
||||
except HTTPException:
|
||||
except HTTPException as e:
|
||||
logger.error("Refund failed", extra={"exception": e})
|
||||
# Re-raise HTTP exceptions (like 400 for balance too small)
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error("Refund failed", extra={"exception": e})
|
||||
# If refund fails, don't modify the database
|
||||
error_msg = str(e)
|
||||
if (
|
||||
@@ -228,7 +227,7 @@ async def refund_wallet_endpoint(
|
||||
):
|
||||
raise HTTPException(status_code=503, detail="Mint service unavailable")
|
||||
else:
|
||||
raise HTTPException(status_code=500, detail="Refund failed")
|
||||
raise HTTPException(status_code=500, detail="Refund failed: " + error_msg)
|
||||
|
||||
await _refund_cache_set(bearer_value, result)
|
||||
|
||||
@@ -251,53 +250,6 @@ async def donate(token: str, ref: str | None = None) -> str:
|
||||
return "Invalid token."
|
||||
|
||||
|
||||
@router.post("/child-key")
|
||||
async def create_child_key(
|
||||
key: ApiKey = Depends(get_key_from_header),
|
||||
session: AsyncSession = Depends(get_session),
|
||||
) -> dict:
|
||||
"""Creates a child API key that uses the parent's balance."""
|
||||
# Check if this is already a child key
|
||||
if key.parent_key_hash:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Cannot create a child key for another child key.",
|
||||
)
|
||||
|
||||
cost = settings.child_key_cost
|
||||
|
||||
if key.total_balance < cost:
|
||||
raise HTTPException(
|
||||
status_code=402,
|
||||
detail=f"Insufficient balance to create child key. {cost} mSats required.",
|
||||
)
|
||||
|
||||
# Deduct cost from parent
|
||||
key.balance -= cost
|
||||
key.total_spent += cost
|
||||
session.add(key)
|
||||
|
||||
# Generate new key
|
||||
import secrets
|
||||
|
||||
new_key_raw = secrets.token_hex(32)
|
||||
new_key_hash = new_key_raw # We use the raw key as the hash for sk- keys
|
||||
|
||||
child_key = ApiKey(
|
||||
hashed_key=new_key_hash,
|
||||
balance=0,
|
||||
parent_key_hash=key.hashed_key,
|
||||
)
|
||||
session.add(child_key)
|
||||
await session.commit()
|
||||
|
||||
return {
|
||||
"api_key": "sk-" + new_key_hash,
|
||||
"cost_msats": cost,
|
||||
"parent_balance": key.balance,
|
||||
}
|
||||
|
||||
|
||||
@router.api_route(
|
||||
"/{path:path}",
|
||||
methods=["GET", "POST", "PUT", "DELETE"],
|
||||
|
||||
@@ -124,7 +124,7 @@ async def partial_apikeys(request: Request) -> str:
|
||||
|
||||
rows = "".join(
|
||||
[
|
||||
f"<tr><td>{key.hashed_key}{' <br><small>(Child of ' + key.parent_key_hash[:8] + '...)</small>' if key.parent_key_hash else ''}</td><td>{key.balance}</td><td>{key.total_spent}</td><td>{key.total_requests}</td><td>{key.refund_address}</td><td>{fmt_time(key.key_expiry_time)}</td></tr>"
|
||||
f"<tr><td>{key.hashed_key}</td><td>{key.balance}</td><td>{key.total_spent}</td><td>{key.total_requests}</td><td>{key.refund_address}</td><td>{fmt_time(key.key_expiry_time)}</td></tr>"
|
||||
for key in api_keys
|
||||
]
|
||||
)
|
||||
@@ -158,7 +158,6 @@ async def get_temporary_balances_api(request: Request) -> list[dict[str, object]
|
||||
"total_requests": key.total_requests,
|
||||
"refund_address": key.refund_address,
|
||||
"key_expiry_time": key.key_expiry_time,
|
||||
"parent_key_hash": key.parent_key_hash,
|
||||
}
|
||||
for key in api_keys
|
||||
]
|
||||
|
||||
@@ -47,9 +47,6 @@ class ApiKey(SQLModel, table=True): # type: ignore
|
||||
default=None,
|
||||
description="Currency of the cashu-token",
|
||||
)
|
||||
parent_key_hash: str | None = Field(
|
||||
default=None, foreign_key="api_keys.hashed_key", index=True
|
||||
)
|
||||
|
||||
@property
|
||||
def total_balance(self) -> int:
|
||||
|
||||
@@ -13,7 +13,10 @@ from starlette.exceptions import HTTPException
|
||||
from ..balance import balance_router, deprecated_wallet_router
|
||||
from ..discovery import providers_cache_refresher, providers_router
|
||||
from ..nip91 import announce_provider
|
||||
from ..payment.models import models_router, update_sats_pricing
|
||||
from ..payment.models import (
|
||||
models_router,
|
||||
update_sats_pricing,
|
||||
)
|
||||
from ..payment.price import update_prices_periodically
|
||||
from ..proxy import initialize_upstreams, proxy_router, refresh_model_maps_periodically
|
||||
from ..wallet import periodic_payout
|
||||
|
||||
@@ -6,7 +6,7 @@ import os
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from pydantic.v1 import BaseModel, BaseSettings, Field, validator
|
||||
from pydantic.v1 import BaseModel, BaseSettings, Field
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
|
||||
@@ -37,13 +37,6 @@ class Settings(BaseSettings):
|
||||
|
||||
# Cashu
|
||||
cashu_mints: list[str] = Field(default_factory=list, env="CASHU_MINTS")
|
||||
|
||||
@validator("cashu_mints", pre=True, each_item=True)
|
||||
def normalize_mint_url(cls, v: str) -> str:
|
||||
if isinstance(v, str):
|
||||
return v.rstrip("/")
|
||||
return v
|
||||
|
||||
receive_ln_address: str = Field(default="", env="RECEIVE_LN_ADDRESS")
|
||||
primary_mint: str = Field(default="", env="PRIMARY_MINT_URL")
|
||||
primary_mint_unit: str = Field(default="sat", env="PRIMARY_MINT_UNIT")
|
||||
@@ -59,7 +52,6 @@ class Settings(BaseSettings):
|
||||
exchange_fee: float = Field(default=1.005, env="EXCHANGE_FEE")
|
||||
upstream_provider_fee: float = Field(default=1.05, env="UPSTREAM_PROVIDER_FEE")
|
||||
tolerance_percentage: float = Field(default=1.0, env="TOLERANCE_PERCENTAGE")
|
||||
child_key_cost: int = Field(default=1000, env="CHILD_KEY_COST")
|
||||
# Minimum per-request charge in millisatoshis when model pricing is free/zero
|
||||
min_request_msat: int = Field(default=1, env="MIN_REQUEST_MSAT")
|
||||
reset_reserved_balance_on_startup: bool = Field(
|
||||
|
||||
@@ -16,7 +16,6 @@ from .core.db import (
|
||||
create_session,
|
||||
get_session,
|
||||
)
|
||||
from .core.settings import settings
|
||||
from .payment.helpers import (
|
||||
calculate_discounted_max_cost,
|
||||
check_token_balance,
|
||||
@@ -26,7 +25,6 @@ from .payment.helpers import (
|
||||
from .payment.models import Model
|
||||
from .upstream import BaseUpstreamProvider
|
||||
from .upstream.helpers import init_upstreams
|
||||
from .wallet import deserialize_token_from_string
|
||||
|
||||
logger = get_logger(__name__)
|
||||
proxy_router = APIRouter()
|
||||
@@ -148,24 +146,6 @@ async def proxy(
|
||||
else:
|
||||
model_id = request_body_dict.get("model", "unknown")
|
||||
|
||||
if "https://testnut.cashu.space" in settings.cashu_mints:
|
||||
try:
|
||||
token_str = None
|
||||
if x_cashu_header := headers.get("x-cashu"):
|
||||
token_str = x_cashu_header
|
||||
elif auth_header := headers.get("authorization"):
|
||||
parts = auth_header.split(" ")
|
||||
if len(parts) > 1 and not parts[1].startswith("sk-"):
|
||||
token_str = parts[1]
|
||||
|
||||
if token_str:
|
||||
token_obj = deserialize_token_from_string(token_str)
|
||||
if token_obj.mint == "https://testnut.cashu.space":
|
||||
model_id = "mock/gpt-420-mock"
|
||||
request_body_dict["model"] = model_id
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
model_obj = get_model_instance(model_id)
|
||||
if not model_obj:
|
||||
return create_error_response(
|
||||
|
||||
@@ -1,265 +0,0 @@
|
||||
import asyncio
|
||||
import json
|
||||
import random
|
||||
from typing import AsyncIterator
|
||||
|
||||
from fastapi import Request
|
||||
from fastapi.responses import Response, StreamingResponse
|
||||
|
||||
from ..core.db import ApiKey, AsyncSession
|
||||
from ..payment.models import Architecture, Model, Pricing
|
||||
from .base import BaseUpstreamProvider
|
||||
|
||||
|
||||
class MockUpstreamProvider(BaseUpstreamProvider):
|
||||
"""Fack Mock Upstream provider specifically for Testing."""
|
||||
|
||||
provider_type = "mock"
|
||||
|
||||
async def forward_request(
|
||||
self,
|
||||
request: Request,
|
||||
path: str,
|
||||
headers: dict,
|
||||
request_body: bytes | None,
|
||||
key: ApiKey,
|
||||
max_cost_for_model: int,
|
||||
session: AsyncSession,
|
||||
model_obj: Model,
|
||||
) -> Response | StreamingResponse:
|
||||
if path.endswith("chat/completions"):
|
||||
is_streaming = False
|
||||
if request_body:
|
||||
request_data = json.loads(request_body)
|
||||
is_streaming = request_data.get("stream", False)
|
||||
|
||||
if is_streaming:
|
||||
|
||||
async def fake_streaming_response(
|
||||
chunk_size: int | None = None,
|
||||
) -> AsyncIterator[bytes]:
|
||||
suffix = random.randint(1000, 9999)
|
||||
req_id = f"gen-mock-stream-{suffix}"
|
||||
created = 1766138895
|
||||
model = "mock/gpt-420-mock"
|
||||
|
||||
def make_chunk(
|
||||
delta: dict,
|
||||
finish_reason: str | None = None,
|
||||
usage: dict | None = None,
|
||||
) -> bytes:
|
||||
chunk = {
|
||||
"id": req_id,
|
||||
"provider": "MockProvider",
|
||||
"model": model,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": created,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": delta,
|
||||
"finish_reason": finish_reason,
|
||||
"native_finish_reason": "completed"
|
||||
if finish_reason
|
||||
else None,
|
||||
"logprobs": None,
|
||||
}
|
||||
],
|
||||
}
|
||||
if usage:
|
||||
chunk["usage"] = usage
|
||||
return f"data: {json.dumps(chunk)}\n\n".encode()
|
||||
|
||||
# 1. Initial chunk
|
||||
yield make_chunk({"role": "assistant", "content": ""})
|
||||
await asyncio.sleep(0.02)
|
||||
|
||||
# 2. Reasoning chunks
|
||||
reasoning_tokens = ["Mock", " reason", "ing", "..."]
|
||||
for token in reasoning_tokens:
|
||||
delta = {
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"reasoning": token,
|
||||
"reasoning_details": [
|
||||
{
|
||||
"type": "reasoning.summary",
|
||||
"summary": token,
|
||||
"format": "openai-responses-v1",
|
||||
"index": 0,
|
||||
}
|
||||
],
|
||||
}
|
||||
yield make_chunk(delta)
|
||||
await asyncio.sleep(0.03)
|
||||
|
||||
# 3. Content chunks
|
||||
content_tokens = ["This", " is", " a", " mock", " stream", "."]
|
||||
for token in content_tokens:
|
||||
yield make_chunk({"role": "assistant", "content": token})
|
||||
await asyncio.sleep(0.03)
|
||||
|
||||
# 4. Finish chunk
|
||||
yield make_chunk(
|
||||
{"role": "assistant", "content": ""}, finish_reason="stop"
|
||||
)
|
||||
|
||||
# 5. Usage chunk
|
||||
usage_data = {
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 20,
|
||||
"total_tokens": 30,
|
||||
"cost": 0.001,
|
||||
"is_byok": False,
|
||||
"prompt_tokens_details": {
|
||||
"cached_tokens": 0,
|
||||
"audio_tokens": 0,
|
||||
"video_tokens": 0,
|
||||
},
|
||||
"cost_details": {
|
||||
"upstream_inference_cost": None,
|
||||
"upstream_inference_prompt_cost": 0,
|
||||
"upstream_inference_completions_cost": 0.001,
|
||||
},
|
||||
"completion_tokens_details": {
|
||||
"reasoning_tokens": 10,
|
||||
"image_tokens": 0,
|
||||
},
|
||||
}
|
||||
|
||||
usage_chunk = {
|
||||
"id": req_id,
|
||||
"provider": "MockProvider",
|
||||
"model": model,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": created,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"role": "assistant", "content": ""},
|
||||
"finish_reason": None,
|
||||
"native_finish_reason": None,
|
||||
"logprobs": None,
|
||||
}
|
||||
],
|
||||
"usage": usage_data,
|
||||
}
|
||||
yield f"data: {json.dumps(usage_chunk)}\n\n".encode()
|
||||
|
||||
# 6. DONE
|
||||
yield b"data: [DONE]\n\n"
|
||||
|
||||
# 7. Cost
|
||||
cost_chunk = {
|
||||
"cost": {
|
||||
"base_msats": 0,
|
||||
"input_msats": 2,
|
||||
"output_msats": 10,
|
||||
"total_msats": 12,
|
||||
}
|
||||
}
|
||||
yield f"data: {json.dumps(cost_chunk)}\n\n".encode()
|
||||
|
||||
return StreamingResponse(
|
||||
fake_streaming_response(),
|
||||
200,
|
||||
)
|
||||
|
||||
else:
|
||||
suffix = random.randint(1000, 9999)
|
||||
content_dict = {
|
||||
"id": f"gen-mock-{suffix}",
|
||||
"provider": "MockProvider",
|
||||
"model": "mock/gpt-5-mini",
|
||||
"object": "chat.completion",
|
||||
"created": 1766138655,
|
||||
"choices": [
|
||||
{
|
||||
"logprobs": None,
|
||||
"finish_reason": "length",
|
||||
"native_finish_reason": "max_output_tokens",
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": f"Mock Content {suffix}",
|
||||
"refusal": None,
|
||||
"reasoning": f"Mock Reasoning {suffix}",
|
||||
"reasoning_details": [
|
||||
{
|
||||
"format": "openai-responses-v1",
|
||||
"index": 0,
|
||||
"type": "reasoning.summary",
|
||||
"summary": f"Mock Summary {suffix}",
|
||||
},
|
||||
{
|
||||
"id": f"rs_mock_{suffix}",
|
||||
"format": "openai-responses-v1",
|
||||
"index": 0,
|
||||
"type": "reasoning.encrypted",
|
||||
"data": "mock_encrypted_data",
|
||||
},
|
||||
],
|
||||
},
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 10,
|
||||
"total_tokens": 20,
|
||||
"cost": 0,
|
||||
"is_byok": False,
|
||||
"prompt_tokens_details": {
|
||||
"cached_tokens": 0,
|
||||
"audio_tokens": 0,
|
||||
"video_tokens": 0,
|
||||
},
|
||||
"cost_details": {
|
||||
"upstream_inference_cost": None,
|
||||
"upstream_inference_prompt_cost": 0,
|
||||
"upstream_inference_completions_cost": 0,
|
||||
},
|
||||
"completion_tokens_details": {
|
||||
"reasoning_tokens": 5,
|
||||
"image_tokens": 0,
|
||||
},
|
||||
},
|
||||
"cost": {
|
||||
"base_msats": 0,
|
||||
"input_msats": 0,
|
||||
"output_msats": 0,
|
||||
"total_msats": 0,
|
||||
},
|
||||
}
|
||||
return Response(json.dumps(content_dict).encode(), 200)
|
||||
|
||||
elif path.endswith("embeddings"):
|
||||
raise NotImplementedError
|
||||
elif path.endswith("responses"):
|
||||
raise NotImplementedError
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
async def fetch_models(self) -> list[Model]:
|
||||
return [
|
||||
Model(
|
||||
id="mock/gpt-420-mock",
|
||||
name="mock/gpt-420-mock",
|
||||
created=0,
|
||||
description="mock model for testing",
|
||||
context_length=8192,
|
||||
architecture=Architecture(
|
||||
modality="text",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["text"],
|
||||
tokenizer="",
|
||||
instruct_type=None,
|
||||
),
|
||||
pricing=Pricing(prompt=0.01, completion=0.01),
|
||||
),
|
||||
]
|
||||
|
||||
def transform_model_name(self, model_id: str) -> str:
|
||||
return "fake-model"
|
||||
|
||||
async def get_balance(self) -> float | None:
|
||||
return 420.69
|
||||
@@ -218,14 +218,6 @@ async def init_upstreams() -> list[BaseUpstreamProvider]:
|
||||
results = await asyncio.gather(*tasks)
|
||||
upstreams = [p for p in results if p is not None]
|
||||
|
||||
if "https://testnut.cashu.space" in settings.cashu_mints:
|
||||
from .fake import MockUpstreamProvider
|
||||
|
||||
mock_provider = MockUpstreamProvider("mock", "mock")
|
||||
await mock_provider.refresh_models_cache()
|
||||
upstreams.append(mock_provider)
|
||||
logger.info("Initialized MockUpstreamProvider for testnut mint")
|
||||
|
||||
return upstreams
|
||||
|
||||
|
||||
|
||||
+8
-3
@@ -40,16 +40,23 @@ async def recieve_token(
|
||||
async def send(amount: int, unit: str, mint_url: str | None = None) -> tuple[int, str]:
|
||||
"""Internal send function - returns amount and serialized token"""
|
||||
wallet: Wallet = await get_wallet(mint_url or settings.primary_mint, unit)
|
||||
logger.info("Sending", extra={"amount": amount, "unit": unit, "mint_url": mint_url})
|
||||
proofs = get_proofs_per_mint_and_unit(
|
||||
wallet, mint_url or settings.primary_mint, unit
|
||||
)
|
||||
|
||||
logger.info("Proofs", extra={"proofs": proofs})
|
||||
logger.info(
|
||||
"Selecting to send",
|
||||
extra={"amount": amount, "unit": unit, "mint_url": mint_url},
|
||||
)
|
||||
send_proofs, _ = await wallet.select_to_send(
|
||||
proofs, amount, set_reserved=True, include_fees=False
|
||||
)
|
||||
logger.info("Send proofs", extra={"send_proofs": send_proofs})
|
||||
token = await wallet.serialize_proofs(
|
||||
send_proofs, include_dleq=False, legacy=False, memo=None
|
||||
)
|
||||
logger.info("Token created", extra={"token": token})
|
||||
return amount, token
|
||||
|
||||
|
||||
@@ -313,8 +320,6 @@ async def periodic_payout() -> None:
|
||||
try:
|
||||
async with db.create_session() as session:
|
||||
for mint_url in settings.cashu_mints:
|
||||
if mint_url == "https://testnut.cashu.space":
|
||||
continue
|
||||
for unit in ["sat", "msat"]:
|
||||
wallet = await get_wallet(mint_url, unit)
|
||||
proofs = get_proofs_per_mint_and_unit(
|
||||
|
||||
@@ -1,137 +0,0 @@
|
||||
import secrets
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from routstr.auth import adjust_payment_for_tokens, pay_for_request
|
||||
from routstr.balance import create_child_key
|
||||
from routstr.core.db import ApiKey
|
||||
from routstr.core.settings import settings
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_child_key_flow(integration_session: AsyncSession) -> None:
|
||||
# 1. Create a parent key with balance
|
||||
parent_raw = "parent_test_key_" + secrets.token_hex(4)
|
||||
parent_key = ApiKey(
|
||||
hashed_key=parent_raw,
|
||||
balance=10000, # 10 sats
|
||||
)
|
||||
integration_session.add(parent_key)
|
||||
await integration_session.commit()
|
||||
await integration_session.refresh(parent_key)
|
||||
|
||||
# Mock settings
|
||||
settings.child_key_cost = 1000 # 1 sat
|
||||
|
||||
# 2. Call create_child_key
|
||||
result = await create_child_key(parent_key, integration_session)
|
||||
|
||||
assert "api_key" in result
|
||||
assert result["cost_msats"] == 1000
|
||||
assert result["parent_balance"] == 9000
|
||||
|
||||
child_key_raw = result["api_key"][3:] # remove sk-
|
||||
|
||||
# 3. Verify child key exists in DB
|
||||
child_key_db = await integration_session.get(ApiKey, child_key_raw)
|
||||
assert child_key_db is not None
|
||||
assert child_key_db.parent_key_hash == parent_key.hashed_key
|
||||
assert child_key_db.balance == 0
|
||||
|
||||
# 4. Test payment with child key
|
||||
cost = 500
|
||||
await pay_for_request(child_key_db, cost, integration_session)
|
||||
|
||||
# Refresh keys
|
||||
await integration_session.refresh(parent_key)
|
||||
await integration_session.refresh(child_key_db)
|
||||
|
||||
# Parent should be charged
|
||||
assert parent_key.reserved_balance == 500
|
||||
assert parent_key.total_requests == 1
|
||||
|
||||
# Child should have total_requests incremented
|
||||
assert child_key_db.total_requests == 1
|
||||
|
||||
# 5. Test adjustment
|
||||
response_data = {"model": "test-model", "usage": {"total_tokens": 10}}
|
||||
|
||||
# Mock calculate_cost
|
||||
import routstr.auth
|
||||
from routstr.payment.cost_calculation import CostData
|
||||
|
||||
async def mock_calculate_cost(*args: Any, **kwargs: Any) -> CostData:
|
||||
return CostData(
|
||||
base_msats=0, input_msats=200, output_msats=200, total_msats=400
|
||||
)
|
||||
|
||||
# Patch calculate_cost
|
||||
original_calculate_cost = routstr.auth.calculate_cost
|
||||
routstr.auth.calculate_cost = mock_calculate_cost
|
||||
|
||||
try:
|
||||
adjustment = await adjust_payment_for_tokens(
|
||||
child_key_db, response_data, integration_session, 500
|
||||
)
|
||||
assert adjustment["total_msats"] == 400
|
||||
|
||||
# Refresh keys
|
||||
await integration_session.refresh(parent_key)
|
||||
await integration_session.refresh(child_key_db)
|
||||
|
||||
# Parent should have updated balance and total_spent
|
||||
assert parent_key.reserved_balance == 0
|
||||
assert parent_key.balance == 9000 - 400
|
||||
assert (
|
||||
parent_key.total_spent == 1400
|
||||
) # 1000 for child key creation + 400 for request
|
||||
|
||||
# Child should also have total_spent updated
|
||||
assert child_key_db.total_spent == 400
|
||||
|
||||
finally:
|
||||
routstr.auth.calculate_cost = original_calculate_cost
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_child_key_insufficient_balance(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
parent_key = ApiKey(
|
||||
hashed_key="poor_parent_" + secrets.token_hex(4),
|
||||
balance=500,
|
||||
)
|
||||
integration_session.add(parent_key)
|
||||
await integration_session.commit()
|
||||
await integration_session.refresh(parent_key)
|
||||
|
||||
settings.child_key_cost = 1000
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await create_child_key(parent_key, integration_session)
|
||||
assert exc.value.status_code == 402
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_child_key_cannot_create_child(integration_session: AsyncSession) -> None:
|
||||
parent_key = ApiKey(
|
||||
hashed_key="parent_" + secrets.token_hex(4),
|
||||
balance=10000,
|
||||
)
|
||||
child_key = ApiKey(
|
||||
hashed_key="child_" + secrets.token_hex(4),
|
||||
balance=0,
|
||||
parent_key_hash=parent_key.hashed_key,
|
||||
)
|
||||
integration_session.add(parent_key)
|
||||
integration_session.add(child_key)
|
||||
await integration_session.commit()
|
||||
await integration_session.refresh(child_key)
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await create_child_key(child_key, integration_session)
|
||||
assert exc.value.status_code == 400
|
||||
assert "Cannot create a child key for another child key" in str(exc.value.detail)
|
||||
@@ -60,11 +60,7 @@ export function TemporaryBalances({
|
||||
let totalRequests = 0;
|
||||
|
||||
balances.forEach((balance) => {
|
||||
// Only count parents for total balance to avoid double counting
|
||||
// since child keys use parent balance
|
||||
if (!balance.parent_key_hash) {
|
||||
totalBalance += balance.balance || 0;
|
||||
}
|
||||
totalBalance += balance.balance || 0;
|
||||
totalSpent += balance.total_spent || 0;
|
||||
totalRequests += balance.total_requests || 0;
|
||||
});
|
||||
@@ -76,34 +72,6 @@ export function TemporaryBalances({
|
||||
? calculateTotals(data)
|
||||
: { totalBalance: 0, totalSpent: 0, totalRequests: 0 };
|
||||
|
||||
// Group parents and children
|
||||
const hierarchicalData = (() => {
|
||||
if (!data) return [];
|
||||
|
||||
const parents = filteredData.filter((item) => !item.parent_key_hash);
|
||||
const result: (TemporaryBalance & { isChild?: boolean })[] = [];
|
||||
|
||||
parents.forEach((parent) => {
|
||||
result.push(parent);
|
||||
const children = data.filter(
|
||||
(item) => item.parent_key_hash === parent.hashed_key
|
||||
);
|
||||
children.forEach((child) => {
|
||||
result.push({ ...child, isChild: true });
|
||||
});
|
||||
});
|
||||
|
||||
// Add children whose parents didn't match the search or aren't in the list
|
||||
const orphans = filteredData.filter(
|
||||
(item) =>
|
||||
item.parent_key_hash &&
|
||||
!result.some((r) => r.hashed_key === item.hashed_key)
|
||||
);
|
||||
result.push(...orphans.map((o) => ({ ...o, isChild: true })));
|
||||
|
||||
return result;
|
||||
})();
|
||||
|
||||
return (
|
||||
<>
|
||||
<Card className='h-full w-full shadow-sm'>
|
||||
@@ -214,37 +182,22 @@ export function TemporaryBalances({
|
||||
<div className='text-right'>Expiry Time</div>
|
||||
</div>
|
||||
|
||||
{hierarchicalData.length > 0 ? (
|
||||
hierarchicalData.map((balance, index) => (
|
||||
{filteredData.length > 0 ? (
|
||||
filteredData.map((balance, index) => (
|
||||
<div
|
||||
key={index}
|
||||
className={cn(
|
||||
'hover:bg-muted/50 border-t p-3 text-sm transition-colors',
|
||||
balance.balance === 0 &&
|
||||
!balance.isChild &&
|
||||
'opacity-60',
|
||||
balance.isChild &&
|
||||
'ml-4 border-l-2 border-l-blue-200 bg-blue-50/30'
|
||||
balance.balance === 0 && 'opacity-60'
|
||||
)}
|
||||
>
|
||||
{/* Desktop Layout */}
|
||||
<div className='hidden grid-cols-6 gap-2 md:grid'>
|
||||
<div className='flex max-w-48 items-center gap-2 truncate font-mono text-xs break-all'>
|
||||
{balance.isChild && (
|
||||
<span className='rounded bg-blue-100 px-1 py-0.5 text-[10px] font-bold text-blue-700 uppercase'>
|
||||
Child
|
||||
</span>
|
||||
)}
|
||||
<div className='max-w-32 truncate font-mono text-xs break-all'>
|
||||
{balance.hashed_key}
|
||||
</div>
|
||||
<div className='text-right font-mono'>
|
||||
{balance.isChild ? (
|
||||
<span className='text-muted-foreground italic'>
|
||||
(Parent)
|
||||
</span>
|
||||
) : (
|
||||
formatBalance(balance.balance)
|
||||
)}
|
||||
{formatBalance(balance.balance)}
|
||||
</div>
|
||||
<div className='text-right font-mono'>
|
||||
{formatBalance(balance.total_spent)}
|
||||
@@ -273,20 +226,13 @@ export function TemporaryBalances({
|
||||
|
||||
{/* Mobile Layout */}
|
||||
<div className='space-y-3 md:hidden'>
|
||||
<div className='flex items-center justify-between'>
|
||||
<div className='space-y-1'>
|
||||
<span className='text-muted-foreground text-xs font-medium'>
|
||||
{balance.isChild ? 'Child Key' : 'Key'}
|
||||
</span>
|
||||
<div className='font-mono text-xs break-all'>
|
||||
{balance.hashed_key}
|
||||
</div>
|
||||
<div className='space-y-1'>
|
||||
<span className='text-muted-foreground text-xs font-medium'>
|
||||
Key
|
||||
</span>
|
||||
<div className='font-mono text-xs break-all'>
|
||||
{balance.hashed_key}
|
||||
</div>
|
||||
{balance.isChild && (
|
||||
<span className='rounded bg-blue-100 px-1.5 py-0.5 text-[10px] font-bold text-blue-700 uppercase'>
|
||||
Child
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div className='grid grid-cols-2 gap-3'>
|
||||
@@ -295,13 +241,7 @@ export function TemporaryBalances({
|
||||
Balance
|
||||
</div>
|
||||
<div className='truncate font-mono text-sm'>
|
||||
{balance.isChild ? (
|
||||
<span className='text-muted-foreground text-xs italic'>
|
||||
(Uses Parent)
|
||||
</span>
|
||||
) : (
|
||||
formatBalance(balance.balance)
|
||||
)}
|
||||
{formatBalance(balance.balance)}
|
||||
</div>
|
||||
</div>
|
||||
<div className='space-y-1'>
|
||||
|
||||
@@ -926,7 +926,6 @@ export const TemporaryBalanceSchema = z.object({
|
||||
total_requests: z.number(),
|
||||
refund_address: z.string().nullable(),
|
||||
key_expiry_time: z.number().nullable(),
|
||||
parent_key_hash: z.string().nullable().optional(),
|
||||
});
|
||||
|
||||
export type TemporaryBalance = z.infer<typeof TemporaryBalanceSchema>;
|
||||
|
||||
Reference in New Issue
Block a user