@modular-prompt/driver 0.14.0 → 0.15.0

package/scripts/runtime-cli.js

@@ -0,0 +1,315 @@
+#!/usr/bin/env node
+
+/**
+ * Python runtime 管理 CLI
+ *
+ * setup mlx     — ~/.modular-prompt/runtimes/mlx に venv を作成
+ * setup pytorch — ~/.modular-prompt/runtimes/pytorch に cpu-minimal venv を作成
+ * setup --status
+ * cleanup mlx [--yes]
+ * cleanup --all [--yes]
+ */
+
+import { execSync } from 'child_process';
+import { existsSync, readFileSync, rmSync, mkdirSync } from 'fs';
+import { createInterface } from 'readline';
+import { dirname, join } from 'path';
+import { fileURLToPath, pathToFileURL } from 'url';
+
+const __dirname = dirname(fileURLToPath(import.meta.url));
+const packageRoot = join(__dirname, '..');
+const driverVersion = readPackageVersion();
+
+function runtimeModuleUrl(name) {
+  const distPath = join(packageRoot, 'dist', 'runtime', name);
+  const srcPath = join(packageRoot, 'src', 'runtime', name);
+  return pathToFileURL(existsSync(distPath) ? distPath : srcPath).href;
+}
+
+const {
+  RUNTIME_PROFILES,
+  getModularPromptHome,
+  getRuntimeDir,
+  getVenvPath,
+  getMlxPythonDir,
+  getPytorchPythonDir,
+  isRuntimeReady,
+} = await import(runtimeModuleUrl('paths-core.mjs'));
+
+const {
+  collectInstalledPackages,
+  readManifest,
+  writeManifest,
+} = await import(runtimeModuleUrl('manifest-core.mjs'));
+
+function readPackageVersion() {
+  try {
+    const pkg = JSON.parse(readFileSync(join(packageRoot, 'package.json'), 'utf8'));
+    return pkg.version ?? 'unknown';
+  } catch {
+    return 'unknown';
+  }
+}
+
+function ensureUv() {
+  try {
+    execSync('uv --version', { stdio: 'ignore' });
+    return;
+  } catch {
+    console.log('⚠️  uv is not installed. Installing uv...');
+    execSync('curl -LsSf https://astral.sh/uv/install.sh | sh', { stdio: 'inherit' });
+  }
+}
+
+function setupMlx() {
+  if (process.platform !== 'darwin') {
+    console.error('❌ MLX runtime is only available on macOS (Apple Silicon).');
+    console.error('   For local inference on this platform, use setup-pytorch or vLLM.');
+    process.exit(1);
+  }
+
+  const pythonDir = getMlxPythonDir(packageRoot);
+  if (!existsSync(pythonDir)) {
+    console.error(`❌ MLX Python project not found: ${pythonDir}`);
+    process.exit(1);
+  }
+
+  const venvPath = getVenvPath('mlx');
+  const runtimeDir = getRuntimeDir('mlx');
+
+  console.log('🚀 Setting up MLX runtime...\n');
+  console.log(`📁 Python project: ${pythonDir}`);
+  console.log(`📁 Runtime venv:  ${venvPath}\n`);
+
+  ensureUv();
+  mkdirSync(runtimeDir, { recursive: true });
+
+  const env = {
+    ...process.env,
+    UV_PROJECT_ENVIRONMENT: venvPath,
+  };
+
+  try {
+    execSync('uv venv --python 3.13', { cwd: pythonDir, stdio: 'inherit', env });
+    execSync('uv pip install -e .', { cwd: pythonDir, stdio: 'inherit', env });
+
+    writeManifest('mlx', {
+      profile: 'mlx',
+      driverVersion,
+      platform: process.platform,
+      pythonVersion: '3.13',
+      createdAt: new Date().toISOString(),
+      packages: collectInstalledPackages(pythonDir, venvPath),
+    });
+
+    console.log('\n✅ MLX runtime setup completed.');
+    console.log(`   Home: ${getModularPromptHome()}`);
+    console.log('   You can now use MlxDriver from @modular-prompt/driver');
+  } catch (error) {
+    const message = error instanceof Error ? error.message : String(error);
+    console.error('❌ Failed to setup MLX runtime:', message);
+    process.exit(1);
+  }
+}
+
+const PYTORCH_CPU_INDEX = 'https://download.pytorch.org/whl/cpu';
+const PYTORCH_PYTHON_VERSION = '3.12';
+
+function setupPytorch() {
+  const pythonDir = getPytorchPythonDir(packageRoot);
+  if (!existsSync(pythonDir)) {
+    console.error(`❌ PyTorch Python project not found: ${pythonDir}`);
+    process.exit(1);
+  }
+
+  const venvPath = getVenvPath('pytorch');
+  const runtimeDir = getRuntimeDir('pytorch');
+
+  console.log('🚀 Setting up PyTorch runtime (cpu-minimal)...\n');
+  console.log(`📁 Python project: ${pythonDir}`);
+  console.log(`📁 Runtime venv:  ${venvPath}`);
+  console.log(`📦 torch index:    ${PYTORCH_CPU_INDEX}\n`);
+
+  ensureUv();
+  mkdirSync(runtimeDir, { recursive: true });
+
+  const env = {
+    ...process.env,
+    UV_PROJECT_ENVIRONMENT: venvPath,
+  };
+
+  try {
+    execSync(`uv venv --python ${PYTORCH_PYTHON_VERSION}`, { cwd: pythonDir, stdio: 'inherit', env });
+    const venvPython =
+      process.platform === 'win32'
+        ? join(venvPath, 'Scripts', 'python.exe')
+        : join(venvPath, 'bin', 'python');
+    execSync(`uv pip install --python "${venvPython}" "torch==2.9.1" --index-url ${PYTORCH_CPU_INDEX}`, {
+      cwd: pythonDir,
+      stdio: 'inherit',
+      env,
+    });
+    execSync(`uv pip install --python "${venvPython}" -e .`, { cwd: pythonDir, stdio: 'inherit', env });
+
+    const packages = collectInstalledPackages(pythonDir, venvPath);
+    writeManifest('pytorch', {
+      profile: 'pytorch',
+      variant: 'cpu-minimal',
+      driverVersion,
+      platform: process.platform,
+      pythonVersion: PYTORCH_PYTHON_VERSION,
+      torchVersion: packages?.torch,
+      createdAt: new Date().toISOString(),
+      packages,
+    });
+
+    console.log('\n✅ PyTorch runtime setup completed (cpu-minimal).');
+    console.log(`   Home: ${getModularPromptHome()}`);
+    console.log('   You can now use PyTorchDriver from @modular-prompt/driver');
+    console.log('   For CUDA / custom environments, see docs/LOCAL_MODEL_SETUP.md');
+  } catch (error) {
+    const message = error instanceof Error ? error.message : String(error);
+    console.error('❌ Failed to setup PyTorch runtime:', message);
+    process.exit(1);
+  }
+}
+
+function formatManifestDetail(manifest) {
+  const parts = [`driver ${manifest.driverVersion}`];
+  if (manifest.variant) {
+    parts.push(`variant ${manifest.variant}`);
+  }
+  const torchVersion = manifest.torchVersion ?? manifest.packages?.torch;
+  if (torchVersion) {
+    parts.push(`torch ${torchVersion}`);
+  }
+  parts.push(manifest.createdAt);
+  return ` (${parts.join(', ')})`;
+}
+
+function printStatus() {
+  console.log(`modular-prompt home: ${getModularPromptHome()}\n`);
+  for (const profile of RUNTIME_PROFILES) {
+    const ready = isRuntimeReady(profile);
+    const manifest = ready ? readManifest(profile) : null;
+    const detail = manifest ? formatManifestDetail(manifest) : '';
+    const icon = ready ? '✅' : '❌';
+    const runtimePath = getRuntimeDir(profile);
+    console.log(`${icon} ${profile}: ${ready ? 'ready' : 'not installed'}${detail}`);
+    console.log(`   ${runtimePath}`);
+  }
+  const setupHints = [];
+  if (!isRuntimeReady('mlx') && process.platform === 'darwin') {
+    setupHints.push('pnpm run setup-mlx -w @modular-prompt/driver');
+  }
+  if (!isRuntimeReady('pytorch')) {
+    setupHints.push('pnpm run setup-pytorch -w @modular-prompt/driver');
+  }
+  if (setupHints.length > 0) {
+    console.log(`\nRun: ${setupHints.join('  or  ')}`);
+  }
+}
+
+async function confirm(message) {
+  if (process.argv.includes('--yes') || process.argv.includes('-y')) {
+    return true;
+  }
+  const rl = createInterface({ input: process.stdin, output: process.stdout });
+  const answer = await new Promise((resolve) => {
+    rl.question(`${message} [y/N] `, resolve);
+  });
+  rl.close();
+  return String(answer).toLowerCase() === 'y' || String(answer).toLowerCase() === 'yes';
+}
+
+async function cleanupProfile(profile) {
+  if (!existsSync(getRuntimeDir(profile))) {
+    console.log(`ℹ️  ${profile}: nothing to clean (${getRuntimeDir(profile)} does not exist)`);
+    return;
+  }
+  const ok = await confirm(`Delete ${profile} runtime at ${getRuntimeDir(profile)}?`);
+  if (!ok) {
+    console.log('Cancelled.');
+    return;
+  }
+  rmSync(getRuntimeDir(profile), { recursive: true, force: true });
+  console.log(`✅ Removed ${profile} runtime.`);
+}
+
+async function cleanupAll() {
+  const home = getModularPromptHome();
+  if (!existsSync(home)) {
+    console.log(`ℹ️  Nothing to clean (${home} does not exist)`);
+    return;
+  }
+  const ok = await confirm(`Delete entire modular-prompt home at ${home}?`);
+  if (!ok) {
+    console.log('Cancelled.');
+    return;
+  }
+  rmSync(home, { recursive: true, force: true });
+  console.log(`✅ Removed ${home}`);
+}
+
+function printUsage() {
+  console.log(`Usage:
+  node scripts/runtime-cli.js setup mlx         Set up MLX Python runtime (macOS only)
+  node scripts/runtime-cli.js setup pytorch     Set up PyTorch runtime (cpu-minimal)
+  node scripts/runtime-cli.js setup --status    Show runtime status
+  node scripts/runtime-cli.js cleanup mlx       Remove MLX runtime
+  node scripts/runtime-cli.js cleanup pytorch   Remove PyTorch runtime
+  node scripts/runtime-cli.js cleanup --all     Remove entire ~/.modular-prompt
+  node scripts/runtime-cli.js cleanup ... --yes   Skip confirmation
+
+  npm scripts: setup-mlx, setup-pytorch, runtime:status, runtime:cleanup`);
+}
+
+async function main() {
+  const [command, target] = process.argv.slice(2);
+
+  if (!command || command === '--help' || command === '-h') {
+    printUsage();
+    process.exit(0);
+  }
+
+  if (command === 'setup') {
+    if (target === '--status' || target === 'status') {
+      printStatus();
+      return;
+    }
+    if (target === 'mlx') {
+      setupMlx();
+      return;
+    }
+    if (target === 'pytorch') {
+      setupPytorch();
+      return;
+    }
+    console.error(`Unknown setup target: ${target ?? '(none)'}`);
+    printUsage();
+    process.exit(1);
+  }
+
+  if (command === 'cleanup') {
+    if (target === '--all' || target === 'all') {
+      await cleanupAll();
+      return;
+    }
+    if (target === 'mlx' || target === 'pytorch') {
+      await cleanupProfile(target);
+      return;
+    }
+    console.error(`Unknown cleanup target: ${target ?? '(none)'}`);
+    printUsage();
+    process.exit(1);
+  }
+
+  console.error(`Unknown command: ${command}`);
+  printUsage();
+  process.exit(1);
+}
+
+main().catch((error) => {
+  console.error(error);
+  process.exit(1);
+});

package/src/mlx-ml/python/__main__.py

@@ -6,7 +6,28 @@ from utils.vlm_utils import detect_model_kind
 from server import Server
 
 model_name = sys.argv[1] if len(sys.argv) > 1 else "mlx-community/gemma-3-270m-it-qat-4bit"
-text_only = "--text-only" in sys.argv
+
+VALID_BACKENDS = {"auto", "lm", "vlm", "optiq"}
+
+
+def parse_backend_mode() -> str:
+    if "--backend" in sys.argv:
+        idx = sys.argv.index("--backend")
+        if idx + 1 < len(sys.argv):
+            mode = sys.argv[idx + 1]
+            if mode not in VALID_BACKENDS:
+                sys.stderr.write(
+                    f"Invalid --backend value: {mode}. "
+                    f"Use one of: {', '.join(sorted(VALID_BACKENDS))}\n"
+                )
+                sys.exit(1)
+            return mode
+    if "--text-only" in sys.argv:
+        return "lm"
+    return "auto"
+
+
+backend_mode = parse_backend_mode()
 
 drafter_model = None
 if "--drafter" in sys.argv:
@@ -25,8 +46,26 @@ if "--draft-block-size" in sys.argv:
             sys.exit(1)
 
 
-def create_backend(model_name: str, text_only: bool = False):
-    model_kind = "lm" if text_only else detect_model_kind(model_name)
+def ensure_optiq_registered() -> None:
+    import optiq  # noqa: F401 - registers vendored architectures into mlx-lm
+
+
+def create_backend(model_name: str, mode: str = "auto"):
+    if mode == "optiq":
+        ensure_optiq_registered()
+        mode = "lm"
+
+    if mode == "lm":
+        backend = MlxLmBackend()
+        backend.load(model_name)
+        return backend, "lm"
+
+    if mode == "vlm":
+        backend = MlxVlmBackend()
+        backend.load(model_name)
+        return backend, "vlm"
+
+    model_kind = detect_model_kind(model_name)
 
     if model_kind == "vlm":
         backend = MlxVlmBackend()
@@ -42,7 +81,7 @@ def create_backend(model_name: str, text_only: bool = False):
 
 
 if __name__ == "__main__":
-    backend, model_kind = create_backend(model_name, text_only)
+    backend, model_kind = create_backend(model_name, backend_mode)
 
     if drafter_model:
         backend.load_drafter(drafter_model)

package/src/mlx-ml/python/handlers/__init__.py

@@ -1,6 +1,7 @@
 from handlers.cache import handle_cache_prefill
 from handlers.capabilities import handle_capabilities
-from handlers.chat import handle_chat
 from handlers.completion import handle_completion
 from handlers.format_test import handle_format_test
+from handlers.generate import handle_generate
+from handlers.render import handle_render
 from handlers.tokenize import handle_tokenize

package/src/mlx-ml/python/handlers/completion.py

@@ -1,11 +1,7 @@
 from __future__ import annotations
 
-import os
-import re
-import sys
-
 from backends.base import ModelBackend
-from handlers.cancel import poll_cancel
+from handlers.generate import handle_generate
 
 
 def handle_completion(
@@ -15,25 +11,5 @@ def handle_completion(
     images: list | None = None,
     max_image_size: int = 768,
 ) -> None:
-    """completion API の処理"""
-    if options is None:
-        options = {}
-
-    final_options = dict(options)
-    if images:
-        final_options["max_image_size"] = max_image_size
-        if os.getenv('MLX_DEBUG'):
-            display_prompt = re.sub(r'(<\|image_pad\|>)+', '<|image_pad|>...', prompt)
-            sys.stderr.write(f"--- vlm completion (images: {len(images)}, max_size: {max_image_size})\n{display_prompt}\n")
-    elif os.getenv('MLX_DEBUG'):
-        if isinstance(prompt, list):
-            sys.stderr.write(f"--- prompt: len={len(prompt)}\n")
-        else:
-            sys.stderr.write(f"--- prompt\n{prompt}\n")
-
-    for response in backend.stream_generate(prompt, final_options, images):
-        if poll_cancel():
-            break
-        print(response.text.replace("\0", ""), end="", flush=True)
-
-    print("\n", end="\0", flush=True)
+    """completion API（後方互換）— generate に委譲"""
+    handle_generate(backend, prompt, options, images, max_image_size)

package/src/mlx-ml/python/handlers/{chat.py → generate.py}

@@ -1,12 +1,12 @@
 from __future__ import annotations
 
 import json
+import os
 import re
 import sys
 
 from backends.base import ModelBackend
 from mlx_lm.models.cache import trim_prompt_cache
-from utils.prompt_builder import generate_merged_prompt, supports_chat_template
 from handlers.cancel import poll_cancel
 
 
@@ -53,65 +53,42 @@ def _stream_to_stdout(
         print("", end="\0", flush=True)
 
 
-def handle_chat(
+def handle_generate(
     backend: ModelBackend,
-    capabilities: dict,
-    messages: list,
-    primer: str | None = None,
+    prompt: str | list[int],
     options: dict | None = None,
-    tools: list | None = None,
     images: list | None = None,
     max_image_size: int = 768,
-    reasoning_effort: str | None = None,
+    primer: str | None = None,
     cache_path: str | None = None,
     cache_trim_tokens: int | None = None,
 ) -> None:
-    """chat API の処理"""
+    """LIP generate: 整形済み prompt のストリーム推論"""
     if options is None:
         options = {}
 
-    tokenizer = backend.get_tokenizer()
-
-    if backend.supports_vision():
-        add_generation_prompt = True
-        fmt_messages = list(messages)
-        if primer is not None:
-            fmt_messages.append({"role": "assistant", "content": primer})
-            add_generation_prompt = False
-
-        try:
-            prompt = tokenizer.apply_chat_template(
-                fmt_messages,
-                tools=tools,
-                add_generation_prompt=add_generation_prompt,
-                tokenize=False,
-            )
-        except TypeError:
-            prompt = tokenizer.apply_chat_template(
-                fmt_messages,
-                add_generation_prompt=add_generation_prompt,
-                tokenize=False,
-            )
-
-        if primer is not None:
-            prompt = primer.join(prompt.split(primer)[0:-1]) + primer
-
-        display_prompt = re.sub(r'(<\|image_pad\|>)+', '<|image_pad|>...', prompt)
-        sys.stderr.write(f"--- vlm prompt (images: {len(images) if images else 0}, max_size: {max_image_size})\n{display_prompt}\n")
-
-        final_options = dict(options)
+    final_options = dict(options)
+    if images:
         final_options["max_image_size"] = max_image_size
-        _stream_to_stdout(
-            backend,
-            prompt,
-            final_options,
-            images=images,
-            primer=primer,
-        )
-        return
+        if os.getenv('MLX_DEBUG'):
+            if isinstance(prompt, str):
+                display_prompt = re.sub(r'(<\|image_pad\|>)+', '<|image_pad|>...', prompt)
+            else:
+                display_prompt = f"<token ids: len={len(prompt)}>"
+            sys.stderr.write(
+                f"--- vlm generate (images: {len(images)}, max_size: {max_image_size})\n{display_prompt}\n"
+            )
+    elif os.getenv('MLX_DEBUG'):
+        if isinstance(prompt, list):
+            sys.stderr.write(f"--- prompt: len={len(prompt)}\n")
+        else:
+            sys.stderr.write(f"--- prompt\n{prompt}\n")
 
-    prompt_cache = backend.load_cache_from_file(cache_path) if cache_path else None
+    # VLM 経路では KV キャッシュを使わない（旧 chat ハンドラと同様）
+    prompt_cache = None
     cache_tokens = 0
+    if cache_path and not images:
+        prompt_cache = backend.load_cache_from_file(cache_path)
     if prompt_cache is not None:
         if cache_trim_tokens is not None:
             current_offset = backend.get_cache_offset(prompt_cache)
@@ -128,7 +105,6 @@ def handle_chat(
             if meta_count is not None:
                 cache_tokens = meta_count
             else:
-                # Legacy cache without meta file - skip it for safety
                 sys.stderr.write(
                     f"WARNING: Cache file exists but no .meta.json found at {cache_path}. "
                     "Ignoring cache for safety (may be from old implementation).\n"
@@ -142,67 +118,11 @@ def handle_chat(
     elif cache_path:
         sys.stderr.write(f"KV cache load FAILED: {cache_path}\n")
 
-    if not supports_chat_template(tokenizer):
-        prompt = generate_merged_prompt(messages, capabilities)
-        if prompt_cache is not None:
-            sys.stderr.write("KV cache ignored: model does not support chat template\n")
-        _stream_to_stdout(backend, prompt, options, primer=primer)
-        return
-
-    add_generation_prompt = True
-    fmt_messages = list(messages)
-    if primer is not None:
-        fmt_messages.append({"role": "assistant", "content": primer})
-        add_generation_prompt = False
-
-    extra_kwargs = {}
-    if tools is not None:
-        extra_kwargs["tools"] = tools
-    if reasoning_effort is not None:
-        extra_kwargs["reasoning_effort"] = reasoning_effort
-
-    trust_remote_code = options.get("trust_remote_code")
-    if trust_remote_code is not None:
-        extra_kwargs["trust_remote_code"] = trust_remote_code
-
-    try:
-        prompt = tokenizer.apply_chat_template(
-            fmt_messages,
-            add_generation_prompt=add_generation_prompt,
-            tokenize=False,
-            **extra_kwargs,
-        )
-    except TypeError:
-        try:
-            fallback_kwargs = {}
-            if tools is not None:
-                fallback_kwargs["tools"] = tools
-            prompt = tokenizer.apply_chat_template(
-                fmt_messages,
-                add_generation_prompt=add_generation_prompt,
-                tokenize=False,
-                **fallback_kwargs,
-            )
-        except TypeError:
-            prompt = tokenizer.apply_chat_template(
-                fmt_messages,
-                add_generation_prompt=add_generation_prompt,
-                tokenize=False,
-            )
-
-    if primer is not None:
-        prompt = primer.join(prompt.split(primer)[0:-1]) + primer
-
-    if isinstance(prompt, list):
-        sys.stderr.write(f"--- prompt: len={len(prompt)}\n")
-    else:
-        sys.stderr.write(f"--- prompt\n{prompt}\n")
-
-    final_options = dict(options)
     final_options.pop("trust_remote_code", None)
 
     effective_prompt = prompt
     if prompt_cache is not None and cache_tokens > 0 and isinstance(prompt, str):
+        tokenizer = backend.get_tokenizer()
         add_special = tokenizer.bos_token is None or not prompt.startswith(
             tokenizer.bos_token
         )
@@ -221,4 +141,11 @@ def handle_chat(
             )
             prompt_cache = None
 
-    _stream_to_stdout(backend, effective_prompt, final_options, primer=primer, prompt_cache=prompt_cache)
+    _stream_to_stdout(
+        backend,
+        effective_prompt,
+        final_options,
+        images=images,
+        primer=primer,
+        prompt_cache=prompt_cache,
+    )

package/src/mlx-ml/python/handlers/render.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+import json
+
+from backends.base import ModelBackend
+from utils.template_render import apply_chat_template_prompt
+
+
+def handle_render(
+    backend: ModelBackend,
+    messages: list,
+    options: dict | None = None,
+    tools: list | None = None,
+    reasoning_effort: str | None = None,
+) -> None:
+    """LIP render: apply_chat_template のみ（推論しない）"""
+    if options is None:
+        options = {}
+
+    result: dict = {
+        "formatted_prompt": None,
+        "error": None,
+    }
+
+    try:
+        trust_remote_code = options.get("trust_remote_code")
+        primer = options.get("primer")
+        prompt = apply_chat_template_prompt(
+            backend,
+            messages,
+            primer=primer,
+            tools=tools,
+            reasoning_effort=reasoning_effort,
+            trust_remote_code=trust_remote_code,
+        )
+        result["formatted_prompt"] = prompt
+    except Exception as e:
+        result["error"] = str(e)
+
+    print(json.dumps(result), end="\0", flush=True)

package/src/mlx-ml/python/pyproject.toml

@@ -2,14 +2,15 @@
 name = "mlx_driver"
 version = "0.1.0"
 description = "MLX driver for moduler-prompt"
-requires-python = ">=3.10,<3.14"
+requires-python = ">=3.11,<3.14"
 dependencies = [
     "flex==6.14.1",
     "hf-xet==1.2.0",
     "jinja2==3.1.6",
     "mlx>=0.31.2; sys_platform == 'darwin'",
     "mlx-lm==0.31.3; sys_platform == 'darwin'",
-    "mlx-vlm==0.6.3; sys_platform == 'darwin'",
+    "mlx-optiq>=0.4.4; sys_platform == 'darwin'",
+    "mlx-vlm==0.6.4; sys_platform == 'darwin'",
     "tokenizers==0.22.2",
     "torch==2.9.1",
     "torchvision==0.24.1",