diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index b506e388..10232749 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -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", diff --git a/tests/integration/test_usage_injection.py b/tests/integration/test_usage_injection.py new file mode 100644 index 00000000..59b672fd --- /dev/null +++ b/tests/integration/test_usage_injection.py @@ -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