diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index a244b1fa..00d7563a 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -2152,6 +2152,17 @@ class BaseUpstreamProvider: }, ) + if cost_data: + for i, line in enumerate(lines): + if line.startswith("data: "): + try: + data_json = json.loads(line[6:]) + if "usage" in data_json and data_json["usage"]: + data_json["usage"]["cost_sats"] = cost_data.total_msats // 1000 + lines[i] = "data: " + json.dumps(data_json) + except json.JSONDecodeError: + pass + async def generate() -> AsyncGenerator[bytes, None]: for line in lines: yield (line + "\n").encode("utf-8") @@ -2196,6 +2207,9 @@ class BaseUpstreamProvider: response_json = json.loads(content_str) cost_data = await self.get_x_cashu_cost(response_json, max_cost_for_model) + if cost_data and "usage" in response_json: + response_json["usage"]["cost_sats"] = cost_data.total_msats // 1000 + if not cost_data: logger.error( "Failed to calculate cost for response", @@ -2262,7 +2276,7 @@ class BaseUpstreamProvider: ) return Response( - content=content_str, + content=json.dumps(response_json), status_code=response.status_code, headers=response_headers, media_type="application/json", @@ -3067,6 +3081,17 @@ class BaseUpstreamProvider: }, ) + if cost_data: + for i, line in enumerate(lines): + if line.startswith("data: "): + try: + data_json = json.loads(line[6:]) + if "usage" in data_json and data_json["usage"]: + data_json["usage"]["cost_sats"] = cost_data.total_msats // 1000 + lines[i] = "data: " + json.dumps(data_json) + except json.JSONDecodeError: + pass + async def generate() -> AsyncGenerator[bytes, None]: for line in lines: yield (line + "\\n").encode("utf-8") @@ -3099,6 +3124,9 @@ class BaseUpstreamProvider: response_json = json.loads(content_str) cost_data = await self.get_x_cashu_cost(response_json, max_cost_for_model) + if cost_data and "usage" in response_json: + response_json["usage"]["cost_sats"] = cost_data.total_msats // 1000 + if not cost_data: logger.error( "Failed to calculate cost for Responses API response", @@ -3165,7 +3193,7 @@ class BaseUpstreamProvider: ) return Response( - content=content_str, + content=json.dumps(response_json), status_code=response.status_code, headers=response_headers, media_type="application/json", diff --git a/tests/unit/test_x_cashu_cost_sats.py b/tests/unit/test_x_cashu_cost_sats.py new file mode 100644 index 00000000..d37d7acc --- /dev/null +++ b/tests/unit/test_x_cashu_cost_sats.py @@ -0,0 +1,216 @@ +import json +import os +from unittest.mock import AsyncMock, patch + +import httpx +import pytest + +os.environ.setdefault("UPSTREAM_BASE_URL", "http://test") +os.environ.setdefault("UPSTREAM_API_KEY", "test") + +from routstr.payment.cost_calculation import CostData # noqa: E402 +from routstr.upstream.base import BaseUpstreamProvider # noqa: E402 + + +def _make_provider() -> BaseUpstreamProvider: + return BaseUpstreamProvider(base_url="http://test", api_key="test-key") + + +def _make_httpx_response(status_code: int = 200) -> httpx.Response: + return httpx.Response(status_code, headers={}) + + +def _make_cost_data(total_msats: int = 5000) -> CostData: + return CostData( + base_msats=0, + input_msats=3000, + output_msats=2000, + total_msats=total_msats, + total_usd=0.00025, + input_tokens=100, + output_tokens=50, + ) + + +# --------------------------------------------------------------------------- +# Non-streaming (chat completions) +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_non_streaming_includes_cost_sats() -> None: + provider = _make_provider() + cost_data = _make_cost_data(total_msats=5000) + + response_body = { + "model": "gpt-4o", + "usage": { + "prompt_tokens": 100, + "completion_tokens": 50, + "total_tokens": 150, + "cost": 0.00025, + }, + } + content_str = json.dumps(response_body) + httpx_response = _make_httpx_response() + + with ( + patch.object(provider, "get_x_cashu_cost", new=AsyncMock(return_value=cost_data)), + patch.object(provider, "send_refund", new=AsyncMock(return_value="cashuA_refund_token")), + ): + response = await provider.handle_x_cashu_non_streaming_response( + content_str=content_str, + response=httpx_response, + amount=10000, + unit="msat", + max_cost_for_model=10000, + mint=None, + payment_token_hash=None, + ) + + body = json.loads(response.body) + assert "cost_sats" in body["usage"] + assert body["usage"]["cost_sats"] == 5 # 5000 msats // 1000 + + +@pytest.mark.asyncio +async def test_non_streaming_cost_sats_value_rounds_down() -> None: + provider = _make_provider() + cost_data = _make_cost_data(total_msats=1999) + + response_body = {"model": "gpt-4o", "usage": {"prompt_tokens": 10}} + content_str = json.dumps(response_body) + + with ( + patch.object(provider, "get_x_cashu_cost", new=AsyncMock(return_value=cost_data)), + patch.object(provider, "send_refund", new=AsyncMock(return_value="cashuA_refund_token")), + ): + response = await provider.handle_x_cashu_non_streaming_response( + content_str=content_str, + response=_make_httpx_response(), + amount=10000, + unit="msat", + max_cost_for_model=10000, + ) + + body = json.loads(response.body) + assert body["usage"]["cost_sats"] == 1 # 1999 // 1000 + + +@pytest.mark.asyncio +async def test_non_streaming_preserves_existing_usage_fields() -> None: + provider = _make_provider() + cost_data = _make_cost_data(total_msats=3000) + + response_body = { + "model": "gpt-4o", + "usage": { + "prompt_tokens": 100, + "completion_tokens": 50, + "total_tokens": 150, + "cost": 0.00015, + }, + } + + with ( + patch.object(provider, "get_x_cashu_cost", new=AsyncMock(return_value=cost_data)), + patch.object(provider, "send_refund", new=AsyncMock(return_value="cashuA_refund_token")), + ): + response = await provider.handle_x_cashu_non_streaming_response( + content_str=json.dumps(response_body), + response=_make_httpx_response(), + amount=10000, + unit="msat", + max_cost_for_model=10000, + ) + + body = json.loads(response.body) + usage = body["usage"] + assert usage["prompt_tokens"] == 100 + assert usage["completion_tokens"] == 50 + assert usage["total_tokens"] == 150 + assert usage["cost"] == 0.00015 + assert usage["cost_sats"] == 3 + + +# --------------------------------------------------------------------------- +# Streaming (chat completions) +# --------------------------------------------------------------------------- + +async def _collect_streaming(response: object) -> list[str]: + chunks: list[str] = [] + async for chunk in response.body_iterator: # type: ignore[attr-defined] + if isinstance(chunk, bytes): + chunks.append(chunk.decode("utf-8")) + else: + chunks.append(str(chunk)) + return chunks + + +@pytest.mark.asyncio +async def test_streaming_includes_cost_sats_in_usage_chunk() -> None: + provider = _make_provider() + cost_data = _make_cost_data(total_msats=7000) + + usage_chunk = { + "id": "chatcmpl-123", + "model": "gpt-4o", + "usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150}, + } + content_str = "\n".join([ + 'data: {"id":"chatcmpl-123","model":"gpt-4o","choices":[]}', + f"data: {json.dumps(usage_chunk)}", + "data: [DONE]", + ]) + + with patch.object(provider, "get_x_cashu_cost", new=AsyncMock(return_value=cost_data)): + response = await provider.handle_x_cashu_streaming_response( + content_str=content_str, + response=_make_httpx_response(), + amount=10000, + unit="msat", + max_cost_for_model=10000, + mint=None, + payment_token_hash=None, + ) + + chunks = await _collect_streaming(response) + + full_output = "".join(chunks) + usage_line = next( + line for line in full_output.split("\n") if '"usage"' in line and "cost_sats" in line + ) + data_json = json.loads(usage_line.lstrip("data: ").strip()) + assert data_json["usage"]["cost_sats"] == 7 # 7000 // 1000 + + +@pytest.mark.asyncio +async def test_streaming_non_usage_chunks_unmodified() -> None: + provider = _make_provider() + cost_data = _make_cost_data(total_msats=2000) + + regular_chunk = {"id": "chatcmpl-123", "model": "gpt-4o", "choices": [{"delta": {"content": "hi"}}]} + usage_chunk = {"id": "chatcmpl-123", "model": "gpt-4o", "usage": {"prompt_tokens": 10}} + content_str = "\n".join([ + f"data: {json.dumps(regular_chunk)}", + f"data: {json.dumps(usage_chunk)}", + "data: [DONE]", + ]) + + with patch.object(provider, "get_x_cashu_cost", new=AsyncMock(return_value=cost_data)): + response = await provider.handle_x_cashu_streaming_response( + content_str=content_str, + response=_make_httpx_response(), + amount=10000, + unit="msat", + max_cost_for_model=10000, + ) + + chunks = await _collect_streaming(response) + + lines = [ + line for line in "".join(chunks).split("\n") + if line.startswith("data: ") and line != "data: [DONE]" + ] + regular_line_data = json.loads(lines[0][6:]) + # regular chunk should not have cost_sats injected + assert "cost_sats" not in regular_line_data.get("usage", {})