mirror of
https://github.com/Routstr/routstr-core.git
synced 2026-07-22 12:22:20 +00:00
Compare commits
4 Commits
v0.2.1
...
optimize-s
| Author | SHA1 | Date | |
|---|---|---|---|
|
dce9f37b0b | ||
|
|
5a4ba60072 | ||
|
|
d41c214d9e | ||
|
|
ec0fcfb48b |
@@ -6,7 +6,7 @@ import os
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from pydantic.v1 import BaseModel, BaseSettings, Field
|
||||
from pydantic.v1 import BaseModel, BaseSettings, Field, validator
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
|
||||
@@ -37,6 +37,13 @@ class Settings(BaseSettings):
|
||||
|
||||
# Cashu
|
||||
cashu_mints: list[str] = Field(default_factory=list, env="CASHU_MINTS")
|
||||
|
||||
@validator("cashu_mints", pre=True, each_item=True)
|
||||
def normalize_mint_url(cls, v: str) -> str:
|
||||
if isinstance(v, str):
|
||||
return v.rstrip("/")
|
||||
return v
|
||||
|
||||
receive_ln_address: str = Field(default="", env="RECEIVE_LN_ADDRESS")
|
||||
primary_mint: str = Field(default="", env="PRIMARY_MINT_URL")
|
||||
primary_mint_unit: str = Field(default="sat", env="PRIMARY_MINT_UNIT")
|
||||
|
||||
@@ -16,6 +16,7 @@ from .core.db import (
|
||||
create_session,
|
||||
get_session,
|
||||
)
|
||||
from .core.settings import settings
|
||||
from .payment.helpers import (
|
||||
calculate_discounted_max_cost,
|
||||
check_token_balance,
|
||||
@@ -25,6 +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
|
||||
|
||||
logger = get_logger(__name__)
|
||||
proxy_router = APIRouter()
|
||||
@@ -146,6 +148,24 @@ async def proxy(
|
||||
else:
|
||||
model_id = request_body_dict.get("model", "unknown")
|
||||
|
||||
if "https://testnut.cashu.space" in settings.cashu_mints:
|
||||
try:
|
||||
token_str = None
|
||||
if x_cashu_header := headers.get("x-cashu"):
|
||||
token_str = x_cashu_header
|
||||
elif auth_header := headers.get("authorization"):
|
||||
parts = auth_header.split(" ")
|
||||
if len(parts) > 1 and not parts[1].startswith("sk-"):
|
||||
token_str = parts[1]
|
||||
|
||||
if token_str:
|
||||
token_obj = deserialize_token_from_string(token_str)
|
||||
if token_obj.mint == "https://testnut.cashu.space":
|
||||
model_id = "mock/gpt-420-mock"
|
||||
request_body_dict["model"] = model_id
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
model_obj = get_model_instance(model_id)
|
||||
if not model_obj:
|
||||
return create_error_response(
|
||||
|
||||
@@ -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
|
||||
@@ -169,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]:
|
||||
@@ -234,7 +254,11 @@ class BaseUpstreamProvider:
|
||||
)
|
||||
|
||||
# Handle model in input field (alternative format)
|
||||
if "input" in data and isinstance(data["input"], dict) and "model" in data["input"]:
|
||||
if (
|
||||
"input" in data
|
||||
and isinstance(data["input"], dict)
|
||||
and "model" in data["input"]
|
||||
):
|
||||
original_model = model_obj.id
|
||||
transformed_model = self.transform_model_name(original_model)
|
||||
data["input"]["model"] = transformed_model
|
||||
@@ -432,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
|
||||
@@ -722,166 +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
|
||||
@@ -933,8 +769,8 @@ class BaseUpstreamProvider:
|
||||
"model": response_json.get("model", "unknown"),
|
||||
"has_usage": "usage" in response_json,
|
||||
"has_reasoning_tokens": "usage" in response_json
|
||||
and isinstance(response_json.get("usage"), dict)
|
||||
and "reasoning_tokens" in response_json["usage"],
|
||||
and isinstance(response_json.get("usage"), dict)
|
||||
and "reasoning_tokens" in response_json["usage"],
|
||||
},
|
||||
)
|
||||
|
||||
@@ -1679,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(
|
||||
@@ -2480,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")),
|
||||
},
|
||||
)
|
||||
|
||||
@@ -2490,25 +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(
|
||||
@@ -2584,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(),
|
||||
|
||||
265
routstr/upstream/fake.py
Normal file
265
routstr/upstream/fake.py
Normal file
@@ -0,0 +1,265 @@
|
||||
import asyncio
|
||||
import json
|
||||
import random
|
||||
from typing import AsyncIterator
|
||||
|
||||
from fastapi import Request
|
||||
from fastapi.responses import Response, StreamingResponse
|
||||
|
||||
from ..core.db import ApiKey, AsyncSession
|
||||
from ..payment.models import Architecture, Model, Pricing
|
||||
from .base import BaseUpstreamProvider
|
||||
|
||||
|
||||
class MockUpstreamProvider(BaseUpstreamProvider):
|
||||
"""Fack Mock Upstream provider specifically for Testing."""
|
||||
|
||||
provider_type = "mock"
|
||||
|
||||
async def forward_request(
|
||||
self,
|
||||
request: Request,
|
||||
path: str,
|
||||
headers: dict,
|
||||
request_body: bytes | None,
|
||||
key: ApiKey,
|
||||
max_cost_for_model: int,
|
||||
session: AsyncSession,
|
||||
model_obj: Model,
|
||||
) -> Response | StreamingResponse:
|
||||
if path.endswith("chat/completions"):
|
||||
is_streaming = False
|
||||
if request_body:
|
||||
request_data = json.loads(request_body)
|
||||
is_streaming = request_data.get("stream", False)
|
||||
|
||||
if is_streaming:
|
||||
|
||||
async def fake_streaming_response(
|
||||
chunk_size: int | None = None,
|
||||
) -> AsyncIterator[bytes]:
|
||||
suffix = random.randint(1000, 9999)
|
||||
req_id = f"gen-mock-stream-{suffix}"
|
||||
created = 1766138895
|
||||
model = "mock/gpt-420-mock"
|
||||
|
||||
def make_chunk(
|
||||
delta: dict,
|
||||
finish_reason: str | None = None,
|
||||
usage: dict | None = None,
|
||||
) -> bytes:
|
||||
chunk = {
|
||||
"id": req_id,
|
||||
"provider": "MockProvider",
|
||||
"model": model,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": created,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": delta,
|
||||
"finish_reason": finish_reason,
|
||||
"native_finish_reason": "completed"
|
||||
if finish_reason
|
||||
else None,
|
||||
"logprobs": None,
|
||||
}
|
||||
],
|
||||
}
|
||||
if usage:
|
||||
chunk["usage"] = usage
|
||||
return f"data: {json.dumps(chunk)}\n\n".encode()
|
||||
|
||||
# 1. Initial chunk
|
||||
yield make_chunk({"role": "assistant", "content": ""})
|
||||
await asyncio.sleep(0.02)
|
||||
|
||||
# 2. Reasoning chunks
|
||||
reasoning_tokens = ["Mock", " reason", "ing", "..."]
|
||||
for token in reasoning_tokens:
|
||||
delta = {
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"reasoning": token,
|
||||
"reasoning_details": [
|
||||
{
|
||||
"type": "reasoning.summary",
|
||||
"summary": token,
|
||||
"format": "openai-responses-v1",
|
||||
"index": 0,
|
||||
}
|
||||
],
|
||||
}
|
||||
yield make_chunk(delta)
|
||||
await asyncio.sleep(0.03)
|
||||
|
||||
# 3. Content chunks
|
||||
content_tokens = ["This", " is", " a", " mock", " stream", "."]
|
||||
for token in content_tokens:
|
||||
yield make_chunk({"role": "assistant", "content": token})
|
||||
await asyncio.sleep(0.03)
|
||||
|
||||
# 4. Finish chunk
|
||||
yield make_chunk(
|
||||
{"role": "assistant", "content": ""}, finish_reason="stop"
|
||||
)
|
||||
|
||||
# 5. Usage chunk
|
||||
usage_data = {
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 20,
|
||||
"total_tokens": 30,
|
||||
"cost": 0.001,
|
||||
"is_byok": False,
|
||||
"prompt_tokens_details": {
|
||||
"cached_tokens": 0,
|
||||
"audio_tokens": 0,
|
||||
"video_tokens": 0,
|
||||
},
|
||||
"cost_details": {
|
||||
"upstream_inference_cost": None,
|
||||
"upstream_inference_prompt_cost": 0,
|
||||
"upstream_inference_completions_cost": 0.001,
|
||||
},
|
||||
"completion_tokens_details": {
|
||||
"reasoning_tokens": 10,
|
||||
"image_tokens": 0,
|
||||
},
|
||||
}
|
||||
|
||||
usage_chunk = {
|
||||
"id": req_id,
|
||||
"provider": "MockProvider",
|
||||
"model": model,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": created,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"role": "assistant", "content": ""},
|
||||
"finish_reason": None,
|
||||
"native_finish_reason": None,
|
||||
"logprobs": None,
|
||||
}
|
||||
],
|
||||
"usage": usage_data,
|
||||
}
|
||||
yield f"data: {json.dumps(usage_chunk)}\n\n".encode()
|
||||
|
||||
# 6. DONE
|
||||
yield b"data: [DONE]\n\n"
|
||||
|
||||
# 7. Cost
|
||||
cost_chunk = {
|
||||
"cost": {
|
||||
"base_msats": 0,
|
||||
"input_msats": 2,
|
||||
"output_msats": 10,
|
||||
"total_msats": 12,
|
||||
}
|
||||
}
|
||||
yield f"data: {json.dumps(cost_chunk)}\n\n".encode()
|
||||
|
||||
return StreamingResponse(
|
||||
fake_streaming_response(),
|
||||
200,
|
||||
)
|
||||
|
||||
else:
|
||||
suffix = random.randint(1000, 9999)
|
||||
content_dict = {
|
||||
"id": f"gen-mock-{suffix}",
|
||||
"provider": "MockProvider",
|
||||
"model": "mock/gpt-5-mini",
|
||||
"object": "chat.completion",
|
||||
"created": 1766138655,
|
||||
"choices": [
|
||||
{
|
||||
"logprobs": None,
|
||||
"finish_reason": "length",
|
||||
"native_finish_reason": "max_output_tokens",
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": f"Mock Content {suffix}",
|
||||
"refusal": None,
|
||||
"reasoning": f"Mock Reasoning {suffix}",
|
||||
"reasoning_details": [
|
||||
{
|
||||
"format": "openai-responses-v1",
|
||||
"index": 0,
|
||||
"type": "reasoning.summary",
|
||||
"summary": f"Mock Summary {suffix}",
|
||||
},
|
||||
{
|
||||
"id": f"rs_mock_{suffix}",
|
||||
"format": "openai-responses-v1",
|
||||
"index": 0,
|
||||
"type": "reasoning.encrypted",
|
||||
"data": "mock_encrypted_data",
|
||||
},
|
||||
],
|
||||
},
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 10,
|
||||
"total_tokens": 20,
|
||||
"cost": 0,
|
||||
"is_byok": False,
|
||||
"prompt_tokens_details": {
|
||||
"cached_tokens": 0,
|
||||
"audio_tokens": 0,
|
||||
"video_tokens": 0,
|
||||
},
|
||||
"cost_details": {
|
||||
"upstream_inference_cost": None,
|
||||
"upstream_inference_prompt_cost": 0,
|
||||
"upstream_inference_completions_cost": 0,
|
||||
},
|
||||
"completion_tokens_details": {
|
||||
"reasoning_tokens": 5,
|
||||
"image_tokens": 0,
|
||||
},
|
||||
},
|
||||
"cost": {
|
||||
"base_msats": 0,
|
||||
"input_msats": 0,
|
||||
"output_msats": 0,
|
||||
"total_msats": 0,
|
||||
},
|
||||
}
|
||||
return Response(json.dumps(content_dict).encode(), 200)
|
||||
|
||||
elif path.endswith("embeddings"):
|
||||
raise NotImplementedError
|
||||
elif path.endswith("responses"):
|
||||
raise NotImplementedError
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
async def fetch_models(self) -> list[Model]:
|
||||
return [
|
||||
Model(
|
||||
id="mock/gpt-420-mock",
|
||||
name="mock/gpt-420-mock",
|
||||
created=0,
|
||||
description="mock model for testing",
|
||||
context_length=8192,
|
||||
architecture=Architecture(
|
||||
modality="text",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["text"],
|
||||
tokenizer="",
|
||||
instruct_type=None,
|
||||
),
|
||||
pricing=Pricing(prompt=0.01, completion=0.01),
|
||||
),
|
||||
]
|
||||
|
||||
def transform_model_name(self, model_id: str) -> str:
|
||||
return "fake-model"
|
||||
|
||||
async def get_balance(self) -> float | None:
|
||||
return 420.69
|
||||
@@ -218,6 +218,14 @@ async def init_upstreams() -> list[BaseUpstreamProvider]:
|
||||
results = await asyncio.gather(*tasks)
|
||||
upstreams = [p for p in results if p is not None]
|
||||
|
||||
if "https://testnut.cashu.space" in settings.cashu_mints:
|
||||
from .fake import MockUpstreamProvider
|
||||
|
||||
mock_provider = MockUpstreamProvider("mock", "mock")
|
||||
await mock_provider.refresh_models_cache()
|
||||
upstreams.append(mock_provider)
|
||||
logger.info("Initialized MockUpstreamProvider for testnut mint")
|
||||
|
||||
return upstreams
|
||||
|
||||
|
||||
|
||||
@@ -313,6 +313,8 @@ async def periodic_payout() -> None:
|
||||
try:
|
||||
async with db.create_session() as session:
|
||||
for mint_url in settings.cashu_mints:
|
||||
if mint_url == "https://testnut.cashu.space":
|
||||
continue
|
||||
for unit in ["sat", "msat"]:
|
||||
wallet = await get_wallet(mint_url, unit)
|
||||
proofs = get_proofs_per_mint_and_unit(
|
||||
|
||||
Reference in New Issue
Block a user