laantungir/routstr-core
1
0
Fork 0
mirror of https://github.com/Routstr/routstr-core.git synced 2026-10-06 04:38:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'v0.4.0' into add-routstr-provider

Browse Source
This commit is contained in:
9qeklajc
2026-03-06 23:08:30 +01:00
parent 97daddb95a 5788d63892
commit 66975ee271
12 changed files with 606 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files
routstr/algorithm.py
+98
View File
@@ -118,6 +118,13 @@ def create_model_mappings(
candidates: dict[str, list[tuple["Model", "BaseUpstreamProvider"]]] = {}
unique_models: dict[str, "Model"] = {}
seen_model_provider: set[tuple[str, str]] = set()
providers_by_db_id: dict[int, "BaseUpstreamProvider"] = {}
for upstream in upstreams:
db_id = getattr(upstream, "db_id", None)
if isinstance(db_id, int):
providers_by_db_id[db_id] = upstream
# Separate OpenRouter from other providers
openrouter: "BaseUpstreamProvider" | None = None
@@ -134,6 +141,16 @@ 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 get_provider_identity(upstream: "BaseUpstreamProvider") -> str:
"""Get a stable provider identity used for deduplication."""
db_id = getattr(upstream, "db_id", None)
if isinstance(db_id, int):
return f"db:{db_id}"
provider_type = str(getattr(upstream, "provider_type", "") or "").lower()
base_url = str(getattr(upstream, "base_url", "") or "").lower()
return f"{provider_type}|{base_url}"
def _add_candidate(
alias: str, model: "Model", provider: "BaseUpstreamProvider"
) -> None:
@@ -148,6 +165,7 @@ def create_model_mappings(
) -> None:
"""Process all models from a given provider."""
upstream_prefix = getattr(upstream, "upstream_name", None)
provider_key = get_provider_identity(upstream)
for model in upstream.get_cached_models():
if not model.enabled or model.id in disabled_model_ids:
@@ -189,6 +207,7 @@ def create_model_mappings(
# Try to set each alias
for alias in aliases:
_add_candidate(alias, model_to_use, upstream)
seen_model_provider.add((model_to_use.id.lower(), provider_key))
# Process non-OpenRouter providers first
for upstream in other_upstreams:
@@ -198,6 +217,85 @@ def create_model_mappings(
if openrouter:
process_provider_models(openrouter, is_openrouter=True)
# Include enabled DB overrides even when provider discovery misses models.
# This is important for deployment-based providers like Azure.
for model_id, override_data in overrides_by_id.items():
if model_id in disabled_model_ids:
continue
override_row, provider_fee = override_data
upstream_provider_id = getattr(override_row, "upstream_provider_id", None)
if not isinstance(upstream_provider_id, int):
continue
upstream_for_override = providers_by_db_id.get(upstream_provider_id)
if upstream_for_override is None:
continue
provider_key = get_provider_identity(upstream_for_override)
dedupe_key = (model_id.lower(), provider_key)
if dedupe_key in seen_model_provider:
continue
try:
model_to_use = _row_to_model(
override_row, apply_provider_fee=True, provider_fee=provider_fee
)
except Exception as exc:
logger.warning(
"Skipping invalid model override while building model mappings",
extra={
"model_id": model_id,
"upstream_provider_id": upstream_provider_id,
"error": str(exc),
"error_type": type(exc).__name__,
},
)
continue
if not model_to_use.enabled:
continue
base_id = get_base_model_id(model_to_use.id)
is_openrouter = (
getattr(upstream_for_override, "base_url", "")
== "https://openrouter.ai/api/v1"
)
if not is_openrouter or base_id not in unique_models:
unique_model = model_to_use.copy(
update={
"id": base_id,
"upstream_provider_id": upstream_for_override.provider_type,
}
)
unique_models[base_id] = unique_model
try:
aliases = resolve_model_alias(
model_to_use.id,
model_to_use.canonical_slug,
alias_ids=model_to_use.alias_ids,
)
except Exception as exc:
logger.warning(
"Skipping model aliases for invalid override model",
extra={
"model_id": model_id,
"upstream_provider_id": upstream_provider_id,
"error": str(exc),
"error_type": type(exc).__name__,
},
)
continue
upstream_prefix = getattr(upstream_for_override, "upstream_name", None)
if upstream_prefix and "/" not in model_to_use.id:
prefixed_id = f"{upstream_prefix}/{model_to_use.id}"
if prefixed_id not in aliases:
aliases.append(prefixed_id)
for alias in aliases:
_add_candidate(alias, model_to_use, upstream_for_override)
seen_model_provider.add(dedupe_key)
# Sort candidates and build final maps
model_instances: dict[str, "Model"] = {}
provider_map: dict[str, list["BaseUpstreamProvider"]] = {}
routstr/auth.py
+30 -22
View File
@@ -590,12 +590,15 @@ async def pay_for_request(
async def revert_pay_for_request(
key: ApiKey, session: AsyncSession, cost_per_request: int
) -> None:
) -> bool:
"""Revert a previously reserved payment. Returns True if revert succeeded,
False if the reservation was already released (prevents negative reserved_balance)."""
billing_key = await get_billing_key(key, session)
stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
.where(col(ApiKey.reserved_balance) >= cost_per_request)
.values(
reserved_balance=col(ApiKey.reserved_balance) - cost_per_request,
total_requests=col(ApiKey.total_requests) - 1,
@@ -609,6 +612,7 @@ async def revert_pay_for_request(
child_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == key.hashed_key)
.where(col(ApiKey.reserved_balance) >= cost_per_request)
.values(
total_requests=col(ApiKey.total_requests) - 1,
reserved_balance=col(ApiKey.reserved_balance) - cost_per_request,
@@ -618,8 +622,8 @@ async def revert_pay_for_request(
await session.commit()
if result.rowcount == 0:
logger.error(
"Failed to revert payment - insufficient reserved balance",
logger.warning(
"Revert skipped - reservation already released (no-op to prevent negative reserved_balance)",
extra={
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
@@ -627,19 +631,11 @@ async def revert_pay_for_request(
"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. {billing_key.balance} available.",
"type": "payment_error",
"code": "payment_error",
}
},
)
return False
await session.refresh(billing_key)
if billing_key.hashed_key != key.hashed_key:
await session.refresh(key)
return True
async def adjust_payment_for_tokens(
@@ -671,17 +667,19 @@ async def adjust_payment_for_tokens(
release_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
.where(col(ApiKey.reserved_balance) >= deducted_max_cost)
.values(
reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost
)
)
await session.exec(release_stmt) # type: ignore[call-overload]
result = await session.exec(release_stmt) # type: ignore[call-overload]
# Also release on child key if it's different
if billing_key.hashed_key != key.hashed_key:
child_release_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == key.hashed_key)
.where(col(ApiKey.reserved_balance) >= deducted_max_cost)
.values(
reserved_balance=col(ApiKey.reserved_balance)
- deducted_max_cost
@@ -690,14 +688,24 @@ async def adjust_payment_for_tokens(
await session.exec(child_release_stmt) # type: ignore[call-overload]
await session.commit()
logger.warning(
"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,
},
)
if result.rowcount == 0: # type: ignore[union-attr]
logger.warning(
"Release reservation skipped - already released (no-op to prevent negative reserved_balance)",
extra={
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"deducted_max_cost": deducted_max_cost,
},
)
else:
logger.warning(
"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",
routstr/balance.py
+3 -2
View File
@@ -205,8 +205,9 @@ async def refund_wallet_endpoint(
key: ApiKey = await validate_bearer_key(bearer_value, session)
if cached := await _refund_cache_get(bearer_value):
return cached
if key.total_balance <= 0:
if cached := await _refund_cache_get(bearer_value):
return cached
if key.parent_key_hash:
raise HTTPException(
routstr/core/admin.py
+1 -3
View File
@@ -747,9 +747,7 @@ async def get_provider_models(provider_id: int) -> dict[str, object]:
)
db_model_ids = {model.id for model in db_models}
filtered_remote_models = [
m for m in upstream_models if m.name not in db_model_ids
]
filtered_remote_models = [m for m in upstream_models if m.id not in db_model_ids]
return {
"provider": {
routstr/proxy.py
+5 -1
View File
@@ -296,9 +296,13 @@ async def proxy(
session,
model_obj,
)
except UpstreamError:
# Let the outer UpstreamError handler manage retry/revert
raise
except Exception as e:
# Unexpected error (not an upstream failure) — revert and propagate
logger.error(
"Upstream request failed, ensuring payment is reverted",
"Unexpected error in upstream request, reverting payment",
extra={
"error": str(e),
"error_type": type(e).__name__,
routstr/upstream/azure.py
+45 -11
View File
@@ -4,6 +4,7 @@ from .base import BaseUpstreamProvider
if TYPE_CHECKING:
from ..core.db import UpstreamProviderRow
from ..payment.models import Model
class AzureUpstreamProvider(BaseUpstreamProvider):
@@ -58,19 +59,52 @@ class AzureUpstreamProvider(BaseUpstreamProvider):
"platform_url": cls.platform_url,
}
def prepare_headers(self, request_headers: dict) -> dict:
"""Prepare headers for Azure OpenAI, adding api-key."""
headers = super().prepare_headers(request_headers)
if self.api_key:
headers["api-key"] = self.api_key
headers.pop("Authorization", None)
headers.pop("authorization", None)
return headers
def prepare_params(
self, path: str, query_params: Mapping[str, str] | None
) -> Mapping[str, str]:
"""Prepare query parameters for Azure OpenAI, adding API version.
Args:
path: Request path
query_params: Original query parameters from the client
Returns:
Query parameters dict with Azure API version added for chat completions
"""
"""Prepare query parameters for Azure OpenAI, adding API version."""
params = dict(query_params or {})
if path.endswith("chat/completions"):
params["api-version"] = self.api_version
version = (self.api_version or "").replace("\ufeff", "").strip()
if not version or version.lower() == "v1":
version = "2024-02-15-preview"
params["api-version"] = version
return params
def normalize_request_path(
self, path: str, model_obj: "Model | None" = None
) -> str:
"""Build Azure deployment-specific request path."""
clean_path = super().normalize_request_path(path, model_obj).lstrip("/")
if model_obj is None:
return clean_path
deployment_id = getattr(
model_obj, "canonical_slug", None
) or self.transform_model_name(model_obj.id)
deployment_id = deployment_id.split("/")[-1]
return f"openai/deployments/{deployment_id}/{clean_path}"
def get_request_base_url(
self, path: str, model_obj: "Model | None" = None
) -> str:
"""Use endpoint root, stripping accidental /openai/v1 suffix if present."""
base_url = self.base_url.rstrip("/")
marker = "/openai/v1"
if marker in base_url:
base_url = base_url.split(marker, 1)[0].rstrip("/")
return base_url
def transform_model_name(self, model_id: str) -> str:
"""Extract deployment name from model ID."""
if "/" in model_id:
return model_id.split("/")[-1]
return model_id
routstr/upstream/base.py
+32 -32
View File
@@ -13,7 +13,7 @@ from fastapi.responses import Response, StreamingResponse
from pydantic import BaseModel
from sqlmodel import select
from ..auth import adjust_payment_for_tokens, revert_pay_for_request
from ..auth import adjust_payment_for_tokens
from ..core import get_logger
from ..core.db import ApiKey, AsyncSession, UpstreamProviderRow, create_session
from ..core.exceptions import UpstreamError
@@ -197,6 +197,23 @@ class BaseUpstreamProvider:
"""
return model_id
def normalize_request_path(self, path: str, model_obj: Model | None = None) -> str:
"""Normalize request path before forwarding to upstream."""
if path.startswith("v1/"):
return path.replace("v1/", "", 1)
return path
def get_request_base_url(
self, path: str, model_obj: Model | None = None
) -> str:
"""Get upstream base URL used when building forwarding URL."""
return self.base_url.rstrip("/")
def build_request_url(self, path: str, model_obj: Model | None = None) -> str:
"""Build full upstream URL from normalized path."""
clean_path = path.lstrip("/")
return f"{self.get_request_base_url(path, model_obj)}/{clean_path}"
def prepare_responses_request_body(
self, body: bytes | None, model_obj: Model
) -> bytes | None:
@@ -1101,10 +1118,8 @@ class BaseUpstreamProvider:
Returns:
Response or StreamingResponse from upstream with cost tracking
"""
if path.startswith("v1/"):
path = path.replace("v1/", "")
url = f"{self.base_url}/{path}"
path = self.normalize_request_path(path, model_obj)
url = self.build_request_url(path, model_obj)
transformed_body = self.prepare_request_body(request_body, model_obj)
@@ -1279,8 +1294,7 @@ class BaseUpstreamProvider:
},
)
await revert_pay_for_request(key, session, max_cost_for_model)
# Don't revert here — proxy.py owns payment revert to avoid double-revert
if isinstance(exc, httpx.ConnectError):
error_message = "Unable to connect to upstream service"
elif isinstance(exc, httpx.TimeoutException):
@@ -1310,13 +1324,9 @@ class BaseUpstreamProvider:
},
)
await revert_pay_for_request(key, session, max_cost_for_model)
return create_error_response(
"internal_error",
"An unexpected server error occurred",
500,
request=request,
# Don't revert here — proxy.py owns payment revert to avoid double-revert
raise UpstreamError(
"An unexpected server error occurred", status_code=500
)
async def forward_responses_request(
@@ -1345,11 +1355,8 @@ class BaseUpstreamProvider:
Returns:
Response or StreamingResponse from upstream with cost tracking
"""
# Remove v1/ prefix if present for Responses API
if path.startswith("v1/"):
path = path.replace("v1/", "")
url = f"{self.base_url}/{path}"
path = self.normalize_request_path(path, model_obj)
url = self.build_request_url(path, model_obj)
transformed_body = self.prepare_responses_request_body(request_body, model_obj)
@@ -1501,8 +1508,7 @@ class BaseUpstreamProvider:
},
)
await revert_pay_for_request(key, session, max_cost_for_model)
# Don't revert here — proxy.py owns payment revert to avoid double-revert
if isinstance(exc, httpx.ConnectError):
error_message = "Unable to connect to upstream service"
elif isinstance(exc, httpx.TimeoutException):
@@ -1532,13 +1538,9 @@ class BaseUpstreamProvider:
},
)
await revert_pay_for_request(key, session, max_cost_for_model)
return create_error_response(
"internal_error",
"An unexpected server error occurred",
500,
request=request,
# Don't revert here — proxy.py owns payment revert to avoid double-revert
raise UpstreamError(
"An unexpected server error occurred", status_code=500
)
async def forward_get_request(
@@ -1557,10 +1559,8 @@ class BaseUpstreamProvider:
Returns:
StreamingResponse from upstream
"""
if path.startswith("v1/"):
path = path.replace("v1/", "")
url = f"{self.base_url}/{path}"
path = self.normalize_request_path(path)
url = self.build_request_url(path)
logger.info(
"Forwarding GET request to upstream",
routstr/upstream/helpers.py
+3
View File
@@ -205,6 +205,9 @@ async def init_upstreams() -> list[BaseUpstreamProvider]:
provider = _instantiate_provider(provider_row)
if provider:
# Keep provider DB id on runtime instance so model mapping can
# bind DB overrides to the correct upstream.
setattr(provider, "db_id", provider_row.id)
await provider.refresh_models_cache()
logger.debug(
f"Initialized {provider_row.provider_type} provider",
routstr/upstream/ollama.py
+6 -35
View File
@@ -3,13 +3,11 @@ from __future__ import annotations
from typing import TYPE_CHECKING
import httpx
from fastapi import Request
from fastapi.responses import Response, StreamingResponse
from .base import BaseUpstreamProvider
if TYPE_CHECKING:
from ..core.db import ApiKey, AsyncSession, UpstreamProviderRow
from ..core.db import UpstreamProviderRow
from ..payment.models import Model
from ..core.logging import get_logger
@@ -67,38 +65,11 @@ class OllamaUpstreamProvider(BaseUpstreamProvider):
"""Strip 'ollama/' prefix for Ollama API compatibility."""
return model_id.removeprefix("ollama/")
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:
"""Override to use OpenAI-compatible endpoint for proxy requests."""
if path.startswith("v1/"):
path = path.replace("v1/", "")
original_base_url = self.base_url
self.base_url = f"{self.base_url}/v1"
try:
result = await super().forward_request(
request,
path,
headers,
request_body,
key,
max_cost_for_model,
session,
model_obj,
)
return result
finally:
self.base_url = original_base_url
def get_request_base_url(
self, path: str, model_obj: Model | None = None
) -> str:
"""Route proxy traffic through Ollama's OpenAI-compatible /v1 endpoint."""
return f"{self.base_url.rstrip('/')}/v1"
async def fetch_models(self) -> list[Model]:
"""Fetch models from Ollama API using /api/tags endpoint."""
tests/integration/test_reserved_balance_negative.py
+209 -12
View File
@@ -133,13 +133,16 @@ async def test_reserved_balance_with_successful_requests(
@pytest.mark.asyncio
async def test_insufficient_reserved_balance_for_revert(
async def test_revert_with_zero_reserved_balance_is_noop(
integration_session: AsyncSession,
) -> None:
"""Test revert_pay_for_request behavior with insufficient reserved balance."""
"""Test that revert_pay_for_request is a no-op when reserved_balance is 0.
Previously this would drive reserved_balance negative. With the floor guard,
it should return False and leave reserved_balance at 0.
"""
from routstr.auth import revert_pay_for_request
# Create key with zero reserved balance
unique_key = f"test_revert_key_{uuid.uuid4().hex[:8]}"
test_key = ApiKey(
hashed_key=unique_key,
@@ -149,17 +152,211 @@ async def test_insufficient_reserved_balance_for_revert(
integration_session.add(test_key)
await integration_session.commit()
# Try to revert more than available
# Note: Current implementation allows reserved_balance to go negative
await revert_pay_for_request(test_key, integration_session, 100)
# Try to revert more than available — should be a no-op
result = await revert_pay_for_request(test_key, integration_session, 100)
# Refresh to get updated values
await integration_session.refresh(test_key)
# Current implementation allows negative reserved balance
assert test_key.reserved_balance == -100, (
f"Expected reserved_balance to be -100, got: {test_key.reserved_balance}"
assert result is False, "Revert should return False when reservation already released"
assert test_key.reserved_balance == 0, (
f"Reserved balance should remain 0, got: {test_key.reserved_balance}"
)
assert test_key.total_requests == -1, (
f"Expected total_requests to be -1, got: {test_key.total_requests}"
assert test_key.total_requests == 0, (
f"Total requests should remain 0, got: {test_key.total_requests}"
)
@pytest.mark.asyncio
async def test_revert_with_sufficient_reserved_balance_succeeds(
integration_session: AsyncSession,
) -> None:
"""Test that revert_pay_for_request works correctly when there is enough reserved balance."""
from routstr.auth import revert_pay_for_request
unique_key = f"test_revert_ok_{uuid.uuid4().hex[:8]}"
test_key = ApiKey(
hashed_key=unique_key,
balance=5000,
reserved_balance=500,
total_requests=3,
)
integration_session.add(test_key)
await integration_session.commit()
result = await revert_pay_for_request(test_key, integration_session, 500)
await integration_session.refresh(test_key)
assert result is True, "Revert should return True on success"
assert test_key.reserved_balance == 0, (
f"Reserved balance should be 0, got: {test_key.reserved_balance}"
)
assert test_key.total_requests == 2, (
f"Total requests should be 2, got: {test_key.total_requests}"
)
assert test_key.balance == 5000, "Balance should not change on revert"
@pytest.mark.asyncio
async def test_revert_partial_reserved_balance_is_noop(
integration_session: AsyncSession,
) -> None:
"""Test that reverting more than the current reserved_balance is a no-op."""
from routstr.auth import revert_pay_for_request
unique_key = f"test_revert_partial_{uuid.uuid4().hex[:8]}"
test_key = ApiKey(
hashed_key=unique_key,
balance=5000,
reserved_balance=50,
total_requests=1,
)
integration_session.add(test_key)
await integration_session.commit()
# Try to revert 500 when only 50 is reserved — should be no-op
result = await revert_pay_for_request(test_key, integration_session, 500)
await integration_session.refresh(test_key)
assert result is False, "Revert should fail when cost > reserved_balance"
assert test_key.reserved_balance == 50, (
f"Reserved balance should stay at 50, got: {test_key.reserved_balance}"
)
assert test_key.total_requests == 1, (
f"Total requests should stay at 1, got: {test_key.total_requests}"
)
@pytest.mark.asyncio
async def test_double_revert_prevented(
integration_session: AsyncSession,
) -> None:
"""Test that calling revert twice doesn't drive reserved_balance negative.
This simulates the double-revert scenario where both upstream/base.py
and proxy.py attempt to revert the same reservation.
"""
from routstr.auth import revert_pay_for_request
unique_key = f"test_double_revert_{uuid.uuid4().hex[:8]}"
test_key = ApiKey(
hashed_key=unique_key,
balance=10000,
reserved_balance=500,
total_requests=5,
)
integration_session.add(test_key)
await integration_session.commit()
# First revert — should succeed
result1 = await revert_pay_for_request(test_key, integration_session, 500)
await integration_session.refresh(test_key)
assert result1 is True
assert test_key.reserved_balance == 0
assert test_key.total_requests == 4
# Second revert of the same amount — should be no-op
result2 = await revert_pay_for_request(test_key, integration_session, 500)
await integration_session.refresh(test_key)
assert result2 is False, "Second revert should be a no-op"
assert test_key.reserved_balance == 0, (
f"Reserved balance should stay 0, got: {test_key.reserved_balance}"
)
assert test_key.total_requests == 4, (
f"Total requests should stay 4, got: {test_key.total_requests}"
)
@pytest.mark.asyncio
async def test_sequential_reverts_never_go_negative(
integration_session: AsyncSession,
) -> None:
"""Test that multiple reverts don't cause negative reserved_balance.
Simulates the double-revert scenario where multiple code paths
attempt to revert the same reservation.
"""
from routstr.auth import revert_pay_for_request
unique_key = f"test_multi_revert_{uuid.uuid4().hex[:8]}"
test_key = ApiKey(
hashed_key=unique_key,
balance=10000,
reserved_balance=500,
total_requests=5,
)
integration_session.add(test_key)
await integration_session.commit()
# Run 5 sequential reverts for the same 500 reservation
results = []
for _ in range(5):
r = await revert_pay_for_request(test_key, integration_session, 500)
results.append(r)
await integration_session.refresh(test_key)
# Exactly one should succeed, rest should be no-ops
success_count = sum(1 for r in results if r is True)
assert success_count == 1, (
f"Exactly one revert should succeed, got {success_count} successes"
)
assert test_key.reserved_balance == 0, (
f"Reserved balance should be 0, got: {test_key.reserved_balance}"
)
assert test_key.reserved_balance >= 0, (
f"Reserved balance went negative: {test_key.reserved_balance}"
)
@pytest.mark.asyncio
async def test_child_key_revert_floor_guard(
integration_session: AsyncSession,
) -> None:
"""Test that child key reserved_balance also has floor guard on revert."""
from routstr.auth import revert_pay_for_request
parent_key_hash = f"test_parent_{uuid.uuid4().hex[:8]}"
child_key_hash = f"test_child_{uuid.uuid4().hex[:8]}"
parent_key = ApiKey(
hashed_key=parent_key_hash,
balance=10000,
reserved_balance=500,
total_requests=3,
)
child_key = ApiKey(
hashed_key=child_key_hash,
balance=0,
reserved_balance=500,
total_requests=3,
parent_key_hash=parent_key_hash,
)
integration_session.add(parent_key)
integration_session.add(child_key)
await integration_session.commit()
# First revert succeeds
result1 = await revert_pay_for_request(child_key, integration_session, 500)
await integration_session.refresh(parent_key)
await integration_session.refresh(child_key)
assert result1 is True
assert parent_key.reserved_balance == 0
assert child_key.reserved_balance == 0
# Second revert is a no-op for both parent and child
result2 = await revert_pay_for_request(child_key, integration_session, 500)
await integration_session.refresh(parent_key)
await integration_session.refresh(child_key)
assert result2 is False
assert parent_key.reserved_balance == 0, (
f"Parent reserved_balance should stay 0, got: {parent_key.reserved_balance}"
)
assert child_key.reserved_balance == 0, (
f"Child reserved_balance should stay 0, got: {child_key.reserved_balance}"
)
tests/unit/test_algorithm.py
+92 -1
View File
@@ -1,14 +1,18 @@
"""Tests for the model prioritization algorithm."""
import os
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
# Set required env vars before importing
os.environ["UPSTREAM_BASE_URL"] = "http://test"
os.environ["UPSTREAM_API_KEY"] = "test"
from routstr.algorithm import ( # noqa: E402
calculate_model_cost_score,
create_model_mappings,
get_provider_penalty,
)
from routstr.payment.models import Architecture, Model, Pricing # noqa: E402
@@ -45,11 +49,21 @@ def create_test_model(
)
def create_test_provider(name: str, base_url: str = "http://test.com") -> Mock:
def create_test_provider(
name: str,
base_url: str = "http://test.com",
*,
db_id: int | None = None,
models: list[Model] | None = None,
upstream_name: str | None = None,
) -> Mock:
"""Helper to create a test provider mock."""
provider = Mock()
provider.provider_type = name
provider.base_url = base_url
provider.db_id = db_id
provider.upstream_name = upstream_name or name
provider.get_cached_models.return_value = models or []
return provider
@@ -99,3 +113,80 @@ 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_create_model_mappings_includes_db_override_for_missing_cached_model(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Model overrides should still map when provider discovery misses the model."""
provider = create_test_provider(
"azure",
"https://example.openai.azure.com/openai/v1",
db_id=7,
models=[],
)
override_model = create_test_model("azure/gpt-4o")
override_model.canonical_slug = "azure-deployment"
def fake_row_to_model(*args, **kwargs) -> Model: # type: ignore[no-untyped-def]
return override_model
monkeypatch.setattr("routstr.payment.models._row_to_model", fake_row_to_model)
override_row = SimpleNamespace(id="azure/gpt-4o", upstream_provider_id=7, enabled=True)
model_instances, provider_map, unique_models = create_model_mappings(
upstreams=[provider],
overrides_by_id={"azure/gpt-4o": (override_row, 1.01)},
disabled_model_ids=set(),
)
assert "azure/gpt-4o" in model_instances
assert provider_map["azure/gpt-4o"] == [provider]
assert "gpt-4o" in unique_models
def test_create_model_mappings_dedupes_with_provider_identity_not_provider_type(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Different provider instances of same type should both survive dedupe."""
provider_a_model = create_test_model(
"azure/gpt-4o", prompt_price=0.01, completion_price=0.01
)
provider_a = create_test_provider(
"azure",
"https://a.openai.azure.com/openai/v1",
db_id=1,
models=[provider_a_model],
upstream_name="azure-a",
)
provider_b = create_test_provider(
"azure",
"https://b.openai.azure.com/openai/v1",
db_id=2,
models=[],
upstream_name="azure-b",
)
override_model = create_test_model(
"azure/gpt-4o", prompt_price=0.001, completion_price=0.001
)
override_model.canonical_slug = "azure-b-deployment"
def fake_row_to_model(*args, **kwargs) -> Model: # type: ignore[no-untyped-def]
return override_model
monkeypatch.setattr("routstr.payment.models._row_to_model", fake_row_to_model)
override_row = SimpleNamespace(id="azure/gpt-4o", upstream_provider_id=2, enabled=True)
_, provider_map, _ = create_model_mappings(
upstreams=[provider_a, provider_b],
overrides_by_id={"azure/gpt-4o": (override_row, 1.01)},
disabled_model_ids=set(),
)
providers_for_alias = provider_map["azure/gpt-4o"]
assert provider_a in providers_for_alias
assert provider_b in providers_for_alias
assert len(providers_for_alias) == 2
tests/unit/test_upstream_azure.py
+82
View File
@@ -0,0 +1,82 @@
"""Tests for Azure upstream provider request normalization."""
from routstr.payment.models import Architecture, Model, Pricing
from routstr.upstream.azure import AzureUpstreamProvider
def create_test_model(model_id: str, canonical_slug: str | None = None) -> Model:
return Model(
id=model_id,
name=model_id,
created=0,
description="test",
context_length=8192,
architecture=Architecture(
modality="text",
input_modalities=["text"],
output_modalities=["text"],
tokenizer="test",
instruct_type=None,
),
pricing=Pricing(
prompt=0.001,
completion=0.001,
request=0.0,
image=0.0,
web_search=0.0,
internal_reasoning=0.0,
),
canonical_slug=canonical_slug,
)
def test_prepare_headers_uses_azure_api_key_header() -> None:
provider = AzureUpstreamProvider(
base_url="https://example.openai.azure.com",
api_key="azure-key",
api_version="2024-02-15-preview",
)
headers = provider.prepare_headers({"Authorization": "Bearer user-token"})
assert headers["api-key"] == "azure-key"
assert "Authorization" not in headers
assert "authorization" not in headers
def test_prepare_params_normalizes_azure_api_version() -> None:
provider = AzureUpstreamProvider(
base_url="https://example.openai.azure.com",
api_key="azure-key",
api_version="\ufeff v1 ",
)
params = provider.prepare_params("chat/completions", {})
assert params["api-version"] == "2024-02-15-preview"
def test_normalize_request_path_includes_deployment_id() -> None:
provider = AzureUpstreamProvider(
base_url="https://example.openai.azure.com/openai/v1",
api_key="azure-key",
api_version="2024-02-15-preview",
)
model = create_test_model("azure/gpt-4o", canonical_slug="deploy-gpt4o")
path = provider.normalize_request_path("v1/chat/completions", model)
assert path == "openai/deployments/deploy-gpt4o/chat/completions"
def test_get_request_base_url_strips_openai_v1_suffix() -> None:
provider = AzureUpstreamProvider(
base_url="https://example.openai.azure.com/openai/v1",
api_key="azure-key",
api_version="2024-02-15-preview",
)
model = create_test_model("azure/gpt-4o")
base_url = provider.get_request_base_url("chat/completions", model)
assert base_url == "https://example.openai.azure.com"