@modular-prompt/driver 0.14.0 → 0.15.0

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

@@ -3,7 +3,7 @@ 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
 
 
@@ -27,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
@@ -93,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,