OpenAI-compatible HTTP API: stdlib-only gateway over SERVE with KV prefix reuse across stateless requests (#21)
* Add OpenAI-compatible HTTP API * Support browser API clients * Handle missing KV cache during rewind
This commit is contained in:
@@ -4,6 +4,7 @@ colibrì — piccolo motore, modello immenso.
|
||||
CLI per far girare GLM-5.2 (744B) in locale, su CPU, in ~15-26 GB di RAM.
|
||||
|
||||
coli chat chat interattiva (carica il modello UNA volta)
|
||||
coli serve API HTTP compatibile OpenAI (motore persistente)
|
||||
coli run "prompt" generazione singola
|
||||
coli info stato: modello, RAM, disco, config
|
||||
coli bench [task...] benchmark di qualità (MMLU/HellaSwag/...)
|
||||
@@ -375,6 +376,12 @@ def cmd_chat(a):
|
||||
except Exception: pass
|
||||
print(f" {C.teal}ciao{C.r} {C.dim}— il colibrì torna al nido{C.r} 🐦\n")
|
||||
|
||||
def cmd_serve(a):
|
||||
need_model(a.model)
|
||||
from openai_server import serve
|
||||
serve(a.model, a.host, a.port, a.model_id, a.api_key,
|
||||
a.cap, a.ngen, GLM, env_for(a), a.cors_origin)
|
||||
|
||||
def cmd_bench(a):
|
||||
need_model(a.model)
|
||||
banner("bench")
|
||||
@@ -427,13 +434,18 @@ def main():
|
||||
sub.add_parser("build", parents=[common]); sub.add_parser("info", parents=[common])
|
||||
pr=sub.add_parser("run", parents=[common]); pr.add_argument("prompt", nargs="*")
|
||||
sub.add_parser("chat", parents=[common])
|
||||
ps=sub.add_parser("serve", parents=[common])
|
||||
ps.add_argument("--host",default="127.0.0.1"); ps.add_argument("--port",type=int,default=8000)
|
||||
ps.add_argument("--model-id",default=os.environ.get("COLI_MODEL_ID","glm-5.2-colibri"))
|
||||
ps.add_argument("--api-key",default=os.environ.get("COLI_API_KEY"))
|
||||
ps.add_argument("--cors-origin",action="append",default=None)
|
||||
pb=sub.add_parser("bench", parents=[common]); pb.add_argument("tasks", nargs="*")
|
||||
pb.add_argument("--limit",type=int,default=40); pb.add_argument("--data",default=os.path.join(HERE,"bench"))
|
||||
pc=sub.add_parser("convert", parents=[common]); pc.add_argument("--repo",default="zai-org/GLM-5.2-FP8")
|
||||
pc.add_argument("--ebits",type=int,default=4); pc.add_argument("--io-bits",type=int,default=8); pc.add_argument("--xbits",type=int,default=0)
|
||||
pc.add_argument("--no-mtp",action="store_true",help="salta la testa MTP (niente draft speculativi)")
|
||||
a=ap.parse_args()
|
||||
{"build":cmd_build,"info":cmd_info,"run":cmd_run,"chat":cmd_chat,"bench":cmd_bench,
|
||||
{"build":cmd_build,"info":cmd_info,"run":cmd_run,"chat":cmd_chat,"serve":cmd_serve,"bench":cmd_bench,
|
||||
"convert":cmd_convert}.get(a.cmd, lambda _:(banner(),print(__doc__)))(a)
|
||||
|
||||
if __name__=="__main__":
|
||||
|
||||
@@ -1828,12 +1828,14 @@ static void kv_hdr(Model *m, int32_t *h, int nrec){
|
||||
h[0]=c->n_layers; h[1]=c->kv_lora; h[2]=c->qk_rope;
|
||||
h[3]=m->has_dsa?c->index_hd:0; h[4]=nic; h[5]=c->vocab; h[6]=nrec; h[7]=0;
|
||||
}
|
||||
static void kv_disk_reset(void){
|
||||
static void kv_disk_truncate(int nrec){
|
||||
if(!g_kvsave) return;
|
||||
FILE *f=fopen(g_kv_path,"r+b"); if(!f) return;
|
||||
int32_t nz=0; fseek(f,8+6*4,SEEK_SET); fwrite(&nz,4,1,f); fclose(f);
|
||||
g_kv_nrec=0;
|
||||
FILE *f=fopen(g_kv_path,"r+b");
|
||||
if(!f){ g_kv_nrec=0; return; }
|
||||
g_kv_nrec=nrec;
|
||||
int32_t nr=nrec; fseek(f,8+6*4,SEEK_SET); fwrite(&nr,4,1,f); fclose(f);
|
||||
}
|
||||
static void kv_disk_reset(void){ kv_disk_truncate(0); }
|
||||
static void kv_disk_append(Model *m, const int *hist, int len){
|
||||
if(!g_kvsave || len<=g_kv_nrec) return;
|
||||
Cfg *c=&m->c;
|
||||
@@ -1930,33 +1932,79 @@ static void run_serve(Model *m, const char *snap){
|
||||
printf("STAT %d %.2f %.1f %.2f\n", prod, prod/tdt, (dh+dm)>0?100.0*dh/(dh+dm):0.0, rss_gb());
|
||||
fflush(stdout); kv_disk_append(m,hist,len); repin_pass(m); continue; } /* RFC: re-pin a caldo tra i turni / live re-pin between turns */
|
||||
if(nr<1){ printf("\x01\x01" "END" "\x01\x01\n"); printf("STAT 0 0.00 0.0 %.2f\n", rss_gb()); fflush(stdout); continue; }
|
||||
int bl=0; /* costruisce il testo del turno (con template) */
|
||||
/* API mode: an exact, length-prefixed prompt. Unlike the interactive
|
||||
* line protocol this accepts newlines. The tokenized prompt is matched
|
||||
* against hist so the common KV prefix survives stateless HTTP turns.
|
||||
* Per-request generation controls follow the byte count:
|
||||
* \x02PROMPT <bytes> <max_tokens> <temperature> <top_p>\n<prompt>\n */
|
||||
char *raw=NULL, *input=line;
|
||||
int input_n=(int)nr, raw_mode=0, req_ngen=ngen, prompt_tokens=0;
|
||||
float base_temp=g_temp, base_nuc=g_nuc;
|
||||
if(!strncmp(line,"\x02PROMPT ",8)){
|
||||
unsigned long long nb=0; double rt=0, rp=0;
|
||||
if(sscanf(line+8,"%llu %d %lf %lf",&nb,&req_ngen,&rt,&rp)!=4 ||
|
||||
nb>(16u<<20) || req_ngen<1 || rt<0 || rt>2 || rp<=0 || rp>1){
|
||||
printf("\x01\x01" "END" "\x01\x01\n"); printf("STAT 0 0.00 0.0 %.2f 0 0\n",rss_gb()); fflush(stdout); continue;
|
||||
}
|
||||
raw=malloc((size_t)nb+1); if(!raw){fprintf(stderr,"OOM raw prompt\n");exit(1);}
|
||||
if(fread(raw,1,(size_t)nb,stdin)!=(size_t)nb){free(raw);break;}
|
||||
int delim=fgetc(stdin); if(delim!='\n' && delim!=EOF) ungetc(delim,stdin);
|
||||
if(memchr(raw,0,(size_t)nb)){free(raw); printf("\x01\x01" "END" "\x01\x01\n");
|
||||
printf("STAT 0 0.00 0.0 %.2f 0 0\n",rss_gb()); fflush(stdout); continue;}
|
||||
raw[nb]=0; input=raw; input_n=(int)nb; raw_mode=1;
|
||||
if(req_ngen>ngen) req_ngen=ngen;
|
||||
g_temp=(float)rt; g_nuc=(float)rp;
|
||||
}
|
||||
int bl=0, k=0; /* costruisce/tokenizza il turno */
|
||||
/* template UFFICIALE GLM-5.2 (chat_template.jinja): niente \n dopo i ruoli, e dopo
|
||||
* <|assistant|> serve SEMPRE il blocco think — <think></think> lo DISATTIVA (nothink):
|
||||
* col template sbagliato il modello farfuglia e non emette mai lo stop. THINK=1 lo abilita. */
|
||||
const char *tk = getenv("THINK")&&atoi(getenv("THINK"))? "<think>" : "<think></think>";
|
||||
if(templ){ if(first) bl+=snprintf(buf+bl,(1<<16)-bl,"[gMASK]<sop>");
|
||||
bl+=snprintf(buf+bl,(1<<16)-bl,"<|user|>%s<|assistant|>%s",line,tk); }
|
||||
else bl+=snprintf(buf+bl,(1<<16)-bl,"%s",line);
|
||||
int k=tok_encode(&T,buf,bl,hist+len,maxctx-len);
|
||||
if(k<1){ printf("\x01\x01" "END" "\x01\x01\n"); printf("STAT 0 0.00 0.0 %.2f\n", rss_gb()); fflush(stdout); continue; }
|
||||
if(len+k+8+g_draft>=maxctx){ len=0; first=1; kv_disk_reset(); /* contesto pieno: azzera e ricomincia */
|
||||
bl=0; if(templ){ bl+=snprintf(buf+bl,(1<<16)-bl,"[gMASK]<sop><|user|>%s<|assistant|>%s",line,tk); }
|
||||
else bl+=snprintf(buf+bl,(1<<16)-bl,"%s",line);
|
||||
k=tok_encode(&T,buf,bl,hist,maxctx); if(k>maxctx-8-g_draft) k=maxctx-8-g_draft; }
|
||||
if(raw_mode){
|
||||
int *tmp=malloc(maxctx*sizeof(int)); if(!tmp){fprintf(stderr,"OOM raw tokens\n");exit(1);}
|
||||
prompt_tokens=tok_encode(&T,input,input_n,tmp,maxctx-8-g_draft);
|
||||
int old_len=len, prefix=0;
|
||||
while(prefix<old_len && prefix<prompt_tokens && hist[prefix]==tmp[prefix]) prefix++;
|
||||
if(prefix<old_len){
|
||||
len=prefix;
|
||||
if(m->has_mtp) m->kv_start[m->c.n_layers]=-1;
|
||||
kv_disk_truncate(len); /* il prossimo append sovrascrive solo la coda */
|
||||
}
|
||||
k=prompt_tokens-len;
|
||||
if(k>0) memcpy(hist+len,tmp+len,k*sizeof(int));
|
||||
fprintf(stderr,"[API] KV prefix %d/%d token, prefill %d\n",len,prompt_tokens,k);
|
||||
free(tmp);
|
||||
} else {
|
||||
if(templ){ if(first) bl+=snprintf(buf+bl,(1<<16)-bl,"[gMASK]<sop>");
|
||||
bl+=snprintf(buf+bl,(1<<16)-bl,"<|user|>%s<|assistant|>%s",input,tk); }
|
||||
else bl+=snprintf(buf+bl,(1<<16)-bl,"%s",input);
|
||||
k=tok_encode(&T,buf,bl,hist+len,maxctx-len); prompt_tokens=k;
|
||||
if(len+k+8+g_draft>=maxctx){ len=0; first=1; kv_disk_reset();
|
||||
bl=0; if(templ){ bl+=snprintf(buf+bl,(1<<16)-bl,"[gMASK]<sop><|user|>%s<|assistant|>%s",input,tk); }
|
||||
else bl+=snprintf(buf+bl,(1<<16)-bl,"%s",input);
|
||||
k=tok_encode(&T,buf,bl,hist,maxctx); if(k>maxctx-8-g_draft) k=maxctx-8-g_draft;
|
||||
prompt_tokens=k;
|
||||
}
|
||||
}
|
||||
if(prompt_tokens<1){ free(raw); g_temp=base_temp; g_nuc=base_nuc;
|
||||
printf("\x01\x01" "END" "\x01\x01\n"); printf("STAT 0 0.00 0.0 %.2f 0 0\n", rss_gb()); fflush(stdout); continue; }
|
||||
first=0;
|
||||
int cur=ngen; if(len+k+cur+g_draft+2>=maxctx) cur=maxctx-len-k-g_draft-2;
|
||||
int cur=req_ngen; if(len+k+cur+g_draft+2>=maxctx) cur=maxctx-len-k-g_draft-2;
|
||||
uint64_t h0=m->hits, ms0=m->miss; double tt0=now_s();
|
||||
float *logit=step(m,hist+len,k,len); len+=k;
|
||||
float *logit;
|
||||
if(k>0){ logit=step(m,hist+len,k,len); len+=k; }
|
||||
else logit=step(m,hist+len-1,1,len-1); /* prompt identico/prefisso: rigenera i logits */
|
||||
EmitStream es={&T,m,now_s(),0,1};
|
||||
int prod=0;
|
||||
if(cur>0) prod=spec_decode(m,hist,len,cur,eos,logit,emit_stream,&es,&len);
|
||||
else free(logit);
|
||||
double tdt=now_s()-tt0; if(tdt<1e-6) tdt=1e-6;
|
||||
double dh=(double)(m->hits-h0), dm=(double)(m->miss-ms0);
|
||||
printf("\n\x01\x01" "END" "\x01\x01\n");
|
||||
printf("STAT %d %.2f %.1f %.2f\n", prod, prod/tdt, (dh+dm)>0?100.0*dh/(dh+dm):0.0, rss_gb());
|
||||
printf("%s\x01\x01" "END" "\x01\x01\n",raw_mode?"":"\n");
|
||||
printf("STAT %d %.2f %.1f %.2f %d %d\n", prod, prod/tdt,
|
||||
(dh+dm)>0?100.0*dh/(dh+dm):0.0, rss_gb(), prompt_tokens, prod>=cur);
|
||||
fflush(stdout);
|
||||
free(raw); g_temp=base_temp; g_nuc=base_nuc;
|
||||
usage_save(m); /* la cache che impara: storia aggiornata a ogni turno */
|
||||
kv_disk_append(m,hist,len); /* KV su disco: il prossimo avvio riparte da qui */
|
||||
}
|
||||
|
||||
@@ -0,0 +1,466 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Dependency-free OpenAI-compatible HTTP gateway for the colibri engine."""
|
||||
|
||||
import argparse
|
||||
import codecs
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
from urllib.parse import unquote, urlsplit
|
||||
|
||||
|
||||
HERE = Path(__file__).resolve().parent
|
||||
END = b"\x01\x01END\x01\x01\n"
|
||||
READY = b"\x01\x01READY\x01\x01\n"
|
||||
MAX_BODY = 4 << 20
|
||||
DEFAULT_CORS_ORIGINS = (
|
||||
"http://127.0.0.1:5173",
|
||||
"http://localhost:5173",
|
||||
"http://tauri.localhost",
|
||||
"tauri://localhost",
|
||||
)
|
||||
|
||||
|
||||
class APIError(Exception):
|
||||
def __init__(self, status, message, param=None, code=None, error_type="invalid_request_error"):
|
||||
super().__init__(message)
|
||||
self.status = status
|
||||
self.message = message
|
||||
self.param = param
|
||||
self.code = code
|
||||
self.error_type = error_type
|
||||
|
||||
|
||||
def error_object(error):
|
||||
return {"error": {"message": error.message, "type": error.error_type,
|
||||
"param": error.param, "code": error.code}}
|
||||
|
||||
|
||||
def content_text(content, param):
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if not isinstance(content, list):
|
||||
raise APIError(400, "Message content must be a string or an array of text parts.", param)
|
||||
parts = []
|
||||
for index, part in enumerate(content):
|
||||
if not isinstance(part, dict) or part.get("type") not in ("text", "input_text"):
|
||||
raise APIError(400, "Colibri currently supports text message content only.",
|
||||
f"{param}.{index}", "unsupported_content_type")
|
||||
if not isinstance(part.get("text"), str):
|
||||
raise APIError(400, "Text content parts require a string `text` field.",
|
||||
f"{param}.{index}.text")
|
||||
parts.append(part["text"])
|
||||
return "".join(parts)
|
||||
|
||||
|
||||
def render_chat(messages, enable_thinking=False, reasoning_effort=None):
|
||||
"""Render the text-only subset of the official GLM-5.2 chat template."""
|
||||
if not isinstance(messages, list) or not messages:
|
||||
raise APIError(400, "`messages` must be a non-empty array.", "messages")
|
||||
prompt = ["[gMASK]<sop>"]
|
||||
if enable_thinking:
|
||||
effort = "High" if reasoning_effort == "high" else "Max"
|
||||
prompt.append(f"<|system|>Reasoning Effort: {effort}")
|
||||
for index, message in enumerate(messages):
|
||||
if not isinstance(message, dict):
|
||||
raise APIError(400, "Each message must be an object.", f"messages.{index}")
|
||||
role = message.get("role")
|
||||
text = content_text(message.get("content"), f"messages.{index}.content")
|
||||
if role in ("system", "developer"):
|
||||
prompt.append(f"<|system|>{text}")
|
||||
elif role == "user":
|
||||
prompt.append(f"<|user|>{text}")
|
||||
elif role == "assistant":
|
||||
prompt.append(f"<|assistant|><think></think>{text.strip()}")
|
||||
else:
|
||||
raise APIError(400, f"Unsupported message role: {role!r}.",
|
||||
f"messages.{index}.role", "unsupported_role")
|
||||
prompt.append("<|assistant|><think>" if enable_thinking else
|
||||
"<|assistant|><think></think>")
|
||||
return "".join(prompt)
|
||||
|
||||
|
||||
def generation_options(body, limit):
|
||||
if body.get("n", 1) != 1:
|
||||
raise APIError(400, "Colibri currently supports `n=1` only.", "n", "unsupported_value")
|
||||
for name in ("tools", "functions"):
|
||||
if body.get(name):
|
||||
raise APIError(400, f"`{name}` is not supported yet.", name, "unsupported_parameter")
|
||||
if body.get("stop") is not None:
|
||||
raise APIError(400, "Custom stop sequences are not supported yet.", "stop", "unsupported_parameter")
|
||||
if body.get("logprobs"):
|
||||
raise APIError(400, "Log probabilities are not supported yet.", "logprobs", "unsupported_parameter")
|
||||
if body.get("frequency_penalty", 0) or body.get("presence_penalty", 0):
|
||||
raise APIError(400, "Token penalties are not supported yet.", None, "unsupported_parameter")
|
||||
if body.get("seed") is not None:
|
||||
raise APIError(400, "Per-request seeds are not supported yet.", "seed", "unsupported_parameter")
|
||||
response_format = body.get("response_format")
|
||||
if response_format not in (None, {"type": "text"}):
|
||||
raise APIError(400, "Only the default text response format is supported.",
|
||||
"response_format", "unsupported_parameter")
|
||||
|
||||
maximum = body.get("max_completion_tokens")
|
||||
maximum_param = "max_completion_tokens"
|
||||
if maximum is None:
|
||||
maximum = body.get("max_tokens")
|
||||
maximum_param = "max_tokens"
|
||||
if maximum is None:
|
||||
maximum = min(256, limit)
|
||||
temperature = body.get("temperature")
|
||||
top_p = body.get("top_p")
|
||||
temperature = 0.7 if temperature is None else temperature
|
||||
top_p = 0.9 if top_p is None else top_p
|
||||
if isinstance(maximum, bool) or not isinstance(maximum, int) or not 1 <= maximum <= limit:
|
||||
raise APIError(400, f"`{maximum_param}` must be an integer between 1 and {limit}.", maximum_param)
|
||||
if isinstance(temperature, bool) or not isinstance(temperature, (int, float)) or not 0 <= temperature <= 2:
|
||||
raise APIError(400, "`temperature` must be between 0 and 2.", "temperature")
|
||||
if isinstance(top_p, bool) or not isinstance(top_p, (int, float)) or not 0 < top_p <= 1:
|
||||
raise APIError(400, "`top_p` must be greater than 0 and at most 1.", "top_p")
|
||||
return maximum, float(temperature), float(top_p)
|
||||
|
||||
|
||||
def read_engine_turn(stream, sentinel, on_bytes):
|
||||
pending = b""
|
||||
while True:
|
||||
byte = stream.read(1)
|
||||
if byte == b"":
|
||||
raise RuntimeError("colibri engine exited unexpectedly")
|
||||
pending += byte
|
||||
if pending.endswith(sentinel):
|
||||
data = pending[:-len(sentinel)]
|
||||
if data:
|
||||
on_bytes(data)
|
||||
break
|
||||
if len(pending) > len(sentinel):
|
||||
on_bytes(pending[:-len(sentinel)])
|
||||
pending = pending[-len(sentinel):]
|
||||
|
||||
fields = stream.readline().decode("utf-8", "replace").strip().split()
|
||||
if len(fields) < 5 or fields[0] != "STAT":
|
||||
raise RuntimeError(f"invalid engine status: {' '.join(fields)}")
|
||||
return {
|
||||
"completion_tokens": int(fields[1]),
|
||||
"tokens_per_second": float(fields[2]),
|
||||
"cache_hit_percent": float(fields[3]),
|
||||
"rss_gb": float(fields[4]),
|
||||
"prompt_tokens": int(fields[5]) if len(fields) > 5 else 0,
|
||||
"length_limited": bool(int(fields[6])) if len(fields) > 6 else False,
|
||||
}
|
||||
|
||||
|
||||
class Engine:
|
||||
def __init__(self, executable, model, cap=8, max_tokens=1024, env=None):
|
||||
child_env = dict(env or os.environ, SNAP=str(model), SERVE="1", NGEN=str(max_tokens))
|
||||
self.process = subprocess.Popen(
|
||||
[str(executable), str(cap)], env=child_env, stdin=subprocess.PIPE,
|
||||
stdout=subprocess.PIPE, bufsize=0,
|
||||
)
|
||||
self.lock = threading.Lock()
|
||||
read_engine_turn(self.process.stdout, READY, lambda _: None)
|
||||
|
||||
def generate(self, prompt, max_tokens, temperature, top_p, on_text):
|
||||
payload = prompt.encode("utf-8")
|
||||
if b"\0" in payload:
|
||||
raise APIError(400, "NUL bytes are not supported in prompts.", "messages")
|
||||
decoder = codecs.getincrementaldecoder("utf-8")("replace")
|
||||
|
||||
def decode(data):
|
||||
text = decoder.decode(data)
|
||||
if text:
|
||||
on_text(text)
|
||||
|
||||
with self.lock:
|
||||
if self.process.poll() is not None:
|
||||
raise RuntimeError("colibri engine is not running")
|
||||
header = f"\x02PROMPT {len(payload)} {max_tokens} {temperature:.8g} {top_p:.8g}\n".encode()
|
||||
self.process.stdin.write(header + payload + b"\n")
|
||||
self.process.stdin.flush()
|
||||
stats = read_engine_turn(self.process.stdout, END, decode)
|
||||
tail = decoder.decode(b"", final=True)
|
||||
if tail:
|
||||
on_text(tail)
|
||||
return stats
|
||||
|
||||
def close(self):
|
||||
if self.process.poll() is None:
|
||||
self.process.terminate()
|
||||
try:
|
||||
self.process.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
self.process.kill()
|
||||
|
||||
|
||||
def model_object(model_id, created):
|
||||
return {"id": model_id, "object": "model", "created": created, "owned_by": "colibri"}
|
||||
|
||||
|
||||
class APIServer(ThreadingHTTPServer):
|
||||
daemon_threads = True
|
||||
|
||||
def __init__(self, address, engine, model_id, api_key=None, max_tokens=1024,
|
||||
cors_origins=DEFAULT_CORS_ORIGINS):
|
||||
super().__init__(address, APIHandler)
|
||||
self.engine = engine
|
||||
self.model_id = model_id
|
||||
self.api_key = api_key
|
||||
self.max_tokens = max_tokens
|
||||
self.cors_origins = tuple(cors_origins)
|
||||
self.created = int(time.time())
|
||||
|
||||
|
||||
class APIHandler(BaseHTTPRequestHandler):
|
||||
protocol_version = "HTTP/1.1"
|
||||
server_version = "colibri"
|
||||
|
||||
def log_message(self, fmt, *args):
|
||||
sys.stderr.write("[api] %s - %s\n" % (self.address_string(), fmt % args))
|
||||
|
||||
def send_json(self, status, body, request_id=None):
|
||||
data = json.dumps(body, ensure_ascii=False, separators=(",", ":")).encode()
|
||||
self.send_response(status)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(data)))
|
||||
if request_id:
|
||||
self.send_header("x-request-id", request_id)
|
||||
self.send_cors_headers()
|
||||
self.end_headers()
|
||||
self.wfile.write(data)
|
||||
|
||||
def send_cors_headers(self):
|
||||
origin = self.headers.get("Origin")
|
||||
if not origin or ("*" not in self.server.cors_origins and origin not in self.server.cors_origins):
|
||||
return
|
||||
self.send_header("Access-Control-Allow-Origin", "*" if "*" in self.server.cors_origins else origin)
|
||||
self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
|
||||
self.send_header("Access-Control-Allow-Headers", "Authorization, Content-Type")
|
||||
self.send_header("Access-Control-Expose-Headers", "x-request-id")
|
||||
self.send_header("Access-Control-Max-Age", "600")
|
||||
if "*" not in self.server.cors_origins:
|
||||
self.send_header("Vary", "Origin")
|
||||
|
||||
def require_auth(self):
|
||||
if self.server.api_key and self.headers.get("Authorization") != f"Bearer {self.server.api_key}":
|
||||
raise APIError(401, "Invalid or missing API key.", None, "invalid_api_key",
|
||||
"authentication_error")
|
||||
|
||||
def read_json(self):
|
||||
try:
|
||||
length = int(self.headers.get("Content-Length", "0"))
|
||||
except ValueError:
|
||||
raise APIError(400, "Invalid Content-Length header.")
|
||||
if length < 1 or length > MAX_BODY:
|
||||
raise APIError(400, f"Request body must be between 1 and {MAX_BODY} bytes.")
|
||||
try:
|
||||
body = json.loads(self.rfile.read(length))
|
||||
except (json.JSONDecodeError, UnicodeDecodeError):
|
||||
raise APIError(400, "Request body must be valid JSON.")
|
||||
if not isinstance(body, dict):
|
||||
raise APIError(400, "Request body must be a JSON object.")
|
||||
return body
|
||||
|
||||
def check_model(self, body):
|
||||
model = body.get("model")
|
||||
if model != self.server.model_id:
|
||||
raise APIError(404, f"The model `{model}` does not exist.", "model", "model_not_found")
|
||||
|
||||
def do_GET(self):
|
||||
request_id = "req_" + uuid.uuid4().hex
|
||||
try:
|
||||
path = urlsplit(self.path).path
|
||||
if path == "/health":
|
||||
self.send_json(200, {"status": "ok"}, request_id)
|
||||
return
|
||||
self.require_auth()
|
||||
if path == "/v1/models":
|
||||
self.send_json(200, {"object": "list", "data": [model_object(
|
||||
self.server.model_id, self.server.created)]}, request_id)
|
||||
elif path.startswith("/v1/models/") and unquote(path[11:]) == self.server.model_id:
|
||||
self.send_json(200, model_object(self.server.model_id, self.server.created), request_id)
|
||||
else:
|
||||
raise APIError(404, "Not found.", None, "not_found")
|
||||
except APIError as error:
|
||||
self.send_json(error.status, error_object(error), request_id)
|
||||
|
||||
def do_OPTIONS(self):
|
||||
self.send_response(204)
|
||||
self.send_header("Content-Length", "0")
|
||||
self.send_cors_headers()
|
||||
self.end_headers()
|
||||
|
||||
def do_POST(self):
|
||||
request_id = "req_" + uuid.uuid4().hex
|
||||
try:
|
||||
self.require_auth()
|
||||
body = self.read_json()
|
||||
self.check_model(body)
|
||||
path = urlsplit(self.path).path
|
||||
if path == "/v1/chat/completions":
|
||||
self.chat_completion(body, request_id)
|
||||
elif path == "/v1/completions":
|
||||
self.completion(body, request_id)
|
||||
else:
|
||||
raise APIError(404, "Not found.", None, "not_found")
|
||||
except APIError as error:
|
||||
self.send_json(error.status, error_object(error), request_id)
|
||||
except (BrokenPipeError, ConnectionResetError):
|
||||
pass
|
||||
except Exception as error:
|
||||
self.log_error("request failed: %s", error)
|
||||
api_error = APIError(500, "The colibri engine failed to process the request.",
|
||||
None, "engine_error", "server_error")
|
||||
try:
|
||||
self.send_json(500, error_object(api_error), request_id)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def generation(self, body, prompt, request_id, chat):
|
||||
maximum, temperature, top_p = generation_options(body, self.server.max_tokens)
|
||||
stream = body.get("stream", False)
|
||||
if not isinstance(stream, bool):
|
||||
raise APIError(400, "`stream` must be a boolean.", "stream")
|
||||
object_name = "chat.completion" if chat else "text_completion"
|
||||
id_prefix = "chatcmpl-" if chat else "cmpl-"
|
||||
completion_id = id_prefix + uuid.uuid4().hex
|
||||
created = int(time.time())
|
||||
|
||||
if not stream:
|
||||
output = []
|
||||
stats = self.server.engine.generate(prompt, maximum, temperature, top_p, output.append)
|
||||
text = "".join(output)
|
||||
finish = "length" if stats["length_limited"] else "stop"
|
||||
choice = ({"index": 0, "message": {"role": "assistant", "content": text,
|
||||
"refusal": None}, "logprobs": None, "finish_reason": finish} if chat else
|
||||
{"index": 0, "text": text, "logprobs": None, "finish_reason": finish})
|
||||
self.send_json(200, {"id": completion_id, "object": object_name, "created": created,
|
||||
"model": self.server.model_id, "choices": [choice], "usage": self.usage(stats)}, request_id)
|
||||
return
|
||||
|
||||
stream_options = body.get("stream_options")
|
||||
if stream_options is not None and not isinstance(stream_options, dict):
|
||||
raise APIError(400, "`stream_options` must be an object.", "stream_options")
|
||||
include_usage = bool((stream_options or {}).get("include_usage"))
|
||||
stream_object = "chat.completion.chunk" if chat else object_name
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "text/event-stream")
|
||||
self.send_header("Cache-Control", "no-cache")
|
||||
self.send_header("X-Accel-Buffering", "no")
|
||||
self.send_header("x-request-id", request_id)
|
||||
self.send_cors_headers()
|
||||
self.end_headers()
|
||||
connected = True
|
||||
|
||||
def event(choices, usage_marker=False):
|
||||
nonlocal connected
|
||||
if not connected:
|
||||
return
|
||||
event_body = {"id": completion_id, "object": stream_object, "created": created,
|
||||
"model": self.server.model_id, "choices": choices}
|
||||
if include_usage:
|
||||
event_body["usage"] = None if not usage_marker else usage_marker
|
||||
try:
|
||||
data = json.dumps(event_body, ensure_ascii=False, separators=(",", ":"))
|
||||
self.wfile.write(f"data: {data}\n\n".encode())
|
||||
self.wfile.flush()
|
||||
except OSError:
|
||||
connected = False
|
||||
|
||||
if chat:
|
||||
event([{"index": 0, "delta": {"role": "assistant", "content": ""},
|
||||
"logprobs": None, "finish_reason": None}])
|
||||
|
||||
def emit(text):
|
||||
choice = ({"index": 0, "delta": {"content": text}, "logprobs": None,
|
||||
"finish_reason": None} if chat else
|
||||
{"index": 0, "text": text, "logprobs": None, "finish_reason": None})
|
||||
event([choice])
|
||||
|
||||
stats = self.server.engine.generate(prompt, maximum, temperature, top_p, emit)
|
||||
finish = "length" if stats["length_limited"] else "stop"
|
||||
final_choice = ({"index": 0, "delta": {}, "logprobs": None, "finish_reason": finish}
|
||||
if chat else {"index": 0, "text": "", "logprobs": None,
|
||||
"finish_reason": finish})
|
||||
event([final_choice])
|
||||
if include_usage:
|
||||
event([], self.usage(stats))
|
||||
if connected:
|
||||
try:
|
||||
self.wfile.write(b"data: [DONE]\n\n")
|
||||
self.wfile.flush()
|
||||
except OSError:
|
||||
pass
|
||||
self.close_connection = True
|
||||
|
||||
@staticmethod
|
||||
def usage(stats):
|
||||
prompt = stats["prompt_tokens"]
|
||||
completion = stats["completion_tokens"]
|
||||
return {"prompt_tokens": prompt, "completion_tokens": completion,
|
||||
"total_tokens": prompt + completion}
|
||||
|
||||
def chat_completion(self, body, request_id):
|
||||
reasoning_effort = body.get("reasoning_effort")
|
||||
efforts = (None, "none", "minimal", "low", "medium", "high", "xhigh")
|
||||
if reasoning_effort not in efforts:
|
||||
raise APIError(400, "`reasoning_effort` must be none, minimal, low, medium, high, or xhigh.",
|
||||
"reasoning_effort")
|
||||
enable_thinking = body.get("enable_thinking", reasoning_effort not in (None, "none"))
|
||||
if not isinstance(enable_thinking, bool):
|
||||
raise APIError(400, "`enable_thinking` must be a boolean.", "enable_thinking")
|
||||
prompt = render_chat(body.get("messages"), enable_thinking, reasoning_effort)
|
||||
self.generation(body, prompt, request_id, True)
|
||||
|
||||
def completion(self, body, request_id):
|
||||
prompt = body.get("prompt")
|
||||
if not isinstance(prompt, str):
|
||||
raise APIError(400, "Colibri currently requires `prompt` to be a string.", "prompt")
|
||||
self.generation(body, prompt, request_id, False)
|
||||
|
||||
|
||||
def serve(model, host="127.0.0.1", port=8000, model_id="glm-5.2-colibri", api_key=None,
|
||||
cap=8, max_tokens=1024, engine=HERE / "glm", env=None, cors_origins=None):
|
||||
if not 1 <= max_tokens:
|
||||
raise ValueError("max_tokens must be positive")
|
||||
if not 1 <= port <= 65535:
|
||||
raise ValueError("port must be between 1 and 65535")
|
||||
if host not in ("127.0.0.1", "localhost", "::1") and not api_key:
|
||||
print("WARNING: API is listening beyond localhost without COLI_API_KEY", file=sys.stderr)
|
||||
runtime = Engine(engine, model, cap, max_tokens, env)
|
||||
origins = DEFAULT_CORS_ORIGINS if cors_origins is None else tuple(cors_origins)
|
||||
server = APIServer((host, port), runtime, model_id, api_key, max_tokens, origins)
|
||||
print(f"OpenAI-compatible API listening on http://{host}:{port}/v1", file=sys.stderr)
|
||||
previous_sigterm = signal.getsignal(signal.SIGTERM)
|
||||
signal.signal(signal.SIGTERM, lambda *_: threading.Thread(target=server.shutdown, daemon=True).start())
|
||||
try:
|
||||
server.serve_forever()
|
||||
finally:
|
||||
signal.signal(signal.SIGTERM, previous_sigterm)
|
||||
server.server_close()
|
||||
runtime.close()
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model", default=os.environ.get("COLI_MODEL"), required=not os.environ.get("COLI_MODEL"))
|
||||
parser.add_argument("--engine", default=str(HERE / "glm"))
|
||||
parser.add_argument("--host", default="127.0.0.1")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--model-id", default=os.environ.get("COLI_MODEL_ID", "glm-5.2-colibri"))
|
||||
parser.add_argument("--api-key", default=os.environ.get("COLI_API_KEY"))
|
||||
parser.add_argument("--cors-origin", action="append", default=None,
|
||||
help="allowed browser origin; repeat as needed (use '*' for any origin)")
|
||||
parser.add_argument("--cap", type=int, default=8)
|
||||
parser.add_argument("--max-tokens", type=int, default=1024)
|
||||
args = parser.parse_args()
|
||||
serve(args.model, args.host, args.port, args.model_id, args.api_key,
|
||||
args.cap, args.max_tokens, args.engine, cors_origins=args.cors_origin)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,152 @@
|
||||
import io
|
||||
import json
|
||||
import threading
|
||||
import unittest
|
||||
from urllib.error import HTTPError
|
||||
from urllib.request import Request, urlopen
|
||||
|
||||
from openai_server import APIError, APIServer, END, generation_options, read_engine_turn, render_chat
|
||||
|
||||
|
||||
class FakeEngine:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def generate(self, prompt, maximum, temperature, top_p, on_text):
|
||||
self.calls.append((prompt, maximum, temperature, top_p))
|
||||
on_text("Hé")
|
||||
on_text("llo")
|
||||
return {"prompt_tokens": 7, "completion_tokens": 2, "length_limited": False}
|
||||
|
||||
|
||||
class TemplateTest(unittest.TestCase):
|
||||
def test_renders_text_subset_of_official_template(self):
|
||||
prompt = render_chat([
|
||||
{"role": "system", "content": "System"},
|
||||
{"role": "developer", "content": "Developer"},
|
||||
{"role": "user", "content": [{"type": "text", "text": "Hi"}]},
|
||||
{"role": "assistant", "content": " Hello "},
|
||||
{"role": "user", "content": "Again"},
|
||||
])
|
||||
self.assertEqual(
|
||||
prompt,
|
||||
"[gMASK]<sop><|system|>System<|system|>Developer<|user|>Hi"
|
||||
"<|assistant|><think></think>Hello<|user|>Again"
|
||||
"<|assistant|><think></think>",
|
||||
)
|
||||
|
||||
def test_rejects_non_text_content(self):
|
||||
with self.assertRaisesRegex(APIError, "text message content only"):
|
||||
render_chat([{"role": "user", "content": [
|
||||
{"type": "image_url", "image_url": {"url": "x"}}
|
||||
]}])
|
||||
|
||||
def test_renders_thinking_prefix(self):
|
||||
self.assertEqual(
|
||||
render_chat([{"role": "user", "content": "Hi"}], True, "high"),
|
||||
"[gMASK]<sop><|system|>Reasoning Effort: High<|user|>Hi<|assistant|><think>",
|
||||
)
|
||||
|
||||
def test_validates_generation_limits(self):
|
||||
self.assertEqual(generation_options({"max_tokens": 4, "temperature": 0, "top_p": 1}, 8),
|
||||
(4, 0.0, 1.0))
|
||||
with self.assertRaises(APIError):
|
||||
generation_options({"max_tokens": 9}, 8)
|
||||
self.assertEqual(generation_options({"temperature": None, "top_p": None}, 8),
|
||||
(8, 0.7, 0.9))
|
||||
|
||||
|
||||
class ProtocolTest(unittest.TestCase):
|
||||
def test_reads_payload_and_extended_status(self):
|
||||
stream = io.BytesIO(b"hello" + END + b"STAT 2 3.5 44 1.2 7 1\n")
|
||||
chunks = []
|
||||
stats = read_engine_turn(stream, END, chunks.append)
|
||||
self.assertEqual(b"".join(chunks), b"hello")
|
||||
self.assertEqual(stats["prompt_tokens"], 7)
|
||||
self.assertTrue(stats["length_limited"])
|
||||
|
||||
|
||||
class HTTPTest(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.engine = FakeEngine()
|
||||
cls.server = APIServer(("127.0.0.1", 0), cls.engine, "test-model", "secret", 16)
|
||||
cls.thread = threading.Thread(target=cls.server.serve_forever, daemon=True)
|
||||
cls.thread.start()
|
||||
cls.base = f"http://127.0.0.1:{cls.server.server_port}"
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
cls.server.shutdown()
|
||||
cls.server.server_close()
|
||||
cls.thread.join(timeout=2)
|
||||
|
||||
def request(self, path, body=None, key="secret"):
|
||||
headers = {"Authorization": f"Bearer {key}"}
|
||||
data = None
|
||||
if body is not None:
|
||||
data = json.dumps(body).encode()
|
||||
headers["Content-Type"] = "application/json"
|
||||
return urlopen(Request(self.base + path, data=data, headers=headers), timeout=2)
|
||||
|
||||
def test_lists_models_and_checks_auth(self):
|
||||
with self.request("/v1/models") as response:
|
||||
self.assertEqual(json.load(response)["data"][0]["id"], "test-model")
|
||||
with self.assertRaises(HTTPError) as caught:
|
||||
self.request("/v1/models", key="wrong")
|
||||
self.assertEqual(caught.exception.code, 401)
|
||||
|
||||
def test_browser_preflight(self):
|
||||
request = Request(self.base + "/v1/chat/completions", method="OPTIONS", headers={
|
||||
"Origin": "http://localhost:5173",
|
||||
"Access-Control-Request-Method": "POST",
|
||||
"Access-Control-Request-Headers": "authorization,content-type",
|
||||
})
|
||||
with urlopen(request, timeout=2) as response:
|
||||
self.assertEqual(response.status, 204)
|
||||
self.assertEqual(response.headers["Access-Control-Allow-Origin"], "http://localhost:5173")
|
||||
self.assertIn("Authorization", response.headers["Access-Control-Allow-Headers"])
|
||||
|
||||
def test_chat_completion(self):
|
||||
with self.request("/v1/chat/completions", {
|
||||
"model": "test-model", "messages": [{"role": "user", "content": "Hi"}],
|
||||
"max_tokens": 4,
|
||||
}) as response:
|
||||
body = json.load(response)
|
||||
self.assertEqual(body["object"], "chat.completion")
|
||||
self.assertEqual(body["choices"][0]["message"]["content"], "Héllo")
|
||||
self.assertEqual(body["usage"], {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9})
|
||||
self.assertIn("<|user|>Hi<|assistant|><think></think>", self.engine.calls[-1][0])
|
||||
|
||||
def test_streaming_chat_completion(self):
|
||||
with self.request("/v1/chat/completions", {
|
||||
"model": "test-model", "messages": [{"role": "user", "content": "Hi"}],
|
||||
"stream": True, "stream_options": {"include_usage": True},
|
||||
}) as response:
|
||||
stream = response.read().decode()
|
||||
self.assertIn('\"delta\":{\"role\":\"assistant\",\"content\":\"\"}', stream)
|
||||
self.assertIn('\"object\":\"chat.completion.chunk\"', stream)
|
||||
self.assertIn('\"content\":\"Hé\"', stream)
|
||||
self.assertIn('\"usage\":{\"prompt_tokens\":7,\"completion_tokens\":2,\"total_tokens\":9}', stream)
|
||||
self.assertTrue(stream.endswith("data: [DONE]\n\n"))
|
||||
|
||||
def test_legacy_completion(self):
|
||||
with self.request("/v1/completions", {
|
||||
"model": "test-model", "prompt": "Complete me", "temperature": 0,
|
||||
}) as response:
|
||||
body = json.load(response)
|
||||
self.assertEqual(body["object"], "text_completion")
|
||||
self.assertEqual(body["choices"][0]["text"], "Héllo")
|
||||
self.assertEqual(self.engine.calls[-1][0], "Complete me")
|
||||
|
||||
def test_rejects_invalid_stream_options(self):
|
||||
with self.assertRaises(HTTPError) as caught:
|
||||
self.request("/v1/chat/completions", {
|
||||
"model": "test-model", "messages": [{"role": "user", "content": "Hi"}],
|
||||
"stream": True, "stream_options": "usage",
|
||||
})
|
||||
self.assertEqual(caught.exception.code, 400)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user