From a1ae6e94e91deecafb6cb1d19290de14db7da4c0 Mon Sep 17 00:00:00 2001 From: 9qeklajc Date: Wed, 15 Apr 2026 19:20:50 +0200 Subject: [PATCH 1/3] match model when versioned --- routstr/proxy.py | 20 ++- tests/integration/test_usage_injection.py | 145 ---------------------- 2 files changed, 19 insertions(+), 146 deletions(-) delete mode 100644 tests/integration/test_usage_injection.py diff --git a/routstr/proxy.py b/routstr/proxy.py index ddab1c31..9d2dbf09 100644 --- a/routstr/proxy.py +++ b/routstr/proxy.py @@ -69,7 +69,25 @@ def get_upstreams() -> list[BaseUpstreamProvider]: def get_model_instance(model_id: str) -> Model | None: """Get Model instance by ID from global cache.""" - return _model_instances.get(model_id.lower()) + if not model_id: + return None + + model_id_lower = model_id.lower() + # Try exact match first + if model := _model_instances.get(model_id_lower): + return model + + # Try stripping common version suffixes (e.g., -20251222) + # This handles cases where upstream returns a specific version + # but we only track the base model name. + import re + + base_model_id = re.sub(r"-\d{8}$", "", model_id_lower) + if base_model_id != model_id_lower: + if model := _model_instances.get(base_model_id): + return model + + return None def get_provider_for_model(model_id: str) -> list[BaseUpstreamProvider] | None: diff --git a/tests/integration/test_usage_injection.py b/tests/integration/test_usage_injection.py deleted file mode 100644 index 59b672fd..00000000 --- a/tests/integration/test_usage_injection.py +++ /dev/null @@ -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 From 45bdfaee58eee4757da70261803f666ec9a1b445 Mon Sep 17 00:00:00 2001 From: 9qeklajc Date: Wed, 15 Apr 2026 20:18:56 +0200 Subject: [PATCH 2/3] best effort to be compatible with legacy code --- routstr/upstream/base.py | 39 ++++++++++++++++++++++----------------- 1 file changed, 22 insertions(+), 17 deletions(-) diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index 10232749..b1827a57 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -120,14 +120,26 @@ 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"]["cost"] = total_usd response_json["usage"]["sats_cost"] = sats_cost response_json["usage"]["remaining_balance_msats"] = key.balance @@ -142,13 +154,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 @@ -716,10 +728,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", @@ -1238,10 +1247,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() @@ -1295,10 +1303,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", From 25f39897d7119dcaaa7d1c3e1089859c74501431 Mon Sep 17 00:00:00 2001 From: 9qeklajc Date: Wed, 15 Apr 2026 20:45:13 +0200 Subject: [PATCH 3/3] mirror logic from chat completion --- routstr/upstream/base.py | 66 ++++++++++++++++++++++++++++++++++------ 1 file changed, 57 insertions(+), 9 deletions(-) diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index b1827a57..d349ab7f 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -140,7 +140,7 @@ class BaseUpstreamProvider: # Inject into top-level usage block (OpenAI/Anthropic style) if "usage" in response_json: response_json["usage"]["cost"] = total_usd - response_json["usage"]["sats_cost"] = sats_cost + response_json["usage"]["cost_sats"] = sats_cost response_json["usage"]["remaining_balance_msats"] = key.balance # Inject into Anthropic nested usage block if present @@ -613,7 +613,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"] = ( @@ -1160,7 +1160,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, @@ -1199,6 +1203,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: @@ -1208,6 +1214,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( @@ -1221,10 +1241,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, @@ -1290,11 +1314,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} @@ -1457,7 +1492,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) @@ -1468,7 +1506,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() @@ -1478,7 +1521,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()