Compare commits

...

16 Commits

Author SHA1 Message Date
Shroominic
8d3c064b29 ruff fix 2026-01-23 09:47:48 +08:00
shroominic
4ac96ade5f Merge branch 'v0.3.0' into provider-fallback 2026-01-23 09:45:14 +08:00
shroominic
180a469399 Merge pull request #304 from Routstr/introduce-child-key
Introduce child key
2026-01-23 09:35:53 +08:00
9qeklajc
db021866d8 fmt u 2026-01-10 21:42:21 +01:00
9qeklajc
d4339287be chore: add type annotations to example and test files 2026-01-10 21:35:16 +01:00
9qeklajc
88bcc0edcb fmt 2026-01-10 21:34:48 +01:00
9qeklajc
917a4d32b1 ui: visualize parent-child relationship in balances page 2026-01-10 20:10:59 +01:00
9qeklajc
daf17f51ab fix: move child-key route before catch-all and fix indentation 2026-01-10 20:10:31 +01:00
9qeklajc
fc042c768c support child keys mapped to parent balance 2026-01-10 19:53:27 +01:00
Shroominic
cb36189db3 fmt 2026-01-09 17:04:09 +08:00
Shroominic
d2487f42b0 fix tests 2026-01-09 17:02:35 +08:00
Shroominic
5255fce7b2 todo comment 2026-01-09 16:55:24 +08:00
Shroominic
ee668ee93b handle specific errors eg deactivate ppq on insufficient balance 2026-01-09 16:49:03 +08:00
Shroominic
cf8b990fc7 fix specify error codes retry 2026-01-09 16:10:14 +08:00
Shroominic
78dd74845b fix merge 2026-01-07 17:44:28 +01:00
Shroominic
e7f4c98475 ranked provider fallback on upstream errors 2026-01-07 17:39:05 +01:00
17 changed files with 860 additions and 326 deletions

View File

@@ -0,0 +1,45 @@
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.")

View File

@@ -0,0 +1,42 @@
"""
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")

View File

@@ -84,93 +84,26 @@ def get_provider_penalty(provider: "BaseUpstreamProvider") -> float:
return penalty
def should_prefer_model(
candidate_model: "Model",
candidate_provider: "BaseUpstreamProvider",
current_model: "Model",
current_provider: "BaseUpstreamProvider",
alias: str,
) -> bool:
"""Determine if candidate model should replace current model for an alias.
This is the core decision function for model prioritization. It considers:
1. Alias matching quality (exact match vs. canonical slug match)
2. Model cost (lower is better)
3. Provider penalties (e.g., slight preference against OpenRouter)
Args:
candidate_model: The new model being considered
candidate_provider: Provider offering the candidate model
current_model: The currently selected model for this alias
current_provider: Provider offering the current model
alias: The model alias being mapped
Returns:
True if candidate should replace current, False otherwise
"""
def get_base_model_id(model_id: str) -> str:
"""Get base model ID by removing provider prefix."""
return model_id.split("/", 1)[1] if "/" in model_id else model_id
def alias_priority(model: "Model") -> int:
"""Rank how strong the mapping of alias->model is.
Highest priority when alias exactly equals the model ID without provider prefix.
Next when alias equals canonical slug without prefix. Otherwise lowest.
"""
model_base = get_base_model_id(model.id)
if model_base == alias:
return 3
if model.canonical_slug:
canonical_base = get_base_model_id(model.canonical_slug)
if canonical_base == alias:
return 2
return 1
candidate_alias_priority = alias_priority(candidate_model)
current_alias_priority = alias_priority(current_model)
# If candidate has better alias match, prefer it regardless of cost
if candidate_alias_priority > current_alias_priority:
return True
# If current has better alias match, keep it regardless of cost
if current_alias_priority > candidate_alias_priority:
return False
# Same alias priority - compare costs
candidate_cost = calculate_model_cost_score(candidate_model)
current_cost = calculate_model_cost_score(current_model)
# Apply provider penalties
candidate_adjusted = candidate_cost * get_provider_penalty(candidate_provider)
current_adjusted = current_cost * get_provider_penalty(current_provider)
# Prefer lower adjusted cost
should_replace = candidate_adjusted < current_adjusted
return should_replace
def create_model_mappings(
upstreams: list["BaseUpstreamProvider"],
overrides_by_id: dict[str, tuple],
disabled_model_ids: set[str],
) -> tuple[dict[str, "Model"], dict[str, "BaseUpstreamProvider"], dict[str, "Model"]]:
) -> tuple[
dict[str, "Model"], dict[str, list["BaseUpstreamProvider"]], dict[str, "Model"]
]:
"""Create optimal model mappings based on cost and provider preferences.
This is the main entry point for the algorithm. It processes all upstream providers
and creates three mappings based on cost optimization:
1. model_instances: alias -> Model (all model aliases mapped to their Model objects)
2. provider_map: alias -> UpstreamProvider (which provider to use for each alias)
2. provider_map: alias -> List[UpstreamProvider] (sorted list of providers for each alias)
3. unique_models: base_id -> Model (unique models without provider prefixes)
The algorithm:
- Processes non-OpenRouter providers first (they're typically cheaper)
- Then processes OpenRouter models (they can still win if cheaper)
- For each model alias, uses should_prefer_model() to select the best provider
- For each model alias, collects all candidates and sorts them by priority and cost.
Args:
upstreams: List of all upstream provider instances
@@ -183,8 +116,7 @@ def create_model_mappings(
from .payment.models import _row_to_model
from .upstream.helpers import resolve_model_alias
model_instances: dict[str, "Model"] = {}
provider_map: dict[str, "BaseUpstreamProvider"] = {}
candidates: dict[str, list[tuple["Model", "BaseUpstreamProvider"]]] = {}
unique_models: dict[str, "Model"] = {}
# Separate OpenRouter from other providers
@@ -202,24 +134,14 @@ def create_model_mappings(
"""Get base model ID by removing provider prefix."""
return model_id.split("/", 1)[1] if "/" in model_id else model_id
def _maybe_set_alias(
def _add_candidate(
alias: str, model: "Model", provider: "BaseUpstreamProvider"
) -> None:
"""Set alias to model/provider if not set or if new model is preferred."""
"""Add candidate model/provider for an alias."""
alias_lower = alias.lower()
existing_model = model_instances.get(alias_lower)
if not existing_model:
# No existing mapping, set it
model_instances[alias_lower] = model
provider_map[alias_lower] = provider
else:
# Check if candidate should replace existing
existing_provider = provider_map[alias_lower]
if should_prefer_model(
model, provider, existing_model, existing_provider, alias
):
model_instances[alias_lower] = model
provider_map[alias_lower] = provider
if alias_lower not in candidates:
candidates[alias_lower] = []
candidates[alias_lower].append((model, provider))
def process_provider_models(
upstream: "BaseUpstreamProvider", is_openrouter: bool = False
@@ -266,21 +188,55 @@ def create_model_mappings(
# Try to set each alias
for alias in aliases:
_maybe_set_alias(alias, model_to_use, upstream)
_add_candidate(alias, model_to_use, upstream)
# Process non-OpenRouter providers first (they're typically cheaper)
# Process non-OpenRouter providers first
for upstream in other_upstreams:
process_provider_models(upstream, is_openrouter=False)
# Process OpenRouter last - models only win if they're cheaper or better matched
# Process OpenRouter last
if openrouter:
process_provider_models(openrouter, is_openrouter=True)
# Log provider distribution
# Sort candidates and build final maps
model_instances: dict[str, "Model"] = {}
provider_map: dict[str, list["BaseUpstreamProvider"]] = {}
def alias_priority(model: "Model", alias: str) -> int:
"""Rank how strong the mapping of alias->model is."""
model_base = get_base_model_id(model.id)
if model_base == alias:
return 3
if model.canonical_slug:
canonical_base = get_base_model_id(model.canonical_slug)
if canonical_base == alias:
return 2
return 1
for alias, items in candidates.items():
# Sort key: (priority DESC, cost ASC)
# Using negative cost for DESC sort overall to keep high priority first
def sort_key(item: tuple["Model", "BaseUpstreamProvider"]) -> tuple[int, float]:
model, provider = item
priority = alias_priority(model, alias)
cost = calculate_model_cost_score(model)
penalty = get_provider_penalty(provider)
adjusted_cost = cost * penalty
return (priority, -adjusted_cost)
items.sort(key=sort_key, reverse=True)
best_model, best_provider = items[0]
model_instances[alias] = best_model
provider_map[alias] = [p for _, p in items]
# Log provider distribution (using top provider for stats)
provider_counts: dict[str, int] = {}
for provider in provider_map.values():
provider_name = getattr(provider, "upstream_name", "unknown")
provider_counts[provider_name] = provider_counts.get(provider_name, 0) + 1
for providers in provider_map.values():
if providers:
provider = providers[0]
provider_name = getattr(provider, "upstream_name", "unknown")
provider_counts[provider_name] = provider_counts.get(provider_name, 0) + 1
logger.debug(
f"Updated model mappings with ({len(unique_models)} unique models and {len(model_instances)} aliases)",

View File

@@ -286,30 +286,55 @@ 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] + "...",
"current_balance": key.balance,
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"current_balance": billing_key.balance,
"required_cost": cost_per_request,
"sufficient_balance": key.balance >= cost_per_request,
"sufficient_balance": billing_key.balance >= cost_per_request,
},
)
if key.total_balance < cost_per_request:
if billing_key.total_balance < cost_per_request:
logger.warning(
"Insufficient balance for request",
extra={
"key_hash": key.hashed_key[:8] + "...",
"balance": key.balance,
"reserved_balance": key.reserved_balance,
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"balance": billing_key.balance,
"reserved_balance": billing_key.reserved_balance,
"required": cost_per_request,
"shortfall": cost_per_request - key.total_balance,
"shortfall": cost_per_request - billing_key.total_balance,
},
)
@@ -317,7 +342,7 @@ async def pay_for_request(
status_code=402,
detail={
"error": {
"message": f"Insufficient balance: {cost_per_request} mSats required. {key.total_balance} available. (reserved: {key.reserved_balance})",
"message": f"Insufficient balance: {cost_per_request} mSats required. {billing_key.total_balance} available. (reserved: {billing_key.reserved_balance})",
"type": "insufficient_quota",
"code": "insufficient_balance",
}
@@ -328,15 +353,16 @@ 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": key.balance,
"balance_before": billing_key.balance,
},
)
# Charge the base cost for the request atomically to avoid race conditions
stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == key.hashed_key)
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
.where(col(ApiKey.balance) - col(ApiKey.reserved_balance) >= cost_per_request)
.values(
reserved_balance=col(ApiKey.reserved_balance) + cost_per_request,
@@ -344,6 +370,16 @@ 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:
@@ -351,8 +387,9 @@ 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": key.balance,
"current_balance": billing_key.balance,
},
)
@@ -361,23 +398,26 @@ async def pay_for_request(
status_code=402,
detail={
"error": {
"message": f"Insufficient balance: {cost_per_request} mSats required. {key.balance} available.",
"message": f"Insufficient balance: {cost_per_request} mSats required. {billing_key.balance} available.",
"type": "insufficient_quota",
"code": "insufficient_balance",
}
},
)
await session.refresh(key)
await session.refresh(billing_key)
if billing_key.hashed_key != key.hashed_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": key.balance,
"total_spent": key.total_spent,
"total_requests": key.total_requests,
"new_balance": billing_key.balance,
"total_spent": billing_key.total_spent,
"total_requests": billing_key.total_requests,
},
)
@@ -387,9 +427,11 @@ 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) == key.hashed_key)
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
.values(
reserved_balance=col(ApiKey.reserved_balance) - cost_per_request,
total_requests=col(ApiKey.total_requests) - 1,
@@ -397,27 +439,40 @@ 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": key.reserved_balance,
"current_reserved_balance": billing_key.reserved_balance,
},
)
raise HTTPException(
status_code=402,
detail={
"error": {
"message": f"failed to revert request payment: {cost_per_request} mSats required. {key.balance} available.",
"message": f"failed to revert request payment: {cost_per_request} mSats required. {billing_key.balance} available.",
"type": "payment_error",
"code": "payment_error",
}
},
)
await session.refresh(key)
await session.refresh(billing_key)
if billing_key.hashed_key != key.hashed_key:
await session.refresh(key)
async def adjust_payment_for_tokens(
@@ -428,15 +483,17 @@ 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": key.balance,
"current_balance": billing_key.balance,
"has_usage": "usage" in response_data,
},
)
@@ -446,8 +503,10 @@ async def adjust_payment_for_tokens(
try:
release_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == key.hashed_key)
.values(reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost)
.where(col(ApiKey.hashed_key) == billing_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()
@@ -455,13 +514,18 @@ 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] + "..."},
extra={
"error": str(e),
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
},
)
match await calculate_cost(response_data, deducted_max_cost, session):
@@ -470,6 +534,7 @@ 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,
},
@@ -477,7 +542,7 @@ async def adjust_payment_for_tokens(
# Finalize by releasing reservation and charging max cost
finalize_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == key.hashed_key)
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
.values(
reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost,
balance=col(ApiKey.balance) - cost.total_msats,
@@ -485,27 +550,41 @@ 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": key.reserved_balance,
"current_reserved_balance": billing_key.reserved_balance,
"total_cost": cost.total_msats,
"model": model,
},
)
await release_reservation_only()
else:
await session.refresh(key)
await session.refresh(billing_key)
if billing_key.hashed_key != key.hashed_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": key.balance,
"new_balance": billing_key.balance,
"model": model,
},
)
@@ -521,6 +600,7 @@ 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,
@@ -533,11 +613,15 @@ async def adjust_payment_for_tokens(
if cost_difference == 0:
logger.debug(
"Finalizing with exact reserved cost",
extra={"key_hash": key.hashed_key[:8] + "...", "model": model},
extra={
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"model": model,
},
)
finalize_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == key.hashed_key)
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
.values(
reserved_balance=col(ApiKey.reserved_balance)
- deducted_max_cost,
@@ -546,8 +630,20 @@ 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(key)
await session.refresh(billing_key)
if billing_key.hashed_key != key.hashed_key:
await session.refresh(key)
return cost.dict()
# this should never happen why do we handle this???
@@ -557,16 +653,17 @@ 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": key.balance,
"sufficient_balance": key.balance >= cost_difference,
"current_balance": billing_key.balance,
"sufficient_balance": billing_key.balance >= cost_difference,
"model": model,
},
)
finalize_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == key.hashed_key)
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
.values(
reserved_balance=col(ApiKey.reserved_balance)
- deducted_max_cost,
@@ -575,18 +672,31 @@ 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(key)
await session.refresh(billing_key)
if billing_key.hashed_key != key.hashed_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": key.balance,
"new_balance": billing_key.balance,
"model": model,
},
)
@@ -595,6 +705,7 @@ 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,
},
@@ -607,15 +718,16 @@ 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": key.balance,
"current_balance": billing_key.balance,
"model": model,
},
)
refund_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == key.hashed_key)
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
.values(
reserved_balance=col(ApiKey.reserved_balance)
- deducted_max_cost,
@@ -624,6 +736,16 @@ 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:
@@ -631,8 +753,9 @@ 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": key.reserved_balance,
"current_reserved_balance": billing_key.reserved_balance,
"total_cost": total_cost_msats,
"model": model,
},
@@ -640,14 +763,17 @@ async def adjust_payment_for_tokens(
await release_reservation_only()
else:
cost.total_msats = total_cost_msats
await session.refresh(key)
await session.refresh(billing_key)
if billing_key.hashed_key != key.hashed_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": key.balance,
"new_balance": billing_key.balance,
"final_cost": cost.total_msats,
"model": model,
},

View File

@@ -32,16 +32,30 @@ async def get_key_from_header(
)
# 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)) -> dict:
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": key.balance,
"reserved": key.reserved_balance,
"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)
# TODO: Implement POST /v1/wallet/create endpoint
# This endpoint should accept:
# - cashu_token (required): The eCash token to deposit
@@ -66,12 +80,11 @@ async def create_balance(
@router.get("/info")
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,
}
async def wallet_info(
key: ApiKey = Depends(get_key_from_header),
session: AsyncSession = Depends(get_session),
) -> dict:
return await get_balance_info(key, session)
class TopupRequest(BaseModel):
@@ -85,6 +98,10 @@ 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:
@@ -94,7 +111,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, key, session)
amount_msats = await credit_balance(cashu_token, billing_key, session)
except ValueError as e:
error_msg = str(e)
if "already spent" in error_msg.lower():
@@ -155,6 +172,12 @@ 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
if key.refund_currency == "sat":
@@ -228,6 +251,53 @@ 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"],

View File

@@ -124,7 +124,7 @@ async def partial_apikeys(request: Request) -> str:
rows = "".join(
[
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>"
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>"
for key in api_keys
]
)
@@ -158,6 +158,7 @@ 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
]

View File

@@ -47,6 +47,9 @@ 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:

View File

@@ -6,6 +6,15 @@ from .logging import get_logger
logger = get_logger(__name__)
class UpstreamError(Exception):
"""Exception raised when an upstream provider fails."""
def __init__(self, message: str, status_code: int = 502):
self.message = message
self.status_code = status_code
super().__init__(message)
async def http_exception_handler(request: Request, exc: Exception) -> JSONResponse:
"""Handle HTTP exceptions and include request ID in response."""
request_id = getattr(request.state, "request_id", "unknown")

View File

@@ -13,10 +13,7 @@ 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

View File

@@ -52,6 +52,7 @@ 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(

View File

@@ -16,6 +16,7 @@ from .core.db import (
create_session,
get_session,
)
from .core.exceptions import UpstreamError
from .payment.helpers import (
calculate_discounted_max_cost,
check_token_balance,
@@ -31,7 +32,9 @@ proxy_router = APIRouter()
_upstreams: list[BaseUpstreamProvider] = []
_model_instances: dict[str, Model] = {} # All aliases -> Model
_provider_map: dict[str, BaseUpstreamProvider] = {} # All aliases -> Provider
_provider_map: dict[
str, list[BaseUpstreamProvider]
] = {} # All aliases -> List[Provider]
_unique_models: dict[str, Model] = {} # Unique model.id -> Model (no duplicates)
@@ -68,8 +71,8 @@ def get_model_instance(model_id: str) -> Model | None:
return _model_instances.get(model_id.lower())
def get_provider_for_model(model_id: str) -> BaseUpstreamProvider | None:
"""Get UpstreamProvider for model ID from global cache."""
def get_provider_for_model(model_id: str) -> list[BaseUpstreamProvider] | None:
"""Get UpstreamProvider list for model ID from global cache."""
return _provider_map.get(model_id.lower())
@@ -154,8 +157,8 @@ async def proxy(
"invalid_model", f"Model '{model_id}' not found", 400, request=request
)
upstream = get_provider_for_model(model_id)
if not upstream:
upstreams = get_provider_for_model(model_id)
if not upstreams:
return create_error_response(
"invalid_model",
f"No provider found for model '{model_id}'",
@@ -163,6 +166,10 @@ async def proxy(
request=request,
)
# todo figure out cost calculation since fallback provider is usually not the same price
# Use first provider for initial checks/cost calculation
# primary_upstream = upstreams[0]
_max_cost_for_model = await get_max_cost_for_model(
model=model_id, session=session, model_obj=model_obj
)
@@ -172,14 +179,31 @@ async def proxy(
check_token_balance(headers, request_body_dict, max_cost_for_model)
if x_cashu := headers.get("x-cashu", None):
if is_responses_api:
return await upstream.handle_x_cashu_responses(
request, x_cashu, path, max_cost_for_model, model_obj
)
else:
return await upstream.handle_x_cashu(
request, x_cashu, path, max_cost_for_model, model_obj
)
last_error = None
for i, upstream in enumerate(upstreams):
try:
if is_responses_api:
return await upstream.handle_x_cashu_responses(
request, x_cashu, path, max_cost_for_model, model_obj
)
else:
return await upstream.handle_x_cashu(
request, x_cashu, path, max_cost_for_model, model_obj
)
except UpstreamError as e:
logger.warning(
f"Upstream {upstream.provider_type} failed (x-cashu): {e}"
)
if i == len(upstreams) - 1:
last_error = e
continue
return create_error_response(
"upstream_error",
str(last_error) if last_error else "All upstreams failed",
502,
request=request,
)
elif auth := headers.get("authorization", None):
key = await get_bearer_token_key(headers, path, session, auth)
@@ -194,56 +218,152 @@ async def proxy(
)
logger.debug("Processing unauthenticated GET request", extra={"path": path})
headers = upstream.prepare_headers(dict(request.headers))
return await upstream.forward_get_request(request, path, headers)
last_error_response = None
for i, upstream in enumerate(upstreams):
try:
headers = upstream.prepare_headers(dict(request.headers))
response = await upstream.forward_get_request(request, path, headers)
if response.status_code in [502, 429] and i < len(upstreams) - 1:
error_message = ""
try:
if hasattr(response, "body"):
body_bytes = response.body
data = json.loads(body_bytes)
if "error" in data:
error_data = data["error"]
if isinstance(error_data, dict):
error_message = error_data.get("message", "")
elif isinstance(error_data, str):
error_message = error_data
except Exception:
pass
await upstream.on_upstream_error_redirect(
response.status_code, error_message
)
logger.warning(
f"Upstream {upstream.provider_type} returned {response.status_code} (GET), trying next provider",
extra={
"status_code": response.status_code,
"upstream": upstream.provider_type,
},
)
continue
return response
except UpstreamError as e:
logger.warning(f"Upstream {upstream.provider_type} failed (GET): {e}")
if i == len(upstreams) - 1:
last_error_response = create_error_response(
"upstream_error", str(e), 502, request=request
)
continue
return last_error_response or create_error_response(
"upstream_error", "All upstreams failed", 502, request=request
)
if request_body_dict:
await pay_for_request(key, max_cost_for_model, session)
headers = upstream.prepare_headers(dict(request.headers))
for i, upstream in enumerate(upstreams):
headers = upstream.prepare_headers(dict(request.headers))
if is_responses_api:
response = await upstream.forward_responses_request(
request,
path,
headers,
request_body,
key,
max_cost_for_model,
session,
model_obj,
)
else:
response = await upstream.forward_request(
request,
path,
headers,
request_body,
key,
max_cost_for_model,
session,
model_obj,
)
try:
if is_responses_api:
response = await upstream.forward_responses_request(
request,
path,
headers,
request_body,
key,
max_cost_for_model,
session,
model_obj,
)
else:
response = await upstream.forward_request(
request,
path,
headers,
request_body,
key,
max_cost_for_model,
session,
model_obj,
)
if response.status_code != 200:
await revert_pay_for_request(key, session, max_cost_for_model)
logger.warning(
"Upstream request failed, revert payment",
extra={
"status_code": response.status_code,
"path": path,
"key_hash": key.hashed_key[:8] + "...",
"key_balance": key.balance,
"max_cost_for_model": max_cost_for_model,
"upstream_headers": response.headers
if hasattr(response, "headers")
else None,
},
)
# Return the mapped error response generated earlier rather than masking with 502
return response
if response.status_code != 200:
# Check if we should retry (502 Upstream Error or 429 Rate Limit)
should_retry = response.status_code in [502, 429, 400, 401, 403, 404]
if should_retry and i < len(upstreams) - 1:
error_message = ""
try:
if hasattr(response, "body"):
body_bytes = response.body
data = json.loads(body_bytes)
if "error" in data:
error_data = data["error"]
if isinstance(error_data, dict):
error_message = error_data.get("message", "")
elif isinstance(error_data, str):
error_message = error_data
except Exception:
pass
return response
await upstream.on_upstream_error_redirect(
response.status_code, error_message
)
logger.warning(
f"Upstream {upstream.provider_type} returned {response.status_code}, trying next provider",
extra={
"status_code": response.status_code,
"upstream": upstream.provider_type,
},
)
continue
# 4xx error (user error), or other non-retryable error, or last provider failed
await revert_pay_for_request(key, session, max_cost_for_model)
logger.warning(
"Upstream request failed, revert payment",
extra={
"status_code": response.status_code,
"path": path,
"key_hash": key.hashed_key[:8] + "...",
"key_balance": key.balance,
"max_cost_for_model": max_cost_for_model,
"upstream_headers": response.headers
if hasattr(response, "headers")
else None,
},
)
return response
return response
except UpstreamError as e:
logger.warning(
f"Upstream {upstream.provider_type} failed: {e}",
extra={"retry": i < len(upstreams) - 1},
)
# If this was the last provider
if i == len(upstreams) - 1:
await revert_pay_for_request(key, session, max_cost_for_model)
return create_error_response(
"upstream_error", str(e), 502, request=request
)
# Otherwise loop continues to next provider
continue
# Should not be reached given logic above
return create_error_response(
"upstream_error", "All upstreams failed", 502, request=request
)
async def get_bearer_token_key(

View File

@@ -15,6 +15,7 @@ from pydantic import BaseModel
from ..auth import adjust_payment_for_tokens
from ..core import get_logger
from ..core.db import ApiKey, AsyncSession, create_session
from ..core.exceptions import UpstreamError
if TYPE_CHECKING:
from ..core.db import UpstreamProviderRow
@@ -340,6 +341,20 @@ class BaseUpstreamProvider:
message = preview[:500]
return message, upstream_code
async def on_upstream_error_redirect(
self, status_code: int, error_message: str
) -> None:
"""Hook called when the proxy redirects to another provider due to an error.
Subclasses can implement this to perform actions like disabling the provider
if it's out of balance.
Args:
status_code: The HTTP status code returned by the upstream
error_message: The error message extracted from the upstream response
"""
pass
async def map_upstream_error_response(
self, request: Request, path: str, upstream_response: httpx.Response
) -> Response:
@@ -1148,6 +1163,14 @@ class BaseUpstreamProvider:
)
if response.status_code != 200:
if response.status_code >= 500:
await response.aclose()
await client.aclose()
raise UpstreamError(
f"Upstream returned status {response.status_code}",
status_code=response.status_code,
)
try:
mapped_error = await self.map_upstream_error_response(
request, path, response
@@ -1238,6 +1261,9 @@ class BaseUpstreamProvider:
background=background_tasks,
)
except UpstreamError:
raise
except httpx.RequestError as exc:
await client.aclose()
error_type = type(exc).__name__
@@ -1265,9 +1291,7 @@ class BaseUpstreamProvider:
else:
error_message = f"Error connecting to upstream service: {error_type}"
return create_error_response(
"upstream_error", error_message, 502, request=request
)
raise UpstreamError(error_message, status_code=502)
except Exception as exc:
await client.aclose()
@@ -1380,6 +1404,14 @@ class BaseUpstreamProvider:
)
if response.status_code != 200:
if response.status_code >= 500:
await response.aclose()
await client.aclose()
raise UpstreamError(
f"Upstream returned status {response.status_code}",
status_code=response.status_code,
)
try:
mapped_error = await self.map_upstream_error_response(
request, path, response
@@ -1447,6 +1479,9 @@ class BaseUpstreamProvider:
background=background_tasks,
)
except UpstreamError:
raise
except httpx.RequestError as exc:
await client.aclose()
error_type = type(exc).__name__
@@ -1474,9 +1509,7 @@ class BaseUpstreamProvider:
else:
error_message = f"Error connecting to upstream service: {error_type}"
return create_error_response(
"upstream_error", error_message, 502, request=request
)
raise UpstreamError(error_message, status_code=502)
except Exception as exc:
await client.aclose()

View File

@@ -196,6 +196,37 @@ class PPQAIUpstreamProvider(BaseUpstreamProvider):
)
return []
async def on_upstream_error_redirect(
self, status_code: int, error_message: str
) -> None:
if "insufficient balance" in error_message.lower():
logger.warning(
f"Disabling PPQ.AI provider ({self.base_url}) due to insufficient balance",
extra={"error": error_message},
)
from sqlmodel import select
from ..core.db import UpstreamProviderRow, create_session
async with create_session() as session:
statement = select(UpstreamProviderRow).where(
UpstreamProviderRow.base_url == self.base_url
and UpstreamProviderRow.api_key == self.api_key
)
result = await session.exec(statement)
provider = result.first()
if provider:
provider.enabled = False
session.add(provider)
await session.commit()
# Trigger re-initialization of providers
# Import here to avoid circular dependency
from ..proxy import reinitialize_upstreams
await reinitialize_upstreams()
async def create_account(self) -> dict[str, object]:
"""Create a new PPQ.AI account.

View File

@@ -0,0 +1,137 @@
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)

View File

@@ -10,7 +10,6 @@ os.environ["UPSTREAM_API_KEY"] = "test"
from routstr.algorithm import ( # noqa: E402
calculate_model_cost_score,
get_provider_penalty,
should_prefer_model,
)
from routstr.payment.models import Architecture, Model, Pricing # noqa: E402
@@ -100,100 +99,3 @@ def test_get_provider_penalty_openrouter() -> None:
provider = create_test_provider("openrouter", "https://openrouter.ai/api/v1")
penalty = get_provider_penalty(provider)
assert penalty == 1.001
def test_should_prefer_model_cheaper_wins() -> None:
"""Test that cheaper model is preferred."""
cheap_model = create_test_model("cheap", prompt_price=0.001, completion_price=0.002)
expensive_model = create_test_model(
"expensive", prompt_price=0.03, completion_price=0.06
)
provider1 = create_test_provider("provider1")
provider2 = create_test_provider("provider2")
# Cheaper model should win
assert should_prefer_model(
cheap_model, provider1, expensive_model, provider2, "test-alias"
)
# More expensive model should not win
assert not should_prefer_model(
expensive_model, provider2, cheap_model, provider1, "test-alias"
)
def test_should_prefer_model_exact_match_wins() -> None:
"""Test that exact alias match beats cheaper price."""
# Make model IDs match the alias differently
exact_match = create_test_model(
"test-model", prompt_price=0.03, completion_price=0.06
)
no_match = create_test_model(
"other-model", prompt_price=0.001, completion_price=0.002
)
provider1 = create_test_provider("provider1")
provider2 = create_test_provider("provider2")
# Exact match should win even though it's more expensive
assert should_prefer_model(
exact_match, provider1, no_match, provider2, "test-model"
)
def test_should_prefer_model_openrouter_slight_penalty() -> None:
"""Test that OpenRouter has slight penalty compared to other providers."""
model1 = create_test_model("model1", prompt_price=0.001, completion_price=0.002)
model2 = create_test_model("model2", prompt_price=0.001, completion_price=0.002)
regular_provider = create_test_provider("regular", "http://provider.com")
openrouter_provider = create_test_provider(
"openrouter", "https://openrouter.ai/api/v1"
)
# Regular provider should be preferred over OpenRouter at same cost
assert should_prefer_model(
model1, regular_provider, model2, openrouter_provider, "test-alias"
)
# OpenRouter should not replace regular provider at same cost
assert not should_prefer_model(
model2, openrouter_provider, model1, regular_provider, "test-alias"
)
def test_should_prefer_model_openrouter_can_win_if_cheaper() -> None:
"""Test that OpenRouter can still win if significantly cheaper."""
cheap_model = create_test_model(
"cheap", prompt_price=0.0001, completion_price=0.0002
)
expensive_model = create_test_model(
"expensive", prompt_price=0.03, completion_price=0.06
)
regular_provider = create_test_provider("regular", "http://provider.com")
openrouter_provider = create_test_provider(
"openrouter", "https://openrouter.ai/api/v1"
)
# OpenRouter should win if it's much cheaper (even with penalty)
assert should_prefer_model(
cheap_model,
openrouter_provider,
expensive_model,
regular_provider,
"test-alias",
)
def test_should_prefer_model_same_cost_first_wins() -> None:
"""Test that when costs are identical, current model is kept."""
model1 = create_test_model("model1", prompt_price=0.001, completion_price=0.002)
model2 = create_test_model("model2", prompt_price=0.001, completion_price=0.002)
provider1 = create_test_provider("provider1")
provider2 = create_test_provider("provider2")
# When costs are equal, should not replace
assert not should_prefer_model(model2, provider2, model1, provider1, "test-alias")

View File

@@ -60,7 +60,11 @@ export function TemporaryBalances({
let totalRequests = 0;
balances.forEach((balance) => {
totalBalance += balance.balance || 0;
// 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;
}
totalSpent += balance.total_spent || 0;
totalRequests += balance.total_requests || 0;
});
@@ -72,6 +76,34 @@ 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'>
@@ -182,22 +214,37 @@ export function TemporaryBalances({
<div className='text-right'>Expiry Time</div>
</div>
{filteredData.length > 0 ? (
filteredData.map((balance, index) => (
{hierarchicalData.length > 0 ? (
hierarchicalData.map((balance, index) => (
<div
key={index}
className={cn(
'hover:bg-muted/50 border-t p-3 text-sm transition-colors',
balance.balance === 0 && 'opacity-60'
balance.balance === 0 &&
!balance.isChild &&
'opacity-60',
balance.isChild &&
'ml-4 border-l-2 border-l-blue-200 bg-blue-50/30'
)}
>
{/* Desktop Layout */}
<div className='hidden grid-cols-6 gap-2 md:grid'>
<div className='max-w-32 truncate font-mono text-xs break-all'>
<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>
)}
{balance.hashed_key}
</div>
<div className='text-right font-mono'>
{formatBalance(balance.balance)}
{balance.isChild ? (
<span className='text-muted-foreground italic'>
(Parent)
</span>
) : (
formatBalance(balance.balance)
)}
</div>
<div className='text-right font-mono'>
{formatBalance(balance.total_spent)}
@@ -226,13 +273,20 @@ export function TemporaryBalances({
{/* Mobile Layout */}
<div className='space-y-3 md:hidden'>
<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 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>
{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'>
@@ -241,7 +295,13 @@ export function TemporaryBalances({
Balance
</div>
<div className='truncate font-mono text-sm'>
{formatBalance(balance.balance)}
{balance.isChild ? (
<span className='text-muted-foreground text-xs italic'>
(Uses Parent)
</span>
) : (
formatBalance(balance.balance)
)}
</div>
</div>
<div className='space-y-1'>

View File

@@ -926,6 +926,7 @@ 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>;