Skip to content

Instantly share code, notes, and snippets.

@kdcyberdude

kdcyberdude/proxy.py

Created May 9, 2026 10:02
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save kdcyberdude/0e624a52d5012e9460cbd7653db1ea02 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save kdcyberdude/0e624a52d5012e9460cbd7653db1ea02 to your computer and use it in GitHub Desktop.
Download ZIP
vLLM preserve_thinking proxy (Anthropic/OpenAI) + start_vllm_and_proxy.sh launcher
Raw
proxy.py
#!/usr/bin/env python3
"""
vLLM preserve_thinking proxy (multi-session, restart-safe)
Behaviour
─────────
1. Forwards requests to UPSTREAM (CLI `--upstream`, env
PRESERVE_THINKING_PROXY_UPSTREAM, or default https://vllm.treowai.com).
2. Injects preserve_thinking=true into chat_template_kwargs when the client did
not set it (OpenAI Chat, Responses API, and Anthropic Messages / Claude Code).
3. Caches assistant reasoning keyed by model + conversation prefix + visible text
so parallel chats with the same short reply do not collide.
4. Optional header X-Preserve-Thinking-Session — strongest per-chat isolation.
5. Claude Code uses POST /v1/messages (Anthropic API); the proxy applies the same
preserve_thinking + reasoning round-trip as for OpenAI clients.
6. Persists reasoning_cache.json and responses_state.json; LRU + optional TTL.
Exact reasoning bytes matter for prefix/KV cache with Qwen-style templates.
"""
from __future__ import annotations
import argparse
import hashlib
import os
import json
import time
from collections import OrderedDict
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
import httpx
import uvicorn
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, StreamingResponse
# ── config ────────────────────────────────────────────────────────────────────
UPSTREAM = "https://vllm.treowai.com"
LOG_FILE = Path(__file__).parent / "proxy_traffic.jsonl"
CACHE_FILE = Path(__file__).parent / "reasoning_cache.json"
RESP_STATE_FILE = Path(__file__).parent / "responses_state.json"
MAX_BODY_LOG = 32 * 1024
SESSION_HEADER = "x-preserve-thinking-session"
CACHE_VERSION = 2
# Defaults; overridden by CLI
MAX_REASONING_ENTRIES = 50_000
MAX_RESPONSES_ENTRIES = 50_000
# 0 = disable TTL checks (entries evicted only by LRU)
CACHE_TTL_SEC = 30 * 24 * 3600
INJECT_PRESERVE_THINKING = True
INJECT_REASONING = True
app = FastAPI(title="vLLM preserve_thinking proxy v4")
_client: httpx.AsyncClient | None = None
# LRU: OrderedDict preserves insertion; move_to_end on access; popitem(False) = LRU evict
_reasoning_entries: OrderedDict[str, dict[str, Any]] = OrderedDict()
_resp_id_cache: OrderedDict[str, dict[str, Any]] = OrderedDict()
# ── lifecycle ─────────────────────────────────────────────────────────────────
@app.on_event("startup")
async def _startup():
global _client
_client = httpx.AsyncClient(
base_url=UPSTREAM,
timeout=httpx.Timeout(600.0),
follow_redirects=True,
limits=httpx.Limits(max_connections=200, max_keepalive_connections=50),
)
_load_reasoning_cache()
_load_responses_state()
_purge_expired_all()
LOG_FILE.write_text("")
print(f"[proxy] Upstream : {UPSTREAM}")
print(f"[proxy] Log : {LOG_FILE}")
print(f"[proxy] Reasoning cache : {CACHE_FILE} ({len(_reasoning_entries)} entries)")
print(f"[proxy] Responses state : {RESP_STATE_FILE} ({len(_resp_id_cache)} entries)")
@app.on_event("shutdown")
async def _shutdown():
if _client:
await _client.aclose()
_flush_reasoning_cache()
_flush_responses_state()
def _now() -> float:
return time.time()
def _purge_expired_all():
if CACHE_TTL_SEC <= 0:
return
t = _now()
dead_r = [
k
for k, v in _reasoning_entries.items()
if t - v.get("updated_at", 0) > CACHE_TTL_SEC
]
for k in dead_r:
del _reasoning_entries[k]
dead_s = [
k
for k, v in _resp_id_cache.items()
if t - v.get("updated_at", 0) > CACHE_TTL_SEC
]
for k in dead_s:
del _resp_id_cache[k]
if dead_r or dead_s:
print(f"[proxy] TTL purge: reasoning {len(dead_r)}, responses {len(dead_s)}")
# ── canonical prefix / session scope ────────────────────────────────────────
def _canonical_any_message(msg: dict) -> dict[str, Any] | None:
"""Stable fingerprint for OpenAI-style and Anthropic-style messages."""
if not isinstance(msg, dict):
return None
role = msg.get("role")
if role is None:
return None
content = msg.get("content")
out: dict[str, Any] = {"role": role}
if isinstance(content, str):
out["content"] = content
elif isinstance(content, list):
parts: list[dict[str, Any]] = []
for p in content:
if not isinstance(p, dict):
continue
t = p.get("type")
if t == "text":
parts.append({"type": "text", "t": p.get("text", "")})
elif t == "thinking":
parts.append(
{
"type": "thinking",
"n": bool((p.get("thinking") or "").strip()),
}
)
elif t == "redacted_thinking":
parts.append({"type": "redacted_thinking"})
elif t == "tool_use":
parts.append(
{"type": "tool_use", "id": p.get("id"), "n": p.get("name")}
)
elif t == "tool_result":
parts.append({"type": "tool_result", "tid": p.get("tool_use_id")})
elif t == "image_url":
url = (p.get("image_url") or {}).get("url", "") if isinstance(p.get("image_url"), dict) else ""
parts.append({"type": "image_url", "u": url[:120]})
elif t == "image":
parts.append({"type": "image"})
else:
parts.append({"type": t})
out["parts"] = parts
else:
out["content"] = str(content)
if msg.get("reasoning"):
out["has_reasoning"] = True
if msg.get("reasoning_content"):
out["has_reasoning_content"] = True
if msg.get("tool_calls"):
out["tool_calls_n"] = len(msg["tool_calls"])
return out
def _canonical_prefix_list(messages: list) -> list[dict[str, Any]]:
out: list[dict[str, Any]] = []
for m in messages:
c = _canonical_any_message(m)
if c is not None:
out.append(c)
return out
def _prefix_hash(messages_prefix: list) -> str:
canon = _canonical_prefix_list(messages_prefix)
blob = json.dumps(canon, sort_keys=True, ensure_ascii=False)
return hashlib.sha256(blob.encode()).hexdigest()
def _session_scope(session_header: str | None, messages_prefix: list) -> str:
if session_header and session_header.strip():
return "h:" + session_header.strip()
return "p:" + _prefix_hash(messages_prefix)
def _reasoning_cache_key(model: str, scope: str, content: str) -> str:
raw = f"{model}|{scope}|{content or ''}"
return hashlib.sha256(raw.encode()).hexdigest()
# ── reasoning cache (LRU + TTL) ─────────────────────────────────────────────
def _load_reasoning_cache():
global _reasoning_entries
_reasoning_entries.clear()
if not CACHE_FILE.exists():
return
try:
data = json.loads(CACHE_FILE.read_text())
except Exception as e:
print(f"[proxy] reasoning cache load failed ({e}), starting fresh")
return
if isinstance(data, dict) and data.get("version") == CACHE_VERSION:
entries = data.get("entries", {})
order = data.get("lru_order", list(entries.keys()))
for k in order:
if k in entries and isinstance(entries[k], dict):
_reasoning_entries[k] = entries[k]
print(f"[proxy] Loaded {len(_reasoning_entries)} reasoning entries (v{CACHE_VERSION})")
return
# Legacy flat str -> str
if isinstance(data, dict) and data and all(isinstance(v, str) for v in data.values()):
print("[proxy] Old reasoning_cache.json format ignored (incompatible keys); starting fresh")
else:
print("[proxy] Unrecognized reasoning_cache.json; starting fresh")
def _flush_reasoning_cache():
try:
payload = {
"version": CACHE_VERSION,
"lru_order": list(_reasoning_entries.keys()),
"entries": dict(_reasoning_entries),
}
CACHE_FILE.write_text(json.dumps(payload, ensure_ascii=False, indent=2))
except Exception as e:
print(f"[proxy] reasoning cache flush failed: {e}")
def _reasoning_evict_lru():
while len(_reasoning_entries) > MAX_REASONING_ENTRIES:
_reasoning_entries.popitem(last=False)
def _reasoning_touch(key: str):
if key in _reasoning_entries:
ent = _reasoning_entries[key]
ent["last_access"] = _now()
_reasoning_entries.move_to_end(key)
def _cache_store(
model: str,
session_header: str | None,
request_messages: list,
content: str,
reasoning: str,
):
if content is None:
return
scope = _session_scope(session_header, request_messages)
t = _now()
variants: list[str] = []
for v in (content, content.strip()):
if v not in variants:
variants.append(v)
for variant in variants:
key = _reasoning_cache_key(model, scope, variant)
_reasoning_entries[key] = {
"reasoning": reasoning,
"updated_at": t,
"last_access": t,
"model": model,
"scope": scope,
}
_reasoning_entries.move_to_end(key)
_reasoning_evict_lru()
_flush_reasoning_cache()
def _cache_get(model: str, session_header: str | None, messages_prefix: list, content: str) -> str | None:
if content is None:
return None
t = _now()
scope = _session_scope(session_header, messages_prefix)
variants: list[str] = []
for v in (content, content.strip() if content else ""):
if v not in variants:
variants.append(v)
for variant in variants:
key = _reasoning_cache_key(model, scope, variant)
ent = _reasoning_entries.get(key)
if ent is None:
continue
if CACHE_TTL_SEC > 0 and t - ent.get("updated_at", 0) > CACHE_TTL_SEC:
del _reasoning_entries[key]
continue
_reasoning_touch(key)
return ent.get("reasoning")
return None
# ── responses state (persisted + LRU + TTL) ───────────────────────────────────
def _load_responses_state():
global _resp_id_cache
_resp_id_cache.clear()
if not RESP_STATE_FILE.exists():
return
try:
data = json.loads(RESP_STATE_FILE.read_text())
except Exception as e:
print(f"[proxy] responses_state load failed ({e})")
return
if not isinstance(data, dict) or data.get("version") != CACHE_VERSION:
print("[proxy] Unrecognized responses_state.json; starting fresh")
return
entries = data.get("entries", {})
order = data.get("lru_order", list(entries.keys()))
for rid in order:
if rid in entries and isinstance(entries[rid], dict):
_resp_id_cache[rid] = entries[rid]
print(f"[proxy] Loaded {len(_resp_id_cache)} response-id entries")
def _flush_responses_state():
try:
payload = {
"version": CACHE_VERSION,
"lru_order": list(_resp_id_cache.keys()),
"entries": dict(_resp_id_cache),
}
RESP_STATE_FILE.write_text(json.dumps(payload, ensure_ascii=False, indent=2))
except Exception as e:
print(f"[proxy] responses_state flush failed: {e}")
def _resp_evict_lru():
while len(_resp_id_cache) > MAX_RESPONSES_ENTRIES:
_resp_id_cache.popitem(last=False)
def _resp_touch(rid: str):
if rid in _resp_id_cache:
ent = _resp_id_cache[rid]
ent["last_access"] = _now()
_resp_id_cache.move_to_end(rid)
def _resp_cache_store(
resp_id: str,
model: str,
reasoning: str,
content: str,
input_messages: list | None = None,
):
t = _now()
_resp_id_cache[resp_id] = {
"reasoning": reasoning,
"content": content,
"model": model,
"input_messages": input_messages or [],
"updated_at": t,
"last_access": t,
}
_resp_id_cache.move_to_end(resp_id)
_resp_evict_lru()
_flush_responses_state()
def _resp_cache_get(resp_id: str) -> dict | None:
ent = _resp_id_cache.get(resp_id)
if ent is None:
return None
if CACHE_TTL_SEC > 0 and _now() - ent.get("updated_at", 0) > CACHE_TTL_SEC:
del _resp_id_cache[resp_id]
_flush_responses_state()
return None
_resp_touch(resp_id)
return ent
# ── request mutation helpers ─────────────────────────────────────────────────
def _inject_preserve_thinking(body: dict) -> tuple[dict, bool]:
if not INJECT_PRESERVE_THINKING:
return body, False
ctk = dict(body.get("chat_template_kwargs") or {})
if "preserve_thinking" in ctk:
return body, False
ctk["preserve_thinking"] = True
body = dict(body)
body["chat_template_kwargs"] = ctk
return body, True
def _extract_text_content(content) -> str:
if content is None:
return ""
if isinstance(content, str):
return content
if isinstance(content, list):
parts = []
for item in content:
if isinstance(item, dict):
parts.append(item.get("text") or item.get("content") or "")
elif isinstance(item, str):
parts.append(item)
return "".join(parts)
return str(content)
def _message_list_from_body(body: dict) -> list:
msgs = body.get("messages") or body.get("input") or []
return msgs if isinstance(msgs, list) else []
def _anthropic_visible_text(content: Any) -> str:
"""Visible assistant text from Anthropic/OpenAI list content (text blocks only)."""
if isinstance(content, str):
return content
if not isinstance(content, list):
return ""
parts: list[str] = []
for b in content:
if isinstance(b, dict) and b.get("type") == "text" and b.get("text"):
parts.append(str(b["text"]))
return "".join(parts)
def _body_looks_like_anthropic_messages(messages: list) -> bool:
for msg in messages:
if not isinstance(msg, dict):
continue
c = msg.get("content")
if not isinstance(c, list):
continue
for b in c:
if not isinstance(b, dict):
continue
t = b.get("type")
if t in ("thinking", "redacted_thinking", "tool_result"):
return True
if t == "tool_use" and isinstance(b.get("input"), dict):
return True
return False
def _anthropic_prepend_thinking(msg: dict, thinking: str) -> dict:
"""Insert a thinking block before the first text block (Anthropic shape)."""
block = {"type": "thinking", "thinking": thinking}
content = msg.get("content")
msg = dict(msg)
if isinstance(content, str):
msg["content"] = [block, {"type": "text", "text": content}]
return msg
if isinstance(content, list):
new_list: list = []
inserted = False
for b in content:
if (
not inserted
and isinstance(b, dict)
and b.get("type") == "text"
):
new_list.append(block)
inserted = True
new_list.append(b)
if not inserted:
new_list.insert(0, block)
msg["content"] = new_list
return msg
msg["content"] = [block]
return msg
def _assistant_has_nonempty_thinking_block(msg: dict) -> bool:
c = msg.get("content")
if not isinstance(c, list):
return False
for b in c:
if isinstance(b, dict) and b.get("type") == "thinking":
if (b.get("thinking") or "").strip():
return True
return False
def _inject_reasoning_into_history(
body: dict,
model: str,
session_header: str | None,
*,
prefer_anthropic_blocks: bool = False,
) -> tuple[dict, int, int, int]:
"""
Returns (body, injected_count, cache_hits, cache_misses).
prefer_anthropic_blocks: True for /v1/messages — inject thinking as content
blocks. If False, use OpenAI top-level ``reasoning`` unless the body already
contains Anthropic-style blocks (tool_result / thinking), then use blocks.
"""
if not INJECT_REASONING:
return body, 0, 0, 0
messages = _message_list_from_body(body)
if not messages:
return body, 0, 0, 0
use_blocks = prefer_anthropic_blocks or _body_looks_like_anthropic_messages(
messages
)
injected = 0
hits = 0
misses = 0
new_messages: list = []
for i, msg in enumerate(messages):
if not isinstance(msg, dict):
new_messages.append(msg)
continue
if msg.get("role") != "assistant":
new_messages.append(msg)
continue
prefix = messages[:i]
if use_blocks:
if msg.get("reasoning") or msg.get("reasoning_content"):
new_messages.append(msg)
continue
if _assistant_has_nonempty_thinking_block(msg):
new_messages.append(msg)
continue
raw_content = msg.get("content")
text_key = _anthropic_visible_text(raw_content).strip()
cached = _cache_get(model, session_header, prefix, text_key)
if cached is None and raw_content is not None:
cached = _cache_get(
model,
session_header,
prefix,
_anthropic_visible_text(raw_content),
)
if cached is not None:
msg = _anthropic_prepend_thinking(msg, cached)
injected += 1
hits += 1
print(
f"[proxy] ✅ injected Anthropic thinking block ({len(cached)} chars) at index {i}"
)
else:
misses += 1
print(
f"[proxy] ⚠️ cache MISS (anthropic) idx={i} "
f"content={repr(text_key[:60])}"
)
new_messages.append(msg)
continue
# OpenAI-style: top-level reasoning field
has_reasoning = bool(msg.get("reasoning") or msg.get("reasoning_content"))
if not has_reasoning:
raw_content = msg.get("content")
content_str = _extract_text_content(raw_content).strip()
cached = _cache_get(model, session_header, prefix, content_str)
if cached is None:
cached = _cache_get(
model, session_header, prefix, _extract_text_content(raw_content)
)
if cached is not None:
msg = dict(msg)
msg["reasoning"] = cached
injected += 1
hits += 1
print(
f"[proxy] ✅ injected reasoning ({len(cached)} chars) at index {i}"
)
else:
misses += 1
print(
f"[proxy] ⚠️ cache MISS idx={i} "
f"scope={_session_scope(session_header, prefix)[:24]}… "
f"content={repr(content_str[:60])}"
)
new_messages.append(msg)
if injected:
body = dict(body)
key = "messages" if "messages" in body else "input"
body[key] = new_messages
return body, injected, hits, misses
def _log_session_scope_for_record(session_header: str | None, body: dict) -> str:
if session_header and session_header.strip():
return f"header:{session_header.strip()[:16]}…"
msgs = _message_list_from_body(body)
if not msgs:
return "auto:empty"
return "auto:" + _prefix_hash(msgs)[:16] + "…"
# ── response parsing ──────────────────────────────────────────────────────────
def _parse_sse_stream(raw: str) -> tuple[str, str]:
reasoning_parts, content_parts = [], []
for line in raw.split("\n"):
if not line.startswith("data:") or "[DONE]" in line:
continue
try:
chunk = json.loads(line[5:].strip())
if "choices" in chunk:
delta = chunk["choices"][0].get("delta", {})
if delta.get("reasoning"):
reasoning_parts.append(delta["reasoning"])
if delta.get("content"):
content_parts.append(delta["content"])
elif chunk.get("type") == "response.output_text.delta":
content_parts.append(chunk.get("delta", ""))
elif chunk.get("type") in (
"response.reasoning_text.delta",
"response.reasoning_summary_text.delta",
"response.thinking.delta",
):
reasoning_parts.append(chunk.get("delta", ""))
elif chunk.get("type") == "response.content_part.added":
part = chunk.get("part", {})
if part.get("type") == "reasoning_text":
reasoning_parts.append(part.get("text", ""))
except Exception:
pass
return "".join(reasoning_parts), "".join(content_parts)
def _parse_anthropic_sse_stream(raw: str) -> tuple[str, str]:
"""Accumulate thinking + text from Claude /v1/messages SSE (event: … data: …)."""
reasoning_parts: list[str] = []
text_parts: list[str] = []
for block in raw.split("\n\n"):
block = block.strip()
if not block:
continue
data_line = None
for line in block.split("\n"):
if line.startswith("data:"):
data_line = line[5:].strip()
break
if not data_line:
continue
try:
obj = json.loads(data_line)
except json.JSONDecodeError:
continue
if obj.get("type") != "content_block_delta":
continue
delta = obj.get("delta") or {}
dt = delta.get("type")
if dt == "thinking_delta" and delta.get("thinking"):
reasoning_parts.append(delta["thinking"])
elif dt == "text_delta" and delta.get("text"):
text_parts.append(delta["text"])
return "".join(reasoning_parts), "".join(text_parts)
def _parse_non_stream(resp_json: dict) -> tuple[str, str]:
# Anthropic Messages non-streaming response
if resp_json.get("type") == "message" and isinstance(resp_json.get("content"), list):
reasoning_parts, content_parts = [], []
for b in resp_json["content"]:
if not isinstance(b, dict):
continue
if b.get("type") == "thinking":
reasoning_parts.append(b.get("thinking") or "")
elif b.get("type") == "text":
content_parts.append(b.get("text") or "")
return "".join(reasoning_parts), "".join(content_parts)
if "choices" in resp_json:
for choice in resp_json["choices"]:
msg = choice.get("message", {})
return (
msg.get("reasoning") or msg.get("reasoning_content") or "",
msg.get("content") or "",
)
if "output" in resp_json:
reasoning_parts, content_parts = [], []
for item in resp_json.get("output", []):
item_type = item.get("type", "")
if item_type == "reasoning":
for c in item.get("content", []):
if c.get("type") == "reasoning_text":
reasoning_parts.append(c.get("text", ""))
elif item_type == "message":
for c in item.get("content", []):
if c.get("type") == "output_text":
content_parts.append(c.get("text", ""))
return "".join(reasoning_parts), "".join(content_parts)
return "", ""
def _log(record: dict):
with LOG_FILE.open("a") as f:
f.write(json.dumps(record, ensure_ascii=False) + "\n")
def _trunc(s: str) -> str:
return s if len(s) <= MAX_BODY_LOG else s[:MAX_BODY_LOG] + f"…[+{len(s)-MAX_BODY_LOG}]"
# ── streaming ─────────────────────────────────────────────────────────────────
async def _proxy_stream(
upstream_resp: httpx.Response,
log_record: dict,
model: str,
session_header: str | None,
request_messages: list,
is_responses_api: bool = False,
req_input: list | None = None,
is_anthropic: bool = False,
):
chunks: list[bytes] = []
try:
async for raw_chunk in upstream_resp.aiter_raw():
chunks.append(raw_chunk)
yield raw_chunk
finally:
full_stream = b"".join(chunks).decode(errors="replace")
if is_anthropic:
reasoning, content = _parse_anthropic_sse_stream(full_stream)
else:
reasoning, content = _parse_sse_stream(full_stream)
if content:
_cache_store(model, session_header, request_messages, content, reasoning)
if is_responses_api:
resp_id = None
for line in full_stream.split("\n"):
if not line.startswith("data:"):
continue
try:
chunk = json.loads(line[5:].strip())
t = chunk.get("type", "")
if t in ("response.created", "response.done", "response.completed"):
resp_id = chunk.get("response", {}).get("id")
if resp_id and t == "response.done":
break
if not resp_id and chunk.get("object") == "response":
resp_id = chunk.get("id")
except Exception:
pass
if resp_id:
_resp_cache_store(
resp_id, model, reasoning, content, input_messages=req_input or []
)
log_record["response_id"] = resp_id
print(
f"[proxy] 📦 cached response id={resp_id[:20]}… "
f"reasoning={len(reasoning)} content={len(content)}"
)
print(
f"[proxy] stream cached: content_len={len(content)} reasoning_len={len(reasoning)}"
)
log_record["response_reasoning_len"] = len(reasoning)
log_record["response_content_snippet"] = _trunc(content[:200])
log_record["response_body"] = _trunc(full_stream)
_log(log_record)
# ── main handler ─────────────────────────────────────────────────────────────
@app.api_route(
"/{path:path}",
methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"],
)
async def proxy(path: str, request: Request):
endpoint = "/" + path
raw_body = await request.body()
fwd_headers = {
k: v
for k, v in request.headers.items()
if k.lower() not in ("host", "content-length", "transfer-encoding")
}
fwd_headers.setdefault(
"user-agent",
"Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 Chrome/124.0 Safari/537.36",
)
try:
body_json: dict = json.loads(raw_body) if raw_body else {}
except Exception:
body_json = {}
session_header = request.headers.get(SESSION_HEADER) or request.headers.get(
SESSION_HEADER.upper()
)
is_anthropic = endpoint.startswith("/v1/messages")
is_responses_api = endpoint.startswith("/v1/responses")
needs_llm_patch = (
endpoint.startswith("/v1/chat/")
or is_responses_api
or is_anthropic
)
model = body_json.get("model", "") if body_json else ""
pt_injected = False
rea_injected = 0
cache_hits = 0
cache_misses = 0
if body_json and needs_llm_patch:
prev_id = body_json.get("previous_response_id")
if prev_id and is_responses_api:
cached_prev = _resp_cache_get(prev_id)
if cached_prev:
prev_reasoning = cached_prev["reasoning"]
prev_content = cached_prev["content"]
prev_input_msgs = cached_prev.get("input_messages", [])
prev_asst: dict = {"role": "assistant", "content": prev_content}
if prev_reasoning:
prev_asst["reasoning"] = prev_reasoning
cur_input = body_json.get("input", "")
if isinstance(cur_input, str):
cur_input = [{"role": "user", "content": cur_input}]
body_json = dict(body_json)
body_json["input"] = prev_input_msgs + [prev_asst] + (
cur_input if isinstance(cur_input, list) else []
)
del body_json["previous_response_id"]
rea_injected += 1 if prev_reasoning else 0
print(
f"[proxy] ✅ expanded previous_response_id={prev_id[:16]}… "
f"history={len(prev_input_msgs)} msgs"
)
body_json, inj, hits, misses = _inject_reasoning_into_history(
body_json,
model,
session_header,
prefer_anthropic_blocks=is_anthropic,
)
rea_injected += inj
cache_hits += hits
cache_misses += misses
body_json, pt_injected = _inject_preserve_thinking(body_json)
raw_body = json.dumps(body_json).encode()
fwd_headers["content-length"] = str(len(raw_body))
is_stream = bool(body_json.get("stream")) if body_json else False
request_messages = _message_list_from_body(body_json) if body_json else []
log_record: dict[str, Any] = {
"ts": datetime.now(timezone.utc).isoformat(),
"endpoint": endpoint,
"method": request.method,
"preserve_thinking_injected": pt_injected,
"reasoning_turns_injected": rea_injected,
"reasoning_cache_hits": cache_hits,
"reasoning_cache_misses": cache_misses,
"session_scope": _log_session_scope_for_record(session_header, body_json)
if body_json
else None,
"chat_template_kwargs": body_json.get("chat_template_kwargs")
if body_json
else None,
"is_stream": is_stream,
"request_body": _trunc(
json.dumps(body_json) if body_json else raw_body.decode(errors="replace")
),
}
try:
if is_stream:
upstream_req = _client.build_request(
request.method, endpoint, content=raw_body, headers=fwd_headers
)
upstream_resp = await _client.send(upstream_req, stream=True)
log_record["response_status"] = upstream_resp.status_code
stream_req_input = body_json.get("input", []) if is_responses_api else None
if isinstance(stream_req_input, str):
stream_req_input = [{"role": "user", "content": stream_req_input}]
return StreamingResponse(
_proxy_stream(
upstream_resp,
log_record,
model,
session_header,
request_messages,
is_responses_api=is_responses_api,
req_input=stream_req_input,
is_anthropic=is_anthropic,
),
status_code=upstream_resp.status_code,
headers={
k: v
for k, v in upstream_resp.headers.items()
if k.lower()
not in ("content-encoding", "transfer-encoding", "content-length")
},
media_type=upstream_resp.headers.get(
"content-type", "text/event-stream"
),
)
upstream_resp = await _client.request(
request.method, endpoint, content=raw_body, headers=fwd_headers
)
resp_bytes = upstream_resp.content
resp_str = resp_bytes.decode(errors="replace")
try:
resp_json = json.loads(resp_bytes)
reasoning, content = _parse_non_stream(resp_json)
if content:
_cache_store(model, session_header, request_messages, content, reasoning)
resp_id = resp_json.get("id") if isinstance(resp_json, dict) else None
if resp_id and is_responses_api:
req_input = body_json.get("input", [])
if isinstance(req_input, str):
req_input = [{"role": "user", "content": req_input}]
_resp_cache_store(
resp_id, model, reasoning, content, input_messages=req_input
)
print(
f"[proxy] 📦 cached response id={resp_id[:20]}… "
f"reasoning={len(reasoning)} content={len(content)}"
)
log_record["response_reasoning_len"] = len(reasoning)
log_record["response_id"] = resp_id
except Exception:
resp_json = None
log_record["response_reasoning_len"] = -1
log_record["response_status"] = upstream_resp.status_code
log_record["response_body"] = _trunc(resp_str)
_log(log_record)
resp_headers = {
k: v
for k, v in upstream_resp.headers.items()
if k.lower() not in ("content-encoding", "transfer-encoding", "content-length")
}
return JSONResponse(
content=resp_json if resp_json is not None else resp_str,
status_code=upstream_resp.status_code,
headers=resp_headers,
)
except Exception as exc:
log_record["error"] = str(exc)
_log(log_record)
return JSONResponse({"error": str(exc)}, status_code=502)
# ── entry point ───────────────────────────────────────────────────────────────
if __name__ == "__main__":
p = argparse.ArgumentParser()
p.add_argument("--port", type=int, default=9000)
p.add_argument("--host", default="0.0.0.0")
p.add_argument(
"--no-inject-pt", action="store_true", help="Disable preserve_thinking injection"
)
p.add_argument(
"--no-inject-rea", action="store_true", help="Disable reasoning re-injection"
)
p.add_argument(
"--max-reasoning-cache",
type=int,
default=50_000,
help="Max reasoning cache entries (LRU)",
)
p.add_argument(
"--max-responses-cache",
type=int,
default=50_000,
help="Max response-id cache entries (LRU)",
)
p.add_argument(
"--cache-ttl-days",
type=float,
default=30.0,
help="TTL in days for cache entries (0 = disable TTL; LRU only)",
)
p.add_argument(
"--upstream",
default=None,
metavar="URL",
help=(
"Upstream vLLM base URL (no trailing path). "
"Default: env PRESERVE_THINKING_PROXY_UPSTREAM or https://vllm.treowai.com"
),
)
args = p.parse_args()
_us = (
args.upstream
or os.environ.get("PRESERVE_THINKING_PROXY_UPSTREAM", "").strip()
or UPSTREAM
)
UPSTREAM = _us.rstrip("/")
if args.no_inject_pt:
INJECT_PRESERVE_THINKING = False
if args.no_inject_rea:
INJECT_REASONING = False
MAX_REASONING_ENTRIES = args.max_reasoning_cache
MAX_RESPONSES_ENTRIES = args.max_responses_cache
CACHE_TTL_SEC = int(args.cache_ttl_days * 24 * 3600) if args.cache_ttl_days > 0 else 0
uvicorn.run(app, host=args.host, port=args.port, log_level="warning")
Raw
start_vllm_and_proxy.sh
#!/usr/bin/env bash
# Start vLLM (OpenAI + Anthropic API) on port 8008, then the preserve_thinking proxy on 8000.
#
# Usage:
# ./start_vllm_and_proxy.sh
#
# First-time vLLM fork install (clone + pip) runs when the clone directory is new,
# or when INSTALL_VLLM=1 is set:
# INSTALL_VLLM=1 ./start_vllm_and_proxy.sh
#
# Environment overrides:
# VLLM_SOURCE_DIR — clone path (default: ~/src/vllm-kdcyberdude)
# VLLM_PORT — default 8008
# PROXY_PORT — default 8000
# CONDA_ENV — default vllm
# SKIP_CONDA — set to 1 to skip conda activate (use current Python)
#
# Required:
# VLLM_API_KEY — same value passed to vLLM --api-key and used for /health checks
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
VLLM_SOURCE_DIR="${VLLM_SOURCE_DIR:-$HOME/src/vllm-kdcyberdude}"
VLLM_PORT="${VLLM_PORT:-8008}"
PROXY_PORT="${PROXY_PORT:-8000}"
CONDA_ENV="${CONDA_ENV:-vllm}"
CHAT_TEMPLATE_URL="https://raw.githubusercontent.com/allanchan339/vLLM-Qwen3-3.5-3.6-chat-template-fix/main/chat-template/qwen3.6-enhanced.jinja"
CHAT_TEMPLATE_FILE="${CHAT_TEMPLATE_FILE:-$HOME/vllm_templates/qwen3.6-enhanced.jinja}"
VLLM_PID=""
cleanup() {
if [[ -n "${VLLM_PID}" ]] && kill -0 "${VLLM_PID}" 2>/dev/null; then
echo "[start] Stopping vLLM (pid ${VLLM_PID})..."
kill "${VLLM_PID}" 2>/dev/null || true
wait "${VLLM_PID}" 2>/dev/null || true
fi
}
trap cleanup EXIT INT TERM
ensure_chat_template() {
mkdir -p "$(dirname "${CHAT_TEMPLATE_FILE}")"
if [[ ! -s "${CHAT_TEMPLATE_FILE}" ]]; then
echo "[start] Downloading chat template → ${CHAT_TEMPLATE_FILE}"
curl -fsSL "${CHAT_TEMPLATE_URL}" -o "${CHAT_TEMPLATE_FILE}"
else
echo "[start] Chat template already present: ${CHAT_TEMPLATE_FILE}"
fi
}
ensure_vllm_fork() {
mkdir -p "$(dirname "${VLLM_SOURCE_DIR}")"
local cloned=0
if [[ ! -d "${VLLM_SOURCE_DIR}/.git" ]]; then
echo "[start] Cloning https://github.com/kdcyberdude/vllm → ${VLLM_SOURCE_DIR}"
git clone https://github.com/kdcyberdude/vllm.git "${VLLM_SOURCE_DIR}"
cloned=1
fi
if [[ "${cloned}" -eq 1 ]] || [[ "${INSTALL_VLLM:-0}" == "1" ]]; then
echo "[start] pip install vLLM fork (VLLM_USE_PRECOMPILED=1)…"
(
cd "${VLLM_SOURCE_DIR}"
VLLM_USE_PRECOMPILED=1 pip install .
)
else
echo "[start] Skipping pip install (set INSTALL_VLLM=1 to force, or delete clone to reinstall)"
fi
}
activate_conda_env() {
if [[ "${SKIP_CONDA:-0}" == "1" ]]; then
echo "[start] SKIP_CONDA=1 — using current Python: $(command -v python3)"
return 0
fi
local base
base="$(conda info --base 2>/dev/null)" || {
echo "[start] ERROR: conda not found. Install Miniconda/Anaconda or set SKIP_CONDA=1." >&2
exit 1
}
# shellcheck source=/dev/null
source "${base}/etc/profile.d/conda.sh"
conda activate "${CONDA_ENV}"
echo "[start] Conda env: ${CONDA_ENV} ($(command -v python3))"
}
wait_for_vllm() {
local auth=(-H "Authorization: Bearer ${VLLM_API_KEY}")
echo "[start] Waiting for vLLM health on http://127.0.0.1:${VLLM_PORT}/health …"
for _ in $(seq 1 360); do
if curl -sf "${auth[@]}" "http://127.0.0.1:${VLLM_PORT}/health" >/dev/null; then
echo "[start] vLLM is up."
return 0
fi
sleep 2
done
echo "[start] ERROR: vLLM did not become healthy in time." >&2
exit 1
}
# ── main ─────────────────────────────────────────────────────────────────────
ensure_chat_template
ensure_vllm_fork
activate_conda_env
if [[ -z "${VLLM_API_KEY:-}" ]]; then
echo "[start] ERROR: VLLM_API_KEY is not set. Export it in your environment before running this script." >&2
exit 1
fi
export CUDA_DEVICE_ORDER=PCI_BUS_ID
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,4}"
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export NCCL_TIMEOUT=3600
export TORCH_NCCL_BLOCKING_WAIT=0
export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
export NCCL_CUMEM_ENABLE=0
export NCCL_P2P_DISABLE=1
export NCCL_IB_DISABLE=1
export NCCL_SHM_DISABLE=0
export NCCL_ALGO=Ring
export NCCL_P2P_LEVEL=LOC
export VLLM_RPC_TIMEOUT=180
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_TEST_FORCE_FP8_MARLIN=1
echo "[start] Launching vLLM on port ${VLLM_PORT} (logs: ${SCRIPT_DIR}/vllm_server.log)"
python3 -m vllm.entrypoints.openai.api_server \
--model Qwen/Qwen3.6-27B-FP8 \
--tensor-parallel-size 1 \
--pipeline-parallel-size 3 \
--max-model-len 950000 \
--reasoning-parser qwen3 \
--enable-auto-tool-choice \
--tool-call-parser qwen3_coder \
--gpu-memory-utilization 0.94 \
--kv-cache-dtype fp8 \
--enable-chunked-prefill \
--enable-prefix-caching \
--trust-remote-code \
--port "${VLLM_PORT}" \
--api-key "${VLLM_API_KEY}" \
--chat-template "${CHAT_TEMPLATE_FILE}" \
--default-chat-template-kwargs '{"preserve_thinking":true}' \
--override-generation-config '{"temperature": 0.6, "top_p":0.95, "top_k":20, "min_p":0.0, "presence_penalty":0.0, "repetition_penalty":1.0}' \
--hf-overrides '{"text_config": {"rope_parameters": {"mrope_interleaved": true, "mrope_section": [11, 11, 10], "rope_type": "yarn", "rope_theta": 10000000, "partial_rotary_factor": 0.25, "factor": 4.0, "original_max_position_embeddings": 262144}}}' \
--max-num-seqs 4 \
>>"${SCRIPT_DIR}/vllm_server.log" 2>&1 &
VLLM_PID=$!
wait_for_vllm
export PRESERVE_THINKING_PROXY_UPSTREAM="http://127.0.0.1:${VLLM_PORT}"
echo "[start] Launching proxy on port ${PROXY_PORT} → ${PRESERVE_THINKING_PROXY_UPSTREAM}"
python3 "${SCRIPT_DIR}/proxy.py" \
--host 0.0.0.0 \
--port "${PROXY_PORT}" \
--upstream "${PRESERVE_THINKING_PROXY_UPSTREAM}"
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment