Buckets:
| """Adaptive thinking-mode router/proxy for an OpenAI-compatible sglang endpoint. | |
| Sits in front of a Qwen3-family sglang server. For every /v1/chat/completions | |
| request, makes a tiny classification call to the SAME upstream server (thinking | |
| forced off, a handful of tokens) to decide whether the real request should run | |
| with thinking enabled, then forwards the real request with | |
| `chat_template_kwargs.enable_thinking` set explicitly to that verdict. | |
| Safety net: the real forwarded call gets a hard wall-clock timeout and a | |
| max_tokens ceiling (when the client didn't already set a tighter one), since an | |
| adaptive "enable thinking" verdict reopens the exact runaway-generation failure | |
| mode that `--default-chat-template-kwargs enable_thinking=false` was built to | |
| avoid (see harbor-hf-benchmark skill, Step 2). | |
| """ | |
| import json | |
| import os | |
| import time | |
| import httpx | |
| from fastapi import FastAPI, Request | |
| from fastapi.responses import JSONResponse, StreamingResponse | |
| UPSTREAM_BASE_URL = os.environ["UPSTREAM_BASE_URL"].rstrip("/") | |
| UPSTREAM_TOKEN = os.environ["UPSTREAM_TOKEN"] | |
| MODEL_NAME = os.environ.get("MODEL_NAME", "Qwen/Qwen3.8-27B-FP8") | |
| # When set ("true"/"false"), skip the classification call entirely and force | |
| # every request's enable_thinking to this value. Used to run the always-on/ | |
| # always-off comparison arms off the same backing server as the adaptive arm, | |
| # instead of standing up a second GPU server. | |
| FORCE_THINKING = os.environ.get("FORCE_THINKING") | |
| if FORCE_THINKING is not None: | |
| FORCE_THINKING = FORCE_THINKING.strip().lower() == "true" | |
| CLASSIFY_MAX_TOKENS = int(os.environ.get("CLASSIFY_MAX_TOKENS", "4")) | |
| CLASSIFY_TIMEOUT_S = float(os.environ.get("CLASSIFY_TIMEOUT_S", "20")) | |
| REAL_CALL_TIMEOUT_S = float(os.environ.get("REAL_CALL_TIMEOUT_S", "300")) | |
| REAL_CALL_MAX_TOKENS_CEILING = int( | |
| os.environ.get("REAL_CALL_MAX_TOKENS_CEILING", "4096") | |
| ) | |
| # Truncate the classification prompt so routing stays cheap even on huge | |
| # agent-history payloads (mini-swe-agent/pi resend full growing history). | |
| CLASSIFY_PROMPT_CHARS = int(os.environ.get("CLASSIFY_PROMPT_CHARS", "4000")) | |
| CLASSIFY_SYSTEM_PROMPT = ( | |
| "You are a fast triage classifier for an AI coding agent. You will see the " | |
| "latest message in an agent's conversation. Decide whether responding well " | |
| "requires careful multi-step reasoning or complex planning (e.g. algorithm " | |
| "design, debugging a subtle failure, multi-step math), as opposed to a " | |
| "quick, mostly-mechanical response (e.g. running a command, reading a " | |
| "file, a simple edit). Answer with exactly one word: YES or NO." | |
| ) | |
| app = FastAPI() | |
| client = httpx.AsyncClient() | |
| _stats = {"total": 0, "thinking_on": 0, "thinking_off": 0, "classify_errors": 0} | |
| def _upstream_headers() -> dict: | |
| return { | |
| "Authorization": f"Bearer {UPSTREAM_TOKEN}", | |
| "Content-Type": "application/json", | |
| } | |
| def _last_user_content(messages: list[dict]) -> str: | |
| for msg in reversed(messages): | |
| if msg.get("role") in ("user", "tool"): | |
| content = msg.get("content") | |
| if isinstance(content, list): | |
| content = " ".join( | |
| part.get("text", "") for part in content if isinstance(part, dict) | |
| ) | |
| return str(content or "")[:CLASSIFY_PROMPT_CHARS] | |
| return "" | |
| async def decide_thinking(messages: list[dict]) -> bool: | |
| prompt = _last_user_content(messages) | |
| if not prompt: | |
| return False | |
| body = { | |
| "model": MODEL_NAME, | |
| "messages": [ | |
| {"role": "system", "content": CLASSIFY_SYSTEM_PROMPT}, | |
| {"role": "user", "content": prompt}, | |
| ], | |
| "max_tokens": CLASSIFY_MAX_TOKENS, | |
| "temperature": 0, | |
| "chat_template_kwargs": {"enable_thinking": False}, | |
| } | |
| try: | |
| resp = await client.post( | |
| f"{UPSTREAM_BASE_URL}/v1/chat/completions", | |
| headers=_upstream_headers(), | |
| json=body, | |
| timeout=CLASSIFY_TIMEOUT_S, | |
| ) | |
| resp.raise_for_status() | |
| data = resp.json() | |
| answer = data["choices"][0]["message"]["content"].strip().upper() | |
| return answer.startswith("YES") | |
| except Exception: | |
| _stats["classify_errors"] += 1 | |
| return False | |
| async def health(): | |
| resp = await client.get( | |
| f"{UPSTREAM_BASE_URL}/health", headers=_upstream_headers(), timeout=10 | |
| ) | |
| return JSONResponse(status_code=resp.status_code, content={}) | |
| async def stats(): | |
| return _stats | |
| async def proxy(path: str, request: Request): | |
| body_bytes = await request.body() | |
| is_chat = path == "v1/chat/completions" and request.method == "POST" | |
| if not is_chat: | |
| # Pass everything else straight through (model listing, etc). | |
| resp = await client.request( | |
| request.method, | |
| f"{UPSTREAM_BASE_URL}/{path}", | |
| headers=_upstream_headers(), | |
| content=body_bytes, | |
| timeout=30, | |
| ) | |
| return JSONResponse(status_code=resp.status_code, content=resp.json()) | |
| body = json.loads(body_bytes) | |
| messages = body.get("messages", []) | |
| if FORCE_THINKING is not None: | |
| enable_thinking = FORCE_THINKING | |
| else: | |
| enable_thinking = await decide_thinking(messages) | |
| _stats["total"] += 1 | |
| _stats["thinking_on" if enable_thinking else "thinking_off"] += 1 | |
| body["chat_template_kwargs"] = { | |
| **body.get("chat_template_kwargs", {}), | |
| "enable_thinking": enable_thinking, | |
| } | |
| if "max_tokens" not in body and "max_completion_tokens" not in body: | |
| body["max_tokens"] = REAL_CALL_MAX_TOKENS_CEILING | |
| is_stream = bool(body.get("stream")) | |
| t0 = time.monotonic() | |
| if is_stream: | |
| async def gen(): | |
| async with client.stream( | |
| "POST", | |
| f"{UPSTREAM_BASE_URL}/v1/chat/completions", | |
| headers=_upstream_headers(), | |
| json=body, | |
| timeout=REAL_CALL_TIMEOUT_S, | |
| ) as resp: | |
| async for chunk in resp.aiter_bytes(): | |
| yield chunk | |
| return StreamingResponse(gen(), media_type="text/event-stream") | |
| try: | |
| resp = await client.post( | |
| f"{UPSTREAM_BASE_URL}/v1/chat/completions", | |
| headers=_upstream_headers(), | |
| json=body, | |
| timeout=REAL_CALL_TIMEOUT_S, | |
| ) | |
| except httpx.TimeoutException: | |
| # Hard safety net: the adaptive "thinking on" verdict reopened the | |
| # runaway-generation risk the server-wide default normally guards | |
| # against. Fail this one call cleanly rather than hang the agent. | |
| return JSONResponse( | |
| status_code=504, | |
| content={ | |
| "error": { | |
| "message": ( | |
| f"router: upstream call exceeded {REAL_CALL_TIMEOUT_S}s " | |
| f"(enable_thinking={enable_thinking})" | |
| ) | |
| } | |
| }, | |
| ) | |
| elapsed = time.monotonic() - t0 | |
| resp_json = resp.json() | |
| resp_json.setdefault("router_meta", {}) | |
| resp_json["router_meta"] = {"enable_thinking": enable_thinking, "elapsed_s": elapsed} | |
| return JSONResponse(status_code=resp.status_code, content=resp_json) | |
Xet Storage Details
- Size:
- 7.33 kB
- Xet hash:
- f1eaaf71f483cfe25b9bc367e7edc2b0fd4f8ae29f99d6dcc1713694210671a9
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.