mirror of
https://github.com/Routstr/routstr-core.git
synced 2026-08-09 02:54:37 +00:00
add cost usages
This commit is contained in:
9qeklajc
+52
-30
@@ -119,6 +119,39 @@ class BaseUpstreamProvider:
|
||||
"can_show_balance": False,
|
||||
}
|
||||
|
||||
def inject_cost_metadata(
|
||||
self, response_json: dict, cost_data: CostData | MaxCostData, key: ApiKey
|
||||
) -> None:
|
||||
"""Unifies the injection of cost and usage metadata across all completion types."""
|
||||
sats_cost = cost_data.total_msats // 1000
|
||||
|
||||
# Inject into top-level usage block (OpenAI/Anthropic style)
|
||||
if "usage" in response_json:
|
||||
response_json["usage"]["cost"] = cost_data.total_usd
|
||||
response_json["usage"]["sats_cost"] = 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_data.dict(),
|
||||
"sats_cost": sats_cost,
|
||||
"remaining_balance_msats": key.balance,
|
||||
}
|
||||
|
||||
# Legacy/Compatibility fields
|
||||
response_json["cost"] = cost_data.dict()
|
||||
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.
|
||||
|
||||
@@ -568,7 +601,7 @@ class BaseUpstreamProvider:
|
||||
usage_chunk_data["usage"]["cost"] = cost_data.get(
|
||||
"total_usd", 0.0
|
||||
)
|
||||
usage_chunk_data["usage"]["cost_sats"] = (
|
||||
usage_chunk_data["usage"]["sats_cost"] = (
|
||||
cost_data.get("total_msats", 0) // 1000
|
||||
)
|
||||
usage_chunk_data["usage"]["remaining_balance_msats"] = (
|
||||
@@ -683,31 +716,10 @@ 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
|
||||
if isinstance(cost_data, (CostData, MaxCostData)):
|
||||
self.inject_cost_metadata(response_json, cost_data, key)
|
||||
else:
|
||||
response_json["cost"] = cost_data
|
||||
|
||||
logger.info(
|
||||
"Payment adjustment completed for non-streaming",
|
||||
@@ -895,7 +907,7 @@ class BaseUpstreamProvider:
|
||||
cost_data.get("total_usd", 0.0)
|
||||
)
|
||||
usage_chunk_data["response"]["usage"][
|
||||
"cost_sats"
|
||||
"sats_cost"
|
||||
] = cost_data.get("total_msats", 0) // 1000
|
||||
usage_chunk_data["response"]["usage"][
|
||||
"remaining_balance_msats"
|
||||
@@ -904,7 +916,7 @@ class BaseUpstreamProvider:
|
||||
usage_chunk_data["usage"]["cost"] = cost_data.get(
|
||||
"total_usd", 0.0
|
||||
)
|
||||
usage_chunk_data["usage"]["cost_sats"] = (
|
||||
usage_chunk_data["usage"]["sats_cost"] = (
|
||||
cost_data.get("total_msats", 0) // 1000
|
||||
)
|
||||
usage_chunk_data["usage"][
|
||||
@@ -1020,7 +1032,7 @@ class BaseUpstreamProvider:
|
||||
# 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"] = (
|
||||
response_json["usage"]["sats_cost"] = (
|
||||
cost_data.get("total_msats", 0) // 1000
|
||||
)
|
||||
response_json["usage"]["remaining_balance_msats"] = (
|
||||
@@ -1225,6 +1237,12 @@ class BaseUpstreamProvider:
|
||||
new_session,
|
||||
max_cost_for_model,
|
||||
)
|
||||
|
||||
if isinstance(cost_data, (CostData, MaxCostData)):
|
||||
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:
|
||||
@@ -1276,7 +1294,11 @@ class BaseUpstreamProvider:
|
||||
cost_data = await adjust_payment_for_tokens(
|
||||
key, response_json, session, deducted_max_cost
|
||||
)
|
||||
response_json["cost"] = cost_data
|
||||
|
||||
if isinstance(cost_data, (CostData, MaxCostData)):
|
||||
self.inject_cost_metadata(response_json, cost_data, key)
|
||||
else:
|
||||
response_json["cost"] = cost_data
|
||||
|
||||
allowed_headers = {
|
||||
"content-type",
|
||||
|
||||
@@ -0,0 +1,145 @@
|
||||
import json
|
||||
from collections.abc import AsyncGenerator
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from httpx import AsyncClient
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_usage_and_metadata_injection_all_endpoints(
|
||||
authenticated_client: AsyncClient,
|
||||
) -> None:
|
||||
"""Test that usage, cost, and metadata are injected into all completion endpoints."""
|
||||
|
||||
endpoints = [
|
||||
(
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": "gpt-3.5-turbo",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
},
|
||||
{
|
||||
"id": "chatcmpl-123",
|
||||
"object": "chat.completion",
|
||||
"model": "gpt-3.5-turbo",
|
||||
"choices": [{"message": {"content": "Hi"}, "finish_reason": "stop"}],
|
||||
"usage": {
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 5,
|
||||
"total_tokens": 15,
|
||||
},
|
||||
},
|
||||
),
|
||||
(
|
||||
"/v1/messages",
|
||||
{
|
||||
"model": "claude-3-opus-20240229",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"max_tokens": 100,
|
||||
},
|
||||
{
|
||||
"id": "msg_123",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-3-opus-20240229",
|
||||
"content": [{"type": "text", "text": "Hi"}],
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5},
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
for path, payload, mock_data in endpoints:
|
||||
with patch("httpx.AsyncClient.send") as mock_send:
|
||||
|
||||
async def mock_iter_bytes(
|
||||
*args: object, **kwargs: object
|
||||
) -> AsyncGenerator[bytes, None]:
|
||||
yield json.dumps(mock_data).encode()
|
||||
|
||||
mock_response = AsyncMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.headers = {"content-type": "application/json"}
|
||||
mock_response.text = json.dumps(mock_data)
|
||||
mock_response.json = AsyncMock(return_value=mock_data)
|
||||
mock_response.aiter_bytes = mock_iter_bytes
|
||||
mock_send.return_value = mock_response
|
||||
|
||||
response = await authenticated_client.post(path, json=payload)
|
||||
assert response.status_code == 200
|
||||
|
||||
data = response.json()
|
||||
if hasattr(data, "__await__"):
|
||||
data = await data
|
||||
|
||||
# Check unified usage injection
|
||||
assert "usage" in data
|
||||
assert "cost" in data["usage"]
|
||||
assert "sats_cost" in data["usage"]
|
||||
assert "remaining_balance_msats" in data["usage"]
|
||||
|
||||
# Check metadata injection
|
||||
assert "metadata" in data
|
||||
assert "routstr" in data["metadata"]
|
||||
assert "cost" in data["metadata"]["routstr"]
|
||||
assert "sats_cost" in data["metadata"]["routstr"]
|
||||
assert "remaining_balance_msats" in data["metadata"]["routstr"]
|
||||
|
||||
# Check legacy cost injection
|
||||
assert "cost" in data
|
||||
assert "sats_cost" in data["cost"]
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_usage_injection(
|
||||
authenticated_client: AsyncClient,
|
||||
) -> None:
|
||||
"""Test usage injection for streaming messages."""
|
||||
|
||||
path = "/v1/messages"
|
||||
payload = {
|
||||
"model": "claude-3-opus-20240229",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"max_tokens": 100,
|
||||
"stream": True,
|
||||
}
|
||||
|
||||
# Final usage chunk for Anthropic
|
||||
streaming_chunks = [
|
||||
b'event: message_start\ndata: {"type": "message_start", "message": {"id": "msg_123", "type": "message", "role": "assistant", "model": "claude-3-opus-20240229", "usage": {"input_tokens": 10, "output_tokens": 0}}}\n\n',
|
||||
b'event: message_delta\ndata: {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": null}, "usage": {"output_tokens": 5}}\n\n',
|
||||
b'event: message_stop\ndata: {"type": "message_stop"}\n\n',
|
||||
]
|
||||
|
||||
with patch("httpx.AsyncClient.send") as mock_send:
|
||||
|
||||
async def mock_iter_bytes(
|
||||
*args: object, **kwargs: object
|
||||
) -> AsyncGenerator[bytes, None]:
|
||||
for chunk in streaming_chunks:
|
||||
yield chunk
|
||||
|
||||
mock_response = AsyncMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.headers = {"content-type": "text/event-stream"}
|
||||
mock_response.aiter_bytes = mock_iter_bytes
|
||||
mock_send.return_value = mock_response
|
||||
|
||||
response = await authenticated_client.post(path, json=payload)
|
||||
assert response.status_code == 200
|
||||
|
||||
# In streaming, we look for the final 'event: cost' which is added by Routstr
|
||||
response_text = response.text
|
||||
if hasattr(response_text, "__await__"):
|
||||
response_text = await response_text
|
||||
lines = response_text.split("\n")
|
||||
|
||||
cost_event = next(
|
||||
(line for line in lines if line.startswith('data: {"cost":')), None
|
||||
)
|
||||
assert cost_event is not None
|
||||
|
||||
cost_data = json.loads(cost_event[6:])["cost"]
|
||||
assert "total_msats" in cost_data
|
||||
Reference in New Issue
Block a user