From b88e0d6fd0b21fcc3f7fdcf8ee3435bfffdaaadc Mon Sep 17 00:00:00 2001 From: 9qeklajc <9qeklajc> Date: Wed, 12 Nov 2025 23:39:35 +0100 Subject: [PATCH] usage calculation --- routstr/payment/cost_caculation.py | 180 ++++++++++++++++++++++++++++- routstr/upstream/base.py | 31 +++-- 2 files changed, 200 insertions(+), 11 deletions(-) diff --git a/routstr/payment/cost_caculation.py b/routstr/payment/cost_caculation.py index f8eb4ffb..f442ae20 100644 --- a/routstr/payment/cost_caculation.py +++ b/routstr/payment/cost_caculation.py @@ -1,4 +1,7 @@ +import base64 import math +import struct +from typing import Optional from pydantic.v1 import BaseModel @@ -25,8 +28,163 @@ class CostDataError(BaseModel): code: str +def get_image_resolution_from_data_url(data_url: str) -> Optional[tuple[int, int]]: + """ + Extract image resolution (width, height) from a base64 data URL without DOM. + Supports PNG and JPEG. Returns None if format unsupported or parsing fails. + """ + try: + if not isinstance(data_url, str) or not data_url.startswith("data:"): + return None + + comma_idx = data_url.find(",") + if comma_idx == -1: + return None + + meta = data_url[5:comma_idx] # e.g. "image/png;base64" + base64_data = data_url[comma_idx + 1:] + + # Decode base64 to binary + try: + binary_data = base64.b64decode(base64_data) + except Exception: + return None + + is_png = "image/png" in meta + is_jpeg = "image/jpeg" in meta or "image/jpg" in meta + + # PNG: width/height are 4-byte big-endian at offsets 16 and 20 + if is_png: + # Validate PNG signature + png_sig = b'\x89PNG\r\n\x1a\n' + if not binary_data.startswith(png_sig): + return None + + if len(binary_data) < 24: + return None + + # Width and height are at bytes 16-19 and 20-23 respectively + width = struct.unpack('>I', binary_data[16:20])[0] + height = struct.unpack('>I', binary_data[20:24])[0] + + if width > 0 and height > 0: + return (width, height) + return None + + # JPEG: parse markers to SOF0/SOF2 for dimensions + if is_jpeg: + offset = 0 + # JPEG SOI 0xFFD8 + if len(binary_data) < 2 or binary_data[0] != 0xFF or binary_data[1] != 0xD8: + return None + offset = 2 + + while offset < len(binary_data): + # Find marker + while offset < len(binary_data) and binary_data[offset] != 0xFF: + offset += 1 + if offset + 1 >= len(binary_data): + break + + # Skip fill bytes 0xFF + while offset < len(binary_data) and binary_data[offset] == 0xFF: + offset += 1 + if offset >= len(binary_data): + break + + marker = binary_data[offset] + offset += 1 + + # Standalone markers without length + if marker == 0xD8 or marker == 0xD9: # SOI/EOI + continue + + if offset + 1 >= len(binary_data): + break + + length = (binary_data[offset] << 8) | binary_data[offset + 1] + offset += 2 + + # SOF0 (0xC0) or SOF2 (0xC2) contain dimensions + if marker == 0xC0 or marker == 0xC2: + if length < 7 or offset + length - 2 > len(binary_data): + return None + # Skip precision byte + height = (binary_data[offset + 1] << 8) | binary_data[offset + 2] + width = (binary_data[offset + 3] << 8) | binary_data[offset + 4] + if width > 0 and height > 0: + return (width, height) + return None + else: + # Skip this segment + offset += length - 2 + + return None + + # Unsupported formats (e.g., webp/gif) - skip for now + return None + except Exception: + return None + + +def calculate_image_tokens_from_messages(messages: list) -> int: + """ + Calculate image tokens from messages using 32px patch method. + """ + image_tokens = 0 + + try: + for msg in messages: + content = msg.get("content") + if isinstance(content, list): + for part in content: + if (isinstance(part, dict) and + part.get("type") == "image_url"): + + image_url = part.get("image_url") + url: Optional[str] = None + if isinstance(image_url, str): + url = image_url + elif isinstance(image_url, dict): + url = image_url.get("url") + else: + continue + + # Expecting a base64 data URL for local image inputs + if url and isinstance(url, str) and url.startswith("data:"): + resolution = get_image_resolution_from_data_url(url) + if resolution: + width, height = resolution + patch_size = 32 + patches_w = (width + patch_size - 1) // patch_size + patches_h = (height + patch_size - 1) // patch_size + tokens_from_image = patches_w * patches_h + image_tokens += tokens_from_image + + logger.debug( + "Calculated image tokens", + extra={ + "width": width, + "height": height, + "tokens_from_image": tokens_from_image, + } + ) + else: + logger.warning( + "Could not determine image resolution", + extra={"url_prefix": url[:50] + "..." if len(url) > 50 else url} + ) + except Exception as e: + logger.error( + "Error calculating image tokens", + extra={"error": str(e), "error_type": type(e).__name__} + ) + + return image_tokens + + async def calculate_cost( # todo: can be sync - response_data: dict, max_cost: int, session: AsyncSession + response_data: dict, max_cost: int, session: AsyncSession, request_data: Optional[dict] = None ) -> CostData | MaxCostData | CostDataError: """ Calculate the cost of an API request based on token usage. @@ -34,6 +192,8 @@ async def calculate_cost( # todo: can be sync Args: response_data: Response data containing usage information max_cost: Maximum cost in millisats + session: Database session + request_data: Original request data containing messages (for image token calculation) Returns: Cost data or error information @@ -132,14 +292,28 @@ async def calculate_cost( # todo: can be sync input_tokens = response_data.get("usage", {}).get("prompt_tokens", 0) output_tokens = response_data.get("usage", {}).get("completion_tokens", 0) - input_msats = round(input_tokens / 1000 * MSATS_PER_1K_INPUT_TOKENS, 3) + # Calculate image tokens from request data if available + image_tokens = 0 + if request_data and "messages" in request_data: + image_tokens = calculate_image_tokens_from_messages(request_data["messages"]) + logger.debug( + "Image tokens calculated", + extra={"image_tokens": image_tokens, "text_input_tokens": input_tokens} + ) + + # Add image tokens to input tokens for cost calculation + total_input_tokens = input_tokens + image_tokens + + input_msats = round(total_input_tokens / 1000 * MSATS_PER_1K_INPUT_TOKENS, 3) output_msats = round(output_tokens / 1000 * MSATS_PER_1K_OUTPUT_TOKENS, 3) token_based_cost = math.ceil(input_msats + output_msats) logger.info( "Calculated token-based cost", extra={ - "input_tokens": input_tokens, + "text_input_tokens": input_tokens, + "image_tokens": image_tokens, + "total_input_tokens": total_input_tokens, "output_tokens": output_tokens, "input_cost_msats": input_msats, "output_cost_msats": output_msats, diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index 2007996e..9c0b663f 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -4,7 +4,7 @@ import json import re import traceback from collections.abc import AsyncGenerator -from typing import Mapping +from typing import Mapping, Optional import httpx from fastapi import BackgroundTasks, HTTPException, Request @@ -875,13 +875,14 @@ class BaseUpstreamProvider: ) async def get_x_cashu_cost( - self, response_data: dict, max_cost_for_model: int + self, response_data: dict, max_cost_for_model: int, request_data: Optional[dict] = None ) -> MaxCostData | CostData | None: """Calculate cost for X-Cashu payment based on response data. Args: response_data: Response data containing model and usage information max_cost_for_model: Maximum cost for the model + request_data: Original request data containing messages (for image token calculation) Returns: Cost data object (MaxCostData or CostData) or None if calculation fails @@ -893,7 +894,7 @@ class BaseUpstreamProvider: ) async with create_session() as session: - match await calculate_cost(response_data, max_cost_for_model, session): + match await calculate_cost(response_data, max_cost_for_model, session, request_data): case MaxCostData() as cost: logger.debug( "Using max cost pricing", @@ -1017,6 +1018,7 @@ class BaseUpstreamProvider: unit: str, max_cost_for_model: int, mint: str | None = None, + request_data: Optional[dict] = None, ) -> StreamingResponse: """Handle streaming response for X-Cashu payment, calculating refund if needed. @@ -1075,7 +1077,7 @@ class BaseUpstreamProvider: response_data = {"usage": usage_data, "model": model} try: cost_data = await self.get_x_cashu_cost( - response_data, max_cost_for_model + response_data, max_cost_for_model, request_data ) if cost_data: if unit == "msat": @@ -1150,6 +1152,7 @@ class BaseUpstreamProvider: unit: str, max_cost_for_model: int, mint: str | None = None, + request_data: Optional[dict] = None, ) -> Response: """Handle non-streaming response for X-Cashu payment, calculating refund if needed. @@ -1170,7 +1173,7 @@ class BaseUpstreamProvider: try: response_json = json.loads(content_str) - cost_data = await self.get_x_cashu_cost(response_json, max_cost_for_model) + cost_data = await self.get_x_cashu_cost(response_json, max_cost_for_model, request_data) if not cost_data: logger.error( @@ -1280,6 +1283,7 @@ class BaseUpstreamProvider: unit: str, max_cost_for_model: int, mint: str | None = None, + request_data: Optional[dict] = None, ) -> StreamingResponse | Response: """Handle chat completion response for X-Cashu payment, detecting streaming vs non-streaming. @@ -1316,11 +1320,11 @@ class BaseUpstreamProvider: if is_streaming: return await self.handle_x_cashu_streaming_response( - content_str, response, amount, unit, max_cost_for_model, mint + content_str, response, amount, unit, max_cost_for_model, mint, request_data ) else: return await self.handle_x_cashu_non_streaming_response( - content_str, response, amount, unit, max_cost_for_model, mint + content_str, response, amount, unit, max_cost_for_model, mint, request_data ) except Exception as e: @@ -1372,6 +1376,17 @@ class BaseUpstreamProvider: request_body = await request.body() transformed_body = self.prepare_request_body(request_body, model_obj) + # Parse request data for image token calculation + request_data = None + try: + if request_body: + request_data = json.loads(request_body.decode('utf-8')) + except Exception as e: + logger.warning( + "Could not parse request body for image token calculation", + extra={"error": str(e)} + ) + logger.debug( "Forwarding request to upstream", extra={ @@ -1457,7 +1472,7 @@ class BaseUpstreamProvider: ) result = await self.handle_x_cashu_chat_completion( - response, amount, unit, max_cost_for_model, mint + response, amount, unit, max_cost_for_model, mint, request_data ) background_tasks = BackgroundTasks() background_tasks.add_task(response.aclose)