clean up

set in/out msat prices correctly

routstr/payment/cost_calculation.py

@@ -109,12 +109,20 @@ async def calculate_cost(  # todo: can be sync
     )
 
     usd_cost = 0.0
+    input_usd = 0.0
+    output_usd = 0.0
 
-    # Prioritize cost_details.upstream_inference_cost
     if "cost_details" in usage_data:
         usd_cost = float(
             usage_data["cost_details"].get("upstream_inference_cost", 0) or 0
         )
+        input_usd = float(
+            usage_data["cost_details"].get("upstream_inference_prompt_cost", 0) or 0
+        )
+        output_usd = float(
+            usage_data["cost_details"].get("upstream_inference_completions_cost", 0)
+            or 0
+        )
 
     # Fallback to cost field if upstream_inference_cost is 0
     if usd_cost == 0 and "cost" in usage_data:
@@ -123,12 +131,34 @@ async def calculate_cost(  # todo: can be sync
         except Exception:
             pass
 
+    MSATS_PER_1K_INPUT_TOKENS: float = (
+        float(settings.fixed_per_1k_input_tokens) * 1000.0
+    )
+    MSATS_PER_1K_OUTPUT_TOKENS: float = (
+        float(settings.fixed_per_1k_output_tokens) * 1000.0
+    )
+
     if usd_cost > 0:
         try:
             sats_per_usd = 1.0 / sats_usd_price()
             cost_in_sats = usd_cost * sats_per_usd
             cost_in_msats = math.ceil(cost_in_sats * 1000)
 
+            input_msats = 0
+            output_msats = 0
+
+            if input_usd > 0 or output_usd > 0:
+                input_msats = int((input_usd * sats_per_usd) * 1000)
+                output_msats = int((output_usd * sats_per_usd) * 1000)
+            else:
+                total_tokens = input_tokens + output_tokens
+                if total_tokens > 0:
+                    input_ratio = input_tokens / total_tokens
+                    input_msats = int(cost_in_msats * input_ratio)
+                    output_msats = cost_in_msats - input_msats
+                else:
+                    output_msats = cost_in_msats
+
             logger.info(
                 "Using cost from usage data/details",
                 extra={
@@ -140,9 +170,9 @@ async def calculate_cost(  # todo: can be sync
             )
 
             return CostData(
-                base_msats=-1,
-                input_msats=-1,  # Cost field doesn't break down by token type
-                output_msats=-1,
+                base_msats=0,
+                input_msats=input_msats,
+                output_msats=output_msats,
                 total_msats=cost_in_msats,
                 total_usd=usd_cost,
                 input_tokens=input_tokens,
@@ -159,13 +189,6 @@ async def calculate_cost(  # todo: can be sync
             )
             # Fall through to token-based calculation
 
-    MSATS_PER_1K_INPUT_TOKENS: float = (
-        float(settings.fixed_per_1k_input_tokens) * 1000.0
-    )
-    MSATS_PER_1K_OUTPUT_TOKENS: float = (
-        float(settings.fixed_per_1k_output_tokens) * 1000.0
-    )
-
     if not settings.fixed_pricing:
         response_model = response_data.get("model", "")
         logger.debug(
@@ -231,10 +254,10 @@ async def calculate_cost(  # todo: can be sync
             output_tokens=output_tokens,
         )
 
-    input_msats = round(input_tokens / 1000 * MSATS_PER_1K_INPUT_TOKENS, 3)
+    calc_input_msats = round(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)
+    calc_output_msats = round(output_tokens / 1000 * MSATS_PER_1K_OUTPUT_TOKENS, 3)
+    token_based_cost = math.ceil(calc_input_msats + calc_output_msats)
     total_usd = (token_based_cost / 1000.0) * sats_usd_price()
 
     logger.info(
@@ -242,8 +265,8 @@ async def calculate_cost(  # todo: can be sync
         extra={
             "input_tokens": input_tokens,
             "output_tokens": output_tokens,
-            "input_cost_msats": input_msats,
-            "output_cost_msats": output_msats,
+            "input_cost_msats": calc_input_msats,
+            "output_cost_msats": calc_output_msats,
             "total_cost_msats": token_based_cost,
             "total_usd": total_usd,
             "model": response_data.get("model", "unknown"),
@@ -252,8 +275,8 @@ async def calculate_cost(  # todo: can be sync
 
     return CostData(
         base_msats=0,
-        input_msats=int(input_msats),
-        output_msats=int(output_msats),
+        input_msats=int(calc_input_msats),
+        output_msats=int(calc_output_msats),
         total_msats=token_based_cost,
         total_usd=total_usd,
         input_tokens=input_tokens,

routstr/upstream/base.py

@@ -452,6 +452,7 @@ class BaseUpstreamProvider:
         key: ApiKey,
         max_cost_for_model: int,
         background_tasks: BackgroundTasks,
+        requested_model: str | None = None,
     ) -> StreamingResponse:
         """Handle streaming chat completion responses with token usage tracking and cost adjustment.
 
@@ -520,11 +521,15 @@ class BaseUpstreamProvider:
                             if isinstance(obj, dict):
                                 if obj.get("model"):
                                     last_model_seen = str(obj.get("model"))
+                                if requested_model:
+                                    obj["model"] = requested_model
 
                                 if isinstance(obj.get("usage"), dict):
                                     # Hold this chunk back to merge cost later
                                     usage_chunk_data = obj
                                     continue
+                                yield b"data: " + json.dumps(obj).encode() + b"\n\n"
+                                continue
                         except json.JSONDecodeError:
                             pass
 
@@ -620,6 +625,7 @@ class BaseUpstreamProvider:
         key: ApiKey,
         session: AsyncSession,
         deducted_max_cost: int,
+        requested_model: str | None = None,
     ) -> Response:
         """Handle non-streaming chat completion responses with token usage tracking and cost adjustment.
 
@@ -668,9 +674,7 @@ class BaseUpstreamProvider:
                 response_json["usage"]["cost_sats"] = (
                     cost_data.get("total_msats", 0) // 1000
                 )
-                response_json["usage"]["remaining_balance_msats"] = (
-                    remaining_balance_msats
-                )
+                response_json["usage"]["remaining_balance_msats"] = remaining_balance_msats
 
             # Keep detailed cost
             response_json["metadata"] = response_json.get("metadata", {})
@@ -714,6 +718,8 @@ class BaseUpstreamProvider:
                 if k.lower() in allowed_headers
             }
 
+            if requested_model:
+                response_json["model"] = requested_model
             return Response(
                 content=json.dumps(response_json).encode(),
                 status_code=response.status_code,
@@ -744,7 +750,7 @@ class BaseUpstreamProvider:
             raise
 
     async def handle_streaming_responses_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:
         """Handle streaming Responses API responses with token usage tracking and cost adjustment.
 
@@ -814,6 +820,8 @@ class BaseUpstreamProvider:
                             if isinstance(obj, dict):
                                 if obj.get("model"):
                                     last_model_seen = str(obj.get("model"))
+                                if requested_model:
+                                    obj["model"] = requested_model
 
                                 # Track reasoning tokens for Responses API
                                 if usage := obj.get("usage", {}):
@@ -934,6 +942,7 @@ class BaseUpstreamProvider:
         key: ApiKey,
         session: AsyncSession,
         deducted_max_cost: int,
+        requested_model: str | None = None,
     ) -> Response:
         """Handle non-streaming Responses API responses with token usage tracking and cost adjustment.
 
@@ -985,9 +994,7 @@ class BaseUpstreamProvider:
                 response_json["usage"]["cost_sats"] = (
                     cost_data.get("total_msats", 0) // 1000
                 )
-                response_json["usage"]["remaining_balance_msats"] = (
-                    remaining_balance_msats
-                )
+                response_json["usage"]["remaining_balance_msats"] = remaining_balance_msats
 
             # Keep detailed cost
             response_json["metadata"] = response_json.get("metadata", {})
@@ -1031,6 +1038,8 @@ class BaseUpstreamProvider:
                 if k.lower() in allowed_headers
             }
 
+            if requested_model:
+                response_json["model"] = requested_model
             return Response(
                 content=json.dumps(response_json).encode(),
                 status_code=response.status_code,
@@ -1294,6 +1303,8 @@ class BaseUpstreamProvider:
         path = self.normalize_request_path(path, model_obj)
         url = self.build_request_url(path, model_obj)
 
+        original_model_id = model_obj.id if model_obj else None
+
         transformed_body = self.prepare_request_body(request_body, model_obj)
 
         logger.info(
@@ -1452,7 +1463,8 @@ class BaseUpstreamProvider:
                         background_tasks.add_task(response.aclose)
                         background_tasks.add_task(client.aclose)
                         result = await self.handle_streaming_chat_completion(
-                            response, key, max_cost_for_model, background_tasks
+                            response, key, max_cost_for_model, background_tasks,
+                            requested_model=original_model_id,
                         )
                         result.background = background_tasks
                         return result
@@ -1461,7 +1473,8 @@ class BaseUpstreamProvider:
                 if response.status_code == 200:
                     try:
                         return await self.handle_non_streaming_chat_completion(
-                            response, key, session, max_cost_for_model
+                            response, key, session, max_cost_for_model,
+                            requested_model=original_model_id,
                         )
                     finally:
                         await response.aclose()
@@ -1576,6 +1589,8 @@ class BaseUpstreamProvider:
         path = self.normalize_request_path(path, model_obj)
         url = self.build_request_url(path, model_obj)
 
+        original_model_id = model_obj.id if model_obj else None
+
         transformed_body = self.prepare_responses_request_body(request_body, model_obj)
 
         logger.info(
@@ -1662,7 +1677,8 @@ class BaseUpstreamProvider:
 
                 if is_streaming and response.status_code == 200:
                     result = await self.handle_streaming_responses_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)
@@ -1673,7 +1689,8 @@ class BaseUpstreamProvider:
                 if response.status_code == 200:
                     try:
                         return await self.handle_non_streaming_responses_completion(
-                            response, key, session, max_cost_for_model
+                            response, key, session, max_cost_for_model,
+                            requested_model=original_model_id,
                         )
                     finally:
                         await response.aclose()