mirror of
https://github.com/Routstr/routstr-core.git
synced 2026-10-06 12:38:23 +00:00
Merge pull request #458 from Routstr/add-cost-usage-to-messages-endpoint
Add cost usage to messages endpoint
This commit is contained in:
9qeklajc
+19
-1
@@ -69,7 +69,25 @@ def get_upstreams() -> list[BaseUpstreamProvider]:
|
||||
|
||||
def get_model_instance(model_id: str) -> Model | None:
|
||||
"""Get Model instance by ID from global cache."""
|
||||
return _model_instances.get(model_id.lower())
|
||||
if not model_id:
|
||||
return None
|
||||
|
||||
model_id_lower = model_id.lower()
|
||||
# Try exact match first
|
||||
if model := _model_instances.get(model_id_lower):
|
||||
return model
|
||||
|
||||
# Try stripping common version suffixes (e.g., -20251222)
|
||||
# This handles cases where upstream returns a specific version
|
||||
# but we only track the base model name.
|
||||
import re
|
||||
|
||||
base_model_id = re.sub(r"-\d{8}$", "", model_id_lower)
|
||||
if base_model_id != model_id_lower:
|
||||
if model := _model_instances.get(base_model_id):
|
||||
return model
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def get_provider_for_model(model_id: str) -> list[BaseUpstreamProvider] | None:
|
||||
|
||||
+107
-8
@@ -119,6 +119,51 @@ class BaseUpstreamProvider:
|
||||
"can_show_balance": False,
|
||||
}
|
||||
|
||||
def inject_cost_metadata(
|
||||
self,
|
||||
response_json: dict,
|
||||
cost_data: CostData | MaxCostData | dict,
|
||||
key: ApiKey,
|
||||
) -> None:
|
||||
"""Unifies the injection of cost and usage metadata across all completion types."""
|
||||
if isinstance(cost_data, dict):
|
||||
total_msats = cost_data.get("total_msats", 0)
|
||||
total_usd = cost_data.get("total_usd", 0.0)
|
||||
cost_dict = cost_data
|
||||
else:
|
||||
total_msats = cost_data.total_msats
|
||||
total_usd = cost_data.total_usd
|
||||
cost_dict = cost_data.dict()
|
||||
|
||||
sats_cost = total_msats // 1000
|
||||
|
||||
# Inject into top-level usage block (OpenAI/Anthropic style)
|
||||
if "usage" in response_json:
|
||||
response_json["usage"]["cost"] = total_usd
|
||||
response_json["usage"]["cost_sats"] = sats_cost
|
||||
response_json["usage"]["remaining_balance_msats"] = key.balance
|
||||
|
||||
# Inject into Anthropic nested usage block if present
|
||||
if (
|
||||
"message" in response_json
|
||||
and isinstance(response_json["message"], dict)
|
||||
and "usage" in response_json["message"]
|
||||
):
|
||||
response_json["message"]["usage"]["sats_cost"] = sats_cost
|
||||
|
||||
# Unified Routstr metadata
|
||||
response_json["metadata"] = response_json.get("metadata", {})
|
||||
response_json["metadata"]["routstr"] = {
|
||||
"cost": cost_dict,
|
||||
"sats_cost": sats_cost,
|
||||
"remaining_balance_msats": key.balance,
|
||||
}
|
||||
|
||||
# Legacy/Compatibility fields
|
||||
response_json["cost"] = cost_dict.copy()
|
||||
response_json["cost"]["sats_cost"] = sats_cost
|
||||
response_json["cost"]["remaining_balance_msats"] = key.balance
|
||||
|
||||
def prepare_headers(self, request_headers: dict) -> dict:
|
||||
"""Prepare headers for upstream request by removing proxy-specific headers and adding authentication.
|
||||
|
||||
@@ -1139,7 +1184,11 @@ class BaseUpstreamProvider:
|
||||
)
|
||||
|
||||
async def handle_streaming_messages_completion(
|
||||
self, response: httpx.Response, key: ApiKey, max_cost_for_model: int
|
||||
self,
|
||||
response: httpx.Response,
|
||||
key: ApiKey,
|
||||
max_cost_for_model: int,
|
||||
requested_model: str | None = None,
|
||||
) -> StreamingResponse:
|
||||
async def stream_with_cost(
|
||||
max_cost_for_model: int,
|
||||
@@ -1178,6 +1227,8 @@ class BaseUpstreamProvider:
|
||||
stored_chunks.append(chunk)
|
||||
try:
|
||||
decoded_chunk = chunk.decode("utf-8", errors="ignore")
|
||||
modified_lines = []
|
||||
changed = False
|
||||
for line in decoded_chunk.split("\n"):
|
||||
if line.startswith("data: "):
|
||||
try:
|
||||
@@ -1187,6 +1238,20 @@ class BaseUpstreamProvider:
|
||||
if msg and msg.get("model"):
|
||||
last_model_seen = str(msg.get("model"))
|
||||
|
||||
if requested_model:
|
||||
# Apply requested_model override
|
||||
model_updated = False
|
||||
if msg:
|
||||
msg["model"] = requested_model
|
||||
model_updated = True
|
||||
if data.get("model"):
|
||||
data["model"] = requested_model
|
||||
model_updated = True
|
||||
|
||||
if model_updated:
|
||||
line = "data: " + json.dumps(data)
|
||||
changed = True
|
||||
|
||||
if usage := msg.get("usage"):
|
||||
input_tokens += usage.get("input_tokens", 0)
|
||||
output_tokens += usage.get(
|
||||
@@ -1200,10 +1265,14 @@ class BaseUpstreamProvider:
|
||||
)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
modified_lines.append(line)
|
||||
|
||||
yield chunk
|
||||
if changed:
|
||||
yield "\n".join(modified_lines).encode("utf-8")
|
||||
else:
|
||||
yield chunk
|
||||
except Exception:
|
||||
yield chunk
|
||||
|
||||
usage_data = {
|
||||
"input_tokens": input_tokens,
|
||||
@@ -1225,6 +1294,11 @@ class BaseUpstreamProvider:
|
||||
new_session,
|
||||
max_cost_for_model,
|
||||
)
|
||||
|
||||
self.inject_cost_metadata(
|
||||
combined_data, cost_data, fresh_key
|
||||
)
|
||||
|
||||
usage_finalized = True
|
||||
yield f"event: cost\ndata: {json.dumps({'cost': cost_data})}\n\n".encode()
|
||||
except Exception:
|
||||
@@ -1264,11 +1338,22 @@ class BaseUpstreamProvider:
|
||||
session: AsyncSession,
|
||||
deducted_max_cost: int,
|
||||
path: str,
|
||||
requested_model: str | None = None,
|
||||
) -> Response:
|
||||
try:
|
||||
content = await response.aread()
|
||||
response_json = json.loads(content)
|
||||
|
||||
if requested_model:
|
||||
if "model" in response_json:
|
||||
response_json["model"] = requested_model
|
||||
if (
|
||||
"message" in response_json
|
||||
and isinstance(response_json["message"], dict)
|
||||
and "model" in response_json["message"]
|
||||
):
|
||||
response_json["message"]["model"] = requested_model
|
||||
|
||||
if path.endswith("count_tokens") and "usage" not in response_json:
|
||||
input_tokens = response_json.get("input_tokens", 0)
|
||||
response_json["usage"] = {"input_tokens": input_tokens}
|
||||
@@ -1276,7 +1361,8 @@ class BaseUpstreamProvider:
|
||||
cost_data = await adjust_payment_for_tokens(
|
||||
key, response_json, session, deducted_max_cost
|
||||
)
|
||||
response_json["cost"] = cost_data
|
||||
|
||||
self.inject_cost_metadata(response_json, cost_data, key)
|
||||
|
||||
allowed_headers = {
|
||||
"content-type",
|
||||
@@ -1430,7 +1516,10 @@ class BaseUpstreamProvider:
|
||||
|
||||
if is_streaming and response.status_code == 200:
|
||||
result = await self.handle_streaming_messages_completion(
|
||||
response, key, max_cost_for_model
|
||||
response,
|
||||
key,
|
||||
max_cost_for_model,
|
||||
requested_model=original_model_id,
|
||||
)
|
||||
background_tasks = BackgroundTasks()
|
||||
background_tasks.add_task(response.aclose)
|
||||
@@ -1441,7 +1530,12 @@ class BaseUpstreamProvider:
|
||||
if response.status_code == 200:
|
||||
try:
|
||||
return await self.handle_non_streaming_messages_completion(
|
||||
response, key, session, max_cost_for_model, path
|
||||
response,
|
||||
key,
|
||||
session,
|
||||
max_cost_for_model,
|
||||
path,
|
||||
requested_model=original_model_id,
|
||||
)
|
||||
finally:
|
||||
await response.aclose()
|
||||
@@ -1451,7 +1545,12 @@ class BaseUpstreamProvider:
|
||||
if response.status_code == 200:
|
||||
try:
|
||||
return await self.handle_non_streaming_messages_completion(
|
||||
response, key, session, max_cost_for_model, path
|
||||
response,
|
||||
key,
|
||||
session,
|
||||
max_cost_for_model,
|
||||
path,
|
||||
requested_model=original_model_id,
|
||||
)
|
||||
finally:
|
||||
await response.aclose()
|
||||
|
||||
Reference in New Issue
Block a user