clean up

set in/out msat prices correctly

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()