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:
ZacharyZcR
2026-07-10 17:04:56 +08:00
committed by GitHub
parent ea5f6fb914
commit 3e4d08b6bf
5 changed files with 737 additions and 19 deletions
+40
View File
@@ -93,6 +93,45 @@ COLI_MODEL=/nvme/glm52_i4 ./coli chat
The engine at runtime is pure C — python is only used by the one-time converter. The engine at runtime is pure C — python is only used by the one-time converter.
### OpenAI-compatible API
`coli serve` keeps one model process loaded and exposes a text-only OpenAI-compatible
HTTP API. The gateway uses only the Python standard library; inference still runs in
the same dependency-free C engine.
```bash
cd c
COLI_MODEL=/nvme/glm52_i4 COLI_API_KEY=local-secret ./coli serve \
--host 127.0.0.1 --port 8000 --model-id glm-5.2-colibri
curl http://127.0.0.1:8000/v1/chat/completions \
-H 'Authorization: Bearer local-secret' \
-H 'Content-Type: application/json' \
-d '{
"model": "glm-5.2-colibri",
"messages": [{"role": "user", "content": "Hello"}],
"stream": true
}'
```
Implemented endpoints are `GET /v1/models`, `GET /v1/models/{model}`,
`POST /v1/chat/completions`, and legacy `POST /v1/completions`. Chat and
completion requests support JSON responses, SSE streaming, usage counts,
`max_tokens`/`max_completion_tokens`, `temperature`, and `top_p`. The extension
`enable_thinking: true` enables GLM-5.2's reasoning block; the standard
`reasoning_effort` field also enables it unless set to `none`.
The first version is deliberately text-only and serves one generation at a time:
the 744B model stays in one persistent process, so concurrent HTTP requests queue
instead of loading duplicate model copies. Tools, image/audio input, custom stop
sequences, log probabilities, and token penalties return an explicit error rather
than being silently ignored. The default bind address is localhost; set
`COLI_API_KEY` before exposing the server beyond the machine.
Browser access from the Vite development server and Tauri local origins is enabled
by default. Repeat `--cors-origin https://your-ui.example` to allow another exact
origin, or use `--cors-origin '*'` only on a trusted local network.
### Experimental resident CUDA backend ### Experimental resident CUDA backend
colibrì includes an opt-in CUDA backend for model-resident tensors. Streaming colibrì includes an opt-in CUDA backend for model-resident tensors. Streaming
@@ -266,6 +305,7 @@ c/
├── backend_cuda.* optional CUDA tier ├── backend_cuda.* optional CUDA tier
├── Makefile build and local checks ├── Makefile build and local checks
├── coli user-facing CLI ├── coli user-facing CLI
├── openai_server.py OpenAI-compatible HTTP gateway
├── setup.sh one-command local setup ├── setup.sh one-command local setup
├── tools/ offline conversion, fixtures and benchmarks ├── tools/ offline conversion, fixtures and benchmarks
├── scripts/ long-running conversion helpers ├── scripts/ long-running conversion helpers
+13 -1
View File
@@ -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. 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 chat chat interattiva (carica il modello UNA volta)
coli serve API HTTP compatibile OpenAI (motore persistente)
coli run "prompt" generazione singola coli run "prompt" generazione singola
coli info stato: modello, RAM, disco, config coli info stato: modello, RAM, disco, config
coli bench [task...] benchmark di qualità (MMLU/HellaSwag/...) coli bench [task...] benchmark di qualità (MMLU/HellaSwag/...)
@@ -375,6 +376,12 @@ def cmd_chat(a):
except Exception: pass except Exception: pass
print(f" {C.teal}ciao{C.r} {C.dim}— il colibrì torna al nido{C.r} 🐦\n") 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): def cmd_bench(a):
need_model(a.model) need_model(a.model)
banner("bench") banner("bench")
@@ -427,13 +434,18 @@ def main():
sub.add_parser("build", parents=[common]); sub.add_parser("info", parents=[common]) sub.add_parser("build", parents=[common]); sub.add_parser("info", parents=[common])
pr=sub.add_parser("run", parents=[common]); pr.add_argument("prompt", nargs="*") pr=sub.add_parser("run", parents=[common]); pr.add_argument("prompt", nargs="*")
sub.add_parser("chat", parents=[common]) 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=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")) 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=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("--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)") pc.add_argument("--no-mtp",action="store_true",help="salta la testa MTP (niente draft speculativi)")
a=ap.parse_args() 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) "convert":cmd_convert}.get(a.cmd, lambda _:(banner(),print(__doc__)))(a)
if __name__=="__main__": if __name__=="__main__":
+65 -17
View File
@@ -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[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; 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; if(!g_kvsave) return;
FILE *f=fopen(g_kv_path,"r+b"); if(!f) return; FILE *f=fopen(g_kv_path,"r+b");
int32_t nz=0; fseek(f,8+6*4,SEEK_SET); fwrite(&nz,4,1,f); fclose(f); if(!f){ g_kv_nrec=0; return; }
g_kv_nrec=0; 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){ static void kv_disk_append(Model *m, const int *hist, int len){
if(!g_kvsave || len<=g_kv_nrec) return; if(!g_kvsave || len<=g_kv_nrec) return;
Cfg *c=&m->c; 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()); 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 */ 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; } 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 /* 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): * <|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. */ * 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>"; const char *tk = getenv("THINK")&&atoi(getenv("THINK"))? "<think>" : "<think></think>";
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>"); 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); } bl+=snprintf(buf+bl,(1<<16)-bl,"<|user|>%s<|assistant|>%s",input,tk); }
else bl+=snprintf(buf+bl,(1<<16)-bl,"%s",line); else bl+=snprintf(buf+bl,(1<<16)-bl,"%s",input);
int k=tok_encode(&T,buf,bl,hist+len,maxctx-len); k=tok_encode(&T,buf,bl,hist+len,maxctx-len); prompt_tokens=k;
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();
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",input,tk); }
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",input);
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;
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; 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(); 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}; EmitStream es={&T,m,now_s(),0,1};
int prod=0; int prod=0;
if(cur>0) prod=spec_decode(m,hist,len,cur,eos,logit,emit_stream,&es,&len); if(cur>0) prod=spec_decode(m,hist,len,cur,eos,logit,emit_stream,&es,&len);
else free(logit); else free(logit);
double tdt=now_s()-tt0; if(tdt<1e-6) tdt=1e-6; double tdt=now_s()-tt0; if(tdt<1e-6) tdt=1e-6;
double dh=(double)(m->hits-h0), dm=(double)(m->miss-ms0); double dh=(double)(m->hits-h0), dm=(double)(m->miss-ms0);
printf("\n\x01\x01" "END" "\x01\x01\n"); printf("%s\x01\x01" "END" "\x01\x01\n",raw_mode?"":"\n");
printf("STAT %d %.2f %.1f %.2f\n", prod, prod/tdt, (dh+dm)>0?100.0*dh/(dh+dm):0.0, rss_gb()); 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); fflush(stdout);
free(raw); g_temp=base_temp; g_nuc=base_nuc;
usage_save(m); /* la cache che impara: storia aggiornata a ogni turno */ 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 */ kv_disk_append(m,hist,len); /* KV su disco: il prossimo avvio riparte da qui */
} }
+466
View File
@@ -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()
+152
View File
@@ -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("")
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\":\"\"', 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()