@modular-prompt/driver 0.14.0 → 0.15.0

package/src/pytorch/python/server.py

@@ -0,0 +1,140 @@
+"""JSON-RPC風サーバー: stdin/stdoutベースのリクエストディスパッチ"""
+import json
+import sys
+
+from backends.base import ModelBackend
+from handlers import 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
+
+
+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
+
+
+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
+        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
+
+        if method == 'cancel':
+            request_cancel()
+            return
+
+        reset_cancel()
+
+        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':
+                self._error_response("cache_prefill is not supported by the PyTorch backend")
+
+            elif method == 'render':
+                messages = req.get('messages')
+                if not messages:
+                    self._error_response("'messages' field is required for render method")
+                    return
+                handle_render(
+                    self.backend,
+                    messages,
+                    options=req.get('options', {}),
+                    tools=req.get('tools'),
+                    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'),
+                )
+
+            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/pytorch/python/utils/chat_template_constraints.py

@@ -0,0 +1,164 @@
+"""
+チャットテンプレートの制約検出
+
+tokenizerのapply_chat_templateを使用して、
+モデルがサポートするメッセージパターンの制約を検出する。
+"""
+
+
+def detect_chat_restrictions(tokenizer) -> dict:
+    """
+    チャットテンプレートの制約を検出
+
+    Args:
+        tokenizer: HuggingFace tokenizer (apply_chat_template対応)
+
+    Returns:
+        dict: chat_restrictions情報
+            {
+                "single_system_at_start": bool,
+                "max_system_messages": int,
+                "alternating_turns": bool,
+                "requires_user_last": bool,
+                "allow_empty_messages": bool
+            }
+    """
+    if not hasattr(tokenizer, 'apply_chat_template'):
+        return None
+
+    # テストパターンを実行
+    test_results = {}
+    for pattern in _get_test_patterns():
+        try:
+            tokenizer.apply_chat_template(
+                pattern['messages'],
+                tokenize=False,
+                add_generation_prompt=False
+            )
+            test_results[pattern['name']] = {'success': True}
+        except Exception as e:
+            test_results[pattern['name']] = {'error': str(e)}
+
+    # テスト結果から制約を推論
+    return _infer_restrictions_from_results(test_results)
+
+
+def _get_test_patterns():
+    """テストパターンの定義"""
+    return [
+        # 基本パターン
+        {
+            'name': 'basic',
+            'messages': [
+                {'role': 'user', 'content': 'Hello'}
+            ]
+        },
+
+        # システムメッセージ付き
+        {
+            'name': 'with-system',
+            'messages': [
+                {'role': 'system', 'content': 'You are a helpful assistant.'},
+                {'role': 'user', 'content': 'Hello'}
+            ]
+        },
+
+        # 複数システムメッセージ
+        {
+            'name': 'multi-system',
+            'messages': [
+                {'role': 'system', 'content': 'First system message.'},
+                {'role': 'system', 'content': 'Second system message.'},
+                {'role': 'user', 'content': 'Hello'}
+            ]
+        },
+
+        # 連続ユーザーメッセージ
+        {
+            'name': 'consecutive-user',
+            'messages': [
+                {'role': 'user', 'content': 'First question'},
+                {'role': 'user', 'content': 'Second question'}
+            ]
+        },
+
+        # アシスタントで終わる
+        {
+            'name': 'assistant-last',
+            'messages': [
+                {'role': 'user', 'content': 'Hello'},
+                {'role': 'assistant', 'content': 'Hi there!'}
+            ]
+        },
+
+        # 交互の会話
+        {
+            'name': 'alternating',
+            'messages': [
+                {'role': 'user', 'content': 'Question 1'},
+                {'role': 'assistant', 'content': 'Answer 1'},
+                {'role': 'user', 'content': 'Question 2'}
+            ]
+        },
+
+        # 空メッセージ
+        {
+            'name': 'empty-message',
+            'messages': [
+                {'role': 'user', 'content': ''}
+            ]
+        },
+
+        # システムメッセージが途中にある
+        {
+            'name': 'system-middle',
+            'messages': [
+                {'role': 'user', 'content': 'First'},
+                {'role': 'system', 'content': 'System in middle'},
+                {'role': 'user', 'content': 'Second'}
+            ]
+        }
+    ]
+
+
+def _infer_restrictions_from_results(test_results: dict) -> dict:
+    """
+    テスト結果から制約を推論
+
+    Args:
+        test_results: テストパターン名をキーとした結果の辞書
+
+    Returns:
+        dict: 検出された制約
+    """
+    restrictions = {}
+
+    # システムメッセージの制約を検出
+    with_system = test_results.get('with-system')
+    multi_system = test_results.get('multi-system')
+
+    if with_system and 'error' in with_system:
+        # 単独のsystemメッセージもエラー → systemロール自体がサポートされていない
+        restrictions['max_system_messages'] = 0
+    elif multi_system and 'error' in multi_system:
+        # 複数はエラーだが単独は成功 → 最大1つまで
+        restrictions['single_system_at_start'] = True
+        restrictions['max_system_messages'] = 1
+    # それ以外（両方成功）→ max_system_messagesキーを設定しない（無制限）
+
+    # 連続ユーザーメッセージのテスト
+    consecutive_user = test_results.get('consecutive-user')
+    if consecutive_user and 'error' in consecutive_user:
+        restrictions['alternating_turns'] = True
+
+    # アシスタントで終わるテスト
+    assistant_last = test_results.get('assistant-last')
+    if assistant_last and 'error' in assistant_last:
+        restrictions['requires_user_last'] = True
+
+    # 空メッセージのテスト
+    empty_message = test_results.get('empty-message')
+    if empty_message and 'error' in empty_message:
+        restrictions['allow_empty_messages'] = False
+
+    return restrictions if restrictions else None

package/src/pytorch/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/pytorch/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