Update proxy_cerebras.py

proxy_cerebras.py

@@ -41,12 +41,15 @@ for idx, k in enumerate(OLLAMA_KEYS, 1):
 
 rr_index = 0
 
+# Global async lock to prevent race condition on rr_index & busy flag
+_key_lock = asyncio.Lock()
+
 
 # =====================================================
 # HELPERS
 # =====================================================
 def log(x):
-    print(f"[{time.strftime('%H:%M:%S')}] {x}")
+    print(f"[{time.strftime('%H:%M:%S')}] {x}", flush=True)
 
 
 def sse(obj):
@@ -58,39 +61,61 @@ def auth_ok(req: Request):
     return token == MASTER_API_KEY
 
 
-def get_key(exclude=None):
+async def get_key(exclude=None):
+    """
+    Thread-safe round-robin key picker.
+    Returns the key string, or None if all are busy/excluded.
+    """
     global rr_index
 
     if exclude is None:
         exclude = set()
 
-    for _ in range(len(OLLAMA_KEYS)):
-        rr_index = (rr_index + 1) % len(OLLAMA_KEYS)
-        k = OLLAMA_KEYS[rr_index]
-
-        st = key_status[k]
+    async with _key_lock:
+        for _ in range(len(OLLAMA_KEYS)):
+            rr_index = (rr_index + 1) % len(OLLAMA_KEYS)
+            k = OLLAMA_KEYS[rr_index]
+            st = key_status[k]
 
-        if st["healthy"] and not st["busy"] and k not in exclude:
-            st["busy"] = True
-            return k
+            if st["healthy"] and not st["busy"] and k not in exclude:
+                st["busy"] = True
+                return k
 
     return None
 
 
-def release_key(k):
-    if k in key_status:
-        key_status[k]["busy"] = False
+async def release_key(k):
+    async with _key_lock:
+        if k in key_status:
+            key_status[k]["busy"] = False
 
 
-def mark_fail(k):
-    if k in key_status:
-        key_status[k]["fail"] += 1
+async def mark_fail(k):
+    async with _key_lock:
+        if k in key_status:
+            key_status[k]["fail"] += 1
 
 
-def mark_ok(k):
-    if k in key_status:
-        key_status[k]["success"] += 1
-        key_status[k]["fail"] = 0
+async def mark_ok(k):
+    async with _key_lock:
+        if k in key_status:
+            key_status[k]["success"] += 1
+            key_status[k]["fail"] = 0
+
+
+async def wait_for_free_key(exclude=None, max_wait=30.0, interval=0.3):
+    """
+    Polls until a free key is available or max_wait seconds pass.
+    Returns the key or None on timeout.
+    """
+    elapsed = 0.0
+    while elapsed < max_wait:
+        key = await get_key(exclude)
+        if key:
+            return key
+        await asyncio.sleep(interval)
+        elapsed += interval
+    return None
 
 
 # =====================================================
@@ -98,17 +123,17 @@ def mark_ok(k):
 # =====================================================
 @app.get("/")
 async def root():
-    safe = {}
-
-    for k, v in key_status.items():
-        masked = k[:4] + "****" + k[-4:]
-        safe[masked] = {
-            "index": v["index"],
-            "healthy": v["healthy"],
-            "busy": v["busy"],
-            "success": v["success"],
-            "fail": v["fail"],
-        }
+    async with _key_lock:
+        safe = {}
+        for k, v in key_status.items():
+            masked = k[:4] + "****" + k[-4:]
+            safe[masked] = {
+                "index": v["index"],
+                "healthy": v["healthy"],
+                "busy": v["busy"],
+                "success": v["success"],
+                "fail": v["fail"],
+            }
 
     return {
         "status": "ok",
@@ -116,6 +141,7 @@ async def root():
         "detail": safe
     }
 
+
 # =====================================================
 # /v1/models
 # =====================================================
@@ -136,9 +162,8 @@ async def models(req: Request):
         return JSONResponse({"error": r.text}, status_code=r.status_code)
 
     data = r.json()
-
-    out = []
     now = int(time.time())
+    out = []
 
     for m in data.get("models", []):
         out.append({
@@ -152,7 +177,7 @@ async def models(req: Request):
 
 
 # =====================================================
-# OPENAI CHAT
+# OPENAI CHAT  /v1/chat/completions
 # =====================================================
 @app.post("/v1/chat/completions")
 async def chat(req: Request):
@@ -161,7 +186,7 @@ async def chat(req: Request):
 
     try:
         body = await req.json()
-    except:
+    except Exception:
         return JSONResponse({"error": "Bad JSON"}, status_code=400)
 
     is_stream = body.get("stream", False)
@@ -173,11 +198,10 @@ async def chat(req: Request):
         tried = set()
 
         for _ in range(len(OLLAMA_KEYS)):
-            key = get_key(tried)
+            key = await wait_for_free_key(exclude=tried)
 
             if not key:
-                await asyncio.sleep(0.3)
-                continue
+                break
 
             tried.add(key)
 
@@ -192,25 +216,23 @@ async def chat(req: Request):
                 txt = r.text.lower()
 
                 if "weekly usage limit" in txt or r.status_code == 429:
-                    mark_fail(key)
+                    log(f"Key {key[:8]}... rate limited (non-stream chat), trying next")
+                    await mark_fail(key)
                     continue
 
-                mark_ok(key)
+                await mark_ok(key)
 
                 return Response(
                     content=r.content,
-                    media_type=r.headers.get(
-                        "content-type",
-                        "application/json"
-                    )
+                    media_type=r.headers.get("content-type", "application/json")
                 )
 
             except Exception as e:
-                log(e)
-                mark_fail(key)
+                log(f"Key {key[:8]}... exception: {e}")
+                await mark_fail(key)
 
             finally:
-                release_key(key)
+                await release_key(key)
 
         return JSONResponse({"error": "All keys failed"}, status_code=500)
 
@@ -221,11 +243,10 @@ async def chat(req: Request):
         tried = set()
 
         for _ in range(len(OLLAMA_KEYS)):
-            key = get_key(tried)
+            key = await wait_for_free_key(exclude=tried)
 
             if not key:
-                await asyncio.sleep(0.3)
-                continue
+                break
 
             tried.add(key)
 
@@ -239,22 +260,37 @@ async def chat(req: Request):
                     ) as r:
 
                         if r.status_code == 429:
-                            mark_fail(key)
+                            log(f"Key {key[:8]}... rate limited (stream chat), trying next")
+                            await mark_fail(key)
                             continue
 
+                        hit_limit_mid_stream = False
+
                         async for line in r.aiter_lines():
-                            if line:
-                                yield line + "\n\n"
+                            if not line:
+                                continue
+
+                            # Detect mid-stream rate limit signal in data payload
+                            if "429" in line or "usage limit" in line.lower():
+                                log(f"Key {key[:8]}... mid-stream limit detected, aborting chunk")
+                                hit_limit_mid_stream = True
+                                break
 
-                        mark_ok(key)
+                            yield line + "\n\n"
+
+                        if hit_limit_mid_stream:
+                            await mark_fail(key)
+                            continue
+
+                        await mark_ok(key)
                         return
 
             except Exception as e:
-                log(e)
-                mark_fail(key)
+                log(f"Key {key[:8]}... stream exception: {e}")
+                await mark_fail(key)
 
             finally:
-                release_key(key)
+                await release_key(key)
 
         yield sse({"error": "All keys failed"})
         yield "data: [DONE]\n\n"
@@ -274,16 +310,16 @@ async def anthropic(req: Request):
         body = await req.json()
     except ClientDisconnect:
         return Response(status_code=499)
+    except Exception:
+        return JSONResponse({"error": "Bad JSON"}, status_code=400)
 
     stream = body.get("stream", False)
 
+    # Build messages list for proxy
     messages = []
 
     if body.get("system"):
-        messages.append({
-            "role": "system",
-            "content": body["system"]
-        })
+        messages.append({"role": "system", "content": body["system"]})
 
     for m in body.get("messages", []):
         content = m.get("content", "")
@@ -295,10 +331,7 @@ async def anthropic(req: Request):
                     txt += x.get("text", "")
             content = txt
 
-        messages.append({
-            "role": m["role"],
-            "content": content
-        })
+        messages.append({"role": m["role"], "content": content})
 
     proxy_body = {
         "model": "minimax-m2.7:cloud",
@@ -313,11 +346,10 @@ async def anthropic(req: Request):
         tried = set()
 
         for _ in range(len(OLLAMA_KEYS)):
-            key = get_key(tried)
+            key = await wait_for_free_key(exclude=tried)
 
             if not key:
-                await asyncio.sleep(0.3)
-                continue
+                break
 
             tried.add(key)
 
@@ -332,11 +364,11 @@ async def anthropic(req: Request):
                 txt = r.text.lower()
 
                 if "weekly usage limit" in txt or r.status_code == 429:
-                    mark_fail(key)
+                    log(f"Key {key[:8]}... rate limited (non-stream anthropic), trying next")
+                    await mark_fail(key)
                     continue
 
                 data = r.json()
-
                 ans = data["choices"][0]["message"]["content"]
 
                 out = {
@@ -344,70 +376,42 @@ async def anthropic(req: Request):
                     "type": "message",
                     "role": "assistant",
                     "model": body.get("model", "claude-opus-4-7"),
-                    "content": [
-                        {
-                            "type": "text",
-                            "text": ans
-                        }
-                    ],
+                    "content": [{"type": "text", "text": ans}],
                     "stop_reason": "end_turn",
                     "stop_sequence": None,
-                    "usage": {
-                        "input_tokens": 0,
-                        "output_tokens": 0
-                    }
+                    "usage": {"input_tokens": 0, "output_tokens": 0}
                 }
 
-                mark_ok(key)
+                await mark_ok(key)
                 return JSONResponse(out)
 
             except Exception as e:
-                log(e)
-                mark_fail(key)
+                log(f"Key {key[:8]}... exception: {e}")
+                await mark_fail(key)
 
             finally:
-                release_key(key)
+                await release_key(key)
 
         return JSONResponse({"error": "All keys failed"}, status_code=500)
 
     # -----------------------------------------
-    # STREAM
+    # STREAM  (Anthropic SSE format)
     # -----------------------------------------
     async def agen():
         tried = set()
         msg_id = "msg_" + uuid.uuid4().hex[:10]
+        sent_any_delta = False
 
-        start_payload = {
-            "type": "message_start",
-            "message": {
-                "id": msg_id,
-                "type": "message",
-                "role": "assistant",
-                "model": body.get("model", "claude-opus-4-7"),
-                "content": [],
-                "stop_reason": None,
-                "stop_sequence": None,
-                "usage": {
-                    "input_tokens": 0,
-                    "output_tokens": 0
-                }
-            }
-        }
-
-        yield sse(start_payload)
-
-        yield sse({
-            "type": "content_block_start",
-            "index": 0,
-            "content_block": {"type": "text"}
-        })
+        # Send Anthropic envelope headers ONCE before first key attempt
+        # We defer these until we have a successful connection to avoid
+        # sending headers before knowing if any key works.
+        # Instead we buffer and yield only on confirmed success.
 
         for _ in range(len(OLLAMA_KEYS)):
-            key = get_key(tried)
+            key = await wait_for_free_key(exclude=tried)
 
             if not key:
-                await asyncio.sleep(0.3)
-                continue
+                break
 
             tried.add(key)
 
@@ -421,9 +425,33 @@ async def anthropic(req: Request):
                     ) as r:
 
                         if r.status_code == 429:
-                            mark_fail(key)
+                            log(f"Key {key[:8]}... rate limited (stream anthropic), trying next")
+                            await mark_fail(key)
                             continue
 
+                        # Only emit Anthropic envelope on first successful key
+                        if not sent_any_delta:
+                            yield sse({
+                                "type": "message_start",
+                                "message": {
+                                    "id": msg_id,
+                                    "type": "message",
+                                    "role": "assistant",
+                                    "model": body.get("model", "claude-opus-4-7"),
+                                    "content": [],
+                                    "stop_reason": None,
+                                    "stop_sequence": None,
+                                    "usage": {"input_tokens": 0, "output_tokens": 0}
+                                }
+                            })
+                            yield sse({
+                                "type": "content_block_start",
+                                "index": 0,
+                                "content_block": {"type": "text"}
+                            })
+
+                        hit_limit_mid_stream = False
+
                         async for line in r.aiter_lines():
                             if not line.startswith("data: "):
                                 continue
@@ -433,55 +461,51 @@ async def anthropic(req: Request):
                             if raw == "[DONE]":
                                 break
 
+                            # Detect mid-stream 429 / limit payload
+                            if "429" in raw or "usage limit" in raw.lower():
+                                log(f"Key {key[:8]}... mid-stream limit in anthropic, aborting chunk")
+                                hit_limit_mid_stream = True
+                                break
+
                             try:
                                 j = json.loads(raw)
-                            except:
+                            except Exception:
                                 continue
 
                             delta = j["choices"][0]["delta"]
                             txt = delta.get("content", "")
 
                             if txt:
+                                sent_any_delta = True
                                 yield sse({
                                     "type": "content_block_delta",
                                     "index": 0,
-                                    "delta": {
-                                        "type": "text_delta",
-                                        "text": txt
-                                    }
+                                    "delta": {"type": "text_delta", "text": txt}
                                 })
 
-                        mark_ok(key)
-                        break
+                        if hit_limit_mid_stream:
+                            await mark_fail(key)
+                            # Continue to next key — stream resumes from where it broke
+                            # Note: client will receive continued deltas seamlessly
+                            continue
+
+                        await mark_ok(key)
+                        break  # Success — exit key retry loop
 
             except Exception as e:
-                log(e)
-                mark_fail(key)
+                log(f"Key {key[:8]}... agen exception: {e}")
+                await mark_fail(key)
 
             finally:
-                release_key(key)
-
-        yield sse({
-            "type": "content_block_stop",
-            "index": 0
-        })
+                await release_key(key)
 
+        # Close Anthropic SSE envelope
+        yield sse({"type": "content_block_stop", "index": 0})
         yield sse({
             "type": "message_delta",
-            "delta": {
-                "stop_reason": "end_turn",
-                "stop_sequence": None
-            },
-            "usage": {
-                "output_tokens": 0
-            }
-        })
-
-        yield sse({
-            "type": "message_stop"
+            "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+            "usage": {"output_tokens": 0}
         })
+        yield sse({"type": "message_stop"})
 
-    return StreamingResponse(
-        agen(),
-        media_type="text/event-stream"
-    )
+    return StreamingResponse(agen(), media_type="text/event-stream")