@unblocklabs/unblock-memory 0.3.30 → 0.5.0

package/src/query-tokenizer/tokenizer_config.json

@@ -0,0 +1,19 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|startoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "is_local": true,
+  "legacy": false,
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<|pad|>",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "TokenizersBackend",
+  "use_default_system_prompt": false,
+  "use_fast": true
+}

package/src/query_contract.py

@@ -0,0 +1,60 @@
+"""One student prompt, pinned tokenizer, serialization and target contract for v2."""
+import hashlib
+import json
+from pathlib import Path
+
+ASSETS = Path(__file__).with_name("query-tokenizer")
+CONTRACT = json.loads((ASSETS / "contract.json").read_text())
+MODEL, REVISION, SYSTEM = (CONTRACT[key] for key in ("model", "revision", "system"))
+
+
+def serialize_conversation(conversation):
+    return json.dumps({"history": conversation["history"], "currentRequest": conversation["currentRequest"]},
+                      ensure_ascii=False, separators=(",", ":")).replace("<", "\\u003c")
+
+
+def validate_conversation(conversation, tokenizer):
+    if not isinstance(conversation, dict) or set(conversation) != {"history", "currentRequest"} or \
+            not isinstance(conversation["currentRequest"], str) or not conversation["currentRequest"] or \
+            not isinstance(conversation["history"], list):
+        raise ValueError("Invalid prepared conversation")
+    for message in conversation["history"]:
+        if not isinstance(message, dict) or set(message) != {"role", "content"} or \
+                message["role"] not in ("user", "assistant") or not isinstance(message["content"], str):
+            raise ValueError("Invalid prepared conversation history")
+    serialized = serialize_conversation(conversation)
+    if len(serialized.encode("utf-8")) > CONTRACT["conversationBytes"] or \
+            len(tokenizer.encode(serialized, add_special_tokens=False)) > CONTRACT["conversationTokens"]:
+        raise ValueError("Prepared conversation exceeds shared window")
+
+
+def validate_query_pair(value):
+    if not isinstance(value, dict) or set(value) != {"lex", "vec"} or \
+            any(not isinstance(value[key], str) or not value[key].strip() for key in ("lex", "vec")):
+        raise ValueError("Expected exactly one nonempty lex/vec query pair")
+    return {key: value[key].strip() for key in ("lex", "vec")}
+
+
+def student_messages(conversation):
+    return [{"role": "system", "content": SYSTEM},
+            {"role": "user", "content": "<conversation_data>\n" + serialize_conversation(conversation) + "\n</conversation_data>"}]
+
+
+def student_prompt_tokens(tokenizer, conversation):
+    validate_conversation(conversation, tokenizer)
+    # Transformers 5 may return BatchEncoding for tokenize=True. Encode the
+    # rendered template explicitly, exactly as dataset preparation does.
+    text = tokenizer.apply_chat_template(student_messages(conversation), tokenize=False, add_generation_prompt=True)
+    return tokenizer.encode(text, add_special_tokens=False)
+
+
+def load_tokenizer():
+    from transformers import AutoTokenizer
+    return AutoTokenizer.from_pretrained(ASSETS, local_files_only=True, trust_remote_code=False)
+
+
+def verify_model_tokenizer(root):
+    # The model package must carry the exact tokenizer/template used for collection.
+    for name in ("tokenizer.json", "tokenizer_config.json", "chat_template.jinja"):
+        if hashlib.sha256((root / name).read_bytes()).digest() != hashlib.sha256((ASSETS / name).read_bytes()).digest():
+            raise ValueError("Incompatible model tokenizer or chat template")

package/dist/src/mlx-query.d.ts

@@ -1,22 +0,0 @@
-import type { UnblockMemoryConfig } from "./config.js";
-type Conversation = {
-    history: {
-        role: "user" | "assistant";
-        content: string;
-    }[];
-    currentRequest: string;
-};
-/** Preserve whole visible messages and the complete current request, never tool/thinking text. */
-export declare function queryConversation(prompt: string, messages: readonly unknown[], historyMessages: number): Conversation;
-export declare class MlxQueryGenerator {
-    #private;
-    private readonly config;
-    /** Gateway hook and tool registries can evaluate/register this plugin separately. */
-    static shared(config: NonNullable<UnblockMemoryConfig["memoryWhisperer"]["mlx"]>): MlxQueryGenerator;
-    constructor(config: NonNullable<UnblockMemoryConfig["memoryWhisperer"]["mlx"]>);
-    get closed(): boolean;
-    start(): Promise<void>;
-    generate(conversation: Conversation, signal: AbortSignal): Promise<string[]>;
-    close(): void;
-}
-export {};

package/dist/src/mlx-query.js

@@ -1,155 +0,0 @@
-import { spawn } from "node:child_process";
-import { randomUUID } from "node:crypto";
-import { existsSync } from "node:fs";
-import { fileURLToPath } from "node:url";
-import { createInterface } from "node:readline";
-import { abortable } from "./abortable.js";
-import { messageText } from "./whisperer-context.js";
-import { conversationUserText } from "./response-text.js";
-/** Preserve whole visible messages and the complete current request, never tool/thinking text. */
-export function queryConversation(prompt, messages, historyMessages) {
-    const currentRequest = conversationUserText(prompt)?.text ?? "";
-    const visible = messages.flatMap(message => {
-        const text = messageText(message);
-        const content = text?.role === "user" ? conversationUserText(text.text)?.text : text?.text;
-        return text && content ? [{ role: text.role, content }] : [];
-    });
-    if (visible.at(-1)?.role === "user" && visible.at(-1)?.content === currentRequest)
-        visible.pop();
-    const history = historyMessages ? visible.slice(-historyMessages) : [];
-    while (history.length && Buffer.byteLength(JSON.stringify({ history, currentRequest })) > 24_000)
-        history.shift();
-    if (!currentRequest)
-        throw new Error("Missing or unparseable current request");
-    if (Buffer.byteLength(JSON.stringify({ history, currentRequest })) > 24_000) {
-        throw new Error("Query input exceeds context budget");
-    }
-    return { history, currentRequest };
-}
-export class MlxQueryGenerator {
-    config;
-    /** Gateway hook and tool registries can evaluate/register this plugin separately. */
-    static shared(config) {
-        const scope = globalThis;
-        const symbol = Symbol.for("unblock-memory.mlx-workers.v1");
-        const workers = (scope[symbol] ??= new Map());
-        const key = JSON.stringify([config.pythonPath, config.modelPath]);
-        let worker = workers.get(key);
-        if (!worker || worker.closed) {
-            worker = new MlxQueryGenerator(config);
-            workers.set(key, worker);
-        }
-        return worker;
-    }
-    #child;
-    #ready;
-    #closed = false;
-    #retryAfter = 0;
-    #pending = new Map();
-    constructor(config) {
-        this.config = config;
-    }
-    get closed() { return this.#closed; }
-    start() {
-        if (this.#closed)
-            return Promise.reject(new Error("Query worker stopped"));
-        if (this.#ready)
-            return this.#ready;
-        if (Date.now() < this.#retryAfter)
-            return Promise.reject(new Error("Query worker restarting"));
-        const adjacent = new URL("./mlx-query-worker.py", import.meta.url);
-        const script = existsSync(adjacent) ? adjacent : new URL("../../src/mlx-query-worker.py", import.meta.url);
-        const child = spawn(this.config.pythonPath, ["-u", fileURLToPath(script), this.config.modelPath], {
-            stdio: "pipe", env: { ...process.env, HF_HUB_OFFLINE: "1", TRANSFORMERS_OFFLINE: "1", TOKENIZERS_PARALLELISM: "false" },
-        });
-        this.#child = child;
-        // No source text from Python warnings/exceptions enters Gateway logs.
-        child.stderr.resume();
-        const lines = createInterface({ input: child.stdout });
-        this.#ready = new Promise((resolve, reject) => {
-            const timer = setTimeout(() => child.kill(), 60_000);
-            const fail = () => {
-                clearTimeout(timer);
-                lines.close();
-                reject(new Error("Query worker unavailable"));
-                if (this.#child !== child)
-                    return;
-                this.#child = undefined;
-                this.#ready = undefined;
-                this.#retryAfter = Date.now() + 2000;
-                for (const request of this.#pending.values())
-                    request.reject(new Error("Query worker unavailable"));
-                this.#pending.clear();
-            };
-            child.once("error", fail);
-            child.once("close", fail);
-            child.stdin.on("error", () => child.kill());
-            lines.on("line", line => {
-                try {
-                    const value = JSON.parse(line);
-                    if (!value || typeof value !== "object")
-                        throw new Error("Invalid worker response");
-                    if ("ready" in value && value.ready === true) {
-                        clearTimeout(timer);
-                        resolve();
-                        return;
-                    }
-                    if (!("id" in value) || typeof value.id !== "string")
-                        throw new Error("Missing worker request ID");
-                    const request = this.#pending.get(value.id);
-                    if (!request)
-                        return; // Cancelled request; never surface late results.
-                    this.#pending.delete(value.id);
-                    if (!("text" in value) || typeof value.text !== "string") {
-                        request.reject(new Error("Query generation failed"));
-                        return;
-                    }
-                    let queries;
-                    try {
-                        queries = JSON.parse(value.text).queries;
-                    }
-                    catch { /* reject below */ }
-                    const usable = Array.isArray(queries) ? [...new Set(queries
-                            .filter((query) => typeof query === "string")
-                            .map(query => query.trim()).filter(Boolean))] : [];
-                    if (!usable.length) {
-                        request.reject(new Error("No usable generated queries"));
-                        return;
-                    }
-                    request.resolve(usable);
-                }
-                catch {
-                    child.kill();
-                }
-            });
-        });
-        return this.#ready;
-    }
-    async generate(conversation, signal) {
-        await abortable(this.start(), signal);
-        signal.throwIfAborted();
-        if (this.#pending.size >= 8)
-            throw new Error("Query worker busy");
-        const id = randomUUID();
-        const child = this.#child;
-        const result = new Promise((resolve, reject) => this.#pending.set(id, { resolve, reject }));
-        const cancel = () => {
-            this.#pending.get(id)?.reject(new Error("Query generation cancelled"));
-            this.#pending.delete(id);
-            if (child.stdin.writable)
-                child.stdin.write(JSON.stringify({ id, cancel: true }) + "\n");
-        };
-        signal.addEventListener("abort", cancel, { once: true });
-        child.stdin.write(JSON.stringify({ id, conversation }) + "\n");
-        try {
-            return await result;
-        }
-        finally {
-            signal.removeEventListener("abort", cancel);
-        }
-    }
-    close() {
-        this.#closed = true;
-        this.#child?.kill();
-    }
-}