From 3e4d08b6bfe67b9bc17e8a0694b7c2604862612d Mon Sep 17 00:00:00 2001 From: ZacharyZcR Date: Fri, 10 Jul 2026 17:04:56 +0800 Subject: [PATCH] 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 --- README.md | 40 +++ c/coli | 14 +- c/glm.c | 84 ++++-- c/openai_server.py | 466 ++++++++++++++++++++++++++++++++++ c/tests/test_openai_server.py | 152 +++++++++++ 5 files changed, 737 insertions(+), 19 deletions(-) create mode 100644 c/openai_server.py create mode 100644 c/tests/test_openai_server.py diff --git a/README.md b/README.md index deaed7a..2bf9eb5 100644 --- a/README.md +++ b/README.md @@ -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. +### 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 colibrì includes an opt-in CUDA backend for model-resident tensors. Streaming @@ -266,6 +305,7 @@ c/ ├── backend_cuda.* optional CUDA tier ├── Makefile build and local checks ├── coli user-facing CLI +├── openai_server.py OpenAI-compatible HTTP gateway ├── setup.sh one-command local setup ├── tools/ offline conversion, fixtures and benchmarks ├── scripts/ long-running conversion helpers diff --git a/c/coli b/c/coli index 5e304b7..052d857 100755 --- a/c/coli +++ b/c/coli @@ -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__": diff --git a/c/glm.c b/c/glm.c index 63cb9d7..95f3531 100644 --- a/c/glm.c +++ b/c/glm.c @@ -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 \n\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 — 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"))? "" : ""; - if(templ){ if(first) bl+=snprintf(buf+bl,(1<<16)-bl,"[gMASK]"); - 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]<|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(prefixhas_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]"); + 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]<|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 */ } diff --git a/c/openai_server.py b/c/openai_server.py new file mode 100644 index 0000000..d97bd64 --- /dev/null +++ b/c/openai_server.py @@ -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]"] + 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|>{text.strip()}") + else: + raise APIError(400, f"Unsupported message role: {role!r}.", + f"messages.{index}.role", "unsupported_role") + prompt.append("<|assistant|>" if enable_thinking else + "<|assistant|>") + 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() diff --git a/c/tests/test_openai_server.py b/c/tests/test_openai_server.py new file mode 100644 index 0000000..06382ea --- /dev/null +++ b/c/tests/test_openai_server.py @@ -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]<|system|>System<|system|>Developer<|user|>Hi" + "<|assistant|>Hello<|user|>Again" + "<|assistant|>", + ) + + 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]<|system|>Reasoning Effort: High<|user|>Hi<|assistant|>", + ) + + 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|>", 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()