Upload 4 files

Dockerfile

@@ -1,7 +1,28 @@
 FROM ghcr.io/ggml-org/llama.cpp:server
 
-EXPOSE 7860
+USER root
+
+RUN apt-get update \
+    && apt-get install -y --no-install-recommends python3 python3-pip ca-certificates \
+    && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /srv
+
+COPY requirements.txt /srv/requirements.txt
+RUN python3 -m pip install --no-cache-dir --break-system-packages -r /srv/requirements.txt
 
-ENTRYPOINT ["/app/llama-server"]
+COPY gateway.py /srv/gateway.py
+COPY start.sh /srv/start.sh
+RUN chmod +x /srv/start.sh \
+    && mkdir -p /data
+
+ENV LLAMA_URL=http://127.0.0.1:8080
+ENV STATE_PATH=/data/agent-state.db
+ENV MODEL_ALIAS=spark-x2.5-1.7b
+ENV DEFAULT_LANGUAGE=en
+ENV MAX_QUEUE_SIZE=32
+ENV REQUEST_TIMEOUT_SECONDS=900
+
+EXPOSE 7860
 
-CMD ["--host", "0.0.0.0", "--port", "7860", "--hf-repo", "XHToken/Spark-X2.5-1.7B-GGUF:Q8_0", "--alias", "spark-x2.5-1.7b", "--ctx-size", "32768", "--parallel", "1", "--threads", "2", "--threads-batch", "2", "--batch-size", "512", "--ubatch-size", "256", "--cache-type-k", "q8_0", "--cache-type-v", "q8_0", "--flash-attn", "auto", "--cont-batching", "--metrics"]
+ENTRYPOINT ["/srv/start.sh"]

gateway.py

@@ -0,0 +1,359 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import sqlite3
+import time
+import uuid
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import Any
+
+import httpx
+from fastapi import Depends, FastAPI, Header, HTTPException, Request
+from fastapi.responses import JSONResponse, StreamingResponse
+
+
+LLAMA_URL = os.getenv("LLAMA_URL", "http://127.0.0.1:8080").rstrip("/")
+STATE_PATH = Path(os.getenv("STATE_PATH", "/data/agent-state.db"))
+MODEL_ALIAS = os.getenv("MODEL_ALIAS", "spark-x2.5-1.7b")
+DEFAULT_LANGUAGE = os.getenv("DEFAULT_LANGUAGE", "en")
+MAX_QUEUE_SIZE = int(os.getenv("MAX_QUEUE_SIZE", "32"))
+REQUEST_TIMEOUT = float(os.getenv("REQUEST_TIMEOUT_SECONDS", "900"))
+API_KEY = os.getenv("API_KEY", "").strip()
+ENGLISH_GUARD = (
+    "Respond entirely in English unless the user explicitly requests another "
+    "language. Preserve code, identifiers, and quoted source material exactly."
+)
+
+inference_lock = asyncio.Lock()
+task_queue: asyncio.PriorityQueue[tuple[int, float, str]] = asyncio.PriorityQueue(
+    maxsize=MAX_QUEUE_SIZE
+)
+worker_task: asyncio.Task[None] | None = None
+
+
+def connect() -> sqlite3.Connection:
+    STATE_PATH.parent.mkdir(parents=True, exist_ok=True)
+    db = sqlite3.connect(STATE_PATH, timeout=30)
+    db.row_factory = sqlite3.Row
+    db.execute("PRAGMA journal_mode=WAL")
+    db.execute("PRAGMA foreign_keys=ON")
+    return db
+
+
+def initialise_database() -> None:
+    with connect() as db:
+        db.executescript(
+            """
+            CREATE TABLE IF NOT EXISTS sessions (
+                id TEXT PRIMARY KEY,
+                language TEXT NOT NULL,
+                summary TEXT NOT NULL DEFAULT '',
+                metadata TEXT NOT NULL DEFAULT '{}',
+                created_at REAL NOT NULL,
+                updated_at REAL NOT NULL
+            );
+            CREATE TABLE IF NOT EXISTS messages (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
+                role TEXT NOT NULL,
+                content TEXT NOT NULL,
+                created_at REAL NOT NULL
+            );
+            CREATE INDEX IF NOT EXISTS messages_session_idx
+                ON messages(session_id, id);
+            CREATE TABLE IF NOT EXISTS jobs (
+                id TEXT PRIMARY KEY,
+                session_id TEXT REFERENCES sessions(id) ON DELETE SET NULL,
+                priority INTEGER NOT NULL,
+                status TEXT NOT NULL,
+                request_json TEXT NOT NULL,
+                response_json TEXT,
+                error TEXT,
+                attempts INTEGER NOT NULL DEFAULT 0,
+                created_at REAL NOT NULL,
+                updated_at REAL NOT NULL
+            );
+            CREATE INDEX IF NOT EXISTS jobs_status_idx
+                ON jobs(status, priority, created_at);
+            """
+        )
+        db.execute(
+            "UPDATE jobs SET status='queued', updated_at=? WHERE status='running'",
+            (time.time(),),
+        )
+
+
+def require_api_key(authorization: str | None = Header(default=None)) -> None:
+    if not API_KEY:
+        return
+    if authorization != f"Bearer {API_KEY}":
+        raise HTTPException(status_code=401, detail="Invalid API key")
+
+
+def ensure_language_guard(payload: dict[str, Any]) -> dict[str, Any]:
+    if DEFAULT_LANGUAGE.lower() != "en":
+        return payload
+    messages = list(payload.get("messages") or [])
+    guard = {"role": "system", "content": ENGLISH_GUARD}
+    if messages and messages[0].get("role") == "system":
+        messages[0] = {
+            **messages[0],
+            "content": f"{messages[0].get('content', '')}\n\n{ENGLISH_GUARD}",
+        }
+    else:
+        messages.insert(0, guard)
+    return {**payload, "model": payload.get("model") or MODEL_ALIAS, "messages": messages}
+
+
+async def wait_for_llama() -> None:
+    deadline = time.monotonic() + REQUEST_TIMEOUT
+    async with httpx.AsyncClient() as client:
+        while time.monotonic() < deadline:
+            try:
+                response = await client.get(f"{LLAMA_URL}/health", timeout=5)
+                if response.status_code == 200:
+                    return
+            except httpx.HTTPError:
+                pass
+            await asyncio.sleep(2)
+    raise RuntimeError("llama.cpp did not become healthy before the startup timeout")
+
+
+async def execute_completion(payload: dict[str, Any]) -> dict[str, Any]:
+    await wait_for_llama()
+    async with inference_lock:
+        async with httpx.AsyncClient(timeout=REQUEST_TIMEOUT) as client:
+            response = await client.post(
+                f"{LLAMA_URL}/v1/chat/completions",
+                json=ensure_language_guard(payload),
+            )
+            response.raise_for_status()
+            return response.json()
+
+
+async def job_worker() -> None:
+    while True:
+        _priority, _created_at, job_id = await task_queue.get()
+        try:
+            with connect() as db:
+                row = db.execute("SELECT * FROM jobs WHERE id=?", (job_id,)).fetchone()
+                if row is None or row["status"] != "queued":
+                    continue
+                db.execute(
+                    "UPDATE jobs SET status='running', attempts=attempts+1, updated_at=? WHERE id=?",
+                    (time.time(), job_id),
+                )
+                payload = json.loads(row["request_json"])
+                session_id = row["session_id"]
+                if session_id:
+                    history = db.execute(
+                        "SELECT role, content FROM messages WHERE session_id=? ORDER BY id DESC LIMIT 20",
+                        (session_id,),
+                    ).fetchall()
+                    historical_messages = [dict(item) for item in reversed(history)]
+                    payload = {
+                        **payload,
+                        "messages": historical_messages + list(payload.get("messages") or []),
+                    }
+            try:
+                result = await execute_completion(payload)
+                with connect() as db:
+                    db.execute(
+                        "UPDATE jobs SET status='completed', response_json=?, updated_at=? WHERE id=?",
+                        (json.dumps(result), time.time(), job_id),
+                    )
+                    if session_id:
+                        for message in json.loads(row["request_json"]).get("messages", []):
+                            if message.get("role") in {"user", "assistant"} and isinstance(
+                                message.get("content"), str
+                            ):
+                                db.execute(
+                                    "INSERT INTO messages(session_id, role, content, created_at) VALUES(?,?,?,?)",
+                                    (session_id, message["role"], message["content"], time.time()),
+                                )
+                        choices = result.get("choices") or []
+                        assistant_content = (
+                            choices[0].get("message", {}).get("content") if choices else None
+                        )
+                        if isinstance(assistant_content, str):
+                            db.execute(
+                                "INSERT INTO messages(session_id, role, content, created_at) VALUES(?,?,?,?)",
+                                (session_id, "assistant", assistant_content, time.time()),
+                            )
+                        db.execute(
+                            "UPDATE sessions SET updated_at=? WHERE id=?",
+                            (time.time(), session_id),
+                        )
+            except Exception as exc:
+                with connect() as db:
+                    db.execute(
+                        "UPDATE jobs SET status='failed', error=?, updated_at=? WHERE id=?",
+                        (str(exc)[:2000], time.time(), job_id),
+                    )
+        finally:
+            task_queue.task_done()
+
+
+@asynccontextmanager
+async def lifespan(_app: FastAPI):
+    global worker_task
+    initialise_database()
+    with connect() as db:
+        queued = db.execute(
+            "SELECT id, priority, created_at FROM jobs WHERE status='queued' ORDER BY priority, created_at"
+        ).fetchall()
+    for row in queued:
+        await task_queue.put((row["priority"], row["created_at"], row["id"]))
+    worker_task = asyncio.create_task(job_worker())
+    yield
+    worker_task.cancel()
+    try:
+        await worker_task
+    except asyncio.CancelledError:
+        pass
+
+
+app = FastAPI(title="Spark X2.5 Agent Gateway", version="1.0.0", lifespan=lifespan)
+
+
+@app.get("/health")
+async def health() -> dict[str, Any]:
+    try:
+        async with httpx.AsyncClient() as client:
+            response = await client.get(f"{LLAMA_URL}/health", timeout=5)
+        model_ready = response.status_code == 200
+    except httpx.HTTPError:
+        model_ready = False
+    return {
+        "status": "ok" if model_ready else "degraded",
+        "model_ready": model_ready,
+        "queue_depth": task_queue.qsize(),
+        "inference_busy": inference_lock.locked(),
+    }
+
+
+@app.get("/v1/models", dependencies=[Depends(require_api_key)])
+async def models() -> JSONResponse:
+    async with httpx.AsyncClient() as client:
+        response = await client.get(f"{LLAMA_URL}/v1/models", timeout=30)
+    return JSONResponse(response.json(), status_code=response.status_code)
+
+
+@app.post("/v1/chat/completions", dependencies=[Depends(require_api_key)])
+async def chat_completions(request: Request):
+    payload = await request.json()
+    if not payload.get("stream", False):
+        try:
+            return await execute_completion(payload)
+        except httpx.HTTPStatusError as exc:
+            raise HTTPException(exc.response.status_code, exc.response.text) from exc
+
+    payload = ensure_language_guard(payload)
+
+    async def stream_response():
+        await wait_for_llama()
+        await inference_lock.acquire()
+        try:
+            async with httpx.AsyncClient(timeout=REQUEST_TIMEOUT) as client:
+                async with client.stream(
+                    "POST", f"{LLAMA_URL}/v1/chat/completions", json=payload
+                ) as response:
+                    if response.status_code >= 400:
+                        body = await response.aread()
+                        yield body
+                        return
+                    async for chunk in response.aiter_bytes():
+                        yield chunk
+        finally:
+            inference_lock.release()
+
+    return StreamingResponse(stream_response(), media_type="text/event-stream")
+
+
+@app.post("/api/sessions", dependencies=[Depends(require_api_key)])
+async def create_session(request: Request) -> dict[str, Any]:
+    body = await request.json()
+    now = time.time()
+    session_id = f"ses_{uuid.uuid4().hex}"
+    language = body.get("language", DEFAULT_LANGUAGE)
+    metadata = body.get("metadata", {})
+    with connect() as db:
+        db.execute(
+            "INSERT INTO sessions(id, language, metadata, created_at, updated_at) VALUES(?,?,?,?,?)",
+            (session_id, language, json.dumps(metadata), now, now),
+        )
+    return {"session_id": session_id, "language": language, "created_at": now}
+
+
+@app.get("/api/sessions/{session_id}", dependencies=[Depends(require_api_key)])
+async def get_session(session_id: str) -> dict[str, Any]:
+    with connect() as db:
+        row = db.execute("SELECT * FROM sessions WHERE id=?", (session_id,)).fetchone()
+    if row is None:
+        raise HTTPException(404, "Session not found")
+    result = dict(row)
+    result["metadata"] = json.loads(result["metadata"])
+    with connect() as db:
+        messages = db.execute(
+            "SELECT role, content, created_at FROM messages WHERE session_id=? ORDER BY id DESC LIMIT 20",
+            (session_id,),
+        ).fetchall()
+    result["recent_messages"] = [dict(item) for item in reversed(messages)]
+    return result
+
+
+@app.post("/api/tasks", status_code=202, dependencies=[Depends(require_api_key)])
+async def create_task(request: Request) -> dict[str, Any]:
+    if task_queue.full():
+        raise HTTPException(429, "Queue is full")
+    body = await request.json()
+    payload = body.get("request")
+    if not isinstance(payload, dict) or not isinstance(payload.get("messages"), list):
+        raise HTTPException(422, "request.messages must be provided")
+    priority = max(0, min(int(body.get("priority", 1)), 3))
+    session_id = body.get("session_id")
+    now = time.time()
+    job_id = f"job_{uuid.uuid4().hex}"
+    with connect() as db:
+        if session_id and db.execute(
+            "SELECT 1 FROM sessions WHERE id=?", (session_id,)
+        ).fetchone() is None:
+            raise HTTPException(404, "Session not found")
+        db.execute(
+            "INSERT INTO jobs(id, session_id, priority, status, request_json, created_at, updated_at) "
+            "VALUES(?,?,?,?,?,?,?)",
+            (job_id, session_id, priority, "queued", json.dumps(payload), now, now),
+        )
+    await task_queue.put((priority, now, job_id))
+    return {"task_id": job_id, "status": "queued", "status_url": f"/api/tasks/{job_id}"}
+
+
+@app.get("/api/tasks/{job_id}", dependencies=[Depends(require_api_key)])
+async def get_task(job_id: str) -> dict[str, Any]:
+    with connect() as db:
+        row = db.execute("SELECT * FROM jobs WHERE id=?", (job_id,)).fetchone()
+    if row is None:
+        raise HTTPException(404, "Task not found")
+    result = dict(row)
+    result["request"] = json.loads(result.pop("request_json"))
+    response_json = result.pop("response_json")
+    result["response"] = json.loads(response_json) if response_json else None
+    return result
+
+
+@app.post("/api/tasks/{job_id}/cancel", dependencies=[Depends(require_api_key)])
+async def cancel_task(job_id: str) -> dict[str, Any]:
+    with connect() as db:
+        row = db.execute("SELECT status FROM jobs WHERE id=?", (job_id,)).fetchone()
+        if row is None:
+            raise HTTPException(404, "Task not found")
+        if row["status"] != "queued":
+            raise HTTPException(409, "Only queued tasks can be cancelled")
+        db.execute(
+            "UPDATE jobs SET status='cancelled', updated_at=? WHERE id=?",
+            (time.time(), job_id),
+        )
+    return {"task_id": job_id, "status": "cancelled"}

requirements.txt

@@ -0,0 +1,3 @@
+fastapi==0.116.1
+httpx==0.28.1
+uvicorn[standard]==0.35.0

start.sh

@@ -0,0 +1,40 @@
+#!/bin/sh
+set -eu
+
+MODEL_REPO="${MODEL_REPO:-XHToken/Spark-X2.5-1.7B-GGUF:Q8_0}"
+MODEL_ALIAS="${MODEL_ALIAS:-spark-x2.5-1.7b}"
+CTX_SIZE="${CTX_SIZE:-32768}"
+THREADS="${THREADS:-2}"
+
+/app/llama-server \
+  --host 127.0.0.1 \
+  --port 8080 \
+  --hf-repo "$MODEL_REPO" \
+  --alias "$MODEL_ALIAS" \
+  --ctx-size "$CTX_SIZE" \
+  --parallel 1 \
+  --threads "$THREADS" \
+  --threads-batch "$THREADS" \
+  --batch-size 512 \
+  --ubatch-size 256 \
+  --cache-type-k q8_0 \
+  --cache-type-v q8_0 \
+  --flash-attn auto \
+  --cont-batching \
+  --metrics \
+  --no-webui &
+
+LLAMA_PID=$!
+
+cleanup() {
+  kill "$LLAMA_PID" 2>/dev/null || true
+  wait "$LLAMA_PID" 2>/dev/null || true
+}
+trap cleanup INT TERM EXIT
+
+exec uvicorn gateway:app \
+  --app-dir /srv \
+  --host 0.0.0.0 \
+  --port 7860 \
+  --proxy-headers \
+  --forwarded-allow-ips='*'