@modular-prompt/driver 0.12.0 → 0.13.2

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

@@ -0,0 +1,70 @@
+import json
+
+from backends.base import ModelBackend
+from utils.prompt_builder import generate_merged_prompt, supports_chat_template
+
+
+def handle_format_test(
+    backend: ModelBackend,
+    capabilities: dict,
+    messages: list,
+    options: dict | None = None,
+    tools: list | None = None,
+) -> None:
+    """フォーマットテスト API の処理（実際に生成せずフォーマットのみ）"""
+    if options is None:
+        options = {}
+
+    tokenizer = backend.get_tokenizer()
+    result = {
+        "formatted_prompt": None,
+        "template_applied": False,
+        "model_specific_processing": None,
+        "error": None,
+    }
+
+    try:
+        if supports_chat_template(tokenizer):
+            result["model_specific_processing"] = messages
+
+            primer = options.get("primer")
+            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:
+                formatted_prompt = tokenizer.apply_chat_template(
+                    fmt_messages,
+                    tools=tools,
+                    add_generation_prompt=add_generation_prompt,
+                    tokenize=False,
+                )
+            except TypeError:
+                formatted_prompt = tokenizer.apply_chat_template(
+                    fmt_messages,
+                    add_generation_prompt=add_generation_prompt,
+                    tokenize=False,
+                )
+
+            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:
+            formatted_prompt = generate_merged_prompt(messages, capabilities)
+            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)

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

@@ -0,0 +1,63 @@
+import json
+
+from backends.base import ModelBackend
+from utils.prompt_builder import generate_merged_prompt, supports_chat_template
+
+
+def handle_tokenize(
+    backend: ModelBackend,
+    capabilities: dict,
+    messages: list,
+    tools: list | None = None,
+    reasoning_effort: str | None = None,
+) -> None:
+    """メッセージをchat template適用後にトークン化して返す"""
+    tokenizer = backend.get_tokenizer()
+
+    result = {
+        "token_ids": None,
+        "token_count": 0,
+        "error": None,
+    }
+
+    try:
+        # apply_chat_templateのfallbackパターン (chat.py L165-188 と同じ)
+        # 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
+
+        if supports_chat_template(tokenizer):
+            # chat.py と同じfallbackチェーン
+            prompt = None
+            for kwargs in [extra_kwargs, {k: v for k, v in extra_kwargs.items() if k == "tools"}, {}]:
+                try:
+                    prompt = tokenizer.apply_chat_template(
+                        messages,
+                        add_generation_prompt=False,
+                        tokenize=False,
+                        **kwargs,
+                    )
+                    break
+                except TypeError:
+                    continue
+
+            if prompt is None:
+                prompt = str(messages)
+        else:
+            prompt = generate_merged_prompt(messages, capabilities)
+
+        # トークン化
+        add_special = tokenizer.bos_token is None or not prompt.startswith(
+            tokenizer.bos_token or ""
+        )
+        token_ids = tokenizer.encode(prompt, add_special_tokens=add_special)
+
+        result["token_ids"] = token_ids
+        result["token_count"] = len(token_ids)
+    except Exception as e:
+        result["error"] = str(e)
+
+    print(json.dumps(result), end="\0", flush=True)

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

@@ -9,16 +9,26 @@ dependencies = [
     "jinja2==3.1.6",
     "mlx>=0.31.2; sys_platform == 'darwin'",
     "mlx-lm==0.31.3; sys_platform == 'darwin'",
-    "mlx-vlm==0.4.4",
+    "mlx-vlm==0.5.0",
     "tokenizers==0.22.2",
     "torch==2.9.1",
     "torchvision==0.24.1",
-    "transformers==5.2.0",
+    "transformers>=5.5.0",
 ]
 
+[dependency-groups]
+dev = ["pytest>=9.0"]
+
 [build-system]
 requires = ["setuptools>=61.0"]
 build-backend = "setuptools.build_meta"
 
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
 [tool.setuptools]
-py-modules = ["__main__", "chat_template_constraints", "token_utils"]
+py-modules = ["__main__", "server"]
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["backends*", "handlers*", "utils*"]

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

@@ -0,0 +1,126 @@
+"""JSON-RPC風サーバー: stdin/stdoutベースのリクエストディスパッチ"""
+import json
+import sys
+
+from backends.base import ModelBackend
+from handlers import handle_cache_prefill, handle_capabilities, handle_chat, handle_completion, handle_format_test, handle_tokenize
+
+
+MAX_READ_LINES = 10000
+
+
+def read():
+    lines = []
+    while True:
+        line = sys.stdin.readline()
+        if not line:
+            return None
+        lines.append(line)
+        if len(lines) > MAX_READ_LINES:
+            sys.stderr.write(f"Error: read buffer exceeded {MAX_READ_LINES} lines, discarding\n")
+            lines.clear()
+            continue
+        try:
+            return json.loads(''.join(lines))
+        except json.JSONDecodeError:
+            continue
+
+
+class Server:
+    def __init__(self, backend: ModelBackend, capabilities: dict):
+        self.backend = backend
+        self.capabilities = capabilities
+
+    def run(self):
+        while True:
+            req = read()
+            if req is None:
+                break
+            self._dispatch(req)
+
+    def _error_response(self, message: str) -> None:
+        sys.stderr.write(f"Error: {message}\n")
+        print(json.dumps({"error": message}), end='\0', flush=True)
+
+    def _dispatch(self, req: dict):
+        method = req.get('method')
+        if not method:
+            self._error_response("'method' field is required")
+            return
+
+        try:
+            if method == 'capabilities':
+                handle_capabilities(self.capabilities)
+
+            elif method == 'format_test':
+                messages = req.get('messages')
+                if not messages:
+                    self._error_response("'messages' field is required for format_test method")
+                    return
+                handle_format_test(self.backend, self.capabilities, messages, req.get('options', {}), req.get('tools'))
+
+            elif method == 'tokenize':
+                messages = req.get('messages')
+                if messages is None:
+                    self._error_response("'messages' field is required for tokenize method")
+                    return
+                handle_tokenize(
+                    self.backend, self.capabilities, messages,
+                    tools=req.get('tools'),
+                    reasoning_effort=req.get('reasoning_effort'),
+                )
+
+            elif method == 'cache_prefill':
+                cache_path = req.get('cache_path')
+                messages = req.get('messages')
+                if not cache_path or not messages:
+                    self._error_response("'cache_path' and 'messages' fields are required for cache_prefill")
+                    return
+                handle_cache_prefill(
+                    self.backend, self.capabilities, cache_path, messages,
+                    base_cache_path=req.get('base_cache_path'),
+                    trim_to_tokens=req.get('trim_to_tokens'),
+                    prefix_offsets=req.get('prefix_offsets'),
+                    prefix_hashes=req.get('prefix_hashes'),
+                    tools=req.get('tools'),
+                    reasoning_effort=req.get('reasoning_effort'),
+                )
+
+            elif method == 'chat':
+                messages = req.get('messages')
+                if not messages:
+                    self._error_response("'messages' field is required for chat method")
+                    return
+                handle_chat(
+                    self.backend,
+                    self.capabilities,
+                    messages,
+                    primer=req.get('primer'),
+                    options=req.get('options', {}),
+                    tools=req.get('tools'),
+                    images=req.get('images', []),
+                    max_image_size=req.get('maxImageSize', 768),
+                    reasoning_effort=req.get('reasoning_effort'),
+                    cache_path=req.get('cache_path'),
+                    cache_trim_tokens=req.get('cache_trim_tokens'),
+                )
+
+            elif method == 'completion':
+                prompt = req.get('prompt')
+                if not prompt:
+                    self._error_response("'prompt' field is required for completion method")
+                    return
+                images = req.get('images', [])
+                handle_completion(
+                    self.backend,
+                    prompt,
+                    options=req.get('options', {}),
+                    images=images if images else None,
+                    max_image_size=req.get('maxImageSize', 768),
+                )
+
+            else:
+                self._error_response(f"Unknown method '{method}'")
+
+        except Exception as e:
+            self._error_response(f"Error processing request: {e}")

package/src/mlx-ml/python/utils/prompt_builder.py

@@ -0,0 +1,54 @@
+"""プロンプト生成ユーティリティ"""
+
+
+def supports_chat_template(tokenizer) -> bool:
+    return (hasattr(tokenizer, 'apply_chat_template') and
+            hasattr(tokenizer, 'chat_template') and
+            tokenizer.chat_template is not None)
+
+
+def generate_merged_prompt(messages, capabilities):
+    """apply_chat_templateがない場合のプロンプト生成"""
+    prompt_parts = []
+    special_tokens = capabilities.get('special_tokens', {})
+
+    for msg in messages:
+        role = msg['role']
+        role_upper = role.upper()
+
+        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:
+            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:
+                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:
+                prompt_parts.extend([
+                    f'<!-- begin of {role_upper} -->',
+                    msg['content'].strip(),
+                    f'<!-- end of {role_upper} -->',
+                    ''
+                ])
+
+    return '\n'.join(prompt_parts[:-1])

package/src/mlx-ml/python/{token_utils.py → utils/token_utils.py}

@@ -1,8 +1,7 @@
 """
 トークン関連のユーティリティ関数
 """
-import sys
-from chat_template_constraints import detect_chat_restrictions
+from utils.chat_template_constraints import detect_chat_restrictions
 
 
 def is_eod_token(response, tokenizer):
@@ -188,50 +187,32 @@ def get_special_tokens(tokenizer):
 def detect_tool_call_format(tokenizer):
     """tokenizer設定からtool call/resultのデリミタを検出する
 
-    tokenizer_config.json / chat_template からtool call関連の情報を抽出:
-    - tool_parser_type: パーサー種別（"json_tools" 等）
-    - chat_template テキスト内の <tool_call></tool_call> 等のタグパターン
+    tokenizerアクセスが必要な情報のみを抽出する:
+    - tool_parser_type: tokenizer_config由来のパーサー種別文字列
+    - chat_templateテキストからのデリミタパターン検出
+
+    パーサー種別→デリミタのマッピングはTS側(detector.ts)に一元化。
+    Python側はtokenizerから生の情報を抽出して渡す役割に専念する。
 
     Returns:
         dict or None: {
             "tool_parser_type": str,  # tokenizer_configのtool_parser_type
-            "call_start": str,        # tool callの開始デリミタ
-            "call_end": str,          # tool callの終了デリミタ
+            "call_start": str,        # chat_templateから検出した開始デリミタ
+            "call_end": str,          # chat_templateから検出した終了デリミタ
             "response_start": str,    # tool responseの開始デリミタ（検出時）
             "response_end": str,      # tool responseの終了デリミタ（検出時）
         } or None
     """
     import re
 
-    # tool_parser_type を取得
     tool_parser_type = None
     if hasattr(tokenizer, 'init_kwargs'):
         tool_parser_type = tokenizer.init_kwargs.get('tool_parser_type')
 
-    # 既知パーサーからの逆引き（最優先）
-    KNOWN_TOOL_PARSERS = {
-        "json_tools": {"call_start": "<tool_call>", "call_end": "</tool_call>"},
-        "pythonic": {"call_start": "<|tool_call_start|>", "call_end": "<|tool_call_end|>"},
-        "function_gemma": {"call_start": "<start_function_call>", "call_end": "<end_function_call>"},
-        "mistral": {"call_start": "[TOOL_CALLS]", "call_end": ""},
-        "kimi_k2": {"call_start": "<|tool_calls_section_begin|>", "call_end": "<|tool_calls_section_end|>"},
-        "longcat": {"call_start": "<longcat_tool_call>", "call_end": "</longcat_tool_call>"},
-        "glm47": {"call_start": "<tool_call>", "call_end": "</tool_call>"},
-        "qwen3_coder": {"call_start": "<tool_call>", "call_end": "</tool_call>"},
-        "minimax_m2": {"call_start": "<minimax:tool_call>", "call_end": "</minimax:tool_call>"},
-    }
-
-    if tool_parser_type and tool_parser_type in KNOWN_TOOL_PARSERS:
-        result = {"tool_parser_type": tool_parser_type}
-        result.update(KNOWN_TOOL_PARSERS[tool_parser_type])
-        return result
-
-    # chat_template テキストを取得
     template = getattr(tokenizer, 'chat_template', None)
     if not template and hasattr(tokenizer, 'init_kwargs'):
         template = tokenizer.init_kwargs.get('chat_template', '')
 
-    # tool_parser_type もテンプレートもなければ非対応
     if not tool_parser_type and not template:
         return None
 
@@ -239,21 +220,13 @@ def detect_tool_call_format(tokenizer):
     if tool_parser_type:
         result["tool_parser_type"] = tool_parser_type
 
-    # テンプレートテキストからデリミタを抽出
     if template:
-        # 複数のtool_call関連パターンを順に試行
         tool_call_patterns = [
-            # </tool_call>, <|/tool_call|>, <tool_call|> (終了タグ専用)
             (r'<\|?tool_call\|?>', r'</tool_call>|<\|/tool_call\|>|<tool_call\|>'),
-            # <|tool_call_start|>...<|tool_call_end|>
             (r'<\|tool_call_start\|>', r'<\|tool_call_end\|>'),
-            # <start_function_call>...<end_function_call>
             (r'<start_function_call>', r'<end_function_call>'),
-            # <|tool_calls_section_begin|>...<|tool_calls_section_end|>
             (r'<\|tool_calls_section_begin\|>', r'<\|tool_calls_section_end\|>'),
-            # <longcat_tool_call>...</longcat_tool_call>
             (r'<longcat_tool_call>', r'</longcat_tool_call>'),
-            # <minimax:tool_call>...</minimax:tool_call>
             (r'<minimax:tool_call>', r'</minimax:tool_call>'),
         ]
 
@@ -265,8 +238,6 @@ def detect_tool_call_format(tokenizer):
                 result["call_end"] = end_match.group(0)
                 break
 
-        # Harmony形式の専用検出
-        # テンプレート内で "functions." と <|call|> が共存する場合
         if "call_start" not in result:
             has_functions = re.search(r'"functions\."', template)
             has_call = re.search(r'<\|call\|>', template)
@@ -275,14 +246,12 @@ def detect_tool_call_format(tokenizer):
                 result["call_start"] = "to=functions."
                 result["call_end"] = "<|call|>"
 
-        # Mistral特殊ケース
         if "call_start" not in result:
             mistral_match = re.search(r'\[TOOL_CALLS\]', template)
             if mistral_match:
                 result["call_start"] = "[TOOL_CALLS]"
                 result["call_end"] = ""
 
-        # tool_response タグの検出
         resp_tags = re.findall(r'<[|/]?tool_response[|]?>', template)
         if len(resp_tags) >= 2:
             open_tags = [t for t in resp_tags if '/' not in t]