laantungir/routstr-core
1
0
Fork 0
mirror of https://github.com/Routstr/routstr-core.git synced 2026-08-09 02:54:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'add-cost-usage-to-messages-endpoint' into dynamic-prices

Browse Source
This commit is contained in:
9qeklajc
2026-04-15 21:22:38 +02:00
parent 1190a09acb 25f39897d7
commit d7467dbad2
2 changed files with 79 additions and 171 deletions
Download Patch File Download Diff File Expand all files Collapse all files
routstr/upstream/base.py
+79 -26
View File
@@ -122,15 +122,27 @@ class BaseUpstreamProvider:
}
def inject_cost_metadata(
self, response_json: dict, cost_data: CostData | MaxCostData, key: ApiKey
self,
response_json: dict,
cost_data: CostData | MaxCostData | dict,
key: ApiKey,
) -> None:
"""Unifies the injection of cost and usage metadata across all completion types."""
sats_cost = cost_data.total_msats // 1000
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"] = cost_data.total_usd
response_json["usage"]["sats_cost"] = sats_cost
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
@@ -144,13 +156,13 @@ class BaseUpstreamProvider:
# Unified Routstr metadata
response_json["metadata"] = response_json.get("metadata", {})
response_json["metadata"]["routstr"] = {
"cost": cost_data.dict(),
"cost": cost_dict,
"sats_cost": sats_cost,
"remaining_balance_msats": key.balance,
}
# Legacy/Compatibility fields
response_json["cost"] = cost_data.dict()
response_json["cost"] = cost_dict.copy()
response_json["cost"]["sats_cost"] = sats_cost
response_json["cost"]["remaining_balance_msats"] = key.balance
@@ -603,7 +615,7 @@ class BaseUpstreamProvider:
usage_chunk_data["usage"]["cost"] = cost_data.get(
"total_usd", 0.0
)
usage_chunk_data["usage"]["sats_cost"] = (
usage_chunk_data["usage"]["cost_sats"] = (
cost_data.get("total_msats", 0) // 1000
)
usage_chunk_data["usage"]["remaining_balance_msats"] = (
@@ -718,10 +730,7 @@ class BaseUpstreamProvider:
key, response_json, session, deducted_max_cost
)
if isinstance(cost_data, (CostData, MaxCostData)):
self.inject_cost_metadata(response_json, cost_data, key)
else:
response_json["cost"] = cost_data
self.inject_cost_metadata(response_json, cost_data, key)
logger.info(
"Payment adjustment completed for non-streaming",
@@ -1153,7 +1162,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,
@@ -1192,6 +1205,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:
@@ -1201,6 +1216,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(
@@ -1214,10 +1243,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,
@@ -1240,10 +1273,9 @@ class BaseUpstreamProvider:
max_cost_for_model,
)
if isinstance(cost_data, (CostData, MaxCostData)):
self.inject_cost_metadata(
combined_data, cost_data, fresh_key
)
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()
@@ -1284,11 +1316,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}
@@ -1297,10 +1340,7 @@ class BaseUpstreamProvider:
key, response_json, session, deducted_max_cost
)
if isinstance(cost_data, (CostData, MaxCostData)):
self.inject_cost_metadata(response_json, cost_data, key)
else:
response_json["cost"] = cost_data
self.inject_cost_metadata(response_json, cost_data, key)
allowed_headers = {
"content-type",
@@ -1454,7 +1494,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)
@@ -1465,7 +1508,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()
@@ -1475,7 +1523,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()
tests/integration/test_usage_injection.py
-145
View File
@@ -1,145 +0,0 @@
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