@modular-prompt/driver 0.11.15 → 0.13.1

package/src/mlx-ml/python/backends/mlx_lm.py

@@ -0,0 +1,202 @@
+from __future__ import annotations
+
+import json
+import os
+import sys
+from typing import Any, Iterator
+
+from mlx_lm import load as mlx_lm_load
+from mlx_lm import stream_generate as mlx_lm_stream_generate
+from mlx_lm.models.cache import make_prompt_cache, save_prompt_cache, load_prompt_cache, trim_prompt_cache
+from mlx_lm.sample_utils import make_sampler
+
+from backends.base import ModelBackend
+from utils.token_utils import is_eod_token
+
+
+class MlxLmBackend(ModelBackend):
+    """`mlx_lm` backend for text-only models."""
+
+    def __init__(self) -> None:
+        self.model: Any | None = None
+        self.tokenizer: Any | None = None
+
+    def load(self, model_name: str) -> None:
+        self.model, self.tokenizer = mlx_lm_load(model_name)
+
+    def get_tokenizer(self) -> Any:
+        return self.tokenizer
+
+    def stream_generate(
+        self,
+        prompt: str | list[int],
+        options: dict,
+        images: list | None = None,
+        prompt_cache: list | None = None,
+    ) -> Iterator[Any]:
+        if self.model is None or self.tokenizer is None:
+            raise RuntimeError("Model is not loaded")
+
+        final_options = {"max_tokens": 1000, **options}
+        temperature = final_options.pop("temperature", 1.0)
+        top_p = final_options.pop("top_p", 0.0)
+        top_k = final_options.pop("top_k", 0)
+        final_options["sampler"] = make_sampler(
+            temp=temperature,
+            top_p=top_p,
+            top_k=top_k,
+        )
+
+        if prompt_cache is not None:
+            final_options["prompt_cache"] = prompt_cache
+
+        for response in mlx_lm_stream_generate(
+            self.model,
+            self.tokenizer,
+            prompt,
+            **final_options,
+        ):
+            if is_eod_token(response, self.tokenizer):
+                break
+            yield response
+
+    # get_cache_offset is inherited from ModelBackend base class
+
+    def _tokenize_prompt(self, prompt: str) -> list[int]:
+        """Tokenize a prompt string using the same logic as stream_generate."""
+        add_special = self.tokenizer.bos_token is None or not prompt.startswith(
+            self.tokenizer.bos_token
+        )
+        return self.tokenizer.encode(prompt, add_special_tokens=add_special)
+
+    @staticmethod
+    def _write_cache_meta(
+        cache_path: str,
+        token_count: int,
+        prefix_offsets: list[int] | None = None,
+        prefix_hashes: list[str] | None = None,
+    ) -> None:
+        meta_path = cache_path + '.meta.json'
+        try:
+            meta: dict[str, Any] = {"token_count": token_count}
+            if prefix_offsets and prefix_hashes:
+                meta["prefix_offsets"] = prefix_offsets
+                meta["prefix_hashes"] = prefix_hashes
+            with open(meta_path, 'w') as f:
+                json.dump(meta, f)
+        except Exception as e:
+            sys.stderr.write(f"Failed to write cache meta: {e}\n")
+
+    @staticmethod
+    def _read_cache_meta(cache_path: str) -> int | None:
+        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 cache_prefill(
+        self,
+        cache_path: str,
+        prompt: str,
+        base_cache_path: str | None = None,
+        trim_to_tokens: int | None = None,
+        prefix_offsets: list[int] | None = None,
+        prefix_hashes: list[str] | None = None,
+    ) -> dict:
+        if self.model is None or self.tokenizer is None:
+            raise RuntimeError("Model is not loaded")
+
+        full_tokens = self._tokenize_prompt(prompt)
+        token_count = len(full_tokens)
+        effective_prompt: str | list[int] = prompt
+
+        if base_cache_path is not None:
+            try:
+                prompt_cache = load_prompt_cache(base_cache_path)
+
+                if trim_to_tokens is not None:
+                    current_offset = self.get_cache_offset(prompt_cache)
+                    if current_offset > trim_to_tokens:
+                        trim_count = current_offset - trim_to_tokens
+                        trim_prompt_cache(prompt_cache, trim_count)
+                        if os.getenv('MLX_DEBUG'):
+                            sys.stderr.write(
+                                f"Trimmed base cache: {current_offset} → {trim_to_tokens} tokens\n"
+                            )
+                        cache_offset = trim_to_tokens
+                    else:
+                        cache_offset = current_offset
+                else:
+                    meta_offset = self._read_cache_meta(base_cache_path)
+                    if meta_offset is None:
+                        # Legacy cache without meta file - create fresh cache
+                        sys.stderr.write(
+                            f"WARNING: Cache file exists but no .meta.json found at {base_cache_path}. "
+                            "Creating fresh cache (may be from old implementation).\n"
+                        )
+                        cache_offset = 0
+                        prompt_cache = make_prompt_cache(self.model)
+                    else:
+                        cache_offset = meta_offset
+
+                if cache_offset > 0 and prompt_cache is not None:
+                    if cache_offset < token_count:
+                        effective_prompt = full_tokens[cache_offset:]
+                        if os.getenv('MLX_DEBUG'):
+                            sys.stderr.write(
+                                f"Incremental prefill from: {base_cache_path} "
+                                f"(skip {cache_offset}/{token_count} tokens)\n"
+                            )
+                    else:
+                        if os.getenv('MLX_DEBUG'):
+                            sys.stderr.write(
+                                f"Base cache covers entire prompt "
+                                f"({cache_offset} >= {token_count}), saving as-is\n"
+                            )
+                        save_prompt_cache(cache_path, prompt_cache)
+                        self._write_cache_meta(cache_path, token_count, prefix_offsets, prefix_hashes)
+                        return {"cache_path": cache_path, "token_count": token_count}
+                else:
+                    if os.getenv('MLX_DEBUG'):
+                        sys.stderr.write(f"Incremental prefill from: {base_cache_path}\n")
+            except Exception as e:
+                sys.stderr.write(f"Base cache load failed, creating fresh: {e}\n")
+                prompt_cache = make_prompt_cache(self.model)
+        else:
+            prompt_cache = make_prompt_cache(self.model)
+
+        if os.getenv('MLX_DEBUG'):
+            sys.stderr.write(f"Prefill prompt: {token_count} tokens\n")
+
+        for _ in mlx_lm_stream_generate(
+            self.model, self.tokenizer, effective_prompt,
+            prompt_cache=prompt_cache, max_tokens=1,
+        ):
+            break
+
+        save_prompt_cache(cache_path, prompt_cache)
+        self._write_cache_meta(cache_path, token_count, prefix_offsets, prefix_hashes)
+        if os.getenv('MLX_DEBUG'):
+            sys.stderr.write(f"Cache created: {cache_path} ({token_count} tokens)\n")
+        return {"cache_path": cache_path, "token_count": token_count}
+
+    def load_cache_from_file(self, cache_path: str) -> list | None:
+        try:
+            return load_prompt_cache(cache_path)
+        except FileNotFoundError:
+            sys.stderr.write(f"Cache file not found: {cache_path}\n")
+            return None
+        except Exception as e:
+            sys.stderr.write(f"Failed to load cache: {e}\n")
+            return None
+
+    def supports_vision(self) -> bool:
+        return False
+
+    @property
+    def model_kind(self) -> str:
+        return "lm"

package/src/mlx-ml/python/backends/mlx_vlm.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+import sys
+from typing import Any, Iterator
+
+from mlx_vlm import load as mlx_vlm_load
+from mlx_vlm import stream_generate as mlx_vlm_stream_generate
+
+from backends.base import ModelBackend
+from utils.vlm_utils import load_and_resize_images
+
+
+class MlxVlmBackend(ModelBackend):
+    """`mlx_vlm` backend for vision-language models."""
+
+    def __init__(self) -> None:
+        self.model: Any | None = None
+        self.processor: Any | None = None
+        self.drafter: Any | None = None
+        self.drafter_kind: str | None = None
+        self.draft_block_size: int | None = None
+
+    def load(self, model_name: str) -> None:
+        self.model, self.processor = mlx_vlm_load(model_name)
+
+    def load_drafter(self, drafter_model: str) -> None:
+        from mlx_vlm.speculative.drafters import load_drafter
+        self.drafter, self.drafter_kind = load_drafter(drafter_model)
+        sys.stderr.write(f"Drafter loaded: {drafter_model} (kind={self.drafter_kind})\n")
+
+    def has_drafter(self) -> bool:
+        return self.drafter is not None
+
+    def get_tokenizer(self) -> Any:
+        return self.processor
+
+    def stream_generate(
+        self, prompt: str | list[int], options: dict, images: list | None = None,
+        prompt_cache: list | None = None,
+    ) -> Iterator[Any]:
+        if self.model is None or self.processor is None:
+            raise RuntimeError("Model is not loaded")
+
+        final_options = dict(options)
+        temperature = final_options.pop("temperature", 1.0)
+        max_tokens = final_options.pop("max_tokens", 1000)
+        top_p = final_options.pop("top_p", 0.0)
+        top_k = final_options.pop("top_k", 0)
+
+        processed_images = None
+        if images:
+            max_image_size = final_options.pop("max_image_size", 768)
+            processed_images = load_and_resize_images(images, max_image_size)
+
+        draft_kwargs = {}
+        if self.drafter:
+            draft_kwargs["draft_model"] = self.drafter
+            draft_kwargs["draft_kind"] = self.drafter_kind
+            if self.draft_block_size is not None:
+                draft_kwargs["draft_block_size"] = self.draft_block_size
+
+        # Generate and collect tokens
+        token_count = 0
+        for result in mlx_vlm_stream_generate(
+            self.model,
+            self.processor,
+            prompt,
+            image=processed_images,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            **draft_kwargs,
+        ):
+            token_count += 1
+            yield result
+
+        # Output speculative decoding stats if drafter was used
+        if self.drafter and hasattr(self.drafter, 'accept_lens'):
+            accept_lens = self.drafter.accept_lens
+            if accept_lens:
+                avg_accepted = sum(accept_lens) / len(accept_lens)
+                # Show stats unless MLX_NO_STATS environment variable is set
+                import os
+                if not os.getenv('MLX_NO_STATS'):
+                    sys.stderr.write(f"\n[Speculative Decoding Stats]\n")
+                    sys.stderr.write(f"  Rounds: {len(accept_lens)}\n")
+                    sys.stderr.write(f"  Average accepted tokens/round: {avg_accepted:.2f}\n")
+                    sys.stderr.write(f"  Total tokens generated: {token_count}\n")
+                    sys.stderr.write(f"  Speedup factor: {avg_accepted:.2f}x (theoretical)\n")
+                # Clear for next generation
+                self.drafter.accept_lens = []
+
+    def supports_vision(self) -> bool:
+        return True
+
+    @property
+    def model_kind(self) -> str:
+        return "vlm"

package/src/mlx-ml/python/examples/example_basic.py

@@ -0,0 +1,93 @@
+# This file contains code to use LLM-jp-4 models with mlx-lm on Apple Silicon.
+
+from mlx_lm import load, stream_generate
+from mlx_lm.sample_utils import make_sampler
+
+
+def main():
+    model, tokenizer = load(
+        # "llm-jp/llm-jp-4-8b-instruct",
+        "llm-jp/llm-jp-4-8b-thinking",
+        tokenizer_config={"trust_remote_code": True},
+    )
+
+    messages = [
+        {"role": "user", "content": "日本語で自己紹介してください。"},
+    ]
+
+    prompt: str = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        reasoning_effort="medium",
+    )
+
+    print("--- Prompt ---")
+    print(prompt)
+
+    input_ids = tokenizer.encode(prompt)
+
+    print("--- Input IDs ---")
+    print(input_ids)
+
+    generated_ids: list[int] = []
+
+    sampler = make_sampler(temp=0.7, top_p=0.9)
+
+    for resp in stream_generate(
+        model, tokenizer, prompt=input_ids,
+        max_tokens=1024, sampler=sampler,
+    ):
+        generated_ids.append(resp.token)
+
+    print("--- Generated IDs ---")
+    print(generated_ids)
+
+    response = tokenizer.decode(generated_ids)
+
+    print("\n--- Response ---")
+    print(response)
+
+    parsed = tokenizer.parse_response(response)
+
+    print("\n--- Parsed Response ---")
+    print("Role:", parsed.get("role"))
+    print("Thinking:", parsed.get("thinking"))
+    print("Content:", parsed.get("content"))
+
+    # Harmony parser is bundled as the parse_harmony_message method of the tokenizer.
+    # This function accepts a list of token IDs (not strings)
+    # and returns a list of Harmony's message objects with split tokens.
+
+    # To correctly parse the response,
+    # we need to include the prefill tokens for the assistant's response.
+    response_prefill = tokenizer.encode("<|start|>assistant")
+    parsed_harmony = tokenizer.parse_harmony_message(response_prefill + generated_ids)
+
+    print("\n--- Parsed Harmony Messages ---")
+    for i, message in enumerate(parsed_harmony, start=1):
+        print(f"Message {i}:")
+
+        # The end type can be "END", "CALL", or "INCOMPLETE".
+        print("  End Type:", message.end)
+
+        if message.role:
+            print("  Role Tokens:", message.role.token_ids)
+            print("  Role Text:", repr(tokenizer.decode(message.role.token_ids)))
+            print("  Role Start Position:", message.role.start)
+        if message.channel:
+            print("  Channel Tokens:", message.channel.token_ids)
+            print("  Channel Text:", repr(tokenizer.decode(message.channel.token_ids)))
+            print("  Channel Start Position:", message.channel.start)
+        if message.constrain:
+            print("  Constrain Tokens:", message.constrain.token_ids)
+            print("  Constrain Text:", repr(tokenizer.decode(message.constrain.token_ids)))
+            print("  Constrain Start Position:", message.constrain.start)
+        if message.content:
+            print("  Content Tokens:", message.content.token_ids)
+            print("  Content Text:", repr(tokenizer.decode(message.content.token_ids)))
+            print("  Content Start Position:", message.content.start)
+
+
+if __name__ == "__main__":
+    main()

package/src/mlx-ml/python/examples/example_tool_call.py

@@ -0,0 +1,165 @@
+# Tool call example using LLM-jp-4 with mlx-lm on Apple Silicon.
+
+from mlx_lm import load, stream_generate
+from mlx_lm.sample_utils import make_sampler
+
+
+def generate_response(model, tokenizer, input_ids, sampler):
+    generated_ids: list[int] = []
+    for resp in stream_generate(
+        model, tokenizer, prompt=input_ids,
+        max_tokens=1024, sampler=sampler,
+    ):
+        generated_ids.append(resp.token)
+    return generated_ids
+
+
+def main():
+    model, tokenizer = load(
+        # "llm-jp/llm-jp-4-8b-thinking",
+        # "llm-jp/llm-jp-4-8b-instruct",
+        # "mlx-community/llm-jp-4-32b-a3b-thinking-4bit",
+        "mlx-community/Qwen3.6-27B-4bit",
+        # tokenizer_config={"trust_remote_code": True},
+    )
+
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "get_current_time",
+                "description": "現在の日時を取得する",
+                "parameters": {
+                    "type": "object",
+                    "properties": {},
+                },
+            },
+        },
+        {
+            "type": "function",
+            "function": {
+                "name": "get_weather",
+                "description": "指定された都市の現在の天気を取得する",
+                "parameters": {
+                    "type": "object",
+                    "required": ["city"],
+                    "properties": {
+                        "city": {
+                            "type": "string",
+                            "description": "都市名（例: 東京、大阪）",
+                        },
+                    },
+                },
+            },
+        },
+    ]
+
+    messages = [
+        {"role": "developer", "content": "必要に応じて応答をTool Callに切り替えてください。functionsで定義されている機能を呼び出すことができます。"},
+        # \nツール実行形式: <|start|>assistant to=functions.get_current_time<|channel|>commentary json<|message|>{"locate": "Asia/Tokyo"}<|call|>
+
+        # few-shot: tool call → tool response の例
+        {"role": "user", "content": "今何時？"},
+        {
+            "role": "assistant",
+            "tool_calls": [{
+                "function": {
+                    "name": "get_current_time",
+                    "arguments": {"locate": "Asia/Tokyo"},
+                },
+            }],
+        },
+        {
+            "role": "tool",
+            "content": {"datetime": "2026-04-24T15:30:00+09:00"},
+        },
+        {
+            "role": "assistant",
+            "content": "現在の時刻は15時30分です。",
+        },
+        # 本番のリクエスト
+        {"role": "user", "content": '東京の天気を教えてください。'},
+    ]
+
+    sampler = make_sampler(temp=0.7, top_p=0.9)
+
+    # --- Turn 1: tool call生成 ---
+    prompt: str = tokenizer.apply_chat_template(
+        messages,
+        tools=tools,
+        tokenize=False,
+        add_generation_prompt=True,
+        trust_remote_code=True,
+        reasoning_effort="middle",
+    )
+
+    print("=== Turn 1: Tool Call ===")
+    print("--- Prompt ---")
+    print(prompt)
+
+    input_ids = tokenizer.encode(prompt)
+    generated_ids = generate_response(model, tokenizer, input_ids, sampler)
+    response = tokenizer.decode(generated_ids)
+
+    print("\n--- Raw Response ---")
+    print(response)
+
+    # Harmony parserでtool callを解析
+    response_prefill = tokenizer.encode("<|start|>assistant")
+    parsed_harmony = tokenizer.parse_harmony_message(response_prefill + generated_ids)
+
+    print("\n--- Parsed Harmony Messages ---")
+    for i, message in enumerate(parsed_harmony, start=1):
+        print(f"Message {i}:")
+        print("  End Type:", message.end)
+        if message.role:
+            print("  Role:", repr(tokenizer.decode(message.role.token_ids)))
+        if message.channel:
+            print("  Channel:", repr(tokenizer.decode(message.channel.token_ids)))
+        if message.constrain:
+            print("  Constrain:", repr(tokenizer.decode(message.constrain.token_ids)))
+        if message.content:
+            print("  Content:", repr(tokenizer.decode(message.content.token_ids)))
+
+    # # --- Turn 2: tool resultを渡して最終応答 ---
+    # messages.append({
+    #     "role": "assistant",
+    #     "tool_calls": [{
+    #         "function": {
+    #             "name": "get_weather",
+    #             "arguments": '{"city": "東京"}',
+    #         },
+    #     }],
+    # })
+    # messages.append({
+    #     "role": "tool",
+    #     "content": '{"city": "東京", "weather": "晴れ", "temperature": 22, "humidity": 45}',
+    # })
+
+    # prompt2: str = tokenizer.apply_chat_template(
+    #     messages,
+    #     tools=tools,
+    #     tokenize=False,
+    #     add_generation_prompt=True,
+    # )
+
+    # print("\n\n=== Turn 2: Final Response ===")
+    # print("--- Prompt ---")
+    # print(prompt2)
+
+    # input_ids2 = tokenizer.encode(prompt2)
+    # generated_ids2 = generate_response(model, tokenizer, input_ids2, sampler)
+    # response2 = tokenizer.decode(generated_ids2)
+
+    # print("\n--- Raw Response ---")
+    # print(response2)
+
+    # parsed = tokenizer.parse_response(response2)
+    # print("\n--- Parsed Response ---")
+    # print("Role:", parsed.get("role"))
+    # print("Thinking:", parsed.get("thinking"))
+    # print("Content:", parsed.get("content"))
+
+
+if __name__ == "__main__":
+    main()

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

@@ -0,0 +1,6 @@
+from handlers.cache import handle_cache_prefill
+from handlers.capabilities import handle_capabilities
+from handlers.chat import handle_chat
+from handlers.completion import handle_completion
+from handlers.format_test import handle_format_test
+from handlers.tokenize import handle_tokenize

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

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+import json
+import sys
+
+from backends.base import ModelBackend
+from utils.prompt_builder import generate_merged_prompt, supports_chat_template
+
+
+def handle_cache_prefill(
+    backend: ModelBackend,
+    capabilities: dict,
+    cache_path: str,
+    messages: list,
+    base_cache_path: str | None = None,
+    trim_to_tokens: int | None = None,
+    prefix_offsets: list[int] | None = None,
+    prefix_hashes: list[str] | None = None,
+    tools: list | None = None,
+    reasoning_effort: str | None = None,
+) -> None:
+    tokenizer = backend.get_tokenizer()
+
+    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):
+        try:
+            prompt = tokenizer.apply_chat_template(
+                messages,
+                add_generation_prompt=False,
+                tokenize=False,
+                **extra_kwargs,
+            )
+        except TypeError:
+            try:
+                fallback_kwargs = {}
+                if tools is not None:
+                    fallback_kwargs["tools"] = tools
+                prompt = tokenizer.apply_chat_template(
+                    messages,
+                    add_generation_prompt=False,
+                    tokenize=False,
+                    **fallback_kwargs,
+                )
+            except TypeError:
+                try:
+                    prompt = tokenizer.apply_chat_template(
+                        messages,
+                        add_generation_prompt=False,
+                        tokenize=False,
+                    )
+                except Exception:
+                    prompt = generate_merged_prompt(messages, capabilities)
+                    sys.stderr.write(
+                        "--- cache_prefill: fallback to generate_merged_prompt\n"
+                    )
+        except Exception:
+            prompt = generate_merged_prompt(messages, capabilities)
+            sys.stderr.write(
+                "--- cache_prefill: fallback to generate_merged_prompt\n"
+            )
+    else:
+        prompt = generate_merged_prompt(messages, capabilities)
+
+    # Only show debug output if MLX_DEBUG environment variable is set
+    import os
+    if os.getenv('MLX_DEBUG'):
+        sys.stderr.write(f"--- cache_prefill {cache_path}\n")
+    result = backend.cache_prefill(
+        cache_path, prompt, base_cache_path,
+        trim_to_tokens=trim_to_tokens,
+        prefix_offsets=prefix_offsets,
+        prefix_hashes=prefix_hashes,
+    )
+    if prefix_offsets and prefix_hashes:
+        result["prefix_offsets"] = prefix_offsets
+        result["prefix_hashes"] = prefix_hashes
+    print(json.dumps(result), end="\0", flush=True)

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

@@ -0,0 +1,6 @@
+import json
+
+
+def handle_capabilities(capabilities: dict) -> None:
+    """capabilities API の処理。JSON出力してnull文字で終端"""
+    print(json.dumps(capabilities), end="\0", flush=True)