@modular-prompt/driver 0.15.0 → 0.17.0

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

@@ -18,6 +18,8 @@ def handle_cache_prefill(
     prefix_hashes: list[str] | None = None,
     tools: list | None = None,
     reasoning_effort: str | None = None,
+    images: list | None = None,
+    max_image_size: int = 768,
 ) -> None:
     tokenizer = backend.get_tokenizer()
 
@@ -74,6 +76,8 @@ def handle_cache_prefill(
         trim_to_tokens=trim_to_tokens,
         prefix_offsets=prefix_offsets,
         prefix_hashes=prefix_hashes,
+        images=images,
+        max_image_size=max_image_size,
     )
     if prefix_offsets and prefix_hashes:
         result["prefix_offsets"] = prefix_offsets

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

@@ -6,7 +6,6 @@ import re
 import sys
 
 from backends.base import ModelBackend
-from mlx_lm.models.cache import trim_prompt_cache
 from handlers.cancel import poll_cancel
 
 
@@ -29,6 +28,7 @@ def _stream_to_stdout(
     images: list | None = None,
     primer: str | None = None,
     prompt_cache: list | None = None,
+    cache_loaded: bool | None = None,
 ) -> None:
     if primer is not None:
         print(primer, end="", flush=True)
@@ -46,6 +46,8 @@ def _stream_to_stdout(
             meta["prompt_tokens"] = last_response.prompt_tokens
         if hasattr(last_response, "generation_tokens"):
             meta["generation_tokens"] = last_response.generation_tokens
+    if cache_loaded is not None:
+        meta["cache_loaded"] = cache_loaded
 
     if meta:
         print(f"\x1e__META__:{json.dumps(meta)}", end="\0", flush=True)
@@ -84,16 +86,22 @@ def handle_generate(
         else:
             sys.stderr.write(f"--- prompt\n{prompt}\n")
 
-    # VLM 経路では KV キャッシュを使わない（旧 chat ハンドラと同様）
     prompt_cache = None
     cache_tokens = 0
-    if cache_path and not images:
-        prompt_cache = backend.load_cache_from_file(cache_path)
+    cache_loaded = None
+    if cache_path:
+        prompt_cache = backend.load_cache_from_file(
+            cache_path,
+            images=images,
+            max_image_size=max_image_size,
+            prompt=prompt,
+        )
+        cache_loaded = prompt_cache is not None
     if prompt_cache is not None:
         if cache_trim_tokens is not None:
             current_offset = backend.get_cache_offset(prompt_cache)
             if current_offset > cache_trim_tokens:
-                trim_prompt_cache(prompt_cache, current_offset - cache_trim_tokens)
+                backend.trim_cache(prompt_cache, current_offset - cache_trim_tokens)
                 sys.stderr.write(
                     f"KV cache trimmed: {current_offset} → {cache_trim_tokens} tokens\n"
                 )
@@ -101,9 +109,21 @@ def handle_generate(
             else:
                 cache_tokens = current_offset
         else:
-            meta_count = _read_cache_token_count(cache_path) if cache_path else None
+            meta_count = (
+                _read_cache_token_count(cache_path)
+                if cache_path
+                else None
+            )
             if meta_count is not None:
                 cache_tokens = meta_count
+            elif (
+                backend.model_kind == "vlm"
+                and cache_path is not None
+                and cache_path.startswith("mlx-vlm-memory://")
+            ):
+                # Keep the Phase 1 in-process ref as a compatibility fallback.
+                # Persistent VLM refs always have the sidecar above.
+                cache_tokens = backend.get_cache_offset(prompt_cache)
             else:
                 sys.stderr.write(
                     f"WARNING: Cache file exists but no .meta.json found at {cache_path}. "
@@ -111,6 +131,7 @@ def handle_generate(
                 )
                 prompt_cache = None
                 cache_tokens = 0
+                cache_loaded = False
         if prompt_cache is not None:
             sys.stderr.write(
                 f"KV cache loaded: {len(prompt_cache)} layers, {cache_tokens} cached tokens\n"
@@ -122,11 +143,11 @@ def handle_generate(
 
     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
+        full_tokens = backend.tokenize_prompt(
+            prompt,
+            images=images,
+            max_image_size=max_image_size,
         )
-        full_tokens = tokenizer.encode(prompt, add_special_tokens=add_special)
 
         if cache_tokens < len(full_tokens):
             effective_prompt = full_tokens[cache_tokens:]
@@ -140,6 +161,7 @@ def handle_generate(
                 f"ignoring prefill state\n"
             )
             prompt_cache = None
+            cache_loaded = False
 
     _stream_to_stdout(
         backend,
@@ -148,4 +170,5 @@ def handle_generate(
         images=images,
         primer=primer,
         prompt_cache=prompt_cache,
+        cache_loaded=cache_loaded,
     )

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

@@ -50,10 +50,7 @@ def handle_tokenize(
             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)
+        token_ids = backend.tokenize_prompt(prompt)
 
         result["token_ids"] = token_ids
         result["token_count"] = len(token_ids)

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

@@ -7,14 +7,20 @@ dependencies = [
     "flex==6.14.1",
     "hf-xet==1.2.0",
     "jinja2==3.1.6",
-    "mlx>=0.31.2; sys_platform == 'darwin'",
+    "mlx>=0.32.2; sys_platform == 'darwin'",
     "mlx-lm==0.31.3; sys_platform == 'darwin'",
     "mlx-optiq>=0.4.4; sys_platform == 'darwin'",
-    "mlx-vlm==0.6.4; sys_platform == 'darwin'",
+    "mlx-vlm==0.7.0; sys_platform == 'darwin'",
     "tokenizers==0.22.2",
     "torch==2.9.1",
     "torchvision==0.24.1",
-    "transformers>=5.5.0",
+    "transformers>=5.14.0",
+]
+
+[tool.uv]
+# mlx-optiq は transformers<5.13 を要求するが、mlx-vlm>=0.6.5 は 5.14+ が必要
+override-dependencies = [
+    "transformers>=5.14.0",
 ]
 
 [dependency-groups]

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

@@ -100,6 +100,8 @@ class Server:
                     prefix_hashes=req.get('prefix_hashes'),
                     tools=req.get('tools'),
                     reasoning_effort=req.get('reasoning_effort'),
+                    images=req.get('images'),
+                    max_image_size=req.get('maxImageSize', 768),
                 )
 
             elif method == 'render':