@modular-prompt/driver 0.12.0 → 0.13.1

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):