File size: 5,892 Bytes
0236599 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | #!/usr/bin/env python3
import json, os, time, uuid, urllib.request, urllib.error
from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler
from transformers import AutoTokenizer
MODEL_DIR=os.environ["DEVSTRAL_MODEL"]
BACKEND=os.environ.get("DEVSTRAL_BACKEND","http://192.0.2.11:8031").rstrip("/")
PUBLIC_MODEL=os.environ.get("DEVSTRAL_PUBLIC_MODEL","devstral-small2-24b")
HOST=os.environ.get("PROXY_HOST","192.0.2.11")
PORT=int(os.environ.get("PROXY_PORT","8025"))
TOKENIZER=AutoTokenizer.from_pretrained(MODEL_DIR, trust_remote_code=True)
def normalize_content(content):
if isinstance(content,str):
return content
if isinstance(content,list):
parts=[]
for item in content:
if isinstance(item,dict) and item.get("type") in ("text","input_text"):
parts.append(str(item.get("text","")))
elif isinstance(item,str):
parts.append(item)
return "\n".join(x for x in parts if x)
if content is None:
return ""
return str(content)
def render_messages(messages):
normalized=[]
for m in messages:
normalized.append({
"role": str(m.get("role","user")),
"content": normalize_content(m.get("content",""))
})
return TOKENIZER.apply_chat_template(
normalized,
tokenize=False,
add_generation_prompt=True
)
def backend_post(path,payload,timeout=600):
req=urllib.request.Request(
BACKEND+path,
data=json.dumps(payload).encode("utf-8"),
headers={"Content-Type":"application/json"},
method="POST"
)
with urllib.request.urlopen(req,timeout=timeout) as r:
return r.status,json.loads(r.read().decode("utf-8"))
def backend_health():
try:
with urllib.request.urlopen(BACKEND+"/health",timeout=3) as r:
return r.status==200
except Exception:
return False
class Handler(BaseHTTPRequestHandler):
server_version="DevstralChatProxyV2/1.0"
def log_message(self,fmt,*args):
print("%s - %s" % (self.address_string(), fmt%args), flush=True)
def send_json(self,status,obj):
raw=json.dumps(obj,ensure_ascii=False).encode("utf-8")
self.send_response(status)
self.send_header("Content-Type","application/json; charset=utf-8")
self.send_header("Content-Length",str(len(raw)))
self.end_headers()
self.wfile.write(raw)
def do_GET(self):
if self.path=="/health":
if backend_health():
self.send_json(200,{"status":"ok","backend":"ok"})
else:
self.send_json(503,{"status":"degraded","backend":"down"})
return
if self.path=="/v1/models":
self.send_json(200,{
"object":"list",
"data":[{
"id":PUBLIC_MODEL,
"object":"model",
"created":int(time.time()),
"owned_by":"local-vllm-proxy",
"root":MODEL_DIR,
"parent":None
}]
})
return
self.send_json(404,{"error":{"message":"not found"}})
def do_POST(self):
try:
length=int(self.headers.get("Content-Length","0"))
body=json.loads(self.rfile.read(length) or b"{}")
except Exception as e:
self.send_json(400,{"error":{"message":"invalid json: "+repr(e)}})
return
if self.path=="/v1/chat/completions":
try:
messages=body.get("messages") or []
prompt=render_messages(messages)
payload={
"model":PUBLIC_MODEL,
"prompt":prompt,
"max_tokens":int(body.get("max_tokens",body.get("max_completion_tokens",256))),
"temperature":float(body.get("temperature",0.0)),
"top_p":float(body.get("top_p",1.0)),
"stream":False
}
if body.get("stop") is not None:
payload["stop"]=body["stop"]
status,out=backend_post("/v1/completions",payload)
text=((out.get("choices") or [{}])[0].get("text") or "")
resp={
"id":"chatcmpl-"+uuid.uuid4().hex[:24],
"object":"chat.completion",
"created":int(time.time()),
"model":PUBLIC_MODEL,
"choices":[{
"index":0,
"message":{"role":"assistant","content":text},
"finish_reason":((out.get("choices") or [{}])[0].get("finish_reason"))
}],
"usage":out.get("usage",{})
}
self.send_json(200 if status==200 else status,resp)
except urllib.error.HTTPError as e:
raw=e.read().decode(errors="replace")
self.send_json(e.code,{"error":{"message":"backend http error","body":raw}})
except Exception as e:
self.send_json(500,{"error":{"message":"proxy render/backend failure","detail":repr(e)}})
return
if self.path=="/v1/completions":
try:
status,out=backend_post("/v1/completions",body)
self.send_json(status,out)
except urllib.error.HTTPError as e:
self.send_json(e.code,{"error":{"message":e.read().decode(errors="replace")}})
except Exception as e:
self.send_json(500,{"error":{"message":repr(e)}})
return
self.send_json(404,{"error":{"message":"not found"}})
print(f"DEVSTRAL_CHAT_PROXY_V2_START host={HOST} port={PORT} backend={BACKEND}",flush=True)
ThreadingHTTPServer((HOST,PORT),Handler).serve_forever()
|