AstraX MLX Runtime

Reusable, OS-agnostic inference engine. MLX on Apple Silicon, llama.cpp (CUDA/Vulkan/CPU) everywhere else — the runtime auto-detects your platform and picks the fast path with no config changes.

This is NOT an AI chatbot. It is the inference foundation for a future agent orchestration platform.

AI assistants: Start with the ai/ context folder before editing this codebase.

Requirements

• Node.js 22+
• macOS + Apple Silicon (M-series): Python 3.10+ with mlx-lm — uses the MLXAdapter
• Linux / Windows / Intel Mac: nothing extra — node-llama-cpp ships prebuilt binaries (CUDA/Vulkan/CPU auto-selected) and installs with npm install, via the LlamaCppAdapter

runtime.config.json's "adapter": "auto" (the default) picks between them based on process.platform/process.arch — see docs/adapters.md.

Quick Start

# Install Node dependencies (also pulls in node-llama-cpp for non-Apple-Silicon platforms)
npm install

# macOS + Apple Silicon only — Python deps for the MLX adapter
pip install -r scripts/requirements.txt

# Download SmolLM3 model: MLX format on Apple Silicon, GGUF everywhere else
bash scripts/download-smollm3.sh        # macOS + Apple Silicon
bash scripts/download-smollm3-gguf.sh   # Linux / Windows / Intel Mac

# Build
npm run build

# Run examples
npm run example:load
npm run example:generate
npm run example:stream
npm run example:cancel
npm run example:chat

Or run all of the above in one shot: bash scripts/setup.sh (auto-detects your platform and runs the matching steps).

Interactive Chat REPL

npm run example:chat

A terminal chat loop (examples/chat-repl.ts) over ChatSession — multi-turn history, streamed token output, and a dim tok/s · latency line after each response.

Command	Effect
/clear	Reset conversation history
/help	List commands
/exit	Quit (also: Ctrl+C at the prompt, or Ctrl+D/EOF)
Ctrl+C while streaming	Cancel that response only — REPL stays open

Run the Playground API Server

A local HTTP server (server/) wraps the runtime for manual testing — it loads a model once at startup and keeps it resident in memory, so every request after the first is fast (no reload cost). It is a dev tool, not part of the published @astrax/mlx-runtime package.

npm run server
# → http://localhost:3000        (API)
# → http://localhost:3000/docs   (Swagger UI — try requests from the browser)

Change the port with PORT=4000 npm run server.

Endpoint	Method	Description
/health	GET	Server status + currently loaded model
/models	GET	List registered models, on-disk + loaded status
/models/load	POST	{ "model": "SmolLM3" } — load/swap the resident model
/models/unload	POST	Unload the current model
/generate	POST	Non-streaming completion — { "prompt": "..." } or { "messages": [...] }
/generate/stream	POST	Same body, streamed back as Server-Sent Events
/cancel	POST	Cancel the in-flight generation

curl -X POST localhost:3000/generate \
  -H 'Content-Type: application/json' \
  -d '{"prompt": "Explain what MLX is in one sentence.", "options": {"maxTokens": 100}}'

Full request/response shapes, including options (temperature, topP, maxTokens, etc.), are documented in Swagger at /docs (raw spec at /openapi.json).

Request Analytics

Every request is logged to a SQLite database at server/data/analytics.db (gitignored — local only) via an Express middleware (server/analyticsMiddleware.ts). Each row has:

• Generic fields (every request): timestamp, method, path, status code, duration, and a full request/response body dump (raw text — the complete prompt and generated text, not just metadata).
• Structured performance fields (generation endpoints): model, prompt/generated token counts, tokens/sec, first-token and total latency, finish reason — pulled straight from the same GenerationResult/GenerationMetrics objects the API returns, not re-parsed from logs.

Query it directly:

sqlite3 server/data/analytics.db \
  "SELECT path, model, generated_tokens, ROUND(tokens_per_second,1) AS tps, ROUND(duration_ms) AS ms
   FROM requests ORDER BY id DESC LIMIT 20;"

Note: streaming (/generate/stream) rows always have prompt_tokens = 0 — the MLX bridge doesn't report it for streamed responses (see ai/architecture.md).

Usage

import { Runtime, generateText, streamText } from '@astrax/mlx-runtime';

const runtime = await Runtime.create();
await runtime.loadModel('SmolLM3');

// Generate
const result = await generateText(runtime, 'Hello, world!');
console.log(result.text);

// Stream
for await (const token of streamText(runtime, 'Tell me a story')) {
  process.stdout.write(token.token);
}

await runtime.dispose();

Configuration

Edit runtime.config.json:

{
  "defaultModel": "SmolLM3",
  "contextLength": 8192,
  "temperature": 0.7,
  "maxTokens": 512,
  "modelsDirectory": "~/models",
  "adapter": "auto"
}

Plugging In a New Model

Models are weights on disk (MLX-format directory or a single GGUF file, depending on which adapter is active) plus a metadata entry in ModelRegistry — no code changes needed to use a model that's already registered (see the pre-registered catalog below).

1. Download weights into your models directory (default ~/models/):

   # MLX format (Apple Silicon) — e.g. from the mlx-community org on Hugging Face
   huggingface-cli download mlx-community/<Model>-mlx --local-dir ~/models/<Model>-mlx
   
   # GGUF format (llama.cpp, everywhere else) — single-file download
   huggingface-cli download <org>/<Model>-GGUF <file>.gguf --local-dir ~/models/<Model>-gguf

2. Register it (skip if it's already one of the pre-registered models below):

   runtime.registry.register({
     name: 'MyModel',
     path: 'MyModel-7B-mlx',   // directory (MLX) or file (GGUF), relative to modelsDirectory
     architecture: 'Llama',
     contextSize: 8192,
     parameters: '7B',
     quantization: '4bit',
     tokenizer: 'Llama',
     description: 'My custom model',
   });

   For a permanent addition, add the entry to the matching catalog (BUILTIN_MLX_MODELS or BUILTIN_GGUF_MODELS) in src/models/ModelRegistry.ts instead. Runtime picks the catalog that matches whichever adapter "auto" (or an explicit config value) resolved to.

3. Load and use it — either in code (await runtime.loadModel('MyModel')) or, without touching code at all, via the playground server: POST /models/load { "model": "MyModel" }.

Pre-registered MLX models (download weights before use): SmolLM3 (3B), Qwen3 (4B), Gemma (2B), Phi (3.8B), DeepSeek (7B). Pre-registered GGUF models: SmolLM3 (3B) — the others don't have a verified GGUF entry yet; register your own the same way. Full field reference: docs/adding-models.md.

Architecture

See docs/architecture.md for the full design.

API Layer → Runtime → ILLM Interface → Adapters (MLX on Apple Silicon, llama.cpp elsewhere)

Documentation

AI Context

Structured context for coding assistants — read these first when working with AI tools:

• AI Context Index
• Project Scope & Boundaries
• Architecture Reference
• Coding Conventions
• File Map

Developer Docs

• Architecture
• Interfaces
• Adapter Design
• Adding Models
• Future Backends

License

MIT