Update gen.py

gen.py

@@ -777,11 +777,9 @@ async def generate_text(
             headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"},
         )
 
-    # ── non-streaming ─────────────────────────
     async with httpx.AsyncClient(timeout=None) as client:
         r = await client.post(url, json=body, headers=headers)
 
-        # navy-vision fallback
         if provider == "navy vision" and r.status_code >= 400:
             print("[FALLBACK] Navy vision failed — switching to fallback")
             fb_url, fb_key = _get_provider_url_and_key(FALLBACK_PROVIDER)
@@ -796,9 +794,6 @@ async def generate_text(
         except Exception:
             payload = {"error": "Upstream returned invalid JSON"}
         else:
-            # Normalize usage: upstream may use prompt_tokens/completion_tokens
-            # (OpenAI/Groq style) — rewrite to a consistent shape and add
-            # router_metadata so callers always see the same fields.
             if "usage" in payload and isinstance(payload["usage"], dict):
                 u = payload["usage"]
                 input_tok  = u.get("prompt_tokens")  or u.get("input_tokens",  0)
@@ -807,7 +802,6 @@ async def generate_text(
                     "prompt_tokens":     input_tok,
                     "completion_tokens": output_tok,
                     "total_tokens":      input_tok + output_tok,
-                    # also include the OpenAI Responses-API names for clients that expect them
                     "input_tokens":  input_tok,
                     "output_tokens": output_tok,
                 }
@@ -821,11 +815,6 @@ async def generate_text(
 
     return JSONResponse(status_code=r.status_code, content=payload)
 
-
-# ──────────────────────────────────────────────
-# PROMPT ANALYZE  (/gen/prompt_analyze)
-# ──────────────────────────────────────────────
-
 @router.post("/prompt_analyze")
 async def analyze_prompt(request: Request):
     body = await request.json()
@@ -840,11 +829,6 @@ async def analyze_prompt(request: Request):
     chosen_model, _ = route_chat(messages, uses_tools=uses_tools)
     return {MODEL_MAP.get(chosen_model, chosen_model)}
 
-
-# ──────────────────────────────────────────────
-# MODELS LIST
-# ──────────────────────────────────────────────
-
 @router.get("/models")
 def return_models_openai():
     return {
@@ -944,7 +928,6 @@ def _build_responses_payload(
         },
     }
 
-
 @router.post("/responses")
 async def create_responses(
     request: Request,
@@ -966,10 +949,12 @@ async def create_responses(
     if not messages:
         raise HTTPException(400, "input could not be parsed")
 
-    # ── shared helper: route + call + return (text, input_tokens, output_tokens) ──
+    # ✅ ROUTE + LOG + RATE LIMIT EARLY
+    chosen_model, provider = route_chat(messages)
+    _log_routing(chosen_model, provider, messages, uses_tools=False)
+    await _check_chat_rate_limit(request, authorization, x_client_id)
+
     async def _generate() -> Tuple[str, int, int]:
-        chosen_model, provider = route_chat(messages)
-        await _check_chat_rate_limit(request, authorization, x_client_id)
         data = await call_chat_completions(messages, chosen_model, provider)
         text = _extract_text_from_response(data)
         input_tokens, output_tokens = _extract_usage(data)
@@ -979,10 +964,20 @@ async def create_responses(
     if stream is False:
         text, input_tokens, output_tokens = await _generate()
         response_id = _resp_id("resp")
-        return JSONResponse(
-            content=_build_responses_payload(model, text, response_id, input_tokens, output_tokens)
+
+        payload = _build_responses_payload(
+            chosen_model,
+            text,
+            response_id,
+            input_tokens,
+            output_tokens,
         )
 
+        payload.setdefault("router_metadata", {})
+        payload["router_metadata"]["provider"] = provider
+
+        return JSONResponse(content=payload)
+
     # ── streaming ─────────────────────────────
     async def event_stream():
         response_id = _resp_id("resp")
@@ -990,18 +985,20 @@ async def create_responses(
         ts          = _resp_ts()
 
         def sse(event_type: str, data: dict) -> str:
-            """Emit a properly-formed SSE frame with both event: and data: lines.
-            The OpenAI SDK dispatches on the `event:` field — without it most
-            events are silently dropped."""
             return f"event: {event_type}\ndata: {json.dumps(data)}\n\n"
 
         # 1. response.created
         yield sse("response.created", {
             "type": "response.created",
             "response": {
-                "id": response_id, "object": "response",
-                "created_at": ts, "status": "in_progress", "model": model,
-                "output": [], "usage": None,
+                "id": response_id,
+                "object": "response",
+                "created_at": ts,
+                "status": "in_progress",
+                "model": chosen_model,
+                "output": [],
+                "usage": None,
+                "router_metadata": {"provider": provider},
             },
         })
 
@@ -1009,8 +1006,11 @@ async def create_responses(
         yield sse("response.in_progress", {
             "type": "response.in_progress",
             "response": {
-                "id": response_id, "object": "response",
-                "created_at": ts, "status": "in_progress", "model": model,
+                "id": response_id,
+                "object": "response",
+                "created_at": ts,
+                "status": "in_progress",
+                "model": chosen_model,
             },
         })
 
@@ -1020,27 +1020,32 @@ async def create_responses(
             yield sse("response.failed", {
                 "type": "response.failed",
                 "response": {
-                    "id": response_id, "object": "response",
-                    "created_at": ts, "status": "failed", "model": model,
+                    "id": response_id,
+                    "object": "response",
+                    "created_at": ts,
+                    "status": "failed",
+                    "model": chosen_model,
                     "error": {"code": "upstream_error", "message": exc.detail},
                 },
             })
             yield "data: [DONE]\n\n"
             return
 
-        # 3. output item added (the assistant message container)
-        output_item = {
-            "id": item_id, "type": "message", "role": "assistant",
-            "status": "in_progress", "content": [],
-        }
+        # 3. output item added
         yield sse("response.output_item.added", {
             "type": "response.output_item.added",
             "response_id": response_id,
             "output_index": 0,
-            "item": output_item,
+            "item": {
+                "id": item_id,
+                "type": "message",
+                "role": "assistant",
+                "status": "in_progress",
+                "content": [],
+            },
         })
 
-        # 4. content part added (the text part within the message)
+        # 4. content part added
         yield sse("response.content_part.added", {
             "type": "response.content_part.added",
             "response_id": response_id,
@@ -1089,18 +1094,27 @@ async def create_responses(
             "response_id": response_id,
             "output_index": 0,
             "item": {
-                "id": item_id, "type": "message", "role": "assistant",
+                "id": item_id,
+                "type": "message",
+                "role": "assistant",
                 "status": "completed",
                 "content": [{"type": "output_text", "text": text, "annotations": []}],
             },
         })
 
-        # 9. response.completed — includes usage, this is what SDK exposes as .usage
+        # 9. response.completed (WITH usage + metadata)
+        final_payload = _build_responses_payload(
+            chosen_model,
+            text,
+            response_id,
+            input_tokens,
+            output_tokens,
+        )
+        final_payload["router_metadata"] = {"provider": provider}
+
         yield sse("response.completed", {
             "type": "response.completed",
-            "response": _build_responses_payload(
-                model, text, response_id, input_tokens, output_tokens
-            ),
+            "response": final_payload,
         })
 
         yield "data: [DONE]\n\n"
@@ -1108,5 +1122,9 @@ async def create_responses(
     return StreamingResponse(
         event_stream(),
         media_type="text/event-stream",
-        headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"},
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+            "X-Accel-Buffering": "no",
+        },
     )