Fix fallback streaming

gen.py

@@ -678,85 +678,91 @@ async def generate_text(
 
     if stream:
         body["stream"] = True
-
-        async def event_generator():
+    
+        async def stream_primary(client, url, body, headers):
+            """
+            Handles the primary provider stream (Navy Vision, Groq, Cerebras, etc.)
+            Returns either:
+            - a StreamingResponse generator, OR
+            - triggers fallback if provider fails
+            """
             try:
-                async with httpx.AsyncClient(timeout=None) as client:
-                    async with client.stream(
-                        "POST",
-                        url,
-                        json=body,
-                        headers=headers,
-                    ) as r:
-        
-                        # ---------- FALLBACK ----------
-                        if r.status_code >= 400:
-        
-                            if provider == "navy vision":
-                                print("[STREAM FALLBACK] Navy vision failed → 17B groq")
-        
-                                groq_keys = os.getenv("GROQ_KEY", "")
-                                groq_keys_list = [k.strip() for k in groq_keys.split(",") if k.strip()]
-        
-                                if groq_keys_list:
-                                    fallback_headers = {
-                                        "Authorization": f"Bearer {random.choice(groq_keys_list)}"
-                                    }
-        
-                                    fallback_body = dict(body)
-                                    fallback_body["model"] = fallback_model
-        
-                                    async with client.stream(
-                                        "POST",
-                                        "https://api.groq.com/openai/v1/chat/completions",
-                                        json=fallback_body,
-                                        headers=fallback_headers,
-                                    ) as fallback_r:
-        
-                                        if fallback_r.status_code < 400:
-                                            async for line in fallback_r.aiter_lines():
-                                                if line:
-                                                    yield line + "\n"
-                                            return
-                                        else:
-                                            print("[FALLBACK FAILED] Groq fallback also failed")
-        
-                            # ---------- ORIGINAL ERROR ----------
-                            error_payload = ""
-                            try:
-                                error_payload = (
-                                    (await r.aread()).decode("utf-8", errors="replace")
-                                )[:800]
-                            except Exception:
-                                error_payload = ""
-        
-                            safe_error_payload = (
-                                error_payload.replace("\\", "\\\\")
-                                .replace('"', '\\"')
-                                .replace("\n", " ")
-                                .replace("\r", " ")
-                            )
-        
-                            yield (
-                                'data: {"error": '
-                                f'"Upstream provider error ({r.status_code}): {safe_error_payload}"'
-                                "}\n\n"
-                            )
+                async with client.stream("POST", url, json=body, headers=headers) as r:
+    
+                    # --- PRIMARY PROVIDER FAILED (HTTP ERROR) ---
+                    if r.status_code >= 400:
+                        print("[STREAM FALLBACK] Primary provider failed → switching to Groq fallback")
+                        async for chunk in stream_fallback(client, body):
+                            yield chunk
+                        return
+    
+                    # --- PRIMARY PROVIDER STREAM ---
+                    async for line in r.aiter_lines():
+                        if not line:
+                            yield "\n"
+                            continue
+    
+                        # Detect SSE-level error inside stream
+                        if '"error"' in line.lower():
+                            print("[STREAM FALLBACK] Error detected inside SSE → switching to Groq fallback")
+                            async for chunk in stream_fallback(client, body):
+                                yield chunk
                             return
-        
-                        # ---------- NORMAL STREAM ----------
-                        async for line in r.aiter_lines():
-                            if line == "":
-                                yield "\n"
-                                continue
-        
-                            yield line + "\n"
-        
-            except asyncio.CancelledError:
-                return
+    
+                        yield line + "\n"
+    
             except Exception as e:
-                yield f'data: {{"error": "{str(e)}"}}\n\n'
-
+                print(f"[STREAM ERROR] {e}")
+                async for chunk in stream_fallback(client, body):
+                    yield chunk
+    
+    
+        async def stream_fallback(client, body):
+            """
+            Clean fallback stream to Groq 17B.
+            This MUST NOT be nested inside another stream.
+            """
+            fallback_body = {
+                "model": fallback_model,
+                "messages": body["messages"],
+                "stream": True,
+            }
+    
+            groq_keys = os.getenv("GROQ_KEY", "")
+            groq_keys_list = [k.strip() for k in groq_keys.split(",") if k.strip()]
+            fallback_headers = {"Authorization": f"Bearer {random.choice(groq_keys_list)}"}
+    
+            print("[FALLBACK] Starting Groq fallback stream")
+    
+            async with client.stream(
+                "POST",
+                "https://api.groq.com/openai/v1/chat/completions",
+                json=fallback_body,
+                headers=fallback_headers,
+            ) as r:
+    
+                if r.status_code >= 400:
+                    err = (await r.aread()).decode("utf-8", errors="replace")
+                    yield f'data: {{"error": "Fallback provider failed: {err[:500]}"}}\n\n'
+                    return
+    
+                async for line in r.aiter_lines():
+                    if not line:
+                        yield "\n"
+                        continue
+    
+                    # Normalize SSE output
+                    if not line.startswith("data:"):
+                        yield f"data: {line}\n\n"
+                    else:
+                        yield line + "\n"
+    
+    
+        async def event_generator():
+            async with httpx.AsyncClient(timeout=None) as client:
+                async for chunk in stream_primary(client, url, body, headers):
+                    yield chunk
+    
         return StreamingResponse(
             event_generator(),
             media_type="text/event-stream",
@@ -766,6 +772,7 @@ async def generate_text(
                 "X-Accel-Buffering": "no",
             },
         )
+
     else:
         async with httpx.AsyncClient(timeout=None) as client:
             r = await client.post(url, json=body, headers=headers)