@modular-prompt/driver 0.11.15 → 0.13.1

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

@@ -0,0 +1,221 @@
+from __future__ import annotations
+
+import json
+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
+
+
+def _read_cache_token_count(cache_path: str) -> int | None:
+    """Read token count from the sidecar .meta.json file."""
+    meta_path = cache_path + '.meta.json'
+    try:
+        with open(meta_path) as f:
+            meta = json.load(f)
+            count = meta.get('token_count')
+            return int(count) if count is not None else None
+    except (FileNotFoundError, json.JSONDecodeError, ValueError, TypeError):
+        return None
+
+
+def _stream_to_stdout(
+    backend: ModelBackend,
+    prompt: str | list[int],
+    options: dict,
+    images: list | None = None,
+    primer: str | None = None,
+    prompt_cache: list | None = None,
+) -> None:
+    if primer is not None:
+        print(primer, end="", flush=True)
+
+    last_response = None
+    for response in backend.stream_generate(prompt, options, images, prompt_cache=prompt_cache):
+        print(response.text.replace("\0", "").replace("\x1e", ""), end="", flush=True)
+        last_response = response
+
+    meta: dict = {}
+    if last_response is not None:
+        if hasattr(last_response, "prompt_tokens"):
+            meta["prompt_tokens"] = last_response.prompt_tokens
+        if hasattr(last_response, "generation_tokens"):
+            meta["generation_tokens"] = last_response.generation_tokens
+
+    if meta:
+        print(f"\x1e__META__:{json.dumps(meta)}", end="\0", flush=True)
+    else:
+        print("", end="\0", flush=True)
+
+
+def handle_chat(
+    backend: ModelBackend,
+    capabilities: dict,
+    messages: list,
+    primer: str | None = None,
+    options: dict | None = None,
+    tools: list | None = None,
+    images: list | None = None,
+    max_image_size: int = 768,
+    reasoning_effort: str | None = None,
+    cache_path: str | None = None,
+    cache_trim_tokens: int | None = None,
+) -> None:
+    """chat API の処理"""
+    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["max_image_size"] = max_image_size
+        _stream_to_stdout(
+            backend,
+            prompt,
+            final_options,
+            images=images,
+            primer=primer,
+        )
+        return
+
+    prompt_cache = backend.load_cache_from_file(cache_path) if cache_path else None
+    cache_tokens = 0
+    if prompt_cache is not None:
+        if cache_trim_tokens is not None:
+            current_offset = backend.get_cache_offset(prompt_cache)
+            if current_offset > cache_trim_tokens:
+                trim_prompt_cache(prompt_cache, current_offset - cache_trim_tokens)
+                sys.stderr.write(
+                    f"KV cache trimmed: {current_offset} → {cache_trim_tokens} tokens\n"
+                )
+                cache_tokens = cache_trim_tokens
+            else:
+                cache_tokens = current_offset
+        else:
+            meta_count = _read_cache_token_count(cache_path) if cache_path else None
+            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"
+                )
+                prompt_cache = None
+                cache_tokens = 0
+        if prompt_cache is not None:
+            sys.stderr.write(
+                f"KV cache loaded: {len(prompt_cache)} layers, {cache_tokens} cached tokens\n"
+            )
+    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):
+        add_special = tokenizer.bos_token is None or not prompt.startswith(
+            tokenizer.bos_token
+        )
+        full_tokens = tokenizer.encode(prompt, add_special_tokens=add_special)
+
+        if cache_tokens < len(full_tokens):
+            effective_prompt = full_tokens[cache_tokens:]
+            sys.stderr.write(
+                f"Prefilled {cache_tokens}/{len(full_tokens)} tokens, "
+                f"generating from {len(effective_prompt)} remaining\n"
+            )
+        else:
+            sys.stderr.write(
+                f"Prefill offset {cache_tokens} >= prompt {len(full_tokens)}, "
+                f"ignoring prefill state\n"
+            )
+            prompt_cache = None
+
+    _stream_to_stdout(backend, effective_prompt, final_options, primer=primer, prompt_cache=prompt_cache)

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

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import os
+import re
+import sys
+
+from backends.base import ModelBackend
+
+
+def handle_completion(
+    backend: ModelBackend,
+    prompt: str | list[int],
+    options: dict | None = None,
+    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)

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

@@ -7,18 +7,28 @@ dependencies = [
     "flex==6.14.1",
     "hf-xet==1.2.0",
     "jinja2==3.1.6",
-    "mlx==0.31.1; sys_platform == 'darwin'",
-    "mlx-lm==0.31.2; sys_platform == 'darwin'",
-    "mlx-vlm==0.4.4",
+    "mlx>=0.31.2; sys_platform == 'darwin'",
+    "mlx-lm==0.31.3; sys_platform == 'darwin'",
+    "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):
@@ -152,7 +151,8 @@ def get_special_tokens(tokenizer):
 
         # ツール関連の単体トークン（追加）
         "tool_calls_marker": "[TOOL_CALLS]",
-        "tool_call_end": "<|call|>",
+        # Harmony形式のcallトークン（tool_call_endとは異なる用途）
+        "harmony_call": "<|call|>",
     }
     
     # VLM processorではconvert_tokens_to_idsがない場合がある
@@ -254,8 +254,6 @@ def detect_tool_call_format(tokenizer):
             (r'<longcat_tool_call>', r'</longcat_tool_call>'),
             # <minimax:tool_call>...</minimax:tool_call>
             (r'<minimax:tool_call>', r'</minimax:tool_call>'),
-            # context-1形式: to=functions.{name}...<|call|>
-            (r'to=functions\.', r'<\|call\|>'),
         ]
 
         for start_pattern, end_pattern in tool_call_patterns:
@@ -266,6 +264,16 @@ 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)
+            if has_functions and has_call:
+                result["tool_parser_type"] = "harmony"
+                result["call_start"] = "to=functions."
+                result["call_end"] = "<|call|>"
+
         # Mistral特殊ケース
         if "call_start" not in result:
             mistral_match = re.search(r'\[TOOL_CALLS\]', template)