Add sats cost and remaining balance to response

This commit is contained in:
redshift
2026-02-21 08:51:59 +00:00
parent 9ce91f58a9
commit 6e4da3d3cf
2 changed files with 74 additions and 0 deletions
+66
View File
@@ -525,10 +525,17 @@ class BaseUpstreamProvider:
session,
max_cost_for_model,
)
remaining_balance_msats = fresh_key.balance
# Merge cost into usage
usage_chunk_data["usage"]["cost"] = cost_data.get(
"total_usd", 0.0
)
usage_chunk_data["usage"]["cost_sats"] = (
cost_data.get("total_msats", 0) // 1000
)
usage_chunk_data["usage"]["remaining_balance_msats"] = (
remaining_balance_msats
)
# Keep detailed cost in metadata
usage_chunk_data["metadata"] = usage_chunk_data.get(
"metadata", {}
@@ -536,6 +543,12 @@ class BaseUpstreamProvider:
usage_chunk_data["metadata"]["routstr"] = {
"cost": cost_data
}
usage_chunk_data["metadata"]["routstr"]["cost"][
"sats_cost"
] = cost_data.get("total_msats", 0) // 1000
usage_chunk_data["metadata"]["routstr"]["cost"][
"remaining_balance_msats"
] = remaining_balance_msats
yield f"data: {json.dumps(usage_chunk_data)}\n\n".encode()
usage_finalized = True
except Exception as e:
@@ -626,14 +639,31 @@ class BaseUpstreamProvider:
key, response_json, session, deducted_max_cost
)
await session.refresh(key)
remaining_balance_msats = key.balance
# Merge cost into usage for OpenCode
if "usage" in response_json:
response_json["usage"]["cost"] = cost_data.get("total_usd", 0.0)
response_json["usage"]["cost_sats"] = (
cost_data.get("total_msats", 0) // 1000
)
response_json["usage"]["remaining_balance_msats"] = (
remaining_balance_msats
)
# Keep detailed cost
response_json["metadata"] = response_json.get("metadata", {})
response_json["metadata"]["routstr"] = {"cost": cost_data}
response_json["metadata"]["routstr"]["cost"]["sats_cost"] = (
cost_data.get("total_msats", 0) // 1000
)
response_json["metadata"]["routstr"]["cost"]["remaining_balance_msats"] = (
remaining_balance_msats
)
response_json["cost"] = cost_data
response_json["cost"]["sats_cost"] = cost_data.get("total_msats", 0) // 1000
response_json["cost"]["remaining_balance_msats"] = remaining_balance_msats
logger.info(
"Payment adjustment completed for non-streaming",
@@ -803,6 +833,7 @@ class BaseUpstreamProvider:
session,
max_cost_for_model,
)
remaining_balance_msats = fresh_key.balance
# Merge cost into usage chunk
if (
"response" in usage_chunk_data
@@ -811,10 +842,22 @@ class BaseUpstreamProvider:
usage_chunk_data["response"]["usage"]["cost"] = (
cost_data.get("total_usd", 0.0)
)
usage_chunk_data["response"]["usage"][
"cost_sats"
] = cost_data.get("total_msats", 0) // 1000
usage_chunk_data["response"]["usage"][
"remaining_balance_msats"
] = remaining_balance_msats
elif "usage" in usage_chunk_data:
usage_chunk_data["usage"]["cost"] = cost_data.get(
"total_usd", 0.0
)
usage_chunk_data["usage"]["cost_sats"] = (
cost_data.get("total_msats", 0) // 1000
)
usage_chunk_data["usage"][
"remaining_balance_msats"
] = remaining_balance_msats
# Keep detailed cost in metadata
usage_chunk_data["metadata"] = usage_chunk_data.get(
@@ -823,6 +866,12 @@ class BaseUpstreamProvider:
usage_chunk_data["metadata"]["routstr"] = {
"cost": cost_data
}
usage_chunk_data["metadata"]["routstr"]["cost"][
"sats_cost"
] = cost_data.get("total_msats", 0) // 1000
usage_chunk_data["metadata"]["routstr"]["cost"][
"remaining_balance_msats"
] = remaining_balance_msats
yield f"data: {json.dumps(usage_chunk_data)}\n\n".encode()
usage_finalized = True
except Exception:
@@ -907,14 +956,31 @@ class BaseUpstreamProvider:
key, response_json, session, deducted_max_cost
)
await session.refresh(key)
remaining_balance_msats = key.balance
# Merge cost into usage for OpenCode
if "usage" in response_json:
response_json["usage"]["cost"] = cost_data.get("total_usd", 0.0)
response_json["usage"]["cost_sats"] = (
cost_data.get("total_msats", 0) // 1000
)
response_json["usage"]["remaining_balance_msats"] = (
remaining_balance_msats
)
# Keep detailed cost
response_json["metadata"] = response_json.get("metadata", {})
response_json["metadata"]["routstr"] = {"cost": cost_data}
response_json["metadata"]["routstr"]["cost"]["sats_cost"] = (
cost_data.get("total_msats", 0) // 1000
)
response_json["metadata"]["routstr"]["cost"]["remaining_balance_msats"] = (
remaining_balance_msats
)
response_json["cost"] = cost_data
response_json["cost"]["sats_cost"] = cost_data.get("total_msats", 0) // 1000
response_json["cost"]["remaining_balance_msats"] = remaining_balance_msats
logger.info(
"Payment adjustment completed for non-streaming Responses API",
+8
View File
@@ -259,7 +259,15 @@ class GeminiUpstreamProvider(BaseUpstreamProvider):
cost_data = await adjust_payment_for_tokens(
key, openai_format_response, session, max_cost_for_model
)
await session.refresh(key)
remaining_balance_msats = key.balance
openai_format_response["cost"] = cost_data
openai_format_response["cost"]["sats_cost"] = (
cost_data.get("total_msats", 0) // 1000
)
openai_format_response["cost"]["remaining_balance_msats"] = (
remaining_balance_msats
)
logger.info(
"Gemini non-streaming payment completed",