usage calculation

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,

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)