server.py

from __future__ import annotations

import json
import os
from http import HTTPStatus
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from urllib.parse import parse_qs, urlparse

from persona_dataset import PersonaDatasetError, dataset_metadata, dataset_status, sample_personas
from silicon_logs import log_status, recent_runs, record_run
from silicon_llm import build_silicon_prompt_preview, resolve_openai_api_key, run_silicon_llm_sampling


APP_DIR = Path(__file__).resolve().parent
STATIC_DIR = APP_DIR / "static" / "immersive"


def read_body(handler: BaseHTTPRequestHandler) -> dict:
    length = int(handler.headers.get("content-length") or 0)
    if length <= 0:
        return {}
    raw = handler.rfile.read(length).decode("utf-8")
    return json.loads(raw) if raw else {}


def write_json(handler: BaseHTTPRequestHandler, payload: dict, status: HTTPStatus = HTTPStatus.OK) -> None:
    data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    handler.send_response(status)
    handler.send_header("content-type", "application/json; charset=utf-8")
    handler.send_header("content-length", str(len(data)))
    handler.end_headers()
    handler.wfile.write(data)


def content_type(path: Path) -> str:
    suffix = path.suffix.lower()
    return {
        ".html": "text/html; charset=utf-8",
        ".js": "text/javascript; charset=utf-8",
        ".css": "text/css; charset=utf-8",
        ".json": "application/json; charset=utf-8",
        ".svg": "image/svg+xml",
        ".png": "image/png",
        ".jpg": "image/jpeg",
        ".jpeg": "image/jpeg",
        ".webp": "image/webp",
    }.get(suffix, "application/octet-stream")


def serve_file(handler: BaseHTTPRequestHandler, path: Path) -> None:
    if not path.exists() or not path.is_file():
        handler.send_error(HTTPStatus.NOT_FOUND)
        return
    data = path.read_bytes()
    handler.send_response(HTTPStatus.OK)
    handler.send_header("content-type", content_type(path))
    handler.send_header("content-length", str(len(data)))
    handler.end_headers()
    handler.wfile.write(data)


class Handler(BaseHTTPRequestHandler):
    def log_message(self, fmt: str, *args: object) -> None:
        print("%s - - %s" % (self.address_string(), fmt % args), flush=True)

    def do_GET(self) -> None:
        parsed = urlparse(self.path)
        path = parsed.path
        if path == "/api/status":
            return write_json(
                self,
                {
                    "ok": True,
                    "openai_api_key_set": bool(resolve_openai_api_key()),
                    "dataset": dataset_status(),
                    "logs": log_status(),
                },
            )
        if path == "/api/logs/status":
            return write_json(self, log_status())
        if path == "/api/logs/recent":
            token = os.getenv("SILICON_LOG_READ_TOKEN", "").strip()
            provided = parse_qs(parsed.query).get("token", [""])[0]
            if not token or provided != token:
                return write_json(self, {"error": "log read token required", "status": log_status()}, HTTPStatus.FORBIDDEN)
            return write_json(self, {"runs": recent_runs(20), "status": log_status()})
        if path == "/api/persona-dataset/metadata":
            try:
                return write_json(self, dataset_metadata())
            except PersonaDatasetError as exc:
                return write_json(self, {"available": False, "error": str(exc), "status": dataset_status()})
        if path in {"/", "/silicon", "/silicon/"}:
            return serve_file(self, STATIC_DIR / "index.html")
        target = (STATIC_DIR / path.lstrip("/")).resolve()
        if STATIC_DIR.resolve() in target.parents or target == STATIC_DIR.resolve():
            return serve_file(self, target)
        self.send_error(HTTPStatus.NOT_FOUND)

    def do_POST(self) -> None:
        parsed = urlparse(self.path)
        payload = read_body(self)
        if parsed.path == "/api/persona-dataset/sample":
            try:
                return write_json(self, sample_personas(payload))
            except PersonaDatasetError as exc:
                return write_json(self, {"error": str(exc), "status": dataset_status()}, HTTPStatus.BAD_REQUEST)
        if parsed.path == "/api/silicon/prompt-preview":
            try:
                return write_json(self, build_silicon_prompt_preview(payload))
            except ValueError as exc:
                return write_json(self, {"error": str(exc)}, HTTPStatus.BAD_REQUEST)
        if parsed.path == "/api/silicon/llm-run":
            try:
                result = run_silicon_llm_sampling(payload)
                try:
                    result["runLog"] = record_run(
                        payload,
                        result,
                        client={
                            "userAgent": self.headers.get("user-agent", "")[:240],
                            "referer": self.headers.get("referer", "")[:240],
                        },
                    )
                except Exception as log_exc:  # noqa: BLE001
                    result["runLogError"] = f"{type(log_exc).__name__}: {log_exc}"
                return write_json(self, result)
            except ValueError as exc:
                return write_json(self, {"error": str(exc)}, HTTPStatus.BAD_REQUEST)
            except RuntimeError as exc:
                return write_json(self, {"error": str(exc)}, HTTPStatus.BAD_REQUEST)
        self.send_error(HTTPStatus.NOT_FOUND)


def main() -> int:
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default=os.getenv("HOST", "127.0.0.1"))
    parser.add_argument("--port", type=int, default=int(os.getenv("PORT", "8765")))
    args = parser.parse_args()
    server = ThreadingHTTPServer((args.host, args.port), Handler)
    print(f"Silicon Sampling Lab running at http://{args.host}:{args.port}", flush=True)
    print(f"OpenAI API key available: {bool(resolve_openai_api_key())}", flush=True)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nShutting down.", flush=True)
    finally:
        server.server_close()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())