mirror of
https://github.com/Routstr/routstr-core.git
synced 2026-07-31 15:56:14 +00:00
Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Shroominic
|
dce9f37b0b |
+5
-41
@@ -1,7 +1,5 @@
|
||||
"""Model prioritization algorithm for selecting cheapest upstream providers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from .core.logging import get_logger
|
||||
@@ -159,21 +157,15 @@ 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"],
|
||||
dict[str, list["BaseUpstreamProvider"]],
|
||||
]:
|
||||
) -> tuple[dict[str, "Model"], dict[str, "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 (the BEST provider to use for each alias)
|
||||
2. provider_map: alias -> UpstreamProvider (which provider to use for each alias)
|
||||
3. unique_models: base_id -> Model (unique models without provider prefixes)
|
||||
4. provider_candidates_map: alias -> list[UpstreamProvider] (all providers offering the model, sorted by preference)
|
||||
|
||||
The algorithm:
|
||||
- Processes non-OpenRouter providers first (they're typically cheaper)
|
||||
@@ -186,7 +178,7 @@ def create_model_mappings(
|
||||
disabled_model_ids: Set of model IDs that should be excluded
|
||||
|
||||
Returns:
|
||||
Tuple of (model_instances, provider_map, unique_models, provider_candidates_map)
|
||||
Tuple of (model_instances, provider_map, unique_models)
|
||||
"""
|
||||
from .payment.models import _row_to_model
|
||||
from .upstream.helpers import resolve_model_alias
|
||||
@@ -194,10 +186,9 @@ def create_model_mappings(
|
||||
model_instances: dict[str, "Model"] = {}
|
||||
provider_map: dict[str, "BaseUpstreamProvider"] = {}
|
||||
unique_models: dict[str, "Model"] = {}
|
||||
provider_candidates_map: dict[str, list["BaseUpstreamProvider"]] = {}
|
||||
|
||||
# Separate OpenRouter from other providers
|
||||
openrouter: BaseUpstreamProvider | None = None
|
||||
openrouter: "BaseUpstreamProvider" | None = None
|
||||
other_upstreams: list["BaseUpstreamProvider"] = []
|
||||
|
||||
for upstream in upstreams:
|
||||
@@ -216,13 +207,6 @@ def create_model_mappings(
|
||||
) -> None:
|
||||
"""Set alias to model/provider if not set or if new model is preferred."""
|
||||
alias_lower = alias.lower()
|
||||
|
||||
# Add to candidates list, to be used later as fallback
|
||||
if alias_lower not in provider_candidates_map:
|
||||
provider_candidates_map[alias_lower] = [provider]
|
||||
else:
|
||||
provider_candidates_map[alias_lower].append(provider)
|
||||
|
||||
existing_model = model_instances.get(alias_lower)
|
||||
if not existing_model:
|
||||
# No existing mapping, set it
|
||||
@@ -292,26 +276,6 @@ def create_model_mappings(
|
||||
if openrouter:
|
||||
process_provider_models(openrouter, is_openrouter=True)
|
||||
|
||||
# Sort and filter provider candidates for each alias using provider_map as reference
|
||||
# We only keep entries that have more than one provider.
|
||||
final_candidates_map: dict[str, list["BaseUpstreamProvider"]] = {}
|
||||
for alias_lower, best_provider in provider_map.items():
|
||||
candidates = provider_candidates_map.get(alias_lower, [])
|
||||
# Remove duplicates
|
||||
unique_candidates = []
|
||||
seen = set()
|
||||
for c in candidates:
|
||||
if c not in seen:
|
||||
unique_candidates.append(c)
|
||||
seen.add(c)
|
||||
|
||||
if len(unique_candidates) > 1:
|
||||
# Keep the best one at the front, others follow.
|
||||
if best_provider in unique_candidates:
|
||||
unique_candidates.remove(best_provider)
|
||||
unique_candidates.insert(0, best_provider)
|
||||
final_candidates_map[alias_lower] = unique_candidates
|
||||
|
||||
# Log provider distribution
|
||||
provider_counts: dict[str, int] = {}
|
||||
for provider in provider_map.values():
|
||||
@@ -323,4 +287,4 @@ def create_model_mappings(
|
||||
extra={"provider_distribution": provider_counts},
|
||||
)
|
||||
|
||||
return model_instances, provider_map, unique_models, provider_candidates_map
|
||||
return model_instances, provider_map, unique_models
|
||||
|
||||
+60
-193
@@ -1,7 +1,6 @@
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request
|
||||
from fastapi.responses import Response, StreamingResponse
|
||||
from sqlmodel import select
|
||||
@@ -27,7 +26,7 @@ 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, recieve_token
|
||||
from .wallet import deserialize_token_from_string
|
||||
|
||||
logger = get_logger(__name__)
|
||||
proxy_router = APIRouter()
|
||||
@@ -35,9 +34,6 @@ proxy_router = APIRouter()
|
||||
_upstreams: list[BaseUpstreamProvider] = []
|
||||
_model_instances: dict[str, Model] = {} # All aliases -> Model
|
||||
_provider_map: dict[str, BaseUpstreamProvider] = {} # All aliases -> Provider
|
||||
_provider_candidates_map: dict[
|
||||
str, list[BaseUpstreamProvider]
|
||||
] = {} # All aliases -> [Providers]
|
||||
_unique_models: dict[str, Model] = {} # Unique model.id -> Model (no duplicates)
|
||||
|
||||
|
||||
@@ -79,21 +75,6 @@ def get_provider_for_model(model_id: str) -> BaseUpstreamProvider | None:
|
||||
return _provider_map.get(model_id.lower())
|
||||
|
||||
|
||||
def get_providers_for_model(model_id: str) -> list[BaseUpstreamProvider]:
|
||||
"""Get list of prioritized UpstreamProviders for model ID from global cache.
|
||||
|
||||
If multiple providers are available, returns the sorted list.
|
||||
Otherwise, returns a list containing only the single best provider.
|
||||
"""
|
||||
candidates = _provider_candidates_map.get(model_id.lower(), [])
|
||||
if candidates:
|
||||
return candidates
|
||||
|
||||
# Fallback to the single best provider if no multi-provider candidates exist
|
||||
best = get_provider_for_model(model_id)
|
||||
return [best] if best else []
|
||||
|
||||
|
||||
def get_unique_models() -> list[Model]:
|
||||
"""Get list of unique models (no duplicates from aliases)."""
|
||||
return list(_unique_models.values())
|
||||
@@ -103,7 +84,7 @@ async def refresh_model_maps() -> None:
|
||||
"""Refresh global model and provider maps using the cost-based algorithm."""
|
||||
from sqlalchemy.orm import selectinload
|
||||
|
||||
global _model_instances, _provider_map, _unique_models, _provider_candidates_map
|
||||
global _model_instances, _provider_map, _unique_models
|
||||
|
||||
async with create_session() as session:
|
||||
# Fetch all providers with their models in a single logical operation
|
||||
@@ -123,12 +104,10 @@ async def refresh_model_maps() -> None:
|
||||
else:
|
||||
disabled_model_ids.add(model.id)
|
||||
|
||||
_model_instances, _provider_map, _unique_models, _provider_candidates_map = (
|
||||
create_model_mappings(
|
||||
upstreams=_upstreams,
|
||||
overrides_by_id=overrides_by_id,
|
||||
disabled_model_ids=disabled_model_ids,
|
||||
)
|
||||
_model_instances, _provider_map, _unique_models = create_model_mappings(
|
||||
upstreams=_upstreams,
|
||||
overrides_by_id=overrides_by_id,
|
||||
disabled_model_ids=disabled_model_ids,
|
||||
)
|
||||
|
||||
|
||||
@@ -193,8 +172,8 @@ async def proxy(
|
||||
"invalid_model", f"Model '{model_id}' not found", 400, request=request
|
||||
)
|
||||
|
||||
upstreams = get_providers_for_model(model_id)
|
||||
if not upstreams:
|
||||
upstream = get_provider_for_model(model_id)
|
||||
if not upstream:
|
||||
return create_error_response(
|
||||
"invalid_model",
|
||||
f"No provider found for model '{model_id}'",
|
||||
@@ -211,80 +190,14 @@ async def proxy(
|
||||
check_token_balance(headers, request_body_dict, max_cost_for_model)
|
||||
|
||||
if x_cashu := headers.get("x-cashu", None):
|
||||
# Redeem token once before trying any providers
|
||||
amount, unit, mint = await recieve_token(x_cashu)
|
||||
|
||||
# Fallback for X-Cashu payments
|
||||
last_exception = None
|
||||
for i, upstream in enumerate(upstreams):
|
||||
try:
|
||||
# Prepare headers for this specific upstream
|
||||
upstream_headers = upstream.prepare_headers(dict(request.headers))
|
||||
|
||||
if is_responses_api:
|
||||
return await upstream.forward_x_cashu_responses_request(
|
||||
request,
|
||||
path,
|
||||
upstream_headers,
|
||||
amount,
|
||||
unit,
|
||||
max_cost_for_model,
|
||||
model_obj,
|
||||
mint,
|
||||
)
|
||||
else:
|
||||
return await upstream.forward_x_cashu_request(
|
||||
request,
|
||||
path,
|
||||
upstream_headers,
|
||||
amount,
|
||||
unit,
|
||||
max_cost_for_model,
|
||||
model_obj,
|
||||
mint,
|
||||
)
|
||||
except (httpx.TimeoutException, httpx.ConnectError) as e:
|
||||
logger.warning(
|
||||
f"Upstream provider {i + 1}/{len(upstreams)} ({upstream.provider_type}) timed out, trying fallback",
|
||||
extra={
|
||||
"model": model_id,
|
||||
"error": str(e),
|
||||
"attempt": i + 1,
|
||||
},
|
||||
)
|
||||
last_exception = e
|
||||
continue
|
||||
|
||||
# If we get here, all providers failed
|
||||
# Since the token was already redeemed, we must issue a refund
|
||||
logger.error(
|
||||
"All providers failed for X-Cashu request, issuing emergency refund",
|
||||
extra={"amount": amount, "unit": unit, "mint": mint},
|
||||
)
|
||||
|
||||
# Try to use the first provider's refund mechanism
|
||||
refund_token = await upstreams[0].send_refund(amount - 60, unit, mint)
|
||||
|
||||
error_message = "All upstream providers timed out"
|
||||
if isinstance(last_exception, httpx.ConnectError):
|
||||
error_message = "Unable to connect to any upstream service"
|
||||
|
||||
error_response = Response(
|
||||
content=json.dumps(
|
||||
{
|
||||
"error": {
|
||||
"message": error_message,
|
||||
"type": "upstream_error",
|
||||
"code": 504,
|
||||
"refund_token": refund_token,
|
||||
}
|
||||
}
|
||||
),
|
||||
status_code=504,
|
||||
media_type="application/json",
|
||||
)
|
||||
error_response.headers["X-Cashu"] = refund_token
|
||||
return error_response
|
||||
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
|
||||
)
|
||||
|
||||
elif auth := headers.get("authorization", None):
|
||||
key = await get_bearer_token_key(headers, path, session, auth)
|
||||
@@ -299,102 +212,56 @@ async def proxy(
|
||||
)
|
||||
|
||||
logger.debug("Processing unauthenticated GET request", extra={"path": path})
|
||||
|
||||
# Try fallback for GET requests too
|
||||
last_exception = None
|
||||
for i, upstream in enumerate(upstreams):
|
||||
try:
|
||||
headers = upstream.prepare_headers(dict(request.headers))
|
||||
return await upstream.forward_get_request(request, path, headers)
|
||||
except (httpx.TimeoutException, httpx.ConnectError) as e:
|
||||
logger.warning(
|
||||
f"Upstream GET provider {i + 1}/{len(upstreams)} ({upstream.provider_type}) timed out, trying fallback",
|
||||
extra={"path": path, "error": str(e)},
|
||||
)
|
||||
last_exception = e
|
||||
continue
|
||||
|
||||
error_message = "Upstream service request timed out"
|
||||
if isinstance(last_exception, httpx.ConnectError):
|
||||
error_message = "Unable to connect to upstream service"
|
||||
|
||||
return create_error_response(
|
||||
"upstream_error", error_message, 502, request=request
|
||||
)
|
||||
headers = upstream.prepare_headers(dict(request.headers))
|
||||
return await upstream.forward_get_request(request, path, headers)
|
||||
|
||||
if request_body_dict:
|
||||
await pay_for_request(key, max_cost_for_model, session)
|
||||
|
||||
# Fallback for API Key payments
|
||||
last_exception = None
|
||||
for i, upstream in enumerate(upstreams):
|
||||
try:
|
||||
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,
|
||||
)
|
||||
headers = upstream.prepare_headers(dict(request.headers))
|
||||
|
||||
if response.status_code != 200:
|
||||
# If it's a 429 (rate limit) or 503 (service unavailable), we might also want to fallback
|
||||
if response.status_code in (429, 503, 502) and i < len(upstreams) - 1:
|
||||
logger.warning(
|
||||
f"Upstream provider {i + 1}/{len(upstreams)} returned {response.status_code}, trying fallback",
|
||||
extra={"model": model_id, "status": response.status_code},
|
||||
)
|
||||
continue
|
||||
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,
|
||||
)
|
||||
|
||||
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
|
||||
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
|
||||
|
||||
return response
|
||||
|
||||
except (httpx.TimeoutException, httpx.ConnectError) as e:
|
||||
logger.warning(
|
||||
f"Upstream provider {i + 1}/{len(upstreams)} ({upstream.provider_type}) timed out, trying fallback",
|
||||
extra={"model": model_id, "error": str(e)},
|
||||
)
|
||||
last_exception = e
|
||||
continue
|
||||
|
||||
# All providers failed with timeout/connect error
|
||||
await revert_pay_for_request(key, session, max_cost_for_model)
|
||||
error_message = "Upstream service request timed out"
|
||||
if isinstance(last_exception, httpx.ConnectError):
|
||||
error_message = "Unable to connect to upstream service"
|
||||
|
||||
return create_error_response("upstream_error", error_message, 502, request=request)
|
||||
return response
|
||||
|
||||
|
||||
async def get_bearer_token_key(
|
||||
@@ -490,7 +357,7 @@ def extract_model_from_responses_request(request_body_dict: dict[str, Any]) -> s
|
||||
|
||||
logger.warning(
|
||||
"No model found in Responses API request",
|
||||
extra={"body_keys": list(request_body_dict.keys())},
|
||||
extra={"body_keys": list(request_body_dict.keys())}
|
||||
)
|
||||
return "unknown"
|
||||
|
||||
|
||||
+130
-344
@@ -2,7 +2,6 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
import traceback
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import TYPE_CHECKING, Mapping
|
||||
@@ -37,8 +36,6 @@ from ..wallet import recieve_token, send_token
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
DEFAULT_PROXY_TIMEOUT = 30.0
|
||||
|
||||
|
||||
class TopupData(BaseModel):
|
||||
"""Universal top-up data schema for Lightning Network invoices."""
|
||||
@@ -171,6 +168,27 @@ class BaseUpstreamProvider:
|
||||
|
||||
return headers
|
||||
|
||||
def _extract_usage_from_tail(
|
||||
self, tail_content: str
|
||||
) -> tuple[dict | None, str | None]:
|
||||
"""Extract usage and model from tail content of a stream."""
|
||||
usage_data = None
|
||||
model = None
|
||||
|
||||
lines = tail_content.strip().split("\n")
|
||||
for line in lines:
|
||||
if line.startswith("data: "):
|
||||
try:
|
||||
data = json.loads(line[6:])
|
||||
if isinstance(data, dict):
|
||||
if "usage" in data:
|
||||
usage_data = data["usage"]
|
||||
if "model" in data:
|
||||
model = data["model"]
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
return usage_data, model
|
||||
|
||||
def prepare_params(
|
||||
self, path: str, query_params: Mapping[str, str] | None
|
||||
) -> Mapping[str, str]:
|
||||
@@ -438,155 +456,64 @@ class BaseUpstreamProvider:
|
||||
async def stream_with_cost(
|
||||
max_cost_for_model: int,
|
||||
) -> AsyncGenerator[bytes, None]:
|
||||
stored_chunks: list[bytes] = []
|
||||
usage_finalized: bool = False
|
||||
last_model_seen: str | None = None
|
||||
|
||||
async def finalize_without_usage() -> bytes | None:
|
||||
nonlocal usage_finalized
|
||||
if usage_finalized:
|
||||
return None
|
||||
async with create_session() as new_session:
|
||||
fresh_key = await new_session.get(key.__class__, key.hashed_key)
|
||||
if not fresh_key:
|
||||
return None
|
||||
try:
|
||||
fallback: dict = {
|
||||
"model": last_model_seen or "unknown",
|
||||
"usage": None,
|
||||
}
|
||||
cost_data = await adjust_payment_for_tokens(
|
||||
fresh_key, fallback, new_session, max_cost_for_model
|
||||
)
|
||||
usage_finalized = True
|
||||
logger.info(
|
||||
"Finalized streaming payment without explicit usage",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"cost_data": cost_data,
|
||||
"balance_after_adjustment": fresh_key.balance,
|
||||
},
|
||||
)
|
||||
return f"data: {json.dumps({'cost': cost_data})}\n\n".encode()
|
||||
except Exception as cost_error:
|
||||
logger.error(
|
||||
"Error finalizing payment without usage",
|
||||
extra={
|
||||
"error": str(cost_error),
|
||||
"error_type": type(cost_error).__name__,
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
},
|
||||
)
|
||||
return None
|
||||
tail_buffer: list[bytes] = []
|
||||
MAX_TAIL_CHUNKS = 5
|
||||
|
||||
try:
|
||||
async for chunk in response.aiter_bytes():
|
||||
stored_chunks.append(chunk)
|
||||
try:
|
||||
for part in re.split(b"data: ", chunk):
|
||||
if not part or part.strip() in (b"[DONE]", b""):
|
||||
continue
|
||||
try:
|
||||
obj = json.loads(part)
|
||||
if isinstance(obj, dict) and obj.get("model"):
|
||||
last_model_seen = str(obj.get("model"))
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
yield chunk
|
||||
tail_buffer.append(chunk)
|
||||
if len(tail_buffer) > MAX_TAIL_CHUNKS:
|
||||
tail_buffer.pop(0)
|
||||
|
||||
logger.debug(
|
||||
"Streaming completed, analyzing usage data",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"chunks_count": len(stored_chunks),
|
||||
},
|
||||
)
|
||||
# Post-stream processing
|
||||
tail_content = b"".join(tail_buffer).decode("utf-8", errors="ignore")
|
||||
usage_data, model = self._extract_usage_from_tail(tail_content)
|
||||
|
||||
for i in range(len(stored_chunks) - 1, -1, -1):
|
||||
chunk = stored_chunks[i]
|
||||
if not chunk:
|
||||
continue
|
||||
try:
|
||||
events = re.split(b"data: ", chunk)
|
||||
for event_data in events:
|
||||
if not event_data or event_data.strip() in (b"[DONE]", b""):
|
||||
continue
|
||||
if usage_data:
|
||||
# Calculate final cost using usage data extracted from stream tail
|
||||
async with create_session() as new_session:
|
||||
fresh_key = await new_session.get(key.__class__, key.hashed_key)
|
||||
if fresh_key:
|
||||
try:
|
||||
data = json.loads(event_data)
|
||||
if isinstance(data, dict) and data.get("model"):
|
||||
last_model_seen = str(data.get("model"))
|
||||
if isinstance(data, dict) and isinstance(
|
||||
data.get("usage"), dict
|
||||
):
|
||||
async with create_session() as new_session:
|
||||
fresh_key = await new_session.get(
|
||||
key.__class__, key.hashed_key
|
||||
)
|
||||
if fresh_key:
|
||||
try:
|
||||
cost_data = (
|
||||
await adjust_payment_for_tokens(
|
||||
fresh_key,
|
||||
data,
|
||||
new_session,
|
||||
max_cost_for_model,
|
||||
)
|
||||
)
|
||||
usage_finalized = True
|
||||
logger.info(
|
||||
"Payment adjustment completed for streaming",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8]
|
||||
+ "...",
|
||||
"cost_data": cost_data,
|
||||
"model": last_model_seen,
|
||||
"balance_after_adjustment": fresh_key.balance,
|
||||
},
|
||||
)
|
||||
yield f"data: {json.dumps({'cost': cost_data})}\n\n".encode()
|
||||
except Exception as cost_error:
|
||||
logger.error(
|
||||
"Error adjusting payment for streaming tokens",
|
||||
extra={
|
||||
"error": str(cost_error),
|
||||
"error_type": type(
|
||||
cost_error
|
||||
).__name__,
|
||||
"key_hash": key.hashed_key[:8]
|
||||
+ "...",
|
||||
},
|
||||
)
|
||||
break
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Error processing streaming response chunk",
|
||||
extra={
|
||||
"error": str(e),
|
||||
"error_type": type(e).__name__,
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
},
|
||||
)
|
||||
# We need to reconstruct a response object for adjust_payment_for_tokens
|
||||
# or call it with usage directly if supported.
|
||||
# adjust_payment_for_tokens expects a dict with "usage" and "model" keys usually.
|
||||
|
||||
if not usage_finalized:
|
||||
maybe_cost_event = await finalize_without_usage()
|
||||
if maybe_cost_event is not None:
|
||||
yield maybe_cost_event
|
||||
data = {
|
||||
"usage": usage_data,
|
||||
"model": model or "unknown",
|
||||
}
|
||||
|
||||
cost_data = await adjust_payment_for_tokens(
|
||||
fresh_key,
|
||||
data,
|
||||
new_session,
|
||||
max_cost_for_model,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Payment adjustment completed for streaming",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"cost_data": cost_data,
|
||||
"model": model,
|
||||
"balance_after_adjustment": fresh_key.balance,
|
||||
},
|
||||
)
|
||||
# We yield the cost data event at the end
|
||||
yield f"data: {json.dumps({'cost': cost_data})}\n\n".encode()
|
||||
except Exception as cost_error:
|
||||
logger.error(
|
||||
"Error adjusting payment for streaming tokens",
|
||||
extra={"error": str(cost_error)},
|
||||
)
|
||||
|
||||
except Exception as stream_error:
|
||||
logger.warning(
|
||||
"Streaming interrupted; finalizing without usage",
|
||||
extra={
|
||||
"error": str(stream_error),
|
||||
"error_type": type(stream_error).__name__,
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
},
|
||||
"Streaming interrupted",
|
||||
extra={"error": str(stream_error)},
|
||||
)
|
||||
await finalize_without_usage()
|
||||
raise
|
||||
|
||||
# Remove inaccurate encoding headers from upstream response
|
||||
@@ -692,7 +619,6 @@ class BaseUpstreamProvider:
|
||||
},
|
||||
)
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Error processing non-streaming chat completion",
|
||||
@@ -729,171 +655,69 @@ class BaseUpstreamProvider:
|
||||
async def stream_with_responses_cost(
|
||||
max_cost_for_model: int,
|
||||
) -> AsyncGenerator[bytes, None]:
|
||||
stored_chunks: list[bytes] = []
|
||||
usage_finalized: bool = False
|
||||
last_model_seen: str | None = None
|
||||
reasoning_tokens: int = 0
|
||||
|
||||
async def finalize_without_usage() -> bytes | None:
|
||||
nonlocal usage_finalized
|
||||
if usage_finalized:
|
||||
return None
|
||||
async with create_session() as new_session:
|
||||
fresh_key = await new_session.get(key.__class__, key.hashed_key)
|
||||
if not fresh_key:
|
||||
return None
|
||||
try:
|
||||
fallback: dict = {
|
||||
"model": last_model_seen or "unknown",
|
||||
"usage": None,
|
||||
}
|
||||
cost_data = await adjust_payment_for_tokens(
|
||||
fresh_key, fallback, new_session, max_cost_for_model
|
||||
)
|
||||
usage_finalized = True
|
||||
logger.info(
|
||||
"Finalized Responses API streaming payment without explicit usage",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"cost_data": cost_data,
|
||||
"balance_after_adjustment": fresh_key.balance,
|
||||
},
|
||||
)
|
||||
return f"data: {json.dumps({'cost': cost_data})}\\n\\n".encode()
|
||||
except Exception as cost_error:
|
||||
logger.error(
|
||||
"Error finalizing Responses API payment without usage",
|
||||
extra={
|
||||
"error": str(cost_error),
|
||||
"error_type": type(cost_error).__name__,
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
},
|
||||
)
|
||||
return None
|
||||
tail_buffer: list[bytes] = []
|
||||
MAX_TAIL_CHUNKS = 5
|
||||
|
||||
try:
|
||||
async for chunk in response.aiter_bytes():
|
||||
stored_chunks.append(chunk)
|
||||
try:
|
||||
for part in re.split(b"data: ", chunk):
|
||||
if not part or part.strip() in (b"[DONE]", b""):
|
||||
continue
|
||||
try:
|
||||
obj = json.loads(part)
|
||||
if isinstance(obj, dict):
|
||||
if obj.get("model"):
|
||||
last_model_seen = str(obj.get("model"))
|
||||
|
||||
# Track reasoning tokens for Responses API
|
||||
if usage := obj.get("usage", {}):
|
||||
if (
|
||||
isinstance(usage, dict)
|
||||
and "reasoning_tokens" in usage
|
||||
):
|
||||
reasoning_tokens += usage.get(
|
||||
"reasoning_tokens", 0
|
||||
)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
yield chunk
|
||||
tail_buffer.append(chunk)
|
||||
if len(tail_buffer) > MAX_TAIL_CHUNKS:
|
||||
tail_buffer.pop(0)
|
||||
|
||||
logger.debug(
|
||||
"Responses API streaming completed, analyzing usage data",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"chunks_count": len(stored_chunks),
|
||||
"reasoning_tokens": reasoning_tokens,
|
||||
},
|
||||
)
|
||||
# Post-stream processing
|
||||
tail_content = b"".join(tail_buffer).decode("utf-8", errors="ignore")
|
||||
usage_data, model = self._extract_usage_from_tail(tail_content)
|
||||
|
||||
# Process final usage data
|
||||
for i in range(len(stored_chunks) - 1, -1, -1):
|
||||
chunk = stored_chunks[i]
|
||||
if not chunk:
|
||||
continue
|
||||
try:
|
||||
events = re.split(b"data: ", chunk)
|
||||
for event_data in events:
|
||||
if not event_data or event_data.strip() in (b"[DONE]", b""):
|
||||
continue
|
||||
if usage_data:
|
||||
async with create_session() as new_session:
|
||||
fresh_key = await new_session.get(key.__class__, key.hashed_key)
|
||||
if fresh_key:
|
||||
try:
|
||||
data = json.loads(event_data)
|
||||
if isinstance(data, dict) and data.get("model"):
|
||||
last_model_seen = str(data.get("model"))
|
||||
if isinstance(data, dict) and isinstance(
|
||||
data.get("usage"), dict
|
||||
):
|
||||
# Include reasoning tokens in usage calculation
|
||||
async with create_session() as new_session:
|
||||
fresh_key = await new_session.get(
|
||||
key.__class__, key.hashed_key
|
||||
)
|
||||
if fresh_key:
|
||||
try:
|
||||
cost_data = (
|
||||
await adjust_payment_for_tokens(
|
||||
fresh_key,
|
||||
data,
|
||||
new_session,
|
||||
max_cost_for_model,
|
||||
)
|
||||
)
|
||||
usage_finalized = True
|
||||
logger.info(
|
||||
"Payment adjustment completed for Responses API streaming",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8]
|
||||
+ "...",
|
||||
"cost_data": cost_data,
|
||||
"model": last_model_seen,
|
||||
"reasoning_tokens": reasoning_tokens,
|
||||
"balance_after_adjustment": fresh_key.balance,
|
||||
},
|
||||
)
|
||||
yield f"data: {json.dumps({'cost': cost_data})}\\n\\n".encode()
|
||||
except Exception as cost_error:
|
||||
logger.error(
|
||||
"Error adjusting payment for Responses API streaming tokens",
|
||||
extra={
|
||||
"error": str(cost_error),
|
||||
"error_type": type(
|
||||
cost_error
|
||||
).__name__,
|
||||
"key_hash": key.hashed_key[:8]
|
||||
+ "...",
|
||||
},
|
||||
)
|
||||
break
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Error processing Responses API streaming response chunk",
|
||||
extra={
|
||||
"error": str(e),
|
||||
"error_type": type(e).__name__,
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
},
|
||||
)
|
||||
# We need to reconstruct a response object for adjust_payment_for_tokens
|
||||
# or call it with usage directly if supported.
|
||||
# adjust_payment_for_tokens expects a dict with "usage" and "model" keys usually.
|
||||
|
||||
if not usage_finalized:
|
||||
maybe_cost_event = await finalize_without_usage()
|
||||
if maybe_cost_event is not None:
|
||||
yield maybe_cost_event
|
||||
data = {
|
||||
"usage": usage_data,
|
||||
"model": model or "unknown",
|
||||
}
|
||||
|
||||
cost_data = await adjust_payment_for_tokens(
|
||||
fresh_key,
|
||||
data,
|
||||
new_session,
|
||||
max_cost_for_model,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Payment adjustment completed for Responses API streaming",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"cost_data": cost_data,
|
||||
"model": model,
|
||||
"balance_after_adjustment": fresh_key.balance,
|
||||
},
|
||||
)
|
||||
# We yield the cost data event at the end
|
||||
yield f"data: {json.dumps({'cost': cost_data})}\\n\\n".encode()
|
||||
except Exception as cost_error:
|
||||
logger.error(
|
||||
"Error adjusting payment for Responses API streaming tokens",
|
||||
extra={
|
||||
"error": str(cost_error),
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
},
|
||||
)
|
||||
|
||||
except Exception as stream_error:
|
||||
logger.warning(
|
||||
"Responses API streaming interrupted; finalizing without usage",
|
||||
"Responses API streaming interrupted",
|
||||
extra={
|
||||
"error": str(stream_error),
|
||||
"error_type": type(stream_error).__name__,
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
},
|
||||
)
|
||||
await finalize_without_usage()
|
||||
raise
|
||||
|
||||
# Remove inaccurate encoding headers from upstream response
|
||||
@@ -1059,7 +883,7 @@ class BaseUpstreamProvider:
|
||||
|
||||
client = httpx.AsyncClient(
|
||||
transport=httpx.AsyncHTTPTransport(retries=1),
|
||||
timeout=DEFAULT_PROXY_TIMEOUT,
|
||||
timeout=None,
|
||||
)
|
||||
|
||||
try:
|
||||
@@ -1202,8 +1026,7 @@ class BaseUpstreamProvider:
|
||||
if isinstance(exc, httpx.ConnectError):
|
||||
error_message = "Unable to connect to upstream service"
|
||||
elif isinstance(exc, httpx.TimeoutException):
|
||||
# Re-raise timeout exception to allow fallback handling in proxy layer
|
||||
raise
|
||||
error_message = "Upstream service request timed out"
|
||||
elif isinstance(exc, httpx.NetworkError):
|
||||
error_message = "Network error while connecting to upstream service"
|
||||
else:
|
||||
@@ -1286,7 +1109,7 @@ class BaseUpstreamProvider:
|
||||
|
||||
client = httpx.AsyncClient(
|
||||
transport=httpx.AsyncHTTPTransport(retries=1),
|
||||
timeout=DEFAULT_PROXY_TIMEOUT,
|
||||
timeout=None,
|
||||
)
|
||||
|
||||
try:
|
||||
@@ -1406,8 +1229,7 @@ class BaseUpstreamProvider:
|
||||
if isinstance(exc, httpx.ConnectError):
|
||||
error_message = "Unable to connect to upstream service"
|
||||
elif isinstance(exc, httpx.TimeoutException):
|
||||
# Re-raise timeout exception to allow fallback handling in proxy layer
|
||||
raise
|
||||
error_message = "Upstream service request timed out"
|
||||
elif isinstance(exc, httpx.NetworkError):
|
||||
error_message = "Network error while connecting to upstream service"
|
||||
else:
|
||||
@@ -1470,7 +1292,7 @@ class BaseUpstreamProvider:
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.AsyncHTTPTransport(retries=1),
|
||||
timeout=DEFAULT_PROXY_TIMEOUT,
|
||||
timeout=None,
|
||||
) as client:
|
||||
try:
|
||||
response = await client.send(
|
||||
@@ -1501,9 +1323,6 @@ class BaseUpstreamProvider:
|
||||
status_code=response.status_code,
|
||||
headers=dict(response.headers),
|
||||
)
|
||||
except (httpx.TimeoutException, httpx.ConnectError):
|
||||
# Re-raise to allow fallback handling in proxy layer
|
||||
raise
|
||||
except Exception as exc:
|
||||
tb = traceback.format_exc()
|
||||
logger.error(
|
||||
@@ -1696,21 +1515,8 @@ class BaseUpstreamProvider:
|
||||
if "content-encoding" in response_headers:
|
||||
del response_headers["content-encoding"]
|
||||
|
||||
usage_data = None
|
||||
model = None
|
||||
|
||||
usage_data, model = self._extract_usage_from_tail(content_str)
|
||||
lines = content_str.strip().split("\n")
|
||||
for line in lines:
|
||||
if line.startswith("data: "):
|
||||
try:
|
||||
data_json = json.loads(line[6:])
|
||||
if "usage" in data_json:
|
||||
usage_data = data_json["usage"]
|
||||
model = data_json.get("model")
|
||||
elif "model" in data_json and not model:
|
||||
model = data_json["model"]
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
if usage_data and model:
|
||||
logger.debug(
|
||||
@@ -2036,7 +1842,7 @@ class BaseUpstreamProvider:
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.AsyncHTTPTransport(retries=1),
|
||||
timeout=DEFAULT_PROXY_TIMEOUT,
|
||||
timeout=None,
|
||||
) as client:
|
||||
try:
|
||||
response = await client.send(
|
||||
@@ -2130,9 +1936,6 @@ class BaseUpstreamProvider:
|
||||
headers=dict(response.headers),
|
||||
background=background_tasks,
|
||||
)
|
||||
except (httpx.TimeoutException, httpx.ConnectError):
|
||||
# Re-raise to allow fallback handling in proxy layer
|
||||
raise
|
||||
except Exception as exc:
|
||||
tb = traceback.format_exc()
|
||||
logger.error(
|
||||
@@ -2300,7 +2103,7 @@ class BaseUpstreamProvider:
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.AsyncHTTPTransport(retries=1),
|
||||
timeout=DEFAULT_PROXY_TIMEOUT,
|
||||
timeout=None,
|
||||
) as client:
|
||||
try:
|
||||
response = await client.send(
|
||||
@@ -2394,9 +2197,6 @@ class BaseUpstreamProvider:
|
||||
headers=dict(response.headers),
|
||||
background=background_tasks,
|
||||
)
|
||||
except (httpx.TimeoutException, httpx.ConnectError):
|
||||
# Re-raise to allow fallback handling in proxy layer
|
||||
raise
|
||||
except Exception as exc:
|
||||
tb = traceback.format_exc()
|
||||
logger.error(
|
||||
@@ -2503,7 +2303,7 @@ class BaseUpstreamProvider:
|
||||
extra={
|
||||
"amount": amount,
|
||||
"unit": unit,
|
||||
"content_lines": len(content_str.strip().split("\\n")),
|
||||
"content_lines": len(content_str.strip().split("\n")),
|
||||
},
|
||||
)
|
||||
|
||||
@@ -2513,28 +2313,14 @@ class BaseUpstreamProvider:
|
||||
if "content-encoding" in response_headers:
|
||||
del response_headers["content-encoding"]
|
||||
|
||||
usage_data = None
|
||||
model = None
|
||||
usage_data, model = self._extract_usage_from_tail(content_str)
|
||||
reasoning_tokens = 0
|
||||
|
||||
lines = content_str.strip().split("\\n")
|
||||
for line in lines:
|
||||
if line.startswith("data: "):
|
||||
try:
|
||||
data_json = json.loads(line[6:])
|
||||
if "usage" in data_json:
|
||||
usage_data = data_json["usage"]
|
||||
model = data_json.get("model")
|
||||
# Track reasoning tokens for Responses API
|
||||
if (
|
||||
isinstance(usage_data, dict)
|
||||
and "reasoning_tokens" in usage_data
|
||||
):
|
||||
reasoning_tokens = usage_data.get("reasoning_tokens", 0)
|
||||
elif "model" in data_json and not model:
|
||||
model = data_json["model"]
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
# Responses API specific: check for reasoning tokens
|
||||
if usage_data and "reasoning_tokens" in usage_data:
|
||||
reasoning_tokens = usage_data.get("reasoning_tokens", 0)
|
||||
|
||||
lines = content_str.strip().split("\n")
|
||||
|
||||
if usage_data and model:
|
||||
logger.debug(
|
||||
@@ -2610,7 +2396,7 @@ class BaseUpstreamProvider:
|
||||
|
||||
async def generate() -> AsyncGenerator[bytes, None]:
|
||||
for line in lines:
|
||||
yield (line + "\\n").encode("utf-8")
|
||||
yield (line + "\n").encode("utf-8")
|
||||
|
||||
return StreamingResponse(
|
||||
generate(),
|
||||
|
||||
Reference in New Issue
Block a user