@modular-prompt/driver 0.16.0 → 0.17.0

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

@@ -52,33 +52,103 @@ class ModelBackend(ABC):
         trim_to_tokens: int | None = None,
         prefix_offsets: list[int] | None = None,
         prefix_hashes: list[str] | None = None,
+        images: list | None = None,
+        max_image_size: int = 768,
     ) -> dict:
         """Build a KV cache from a prompt prefix."""
         raise NotImplementedError(
             f"{type(self).__name__} does not support prompt caching"
         )
 
-    def load_cache_from_file(self, cache_path: str) -> list | None:
-        """Load a prompt cache from file, or None."""
+    def tokenize_prompt(
+        self,
+        prompt: str,
+        images: list | None = None,
+        max_image_size: int = 768,
+    ) -> list[int]:
+        """Tokenize a rendered prompt using this backend's prompt rules.
+
+        A VLM backend returns a processor from ``get_tokenizer()``.  Keeping
+        this operation on the backend prevents shared handlers from assuming
+        that the returned object is a plain tokenizer.  ``images`` and
+        ``max_image_size`` let multimodal backends reproduce the same image
+        placeholder expansion as their generation path; text-only backends
+        ignore them.
+        """
+        tokenizer = self.get_tokenizer()
+        tokenizer = getattr(tokenizer, "tokenizer", tokenizer)
+        bos_token = getattr(tokenizer, "bos_token", None)
+        add_special = bos_token is None or not prompt.startswith(bos_token or "")
+        return list(tokenizer.encode(prompt, add_special_tokens=add_special))
+
+    def trim_cache(self, prompt_cache: list, tokens: int) -> None:
+        """Trim a backend-owned prompt cache in place.
+
+        Cache implementations are backend-specific.  The default handles the
+        cache objects exposed by mlx-vlm without importing mlx-lm into the
+        shared generate handler.
+        """
+        if tokens <= 0:
+            return
+
+        for entry in prompt_cache:
+            trim = getattr(entry, "trim", None)
+            if callable(trim):
+                trim(tokens)
+                continue
+
+            children = getattr(entry, "caches", None)
+            if children is not None:
+                for child in children:
+                    child_trim = getattr(child, "trim", None)
+                    if callable(child_trim):
+                        child_trim(tokens)
+
+    def load_cache_from_file(
+        self,
+        cache_path: str,
+        images: list | None = None,
+        max_image_size: int = 768,
+        prompt: str | list[int] | None = None,
+    ) -> list | None:
+        """Load a prompt cache from file, or None.
+
+        ``prompt`` is optional so backends that persist prompt identities can
+        validate the loaded cache against the current request.
+        """
         return None
 
     def get_cache_offset(self, prompt_cache: list) -> int:
         """Get the number of tokens stored in a loaded prompt cache."""
         if not prompt_cache:
             return 0
-        layer0 = prompt_cache[0]
-        if hasattr(layer0, 'offset'):
-            off = layer0.offset
-            return int(off.item() if hasattr(off, 'item') else off)
-        if hasattr(layer0, 'caches'):
-            for c in layer0.caches:
-                if hasattr(c, 'offset'):
-                    off = c.offset
-                    return int(off.item() if hasattr(off, 'item') else off)
-        try:
-            return int(layer0[0].shape[2])
-        except Exception:
-            pass
-        if hasattr(layer0, 'keys') and layer0.keys is not None:
-            return int(layer0.keys.shape[2])
-        return 0
+
+        def offset_of(cache: Any) -> int:
+            if hasattr(cache, "offset"):
+                off = cache.offset
+                return int(off.item() if hasattr(off, "item") else off)
+
+            size = getattr(cache, "size", None)
+            if callable(size):
+                try:
+                    return int(size())
+                except Exception:
+                    pass
+
+            children = getattr(cache, "caches", None)
+            if children is not None:
+                return max((offset_of(child) for child in children), default=0)
+
+            keys = getattr(cache, "keys", None)
+            if keys is not None:
+                try:
+                    return int(keys.shape[-2])
+                except Exception:
+                    pass
+
+            try:
+                return int(cache[0].shape[2])
+            except Exception:
+                return 0
+
+        return max((offset_of(layer) for layer in prompt_cache), default=0)

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

@@ -78,12 +78,29 @@ class MlxLmBackend(ModelBackend):
 
     # get_cache_offset is inherited from ModelBackend base class
 
-    def _tokenize_prompt(self, prompt: str) -> list[int]:
+    def _tokenize_prompt(
+        self,
+        prompt: str,
+        images: list | None = None,
+        max_image_size: int = 768,
+    ) -> 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)
+        return list(self.tokenizer.encode(prompt, add_special_tokens=add_special))
+
+    def tokenize_prompt(
+        self,
+        prompt: str,
+        images: list | None = None,
+        max_image_size: int = 768,
+    ) -> list[int]:
+        return self._tokenize_prompt(prompt, images, max_image_size)
+
+    def trim_cache(self, prompt_cache: list, tokens: int) -> None:
+        if tokens > 0:
+            trim_prompt_cache(prompt_cache, tokens)
 
     @staticmethod
     def _write_cache_meta(
@@ -122,6 +139,8 @@ class MlxLmBackend(ModelBackend):
         trim_to_tokens: int | None = None,
         prefix_offsets: list[int] | None = None,
         prefix_hashes: list[str] | None = None,
+        images: list | None = None,
+        max_image_size: int = 768,
     ) -> dict:
         if self.model is None or self.tokenizer is None:
             raise RuntimeError("Model is not loaded")
@@ -200,7 +219,13 @@ class MlxLmBackend(ModelBackend):
             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:
+    def load_cache_from_file(
+        self,
+        cache_path: str,
+        images: list | None = None,
+        max_image_size: int = 768,
+        prompt: str | list[int] | None = None,
+    ) -> list | None:
         try:
             return load_prompt_cache(cache_path)
         except FileNotFoundError: