@modular-prompt/driver 0.13.5 → 0.15.0

package/src/mlx-ml/python/handlers/{chat.py → generate.py}

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
 import json
+import os
 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
+from handlers.cancel import poll_cancel
 
 
 def _read_cache_token_count(cache_path: str) -> int | None:
@@ -34,6 +35,8 @@ def _stream_to_stdout(
 
     last_response = None
     for response in backend.stream_generate(prompt, options, images, prompt_cache=prompt_cache):
+        if poll_cancel():
+            break
         print(response.text.replace("\0", "").replace("\x1e", ""), end="", flush=True)
         last_response = response
 
@@ -50,65 +53,42 @@ def _stream_to_stdout(
         print("", end="\0", flush=True)
 
 
-def handle_chat(
+def handle_generate(
     backend: ModelBackend,
-    capabilities: dict,
-    messages: list,
-    primer: str | None = None,
+    prompt: str | list[int],
     options: dict | None = None,
-    tools: list | None = None,
     images: list | None = None,
     max_image_size: int = 768,
-    reasoning_effort: str | None = None,
+    primer: str | None = None,
     cache_path: str | None = None,
     cache_trim_tokens: int | None = None,
 ) -> None:
-    """chat API の処理"""
+    """LIP generate: 整形済み prompt のストリーム推論"""
     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 = dict(options)
+    if images:
         final_options["max_image_size"] = max_image_size
-        _stream_to_stdout(
-            backend,
-            prompt,
-            final_options,
-            images=images,
-            primer=primer,
-        )
-        return
+        if os.getenv('MLX_DEBUG'):
+            if isinstance(prompt, str):
+                display_prompt = re.sub(r'(<\|image_pad\|>)+', '<|image_pad|>...', prompt)
+            else:
+                display_prompt = f"<token ids: len={len(prompt)}>"
+            sys.stderr.write(
+                f"--- vlm generate (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")
 
-    prompt_cache = backend.load_cache_from_file(cache_path) if cache_path else None
+    # VLM 経路では KV キャッシュを使わない（旧 chat ハンドラと同様）
+    prompt_cache = None
     cache_tokens = 0
+    if cache_path and not images:
+        prompt_cache = backend.load_cache_from_file(cache_path)
     if prompt_cache is not None:
         if cache_trim_tokens is not None:
             current_offset = backend.get_cache_offset(prompt_cache)
@@ -125,7 +105,6 @@ def handle_chat(
             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"
@@ -139,67 +118,11 @@ def handle_chat(
     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):
+        tokenizer = backend.get_tokenizer()
         add_special = tokenizer.bos_token is None or not prompt.startswith(
             tokenizer.bos_token
         )
@@ -218,4 +141,11 @@ def handle_chat(
             )
             prompt_cache = None
 
-    _stream_to_stdout(backend, effective_prompt, final_options, primer=primer, prompt_cache=prompt_cache)
+    _stream_to_stdout(
+        backend,
+        effective_prompt,
+        final_options,
+        images=images,
+        primer=primer,
+        prompt_cache=prompt_cache,
+    )

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

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+import json
+
+from backends.base import ModelBackend
+from utils.template_render import apply_chat_template_prompt
+
+
+def handle_render(
+    backend: ModelBackend,
+    messages: list,
+    options: dict | None = None,
+    tools: list | None = None,
+    reasoning_effort: str | None = None,
+) -> None:
+    """LIP render: apply_chat_template のみ（推論しない）"""
+    if options is None:
+        options = {}
+
+    result: dict = {
+        "formatted_prompt": None,
+        "error": None,
+    }
+
+    try:
+        trust_remote_code = options.get("trust_remote_code")
+        primer = options.get("primer")
+        prompt = apply_chat_template_prompt(
+            backend,
+            messages,
+            primer=primer,
+            tools=tools,
+            reasoning_effort=reasoning_effort,
+            trust_remote_code=trust_remote_code,
+        )
+        result["formatted_prompt"] = prompt
+    except Exception as e:
+        result["error"] = str(e)
+
+    print(json.dumps(result), end="\0", flush=True)

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

@@ -2,14 +2,15 @@
 name = "mlx_driver"
 version = "0.1.0"
 description = "MLX driver for moduler-prompt"
-requires-python = ">=3.10,<3.14"
+requires-python = ">=3.11,<3.14"
 dependencies = [
     "flex==6.14.1",
     "hf-xet==1.2.0",
     "jinja2==3.1.6",
     "mlx>=0.31.2; sys_platform == 'darwin'",
     "mlx-lm==0.31.3; sys_platform == 'darwin'",
-    "mlx-vlm==0.6.3; sys_platform == 'darwin'",
+    "mlx-optiq>=0.4.4; sys_platform == 'darwin'",
+    "mlx-vlm==0.6.4; sys_platform == 'darwin'",
     "tokenizers==0.22.2",
     "torch==2.9.1",
     "torchvision==0.24.1",

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

@@ -3,7 +3,8 @@ 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
+from handlers import handle_cache_prefill, handle_capabilities, handle_completion, handle_format_test, handle_generate, handle_render, handle_tokenize
+from handlers.cancel import request_cancel, reset_cancel
 
 
 MAX_READ_LINES = 10000
@@ -26,6 +27,15 @@ def read():
             continue
 
 
+def _is_valid_generate_prompt(prompt) -> bool:
+    """prompt は非空文字列または非空のトークン ID リスト"""
+    if isinstance(prompt, str):
+        return bool(prompt)
+    if isinstance(prompt, list):
+        return len(prompt) > 0
+    return False
+
+
 class Server:
     def __init__(self, backend: ModelBackend, capabilities: dict):
         self.backend = backend
@@ -48,6 +58,12 @@ class Server:
             self._error_response("'method' field is required")
             return
 
+        if method == 'cancel':
+            request_cancel()
+            return
+
+        reset_cancel()
+
         try:
             if method == 'capabilities':
                 handle_capabilities(self.capabilities)
@@ -86,21 +102,32 @@ class Server:
                     reasoning_effort=req.get('reasoning_effort'),
                 )
 
-            elif method == 'chat':
+            elif method == 'render':
                 messages = req.get('messages')
                 if not messages:
-                    self._error_response("'messages' field is required for chat method")
+                    self._error_response("'messages' field is required for render method")
                     return
-                handle_chat(
+                handle_render(
                     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'),
+                )
+
+            elif method == 'generate':
+                prompt = req.get('prompt')
+                if not _is_valid_generate_prompt(prompt):
+                    self._error_response("'prompt' field is required for generate method")
+                    return
+                images = req.get('images', [])
+                handle_generate(
+                    self.backend,
+                    prompt,
+                    options=req.get('options', {}),
+                    images=images if images else None,
+                    max_image_size=req.get('maxImageSize', 768),
+                    primer=req.get('primer'),
                     cache_path=req.get('cache_path'),
                     cache_trim_tokens=req.get('cache_trim_tokens'),
                 )

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

@@ -0,0 +1,80 @@
+"""apply_chat_template によるプロンプト整形（推論なし）"""
+
+from __future__ import annotations
+
+from utils.prompt_builder import supports_chat_template
+
+
+def apply_chat_template_prompt(
+    backend,
+    messages: list,
+    *,
+    primer: str | None = None,
+    tools: list | None = None,
+    reasoning_effort: str | None = None,
+    trust_remote_code: bool | None = None,
+) -> str | list:
+    """messages を chat template で整形して prompt を返す"""
+    tokenizer = backend.get_tokenizer()
+
+    if not supports_chat_template(tokenizer):
+        raise ValueError("chat_template_not_available")
+
+    add_generation_prompt = True
+    fmt_messages = list(messages)
+    if primer is not None:
+        fmt_messages.append({"role": "assistant", "content": primer})
+        add_generation_prompt = False
+
+    if backend.supports_vision():
+        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,
+            )
+    else:
+        extra_kwargs: dict = {}
+        if tools is not None:
+            extra_kwargs["tools"] = tools
+        if reasoning_effort is not None:
+            extra_kwargs["reasoning_effort"] = reasoning_effort
+        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: dict = {}
+                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 and isinstance(prompt, str):
+        prompt = primer.join(prompt.split(primer)[0:-1]) + primer
+
+    return prompt

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

@@ -342,11 +342,11 @@ def get_capabilities(tokenizer):
         dict: capabilities情報
     """
     # 基本メソッド
-    methods = ["capabilities", "completion", "format_test"]
+    methods = ["capabilities", "completion", "generate", "format_test"]
 
     # apply_chat_templateがある場合はchatメソッドを追加
     if hasattr(tokenizer, 'apply_chat_template'):
-        methods.append("chat")
+        methods.append("render")
 
     capabilities = {
         "methods": methods,