Update gen.py

gen.py

@@ -66,6 +66,10 @@ MODEL_MAP = {
 FALLBACK_MODEL = "meta-llama/llama-4-scout-17b-16e-instruct"
 FALLBACK_PROVIDER = "groq"
 
+# Header that API-key authenticated clients send so we know to stream
+# thinking tokens back to them.
+API_KEY_HEADER = "x-api-key"
+
 
 # ──────────────────────────────────────────────
 # CENTRAL ROUTING LOGIC
@@ -253,6 +257,65 @@ def is_cinematic_image_prompt(prompt: str) -> bool:
     return False
 
 
+def _is_api_key_request(request: Request) -> bool:
+    """
+    Return True when the caller authenticated with an API key rather than a
+    session cookie / browser auth.  We use this to decide whether to forward
+    think-tag / reasoning_content tokens to the client.
+    """
+    return bool(
+        request.headers.get(API_KEY_HEADER)
+        or request.headers.get("authorization", "").lower().startswith("bearer ")
+    )
+
+
+def _inject_reasoning_into_chunk(obj: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Some navy models return thinking tokens in a non-standard
+    ``reasoning_content`` field inside each delta.  When that field is
+    present we wrap it in <think>…</think> and prepend it to the regular
+    ``content`` delta so that every SSE-speaking client sees a single,
+    unified text stream.
+
+    The original ``reasoning_content`` field is preserved so clients that
+    know about it can still use it directly.
+    """
+    try:
+        delta = obj["choices"][0]["delta"]
+    except (KeyError, IndexError, TypeError):
+        return obj
+
+    reasoning = delta.get("reasoning_content") or delta.get("reasoning") or ""
+    content   = delta.get("content") or ""
+
+    if reasoning and isinstance(reasoning, str):
+        # Wrap in <think> tags and prepend to the visible content delta.
+        wrapped = f"<think>{reasoning}</think>"
+        delta["content"] = wrapped + content
+        # Keep the raw field so native clients can parse it too.
+        delta["reasoning_content"] = reasoning
+        obj["choices"][0]["delta"] = delta
+
+    return obj
+
+
+def _normalize_usage_block(obj: Dict[str, Any]) -> Dict[str, Any]:
+    """Rewrite the usage block to a canonical shape (in-place, returns obj)."""
+    if "usage" not in obj or not isinstance(obj.get("usage"), dict):
+        return obj
+    u = obj["usage"]
+    input_tok  = u.get("prompt_tokens")     or u.get("input_tokens",  0)
+    output_tok = u.get("completion_tokens") or u.get("output_tokens", 0)
+    obj["usage"] = {
+        "prompt_tokens":     input_tok,
+        "completion_tokens": output_tok,
+        "total_tokens":      input_tok + output_tok,
+        "input_tokens":      input_tok,
+        "output_tokens":     output_tok,
+    }
+    return obj
+
+
 # ──────────────────────────────────────────────
 # IMAGE GENERATION
 # ──────────────────────────────────────────────
@@ -682,6 +745,10 @@ async def generate_text(
 
     await _check_chat_rate_limit(request, authorization, x_client_id)
 
+    # Determine whether the caller is an API-key client that should receive
+    # raw thinking tokens.
+    forward_thinking = _is_api_key_request(request)
+
     body["model"] = chosen_model
     stream = body.get("stream", False)
 
@@ -744,39 +811,79 @@ async def generate_text(
             sent_metadata = False
             async with httpx.AsyncClient(timeout=None) as client:
                 async for chunk in stream_primary(client):
+                    # ── emit router metadata once as the very first SSE frame ──
                     if not sent_metadata:
-                        meta = {"router_metadata": {"model_name": MODEL_MAP.get(chosen_model, chosen_model)}}
+                        meta = {
+                            "router_metadata": {
+                                "model_name": MODEL_MAP.get(chosen_model, chosen_model)
+                            }
+                        }
                         yield f"data: {json.dumps(meta)}\n\n"
                         sent_metadata = True
 
-                    # Intercept the final non-[DONE] data chunk and normalize
-                    # the usage block so callers always see consistent field names.
-                    if chunk.startswith("data:") and "[DONE]" not in chunk:
+                    # ── pass [DONE] straight through ──────────────────────────
+                    if "data: [DONE]" in chunk:
+                        yield chunk
+                        continue
+
+                    # ── process data: … lines ─────────────────────────────────
+                    if chunk.startswith("data:"):
                         raw = chunk[5:].strip()
                         try:
                             obj = json.loads(raw)
-                            if isinstance(obj, dict) and "usage" in obj and isinstance(obj["usage"], dict):
-                                u = obj["usage"]
-                                input_tok  = u.get("prompt_tokens")     or u.get("input_tokens",  0)
-                                output_tok = u.get("completion_tokens") or u.get("output_tokens", 0)
-                                obj["usage"] = {
-                                    "prompt_tokens":     input_tok,
-                                    "completion_tokens": output_tok,
-                                    "total_tokens":      input_tok + output_tok,
-                                    "input_tokens":      input_tok,
-                                    "output_tokens":     output_tok,
-                                }
-                                yield f"data: {json.dumps(obj)}\n\n"
-                                continue
                         except Exception:
-                            pass
+                            # Not valid JSON — forward verbatim (keeps partial
+                            # chunks from blocking the stream).
+                            yield chunk
+                            continue
+
+                        if not isinstance(obj, dict):
+                            yield chunk
+                            continue
+
+                        # Normalize usage block whenever it appears.
+                        _normalize_usage_block(obj)
+
+                        # ── thinking / reasoning tokens ───────────────────────
+                        # Navy models may embed thinking in two ways:
+                        #
+                        #   1. As delta.reasoning_content (separate field)
+                        #   2. Inline inside delta.content wrapped in <think>…</think>
+                        #
+                        # For API-key callers we always surface both forms.
+                        # For browser/session callers we strip reasoning_content
+                        # so it doesn't confuse UI clients that don't expect it,
+                        # but <think> tags already present in content are left
+                        # alone (they arrived that way from upstream).
+                        if forward_thinking:
+                            # Merge reasoning_content into content as
+                            # <think>…</think> and keep the raw field.
+                            obj = _inject_reasoning_into_chunk(obj)
+                        else:
+                            # Strip the non-standard field so browser clients
+                            # don't see unexpected keys.
+                            try:
+                                delta = obj["choices"][0]["delta"]
+                                delta.pop("reasoning_content", None)
+                                delta.pop("reasoning", None)
+                                obj["choices"][0]["delta"] = delta
+                            except (KeyError, IndexError, TypeError):
+                                pass
 
+                        yield f"data: {json.dumps(obj)}\n\n"
+                        continue
+
+                    # ── any other line (comments, keep-alives, …) ─────────────
                     yield chunk
 
         return StreamingResponse(
             event_generator(),
             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",
+            },
         )
 
     # ── non-streaming ─────────────────────────
@@ -789,7 +896,11 @@ async def generate_text(
             fb_url, fb_key = _get_provider_url_and_key(FALLBACK_PROVIDER)
             fallback_body = dict(body)
             fallback_body["model"] = FALLBACK_MODEL
-            r = await client.post(fb_url, json=fallback_body, headers={"Authorization": f"Bearer {fb_key}"})
+            r = await client.post(
+                fb_url,
+                json=fallback_body,
+                headers={"Authorization": f"Bearer {fb_key}"},
+            )
 
     content_type = (r.headers.get("content-type") or "").lower()
     if "application/json" in content_type:
@@ -798,22 +909,35 @@ 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)
-                output_tok = u.get("completion_tokens") or u.get("output_tokens", 0)
-                payload["usage"] = {
-                    "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,
-                }
-            payload.setdefault("router_metadata", {})["model_name"] = MODEL_MAP.get(chosen_model, chosen_model)
+            # Normalize usage fields.
+            _normalize_usage_block(payload)
+
+            # ── thinking tokens in non-streaming responses ────────────────────
+            # Some navy models put thinking content in
+            # message.reasoning_content.  For API-key callers we prepend it to
+            # message.content wrapped in <think>…</think>; for others we drop
+            # the non-standard field.
+            try:
+                message = payload["choices"][0]["message"]
+                reasoning = (
+                    message.pop("reasoning_content", None)
+                    or message.pop("reasoning", None)
+                    or ""
+                )
+                if reasoning and isinstance(reasoning, str):
+                    if forward_thinking:
+                        existing = message.get("content") or ""
+                        message["content"] = f"<think>{reasoning}</think>{existing}"
+                        # Restore the raw field for clients that want it.
+                        message["reasoning_content"] = reasoning
+                    # else: already popped — nothing to do.
+                    payload["choices"][0]["message"] = message
+            except (KeyError, IndexError, TypeError):
+                pass
+
+            payload.setdefault("router_metadata", {})["model_name"] = MODEL_MAP.get(
+                chosen_model, chosen_model
+            )
     else:
         payload = {
             "error": "Upstream returned non-JSON response",