diff --git a/README.md b/README.md index e7a99102..6c89b398 100644 --- a/README.md +++ b/README.md @@ -465,6 +465,34 @@ Advanced Options: ``` +## 🔌 OpenAI-compatible server (`--server`) + +Serve the model behind the HTTP API OpenAI clients already speak — no external dependencies +(JDK `HttpServer`), streaming (SSE) and non-streaming. + +```bash +llama-tornado --gpu --cuda --model model.gguf --server --port 8080 +# or directly: +java ... org.beehive.gpullama3.server.OpenAIServer --model model.gguf --port 8080 --gpu +``` + +Endpoints: `POST /v1/chat/completions`, `POST /v1/completions`, `GET /v1/models`, `GET /health`. + +```bash +curl http://localhost:8080/ + +curl http://localhost:8080/v1/chat/completions -H 'Content-Type: application/json' \ + -d '{"messages":[{"role":"user","content":"Capital of France?"}],"max_tokens":16}' + +# streaming +curl -N http://localhost:8080/v1/chat/completions -H 'Content-Type: application/json' \ + -d '{"messages":[{"role":"user","content":"Tell me a joke"}],"stream":true}' +``` + +Any OpenAI SDK works by pointing `base_url` at the server. Generation is serialized on one GPU +context (requests queue); the reusable core is `server/InferenceService`. Smoke test: +`scripts/server-smoke-test.sh http://localhost:8080`. + ## Debug & Profiling Options View TornadoVM's internal behavior: ```bash diff --git a/llama-tornado b/llama-tornado index 78388295..581caba6 100755 --- a/llama-tornado +++ b/llama-tornado @@ -200,7 +200,9 @@ class LlamaRunner: [ "-cp", self._find_llama_jar(), - "org.beehive.gpullama3.LlamaApp", + "org.beehive.gpullama3.server.OpenAIServer" + if getattr(args, "server", False) + else "org.beehive.gpullama3.LlamaApp", ] ) cmd.extend(module_config) @@ -227,6 +229,12 @@ class LlamaRunner: def _add_llama_args(self, cmd: List[str], args: argparse.Namespace) -> List[str]: """Add LLaMA-specific arguments to the command.""" + if getattr(args, "server", False): + # OpenAI server mode: the server parses only --model / --port / --gpu. + srv = ["--model", args.model_path, "--port", str(getattr(args, "port", 8080))] + if args.use_gpu: + srv.append("--gpu") + return cmd + srv llama_args = [ "-m", args.model_path, @@ -467,6 +475,11 @@ def create_parser() -> argparse.ArgumentParser: help="Directory for profiler output", ) + # OpenAI-compatible server + server_group = parser.add_argument_group("OpenAI-compatible server") + server_group.add_argument("--server", action="store_true", help="Run the OpenAI-compatible HTTP server instead of inference") + server_group.add_argument("--port", type=int, default=8080, help="Server port (default 8080)") + # TornadoVM Execution Verbose options verbose_group = parser.add_argument_group("TornadoVM Execution Verbose") verbose_group.add_argument( diff --git a/scripts/server-smoke-test.sh b/scripts/server-smoke-test.sh new file mode 100755 index 00000000..73d6ad63 --- /dev/null +++ b/scripts/server-smoke-test.sh @@ -0,0 +1,51 @@ +#!/usr/bin/env bash +# Smoke test for the GPULlama3 OpenAI-compatible server. +# Usage: scripts/server-smoke-test.sh [host] (default http://localhost:8080) +# Exits non-zero on the first failed check. +set -u +H="${1:-http://localhost:8080}" +pass=0; fail=0 +check() { # name, condition(0/1) + if [ "$2" -eq 0 ]; then echo " ok $1"; pass=$((pass+1)); else echo " FAIL $1"; fail=$((fail+1)); fi +} +j() { python3 -c "import json,sys; d=json.load(sys.stdin); print($1)"; } + +echo "== health ==" +curl -sf "$H/health" | grep -q '"status":"ok"'; check "GET /health" $? + +echo "== models ==" +curl -sf "$H/v1/models" | grep -q '"object":"list"'; check "GET /v1/models" $? + +echo "== chat completion (non-stream) ==" +R=$(curl -sf "$H/v1/chat/completions" -H 'Content-Type: application/json' \ + -d '{"messages":[{"role":"user","content":"Capital of France? One word."}],"max_tokens":10,"temperature":0}') +echo "$R" | j "d['choices'][0]['message']['content']" | grep -qi paris; check "chat returns Paris" $? +echo "$R" | grep -q '"total_tokens"'; check "usage present" $? + +echo "== text completion ==" +curl -sf "$H/v1/completions" -H 'Content-Type: application/json' \ + -d '{"prompt":"The capital of France is","max_tokens":8,"temperature":0}' \ + | j "d['choices'][0]['text']" | grep -qi paris; check "completion returns Paris" $? + +echo "== streaming (SSE) ==" +S=$(curl -sfN "$H/v1/chat/completions" -H 'Content-Type: application/json' \ + -d '{"messages":[{"role":"user","content":"Say hi."}],"max_tokens":6,"temperature":0,"stream":true}') +echo "$S" | grep -q 'chat.completion.chunk'; check "SSE emits chunks" $? +echo "$S" | grep -q 'data: \[DONE\]'; check "SSE terminates with [DONE]" $? + +echo "== errors ==" +[ "$(curl -s -o /dev/null -w '%{http_code}' "$H/v1/chat/completions" -d '{bad')" = 400 ]; check "400 on bad JSON" $? +[ "$(curl -s -o /dev/null -w '%{http_code}' "$H/v1/chat/completions" -d '{}')" = 400 ]; check "400 on missing messages" $? +[ "$(curl -s -o /dev/null -w '%{http_code}' "$H/v1/chat/completions")" = 405 ]; check "405 on GET" $? + +echo "== concurrency (5 parallel) ==" +rm -f /tmp/smoke_*.done +for i in 1 2 3 4 5; do + ( curl -sf "$H/v1/chat/completions" -d '{"messages":[{"role":"user","content":"Hi"}],"max_tokens":4,"temperature":0}' \ + -o "/tmp/smoke_$i.json" && touch "/tmp/smoke_$i.done" ) & +done; wait +[ "$(ls /tmp/smoke_*.done 2>/dev/null | wc -l)" -eq 5 ]; check "5 concurrent requests all succeed" $? + +echo +echo "passed=$pass failed=$fail" +[ "$fail" -eq 0 ] diff --git a/src/main/java/org/beehive/gpullama3/server/InferenceService.java b/src/main/java/org/beehive/gpullama3/server/InferenceService.java new file mode 100644 index 00000000..2d98e2ef --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/server/InferenceService.java @@ -0,0 +1,108 @@ +package org.beehive.gpullama3.server; + +import org.beehive.gpullama3.inference.sampler.Sampler; +import org.beehive.gpullama3.inference.state.State; +import org.beehive.gpullama3.model.Model; +import org.beehive.gpullama3.model.format.ChatFormat; +import org.beehive.gpullama3.tornadovm.TornadoVMMasterPlan; + +import java.util.ArrayList; +import java.util.List; +import java.util.Set; +import java.util.function.IntConsumer; + +/** + * Reusable, thread-safe inference wrapper over a single loaded {@link Model}. + * + *

Holds one {@link State} and one {@link TornadoVMMasterPlan} (the expensive, compile-once GPU + * task graph) and serializes generation on them — the GPU is a single context, so requests run + * one at a time. This is the reuse seam shared by the OpenAI server and any other front-end: + * build the prompt, stream decoded tokens through a callback, return the full text. The single + * proven single-token decode path is used verbatim (KV cache is overwritten from position 0 each + * request, so state is safely reused without reallocation).

+ */ +public final class InferenceService { + + private final Model model; + private final State state; + private final TornadoVMMasterPlan plan; + private final boolean gpu; + private final int initialToken; // model start/BOS token that createNewState() seeds + private final Object lock = new Object(); + + public InferenceService(Model model, boolean gpu) { + this.model = model; + this.gpu = gpu; + this.state = model.createNewState(); + // createNewState() seeds latestToken with the model's start token (BOS / header) — the + // first forward pass consumes it. Capture it so each request can reset to a clean start. + this.initialToken = state.latestToken; + this.plan = gpu ? TornadoVMMasterPlan.initializeTornadoVMPlan(state, model) : null; + } + + public Model model() { + return model; + } + + /** A single generation request. {@code messages} are role/content turns (chat); a lone + * user turn covers the completions endpoint. */ + public record Request(List messages, int maxTokens, float temperature, float topP, long seed) {} + + /** Result of a generation: the text and the token counts (for {@code usage}). */ + public record Result(String text, int promptTokens, int completionTokens, boolean stopped) {} + + /** + * Generate a completion. If {@code onToken} is non-null each decoded token's text is streamed + * to it as it is produced; the full text is always returned. Thread-safe (serialized). + */ + public Result generate(Request req, java.util.function.Consumer onToken) { + synchronized (lock) { + ChatFormat cf = model.chatFormat(); + List promptTokens = new ArrayList<>(); + if (model.shouldAddBeginOfText()) { + promptTokens.add(cf.getBeginOfText()); + } + for (ChatFormat.Message m : req.messages()) { + promptTokens.addAll(cf.encodeMessage(m)); + } + promptTokens.addAll(cf.encodeHeader(new ChatFormat.Message(ChatFormat.Role.ASSISTANT, ""))); + if (model.shouldIncludeReasoning()) { + promptTokens.addAll(model.tokenizer().encode("\n", model.tokenizer().getSpecialTokens().keySet())); + } + + int maxTokens = req.maxTokens() > 0 ? promptTokens.size() + req.maxTokens() : 0; + Sampler sampler = Sampler.selectSampler(model.configuration().vocabularySize(), + req.temperature(), req.topP(), req.seed()); + Set stopTokens = cf.getStopTokens(); + + StringBuilder full = new StringBuilder(); + IntConsumer tokenConsumer = token -> { + if (model.tokenizer().shouldDisplayToken(token)) { + String piece = model.tokenizer().decode(List.of(token)); + full.append(piece); + if (onToken != null) { + onToken.accept(piece); + } + } + }; + + // Reset to the fresh-state convention: model start token + KV overwritten from pos 0. + state.latestToken = initialToken; + + List responseTokens = gpu + ? model.generateTokensGPU(state, 0, promptTokens, stopTokens, maxTokens, sampler, false, tokenConsumer, plan) + : model.generateTokens(state, 0, promptTokens, stopTokens, maxTokens, sampler, false, tokenConsumer); + + boolean stopped = !responseTokens.isEmpty() && stopTokens.contains(responseTokens.getLast()); + int completion = responseTokens.size() - (stopped ? 1 : 0); + return new Result(full.toString(), promptTokens.size(), Math.max(0, completion), stopped); + } + } + + /** Free the GPU plan (call on server shutdown). */ + public void close() { + if (plan != null) { + plan.freeTornadoExecutionPlan(); + } + } +} diff --git a/src/main/java/org/beehive/gpullama3/server/Json.java b/src/main/java/org/beehive/gpullama3/server/Json.java new file mode 100644 index 00000000..a0b06b8c --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/server/Json.java @@ -0,0 +1,268 @@ +package org.beehive.gpullama3.server; + +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +/** + * Tiny dependency-free JSON reader/writer — just enough for the OpenAI-compatible request/response + * bodies. Parses into {@code Map} / {@code List} / String / Double / Boolean + * / null; serializes the same. Reusable and self-contained (GPULlama3 pulls in no JSON library). + */ +public final class Json { + + private Json() {} + + // ── Parsing ─────────────────────────────────────────────────────────────── + + public static Object parse(String s) { + Parser p = new Parser(s); + p.ws(); + Object v = p.value(); + p.ws(); + if (!p.eof()) { + throw new IllegalArgumentException("Trailing characters at " + p.i); + } + return v; + } + + @SuppressWarnings("unchecked") + public static Map parseObject(String s) { + Object v = parse(s); + if (!(v instanceof Map)) { + throw new IllegalArgumentException("Expected a JSON object"); + } + return (Map) v; + } + + private static final class Parser { + final String s; + int i; + + Parser(String s) { this.s = s; } + + boolean eof() { return i >= s.length(); } + + void ws() { + while (i < s.length()) { + char c = s.charAt(i); + if (c == ' ' || c == '\t' || c == '\n' || c == '\r') { + i++; + } else { + break; + } + } + } + + char peek() { + if (i >= s.length()) { + throw new IllegalArgumentException("Unexpected end of JSON"); + } + return s.charAt(i); + } + + Object value() { + char c = peek(); + return switch (c) { + case '{' -> object(); + case '[' -> array(); + case '"' -> string(); + case 't', 'f' -> bool(); + case 'n' -> nul(); + default -> number(); + }; + } + + Map object() { + Map m = new LinkedHashMap<>(); + i++; // { + ws(); + if (peek() == '}') { i++; return m; } + while (true) { + ws(); + String key = string(); + ws(); + if (peek() != ':') { + throw new IllegalArgumentException("Expected ':' at " + i); + } + i++; + ws(); + m.put(key, value()); + ws(); + char c = peek(); + if (c == ',') { i++; continue; } + if (c == '}') { i++; break; } + throw new IllegalArgumentException("Expected ',' or '}' at " + i); + } + return m; + } + + List array() { + List a = new ArrayList<>(); + i++; // [ + ws(); + if (peek() == ']') { i++; return a; } + while (true) { + ws(); + a.add(value()); + ws(); + char c = peek(); + if (c == ',') { i++; continue; } + if (c == ']') { i++; break; } + throw new IllegalArgumentException("Expected ',' or ']' at " + i); + } + return a; + } + + String string() { + if (peek() != '"') { + throw new IllegalArgumentException("Expected string at " + i); + } + i++; + StringBuilder sb = new StringBuilder(); + while (true) { + char c = s.charAt(i++); + if (c == '"') { + break; + } + if (c == '\\') { + char e = s.charAt(i++); + switch (e) { + case '"' -> sb.append('"'); + case '\\' -> sb.append('\\'); + case '/' -> sb.append('/'); + case 'b' -> sb.append('\b'); + case 'f' -> sb.append('\f'); + case 'n' -> sb.append('\n'); + case 'r' -> sb.append('\r'); + case 't' -> sb.append('\t'); + case 'u' -> { + sb.append((char) Integer.parseInt(s.substring(i, i + 4), 16)); + i += 4; + } + default -> throw new IllegalArgumentException("Bad escape \\" + e); + } + } else { + sb.append(c); + } + } + return sb.toString(); + } + + Boolean bool() { + if (s.startsWith("true", i)) { i += 4; return Boolean.TRUE; } + if (s.startsWith("false", i)) { i += 5; return Boolean.FALSE; } + throw new IllegalArgumentException("Bad literal at " + i); + } + + Object nul() { + if (s.startsWith("null", i)) { i += 4; return null; } + throw new IllegalArgumentException("Bad literal at " + i); + } + + Double number() { + int start = i; + while (i < s.length()) { + char c = s.charAt(i); + if ((c >= '0' && c <= '9') || c == '-' || c == '+' || c == '.' || c == 'e' || c == 'E') { + i++; + } else { + break; + } + } + return Double.parseDouble(s.substring(start, i)); + } + } + + // ── Serialization ───────────────────────────────────────────────────────── + + public static String write(Object v) { + StringBuilder sb = new StringBuilder(); + writeValue(sb, v); + return sb.toString(); + } + + @SuppressWarnings("unchecked") + private static void writeValue(StringBuilder sb, Object v) { + if (v == null) { + sb.append("null"); + } else if (v instanceof String str) { + writeString(sb, str); + } else if (v instanceof Boolean || v instanceof Integer || v instanceof Long) { + sb.append(v); + } else if (v instanceof Double d) { + if (d == Math.rint(d) && !d.isInfinite()) { + sb.append(d.longValue()); + } else { + sb.append(d); + } + } else if (v instanceof Map m) { + sb.append('{'); + boolean first = true; + for (var e : ((Map) m).entrySet()) { + if (!first) { sb.append(','); } + first = false; + writeString(sb, e.getKey()); + sb.append(':'); + writeValue(sb, e.getValue()); + } + sb.append('}'); + } else if (v instanceof List list) { + sb.append('['); + for (int k = 0; k < list.size(); k++) { + if (k > 0) { sb.append(','); } + writeValue(sb, list.get(k)); + } + sb.append(']'); + } else { + writeString(sb, v.toString()); + } + } + + public static void writeString(StringBuilder sb, String s) { + sb.append('"'); + for (int k = 0; k < s.length(); k++) { + char c = s.charAt(k); + switch (c) { + case '"' -> sb.append("\\\""); + case '\\' -> sb.append("\\\\"); + case '\n' -> sb.append("\\n"); + case '\r' -> sb.append("\\r"); + case '\t' -> sb.append("\\t"); + case '\b' -> sb.append("\\b"); + case '\f' -> sb.append("\\f"); + default -> { + if (c < 0x20) { + sb.append(String.format("\\u%04x", (int) c)); + } else { + sb.append(c); + } + } + } + } + sb.append('"'); + } + + // ── Typed accessors (lenient) ───────────────────────────────────────────── + + public static String str(Map m, String k, String def) { + Object v = m.get(k); + return v instanceof String s ? s : def; + } + + public static double num(Map m, String k, double def) { + Object v = m.get(k); + return v instanceof Number n ? n.doubleValue() : def; + } + + public static int intVal(Map m, String k, int def) { + Object v = m.get(k); + return v instanceof Number n ? n.intValue() : def; + } + + public static boolean bool(Map m, String k, boolean def) { + Object v = m.get(k); + return v instanceof Boolean b ? b : def; + } +} diff --git a/src/main/java/org/beehive/gpullama3/server/OpenAIServer.java b/src/main/java/org/beehive/gpullama3/server/OpenAIServer.java new file mode 100644 index 00000000..1b416bee --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/server/OpenAIServer.java @@ -0,0 +1,388 @@ +package org.beehive.gpullama3.server; + +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpServer; +import org.beehive.gpullama3.Options; +import org.beehive.gpullama3.model.Model; +import org.beehive.gpullama3.model.format.ChatFormat; + +import java.io.IOException; +import java.io.OutputStream; +import java.net.InetSocketAddress; +import java.nio.charset.StandardCharsets; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicLong; + +import static org.beehive.gpullama3.model.loader.ModelLoader.loadModel; + +/** + * OpenAI-compatible HTTP server for GPULlama3, built on the JDK {@link HttpServer} (no external + * dependencies). Exposes the loaded model behind the endpoints an OpenAI client already speaks: + * + *
    + *
  • {@code POST /v1/chat/completions} — chat, streaming (SSE) or full JSON.
    • + *
  • {@code POST /v1/completions} — text completion (prompt as a single user turn).
    • + *
  • {@code GET /v1/models} — the one served model.
    • + *
  • {@code GET /health} — liveness.
    • + *
+ * + *

Generation is serialized on a single {@link InferenceService} (the GPU is one context); HTTP + * accept is multi-threaded so clients queue cleanly. Run:

+ * 
+ *   java ... org.beehive.gpullama3.server.OpenAIServer --model model.gguf --port 8080 --gpu
+ *
+ */ +public final class OpenAIServer { + + private final InferenceService service; + private final String servedModel; + private final boolean gpu; + private int port; + private final AtomicLong seq = new AtomicLong(); + + public OpenAIServer(InferenceService service, String servedModel, boolean gpu) { + this.service = service; + this.servedModel = servedModel; + this.gpu = gpu; + } + + public static void main(String[] args) throws IOException { + String modelPath = null; + int port = 8080; + boolean gpu = false; + for (int i = 0; i < args.length; i++) { + switch (args[i]) { + case "--model", "-m" -> modelPath = args[++i]; + case "--port", "-p" -> port = Integer.parseInt(args[++i]); + case "--gpu" -> gpu = true; + default -> { } + } + } + if (modelPath == null) { + System.err.println("usage: OpenAIServer --model [--port 8080] [--gpu]"); + System.exit(1); + } + System.setProperty("llama.enableTornadoVM", String.valueOf(gpu)); + + Path path = Paths.get(modelPath); + // interactive=true bypasses the --prompt-required check; the server never uses it. + Options options = new Options(path, "server", null, null, true, 0.0f, 0.95f, 1234L, 512, false, false, gpu, false, 1); + System.err.println("[server] loading " + path.getFileName() + " (gpu=" + gpu + ") ..."); + Model model = loadModel(options); + InferenceService service = new InferenceService(model, gpu); + String served = path.getFileName().toString().replaceAll("\\.gguf$", ""); + + OpenAIServer server = new OpenAIServer(service, served, gpu); + server.start(port); + } + + public void start(int port) throws IOException { + this.port = port; + HttpServer http = HttpServer.create(new InetSocketAddress(port), 0); + http.createContext("/", this::handleIndex); + http.createContext("/health", this::handleHealth); + http.createContext("/v1/models", this::handleModels); + http.createContext("/v1/chat/completions", ex -> handleCompletion(ex, true)); + http.createContext("/v1/completions", ex -> handleCompletion(ex, false)); + // Accept concurrently; generation itself is serialized inside InferenceService. + http.setExecutor(Executors.newFixedThreadPool(8)); + Runtime.getRuntime().addShutdownHook(new Thread(service::close)); + http.start(); + System.err.println("[server] listening on http://localhost:" + port + " model=" + servedModel); + } + + // ── Endpoints ───────────────────────────────────────────────────────────── + + private void handleIndex(HttpExchange ex) throws IOException { + if (!"/".equals(ex.getRequestURI().getPath())) { + sendError(ex, 404, "Not found"); + return; + } + byte[] bytes = INDEX_HTML + .replace("{{model}}", servedModel) + .replace("{{backend}}", gpu ? "GPU (TornadoVM)" : "CPU") + .replace("{{port}}", String.valueOf(port)) + .getBytes(StandardCharsets.UTF_8); + ex.getResponseHeaders().set("Content-Type", "text/html; charset=utf-8"); + ex.sendResponseHeaders(200, bytes.length); + try (OutputStream os = ex.getResponseBody()) { + os.write(bytes); + } + } + + private static final String INDEX_HTML = """ + + + + + GPULlama3.java server + + + + +

GPULlama3.java

+
Local OpenAI-compatible inference server
+ + + + + +
Model{{model}}
Backend{{backend}}
Port{{port}}
+ +

+ This is a local instance of + GPULlama3.java, + a Llama3-family inference engine written in native Java and automatically accelerated on + GPUs with TornadoVM. + It supports Llama3, Mistral, Devstral 2, Qwen2.5, Qwen3, Phi-3, IBM Granite 3.2+, and + IBM Granite 4.0 models in GGUF format, and is also used as the GPU inference engine behind + the Quarkus + and LangChain4j + integrations. +

+ +
+ +

API endpoints

+
    +
  • GET /health — liveness check
    • +
  • GET /v1/models — list the served model
    • +
  • POST /v1/chat/completions — chat completions (supports "stream": true SSE)
    • +
  • POST /v1/completions — text completions (supports "stream": true SSE)
    • +
+ +

Quick test

+ 
curl http://localhost:{{port}}/v1/chat/completions \\
+      -H "Content-Type: application/json" \\
+      -d '{
+            "model": "{{model}}",
+            "messages": [{"role": "user", "content": "Hello!"}]
+          }'
+ +

+ Any OpenAI-compatible client (e.g. the openai Python/Node SDK) can point its + base URL at this server. Generation runs on a single serialized GPU/CPU context, so + concurrent requests are queued and processed one at a time. +

+ + + """; + + private void handleHealth(HttpExchange ex) throws IOException { + sendJson(ex, 200, Map.of("status", "ok")); + } + + private void handleModels(HttpExchange ex) throws IOException { + Map entry = new LinkedHashMap<>(); + entry.put("id", servedModel); + entry.put("object", "model"); + entry.put("created", 0); + entry.put("owned_by", "gpullama3"); + sendJson(ex, 200, Map.of("object", "list", "data", List.of(entry))); + } + + @SuppressWarnings("unchecked") + private void handleCompletion(HttpExchange ex, boolean chat) throws IOException { + if (!"POST".equals(ex.getRequestMethod())) { + sendError(ex, 405, "Method not allowed — use POST"); + return; + } + Map body; + try { + String raw = new String(ex.getRequestBody().readAllBytes(), StandardCharsets.UTF_8); + body = Json.parseObject(raw); + } catch (Exception e) { + sendError(ex, 400, "Invalid JSON body: " + e.getMessage()); + return; + } + + List messages = new ArrayList<>(); + try { + if (chat) { + Object msgs = body.get("messages"); + if (!(msgs instanceof List) || ((List) msgs).isEmpty()) { + sendError(ex, 400, "'messages' must be a non-empty array"); + return; + } + for (Object o : (List) msgs) { + Map m = (Map) o; + String role = Json.str(m, "role", "user"); + String content = Json.str(m, "content", ""); + messages.add(new ChatFormat.Message(new ChatFormat.Role(role), content)); + } + } else { + Object prompt = body.get("prompt"); + String text = prompt instanceof String s ? s : prompt == null ? "" : prompt.toString(); + if (text.isEmpty()) { + sendError(ex, 400, "'prompt' must be a non-empty string"); + return; + } + messages.add(new ChatFormat.Message(ChatFormat.Role.USER, text)); + } + } catch (Exception e) { + sendError(ex, 400, "Malformed request: " + e.getMessage()); + return; + } + + int maxTokens = Json.intVal(body, "max_tokens", chat ? Json.intVal(body, "max_completion_tokens", 256) : 256); + float temperature = (float) Json.num(body, "temperature", 0.0); + float topP = (float) Json.num(body, "top_p", 0.95); + long seed = (long) Json.num(body, "seed", 1234); + boolean stream = Json.bool(body, "stream", false); + + var req = new InferenceService.Request(messages, maxTokens, temperature, topP, seed); + String id = (chat ? "chatcmpl-" : "cmpl-") + seq.incrementAndGet(); + long created = System.currentTimeMillis() / 1000; + + if (stream) { + streamResponse(ex, req, id, created, chat); + } else { + fullResponse(ex, req, id, created, chat); + } + } + + // ── Non-streaming ───────────────────────────────────────────────────────── + + private void fullResponse(HttpExchange ex, InferenceService.Request req, String id, long created, boolean chat) throws IOException { + InferenceService.Result r; + try { + r = service.generate(req, null); + } catch (Exception e) { + sendError(ex, 500, "Generation failed: " + e); + return; + } + String finish = r.stopped() ? "stop" : "length"; + Map choice = new LinkedHashMap<>(); + choice.put("index", 0); + if (chat) { + choice.put("message", Map.of("role", "assistant", "content", r.text())); + } else { + choice.put("text", r.text()); + } + choice.put("finish_reason", finish); + + Map resp = new LinkedHashMap<>(); + resp.put("id", id); + resp.put("object", chat ? "chat.completion" : "text_completion"); + resp.put("created", created); + resp.put("model", servedModel); + resp.put("choices", List.of(choice)); + resp.put("usage", Map.of( + "prompt_tokens", r.promptTokens(), + "completion_tokens", r.completionTokens(), + "total_tokens", r.promptTokens() + r.completionTokens())); + sendJson(ex, 200, resp); + } + + // ── Streaming (Server-Sent Events) ──────────────────────────────────────── + + private void streamResponse(HttpExchange ex, InferenceService.Request req, String id, long created, boolean chat) throws IOException { + ex.getResponseHeaders().set("Content-Type", "text/event-stream; charset=utf-8"); + ex.getResponseHeaders().set("Cache-Control", "no-cache"); + ex.getResponseHeaders().set("Connection", "keep-alive"); + ex.sendResponseHeaders(200, 0); + OutputStream os = ex.getResponseBody(); + + String object = chat ? "chat.completion.chunk" : "text_completion"; + try { + if (chat) { + // First chunk carries the assistant role. + writeSse(os, chunk(id, object, created, roleDelta(), null)); + } + InferenceService.Result r = service.generate(req, piece -> { + try { + writeSse(os, chunk(id, object, created, chat ? contentDelta(piece) : textField(piece), null)); + } catch (IOException io) { + throw new RuntimeException(io); // client disconnected — abort generation + } + }); + String finish = r.stopped() ? "stop" : "length"; + writeSse(os, chunk(id, object, created, chat ? Map.of() : textField(""), finish)); + os.write("data: [DONE]\n\n".getBytes(StandardCharsets.UTF_8)); + os.flush(); + } catch (Exception e) { + // Best-effort: nothing more we can send on a half-written SSE stream. + } finally { + os.close(); + } + } + + private Map roleDelta() { + Map d = new LinkedHashMap<>(); + d.put("role", "assistant"); + d.put("content", ""); + return d; + } + + private Map contentDelta(String piece) { + return Map.of("content", piece); + } + + private Map textField(String piece) { + return Map.of("text", piece); + } + + /** One SSE data object. {@code delta} is the chat delta map or the completion text map. */ + private String chunk(String id, String object, long created, Map delta, String finish) { + Map choice = new LinkedHashMap<>(); + choice.put("index", 0); + if (object.startsWith("chat")) { + choice.put("delta", delta); + } else { + choice.putAll(delta); // {"text": ...} + } + choice.put("finish_reason", finish); + Map obj = new LinkedHashMap<>(); + obj.put("id", id); + obj.put("object", object); + obj.put("created", created); + obj.put("model", servedModel); + obj.put("choices", List.of(choice)); + return Json.write(obj); + } + + private void writeSse(OutputStream os, String jsonData) throws IOException { + os.write(("data: " + jsonData + "\n\n").getBytes(StandardCharsets.UTF_8)); + os.flush(); + } + + // ── Wire helpers ────────────────────────────────────────────────────────── + + private void sendJson(HttpExchange ex, int status, Map body) throws IOException { + byte[] bytes = Json.write(body).getBytes(StandardCharsets.UTF_8); + ex.getResponseHeaders().set("Content-Type", "application/json; charset=utf-8"); + ex.sendResponseHeaders(status, bytes.length); + try (OutputStream os = ex.getResponseBody()) { + os.write(bytes); + } + } + + private void sendError(HttpExchange ex, int status, String message) throws IOException { + Map err = Map.of("error", Map.of( + "message", message, + "type", "invalid_request_error")); + sendJson(ex, status, err); + } +}