@modular-prompt/driver 0.13.5 → 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/skills/driver-usage/SKILL.md

@@ -54,6 +54,8 @@ interface QueryOptions {
   stream?: boolean;
   tools?: ToolDefinition[];
   toolChoice?: ToolChoice;
+  cache?: boolean | 'read-only';  // プロンプトキャッシュ（ドライバー依存）
+  signal?: AbortSignal;           // 推論キャンセル（未対応ドライバーは無視）
 }
 ```
 
@@ -67,6 +69,8 @@ interface QueryResult {
     promptTokens: number;
     completionTokens: number;
     totalTokens: number;
+    cacheReadTokens?: number;    // キャッシュヒット分（MLX 等）
+    cacheWriteTokens?: number;   // 新規キャッシュ書き込み分（MLX 等）
   };
   toolCalls?: ToolCall[];        // ツール呼び出し
   finishReason?: FinishReason;   // 'stop' | 'length' | 'error' | 'tool_calls'
@@ -78,10 +82,47 @@ interface QueryResult {
 ```typescript
 interface StreamResult {
   stream: AsyncIterable<string>;  // テキストチャンクのストリーム
-  result: Promise<QueryResult>;   // 最終結果（ストリーム完了後に解決）
+  result: Promise<QueryResult>;   // 最終結果（ストリーム完了または abort 後に解決）
 }
 ```
 
+- usage は `stream` チャンクではなく **`result.usage` のみ**
+- `signal` でキャンセルした場合も `result` は reject しない（`finishReason: 'error'` + `signal.aborted` で判定）
+
+## 推論キャンセル（AbortSignal）
+
+```typescript
+const controller = new AbortController();
+const { stream, result } = await driver.streamQuery(compiled, {
+  signal: controller.signal,
+});
+
+for await (const chunk of stream) {
+  process.stdout.write(chunk);
+  if (userPressedEsc) controller.abort();
+}
+
+const final = await result;
+```
+
+| ドライバー | 対応 |
+|---|---|
+| MlxDriver | ✅ stdin cancel + Python `poll_cancel()` |
+| その他 | 未実装（`signal` 無視） |
+
+ヘルパー: `watchAbortSignal`, `createAbortedStreamResult`, `isAborted`（`@modular-prompt/driver`）
+
+## トークン使用量とキャッシュ
+
+```typescript
+const final = await result;
+final.usage?.promptTokens;      // プロバイダ報告の prompt 側総数
+final.usage?.cacheReadTokens;   // MLX: KV キャッシュヒット分
+final.usage?.cacheWriteTokens;  // MLX: 同一クエリ内の新規 prefill 分
+```
+
+`buildQueryUsage()` でカスタムドライバーから同形式の usage を組み立てられます。
+
 ## 各ドライバーのConfig
 
 ### OpenAIDriver
@@ -219,6 +260,20 @@ await driver.close();
 
 Apple Silicon専用。Python 3.11以上が必要。
 
+#### AbortSignal（推論キャンセル）
+
+`QueryOptions.signal` を渡すと、ストリーム中の推論をキャンセルできます。
+
+```typescript
+const controller = new AbortController();
+const { stream, result } = await driver.streamQuery(compiled, {
+  signal: controller.signal,
+});
+// controller.abort() → Python 子プロセスに cancel 送信、部分 content 保持
+```
+
+実装: TS が stdin に `{"method":"cancel"}` を送り、Python の `stream_generate` ループが `poll_cancel()` で協調的に終了します。
+
 #### textOnlyオプション
 
 VLM（Vision Language Model）対応モデルを画像なしのテキストのみで使用する場合に`textOnly: true`を指定します。

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/cancel.py

@@ -0,0 +1,53 @@
+"""Cancel request handling for in-flight streaming generation."""
+
+from __future__ import annotations
+
+import json
+import select
+import sys
+
+_cancel_requested = False
+
+
+def request_cancel() -> None:
+    global _cancel_requested
+    _cancel_requested = True
+
+
+def reset_cancel() -> None:
+    global _cancel_requested
+    _cancel_requested = False
+
+
+def is_cancel_requested() -> bool:
+    return _cancel_requested
+
+
+def poll_cancel() -> bool:
+    """Non-blocking check for a cancel command on stdin during streaming."""
+    global _cancel_requested
+    if _cancel_requested:
+        return True
+
+    try:
+        ready, _, _ = select.select([sys.stdin], [], [], 0)
+    except (ValueError, OSError):
+        return False
+
+    if not ready:
+        return False
+
+    line = sys.stdin.readline()
+    if not line:
+        return False
+
+    try:
+        req = json.loads(line)
+    except json.JSONDecodeError:
+        return False
+
+    if req.get("method") == "cancel":
+        _cancel_requested = True
+        return True
+
+    return False

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

@@ -1,10 +1,7 @@
 from __future__ import annotations
 
-import os
-import re
-import sys
-
 from backends.base import ModelBackend
+from handlers.generate import handle_generate
 
 
 def handle_completion(
@@ -14,23 +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):
-        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)