@modular-prompt/driver 0.4.5

package/scripts/setup-mlx.js

@@ -0,0 +1,53 @@
+#!/usr/bin/env node
+
+import { execSync } from 'child_process';
+import { existsSync } from 'fs';
+import { join, dirname } from 'path';
+import { fileURLToPath } from 'url';
+
+const __dirname = dirname(fileURLToPath(import.meta.url));
+const pythonDir = join(__dirname, '..', 'src', 'mlx-ml', 'python');
+const distPythonDir = join(__dirname, '..', 'dist', 'mlx-ml', 'python');
+
+console.log('🚀 Setting up MLX driver dependencies...\n');
+
+// Check Python directory
+const targetDir = existsSync(distPythonDir) ? distPythonDir : pythonDir;
+
+if (!existsSync(targetDir)) {
+  console.log('⚠️  MLX Python directory not found. Skipping setup.');
+  console.log('   MLX driver will not be available.');
+  process.exit(0);
+}
+
+console.log(`📁 Working directory: ${targetDir}`);
+
+// Check if uv is installed
+try {
+  execSync('uv --version', { stdio: 'ignore' });
+  console.log('✅ uv is installed');
+} catch {
+  console.log('⚠️  uv is not installed. Installing uv...');
+  try {
+    execSync('curl -LsSf https://astral.sh/uv/install.sh | sh', { stdio: 'inherit' });
+    console.log('✅ uv installed successfully');
+  } catch (error) {
+    console.error('❌ Failed to install uv. Please install it manually:');
+    console.error('   curl -LsSf https://astral.sh/uv/install.sh | sh');
+    console.error('\n   MLX driver will not be available without uv.');
+    process.exit(0);
+  }
+}
+
+// Setup Python environment
+console.log('\n📦 Setting up Python environment...');
+try {
+  execSync('uv venv --python 3.13', { cwd: targetDir, stdio: 'inherit' });
+  execSync('uv pip install -e .', { cwd: targetDir, stdio: 'inherit' });
+  console.log('\n✅ MLX driver setup completed successfully!');
+  console.log('   You can now use MlxDriver from @moduler-prompt/driver');
+} catch (error) {
+  console.error('❌ Failed to setup Python environment:', error.message);
+  console.error('   MLX driver will not be available.');
+  process.exit(0);
+}

package/src/mlx-ml/python/.python-version

@@ -0,0 +1 @@
+3.13

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

@@ -0,0 +1,312 @@
+import sys
+import json
+from mlx_lm import load, stream_generate
+from mlx_lm.sample_utils import make_sampler
+from token_utils import get_capabilities, is_eod_token
+
+model_name = sys.argv[1] if len(sys.argv) > 1 else "mlx-community/gemma-3-270m-it-qat-4bit"
+
+model, tokenizer = load(model_name)
+
+# Capabilities情報の取得
+capabilities = get_capabilities(tokenizer)
+
+def read():
+    lines = []
+    data = None
+    eof = False
+    while not eof:
+        line = sys.stdin.readline()
+        # sys.stderr.write('line:' + line + '\n')
+        if not line:
+            eof = True
+        else:
+            lines.append(line)
+        try:
+            data = json.loads(''.join(lines))
+        except json.JSONDecodeError as e:
+            data = None
+            continue
+        break
+    return data
+
+
+def supports_chat_template():
+    """
+    チャットテンプレートがサポートされているかを判定
+    
+    apply_chat_templateメソッドの存在と、tokenizer.chat_templateの両方を確認する。
+    tokenizer.chat_templateが設定されていない場合、apply_chat_templateを呼んでも
+    エラーになるため、両方の条件をチェックする必要がある。
+    
+    Returns:
+        bool: チャットテンプレートがサポートされている場合True
+    """
+    return (hasattr(tokenizer, 'apply_chat_template') and 
+            hasattr(tokenizer, 'chat_template') and 
+            tokenizer.chat_template is not None)
+
+def handle_capabilities():
+    """capabilities API の処理"""
+    print(json.dumps(capabilities), end='\0', flush=True)
+
+
+def handle_format_test(messages, options=None):
+    """フォーマットテスト API の処理（実際に生成せずフォーマットのみ）"""
+    if options is None:
+        options = {}
+    
+    result = {
+        "formatted_prompt": None,
+        "template_applied": False,
+        "model_specific_processing": None,
+        "error": None
+    }
+    
+    try:
+        # チャットテンプレートが利用可能かチェック
+        if supports_chat_template():
+            # messagesはTypeScript側で既にモデル固有処理済み
+            result["model_specific_processing"] = messages
+            
+            # プロンプト生成（フォーマットのみ）
+            primer = options.get('primer')
+            add_generation_prompt = True
+            tokenize = False  # 常にテキストで返す
+            
+            if primer is not None:
+                messages.append({'role': 'assistant', 'content': primer})
+                add_generation_prompt = False
+
+            formatted_prompt = tokenizer.apply_chat_template(
+                messages,
+                add_generation_prompt=add_generation_prompt,
+                tokenize=tokenize,
+            )
+
+            if primer is not None:
+                formatted_prompt = primer.join(formatted_prompt.split(primer)[0:-1]) + primer
+            
+            result["formatted_prompt"] = formatted_prompt
+            result["template_applied"] = True
+        else:
+            # チャットテンプレートがない場合はcompletionフォーマット
+            formatted_prompt = generate_merged_prompt(messages)
+            primer = options.get('primer')
+            if primer is not None:
+                formatted_prompt += primer
+            
+            result["formatted_prompt"] = formatted_prompt
+            result["template_applied"] = False
+        
+    except Exception as e:
+        result["error"] = str(e)
+    
+    print(json.dumps(result), end='\0', flush=True)
+
+def handle_chat(messages, primer=None, options=None):
+    """chat API の処理"""
+    if options is None:
+        options = {}
+    
+    # チャットテンプレートが利用可能かチェック
+    if not supports_chat_template():
+        # チャットテンプレートがない場合はcompletionフォーマットに変換
+        # 注意: TypeScript側でAPIを決定するため、ここに来る場合は
+        # TypeScript側でchatが選択されたが、実際にはテンプレートがないケース
+        prompt = generate_merged_prompt(messages)
+        # primerはTypeScript側で既に追加されている場合があるので追加しない
+        # （TypeScript側でcompletion APIへの変換時に追加済み）
+        if primer is not None:
+            print(primer, end='', flush=True)
+        generate_text(prompt, options)
+        return
+    
+    # messagesはTypeScript側で既にモデル固有処理済み
+    
+    # プロンプト生成
+    add_generation_prompt = True
+    tokenize = False
+    
+    if primer is not None:
+        messages.append({'role': 'assistant', 'content': primer})
+        add_generation_prompt = False
+        tokenize = False
+
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=add_generation_prompt,
+        tokenize=tokenize,
+    )
+
+    if primer is not None:
+        prompt = primer.join(prompt.split(primer)[0:-1]) + primer
+        print(primer, end='', flush=True)
+
+    generate_text(prompt, options)
+
+
+def generate_merged_prompt(messages):
+    """apply_chat_templateがない場合のプロンプト生成"""
+    # messagesはTypeScript側で既にmergeSystemMessages処理済み
+    # TypeScript側のformatterと同じフォーマットを維持
+
+    prompt_parts = []
+    special_tokens = capabilities.get('special_tokens', {})
+
+    for msg in messages:
+        role = msg['role']  # 小文字のまま
+        role_upper = role.upper()
+
+        # 1. 専用のspecial_tokenを探す
+        role_token = special_tokens.get(role)
+
+        if role_token and isinstance(role_token, dict) and 'start' in role_token:
+            # 専用トークンがある場合
+            start_token = role_token['start']['text']
+            end_token = role_token['end']['text']
+            prompt_parts.extend([
+                start_token,
+                msg['content'].strip(),
+                end_token,
+                ''  # 空行で区切る
+            ])
+        else:
+            # 2. 専用トークンがない場合、汎用blockトークンを探す
+            # blockやcontextなどの汎用的なペアトークンを探す
+            block_token = None
+            for candidate in ['block', 'context', 'quote', 'section']:
+                token = special_tokens.get(candidate)
+                if token and isinstance(token, dict) and 'start' in token:
+                    block_token = token
+                    break
+
+            if block_token:
+                # 汎用blockトークンがある場合: {block_begin}{role}:\n...{block_end}
+                start_token = block_token['start']['text']
+                end_token = block_token['end']['text']
+                prompt_parts.extend([
+                    f'{start_token}{role_upper}:\n{msg["content"].strip()}',
+                    end_token,
+                    ''  # 空行で区切る
+                ])
+            else:
+                # 3. どちらもない場合は、HTMLコメント形式（フォールバック）
+                prompt_parts.extend([
+                    f'<!-- begin of {role_upper} -->',
+                    msg['content'].strip(),
+                    f'<!-- end of {role_upper} -->',
+                    ''  # 空行で区切る
+                ])
+
+    # 最後の空行を削除して、ダブル改行で結合
+    return '\n'.join(prompt_parts[:-1])
+
+
+def handle_completion(prompt, options=None):
+    """completion API の処理"""
+    if options is None:
+        options = {}
+    
+    # promptはTypeScript側で既にモデル固有処理済み
+    
+    generate_text(prompt, options)
+
+
+def generate_text(prompt, options):
+    """テキスト生成の共通処理
+
+    注意: optionsはTypeScript側で事前にバリデーション済み
+    - temperatureパラメータはsamplerオブジェクトに変換
+    - サポートされていないパラメータはTS側でフィルタリング
+    """
+    # デフォルトオプションの設定
+    default_options = {'max_tokens': 1000}
+
+    # temperatureパラメータを抽出してsamplerを作成
+    temperature = options.pop('temperature', 1.0) if 'temperature' in options else 1.0
+    top_p = options.pop('top_p', 0.0) if 'top_p' in options else 0.0
+    top_k = options.pop('top_k', 0) if 'top_k' in options else 0
+
+    # samplerオブジェクトを作成
+    sampler = make_sampler(temp=temperature, top_p=top_p, top_k=top_k)
+
+    # 残りのオプションとマージ
+    final_options = {**default_options, **options, 'sampler': sampler}
+
+    if isinstance(prompt, list):  # tokenized
+        sys.stderr.write(f"--- prompt: len={len(prompt)}\n")
+    else:
+        sys.stderr.write(f"--- prompt\n{prompt}\n")
+
+    eos_detected = False
+    for response in stream_generate(model, tokenizer, prompt, **final_options):
+        # トークンIDによるEOS判定（より確実）
+        if is_eod_token(response, tokenizer):
+            eos_detected = True
+            print('\n', end='\0', flush=True)
+            break
+        if not eos_detected:
+            print(response.text.replace('\0', ''), end='', flush=True)
+    
+    if not eos_detected:
+        print('\n', end='\0', flush=True)
+
+def main():
+    while True:
+        req = read()
+        if req is None:
+            break
+        
+        method = req.get('method')
+        if not method:
+            sys.stderr.write("Error: 'method' field is required\n")
+            print('\n', end='\0', flush=True)
+            continue
+        
+        try:
+            if method == 'capabilities':
+                handle_capabilities()
+            
+            elif method == 'format_test':
+                messages = req.get('messages')
+                if not messages:
+                    sys.stderr.write("Error: 'messages' field is required for format_test method\n")
+                    print('\n', end='\0', flush=True)
+                    continue
+                
+                options = req.get('options', {})
+                handle_format_test(messages, options)
+            
+            elif method == 'chat':
+                messages = req.get('messages')
+                if not messages:
+                    sys.stderr.write("Error: 'messages' field is required for chat method\n")
+                    print('\n', end='\0', flush=True)
+                    continue
+                
+                primer = req.get('primer')
+                options = req.get('options', {})
+                handle_chat(messages, primer, options)
+            
+            elif method == 'completion':
+                prompt = req.get('prompt')
+                if not prompt:
+                    sys.stderr.write("Error: 'prompt' field is required for completion method\n")
+                    print('\n', end='\0', flush=True)
+                    continue
+                
+                options = req.get('options', {})
+                handle_completion(prompt, options)
+            
+            else:
+                sys.stderr.write(f"Error: Unknown method '{method}'\n")
+                print('\n', end='\0', flush=True)
+        
+        except Exception as e:
+            sys.stderr.write(f"Error processing request: {e}\n")
+            print('\n', end='\0', flush=True)
+
+
+if __name__ == "__main__":
+    main()

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

@@ -0,0 +1,164 @@
+"""
+チャットテンプレートの制約検出
+
+tokenizerのapply_chat_templateを使用して、
+モデルがサポートするメッセージパターンの制約を検出する。
+"""
+
+
+def detect_chat_restrictions(tokenizer) -> dict:
+    """
+    チャットテンプレートの制約を検出
+
+    Args:
+        tokenizer: HuggingFace tokenizer (apply_chat_template対応)
+
+    Returns:
+        dict: chat_restrictions情報
+            {
+                "single_system_at_start": bool,
+                "max_system_messages": int,
+                "alternating_turns": bool,
+                "requires_user_last": bool,
+                "allow_empty_messages": bool
+            }
+    """
+    if not hasattr(tokenizer, 'apply_chat_template'):
+        return None
+
+    # テストパターンを実行
+    test_results = {}
+    for pattern in _get_test_patterns():
+        try:
+            tokenizer.apply_chat_template(
+                pattern['messages'],
+                tokenize=False,
+                add_generation_prompt=False
+            )
+            test_results[pattern['name']] = {'success': True}
+        except Exception as e:
+            test_results[pattern['name']] = {'error': str(e)}
+
+    # テスト結果から制約を推論
+    return _infer_restrictions_from_results(test_results)
+
+
+def _get_test_patterns():
+    """テストパターンの定義"""
+    return [
+        # 基本パターン
+        {
+            'name': 'basic',
+            'messages': [
+                {'role': 'user', 'content': 'Hello'}
+            ]
+        },
+
+        # システムメッセージ付き
+        {
+            'name': 'with-system',
+            'messages': [
+                {'role': 'system', 'content': 'You are a helpful assistant.'},
+                {'role': 'user', 'content': 'Hello'}
+            ]
+        },
+
+        # 複数システムメッセージ
+        {
+            'name': 'multi-system',
+            'messages': [
+                {'role': 'system', 'content': 'First system message.'},
+                {'role': 'system', 'content': 'Second system message.'},
+                {'role': 'user', 'content': 'Hello'}
+            ]
+        },
+
+        # 連続ユーザーメッセージ
+        {
+            'name': 'consecutive-user',
+            'messages': [
+                {'role': 'user', 'content': 'First question'},
+                {'role': 'user', 'content': 'Second question'}
+            ]
+        },
+
+        # アシスタントで終わる
+        {
+            'name': 'assistant-last',
+            'messages': [
+                {'role': 'user', 'content': 'Hello'},
+                {'role': 'assistant', 'content': 'Hi there!'}
+            ]
+        },
+
+        # 交互の会話
+        {
+            'name': 'alternating',
+            'messages': [
+                {'role': 'user', 'content': 'Question 1'},
+                {'role': 'assistant', 'content': 'Answer 1'},
+                {'role': 'user', 'content': 'Question 2'}
+            ]
+        },
+
+        # 空メッセージ
+        {
+            'name': 'empty-message',
+            'messages': [
+                {'role': 'user', 'content': ''}
+            ]
+        },
+
+        # システムメッセージが途中にある
+        {
+            'name': 'system-middle',
+            'messages': [
+                {'role': 'user', 'content': 'First'},
+                {'role': 'system', 'content': 'System in middle'},
+                {'role': 'user', 'content': 'Second'}
+            ]
+        }
+    ]
+
+
+def _infer_restrictions_from_results(test_results: dict) -> dict:
+    """
+    テスト結果から制約を推論
+
+    Args:
+        test_results: テストパターン名をキーとした結果の辞書
+
+    Returns:
+        dict: 検出された制約
+    """
+    restrictions = {}
+
+    # システムメッセージの制約を検出
+    with_system = test_results.get('with-system')
+    multi_system = test_results.get('multi-system')
+
+    if with_system and 'error' in with_system:
+        # 単独のsystemメッセージもエラー → systemロール自体がサポートされていない
+        restrictions['max_system_messages'] = 0
+    elif multi_system and 'error' in multi_system:
+        # 複数はエラーだが単独は成功 → 最大1つまで
+        restrictions['single_system_at_start'] = True
+        restrictions['max_system_messages'] = 1
+    # それ以外（両方成功）→ max_system_messagesキーを設定しない（無制限）
+
+    # 連続ユーザーメッセージのテスト
+    consecutive_user = test_results.get('consecutive-user')
+    if consecutive_user and 'error' in consecutive_user:
+        restrictions['alternating_turns'] = True
+
+    # アシスタントで終わるテスト
+    assistant_last = test_results.get('assistant-last')
+    if assistant_last and 'error' in assistant_last:
+        restrictions['requires_user_last'] = True
+
+    # 空メッセージのテスト
+    empty_message = test_results.get('empty-message')
+    if empty_message and 'error' in empty_message:
+        restrictions['allow_empty_messages'] = False
+
+    return restrictions if restrictions else None

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

@@ -0,0 +1,19 @@
+[project]
+name = "mlx_driver"
+version = "0.1.0"
+description = "MLX driver for moduler-prompt"
+requires-python = ">=3.10,<3.14"
+dependencies = [
+    "flex>=6.14.1",
+    "hf-xet>=1.1.8",
+    "jinja2>=3.1.6",
+    "mlx-lm>=0.28.3",
+    "torch>=2.9.0",
+]
+
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools]
+py-modules = ["__main__", "chat_template_constraints", "token_utils"]