@glassly/engine 0.3.5-beta.710 → 0.3.5-beta.760

package/src/services/LocalLlmModelManager.ts

@@ -24,8 +24,10 @@ export interface LocalLlmModelConfig {
   id: LocalLlmModelId
   /** Human label for the Settings dropdown. */
   label: string
-  /** Approximate download size in bytes (the .pte dominates). */
+  /** Total download size in bytes (the .pte dominates). */
   sizeBytes: number
+  /** Exact size of the published .pte; the availability check compares against it. */
+  modelBytes: number
   /** Rough resident memory while loaded; the Settings copy warns above ~half of device RAM. */
   ramBytes: number
   urls: {
@@ -54,8 +56,10 @@ export interface LocalLlmDownloadStatus {
 
 // XNNPACK 8da4w exports from software-mansion/react-native-executorch-qwen-3,
 // pinned to the v0.10.0 release tag that react-native-executorch 0.10.x
-// expects. Sizes are the .pte file sizes on the Hub (tokenizer adds ~11 MB).
+// expects. modelBytes are the .pte sizes on the Hub; sizeBytes adds the
+// tokenizer files (~11 MB).
 const QWEN3_BASE = "https://huggingface.co/software-mansion/react-native-executorch-qwen-3/resolve/v0.10.0"
+const TOKENIZER_BYTES = 11_431_080
 const MB = 1024 * 1024
 const GB = 1024 * MB
 
@@ -63,7 +67,8 @@ const MODELS: Record<LocalLlmModelId, LocalLlmModelConfig> = {
   "qwen3-0.6b": {
     id: "qwen3-0.6b",
     label: "Qwen3 0.6B",
-    sizeBytes: 517 * MB,
+    sizeBytes: 505_686_400 + TOKENIZER_BYTES,
+    modelBytes: 505_686_400,
     ramBytes: 0.9 * GB,
     urls: {
       model: `${QWEN3_BASE}/0_6b/xnnpack/qwen_3_0_6b_xnnpack_8da4w.pte`,
@@ -74,7 +79,8 @@ const MODELS: Record<LocalLlmModelId, LocalLlmModelConfig> = {
   "qwen3-1.7b": {
     id: "qwen3-1.7b",
     label: "Qwen3 1.7B",
-    sizeBytes: 1315 * MB,
+    sizeBytes: 1_303_738_496 + TOKENIZER_BYTES,
+    modelBytes: 1_303_738_496,
     ramBytes: 1.8 * GB,
     urls: {
       model: `${QWEN3_BASE}/1_7b/xnnpack/qwen_3_1_7b_xnnpack_8da4w.pte`,
@@ -85,7 +91,8 @@ const MODELS: Record<LocalLlmModelId, LocalLlmModelConfig> = {
   "qwen3-4b": {
     id: "qwen3-4b",
     label: "Qwen3 4B",
-    sizeBytes: 2693 * MB,
+    sizeBytes: 2_681_646_976 + TOKENIZER_BYTES,
+    modelBytes: 2_681_646_976,
     ramBytes: 3.4 * GB,
     urls: {
       model: `${QWEN3_BASE}/4b/xnnpack/qwen_3_4b_xnnpack_8da4w.pte`,
@@ -161,7 +168,7 @@ class LocalLlmModelManager {
       // most of the published size. Downloads land in a .part file and are
       // renamed only once whole, so this only catches external tampering.
       const stat = await RNFS.stat(paths.model)
-      return Number(stat.size) >= MODELS[id].sizeBytes * 0.95
+      return Number(stat.size) >= MODELS[id].modelBytes * 0.95
     } catch (error) {
       console.warn("LOCAL_LLM: availability check failed:", error)
       return false
@@ -243,6 +250,9 @@ class LocalLlmModelManager {
       await this.fetchFile(config.urls.model, paths.model, (percent) =>
         this.setStatus({modelId: id, stage: "downloading", percent}),
       )
+      if (!(await this.isModelAvailable(id))) {
+        throw new Error(`downloaded model is smaller than the expected ${config.modelBytes} bytes`)
+      }
       this.setStatus({modelId: id, stage: "complete", percent: 100})
       for (const listener of this.modelReadyListeners) {
         try {

package/src/services/LocalLlmPrompt.ts

@@ -85,55 +85,74 @@ export function toTemplateTools(tools?: LlmTool[]): TemplateTool[] | undefined {
   }))
 }
 
-/** Rough tokens-per-character for a BPE vocabulary over mixed prose/JSON. */
-const CHARS_PER_TOKEN = 3.2
-/** Template framing per message (`<|im_start|>role … <|im_end|>`). */
-const TOKENS_PER_MESSAGE = 6
-
-export function estimateTokens(text: string): number {
-  return Math.ceil(text.length / CHARS_PER_TOKEN)
-}
-
-export function estimateMessageTokens(message: TemplateMessage): number {
-  let total = TOKENS_PER_MESSAGE + estimateTokens(message.content)
-  if (message.role === "assistant" && message.tool_calls) {
-    for (const call of message.tool_calls) total += estimateTokens(JSON.stringify(call.function)) + 8
-  }
-  return total
+export interface FittedLocalPrompt {
+  prompt: string
+  promptTokens: number
+  truncated: boolean
 }
 
-const MIN_TOOL_CONTENT_CHARS = 1_200
-
 /**
- * Fit the conversation into `budgetTokens`. History is grouped into turns
- * (a user message plus everything up to the next user message, so a tool
- * round stays with the question it answers); the oldest turns go first and
- * the last turn is never dropped — the model must at least see the question.
- * When that last turn alone still does not fit, tool outputs are cut to
- * their head.
+ * Measure the complete rendered prompt with the model's tokenizer, including
+ * template framing, system instructions and tool schemas. Drop old turns first,
+ * then shorten tool results / message text in the final turn. Never cut template
+ * delimiters, system instructions, or structured tool calls to make them fit.
  */
-export function fitToBudget(messages: TemplateMessage[], budgetTokens: number, toolTokens = 0): TemplateMessage[] {
-  const total = (list: TemplateMessage[]) => list.reduce((sum, m) => sum + estimateMessageTokens(m), toolTokens)
+export async function fitPromptToContext(
+  messages: TemplateMessage[],
+  budgetTokens: number,
+  render: (messages: TemplateMessage[]) => string,
+  countTokens: (prompt: string) => Promise<number>,
+): Promise<FittedLocalPrompt> {
+  if (!Number.isInteger(budgetTokens) || budgetTokens < 1) {
+    throw new Error("On-device model has no room for a prompt after reserving output tokens")
+  }
   const system = messages[0]?.role === "system" ? [messages[0]] : []
   const turns: TemplateMessage[][] = []
   for (const message of messages.slice(system.length)) {
     if (message.role === "user" || turns.length === 0) turns.push([message])
     else turns[turns.length - 1]!.push(message)
   }
-
-  while (turns.length > 1 && total([...system, ...turns.flat()]) > budgetTokens) {
-    turns.shift()
-  }
-
   let fitted = [...system, ...turns.flat()]
-  if (total(fitted) > budgetTokens) {
-    fitted = fitted.map((m) =>
-      m.role === "tool" && m.content.length > MIN_TOOL_CONTENT_CHARS
-        ? {...m, content: `${m.content.slice(0, MIN_TOOL_CONTENT_CHARS)}\n…[truncated]`}
-        : m,
-    )
+  let truncated = false
+  const marker = "\n…[truncated]…\n"
+  while (true) {
+    const prompt = render(fitted)
+    const promptTokens = await countTokens(prompt)
+    if (!Number.isInteger(promptTokens) || promptTokens < 1) {
+      throw new Error("On-device tokenizer returned an invalid prompt token count")
+    }
+    if (promptTokens <= budgetTokens) return {prompt, promptTokens, truncated}
+    truncated = true
+    if (turns.length > 1) {
+      turns.shift()
+      fitted = [...system, ...turns.flat()]
+      continue
+    }
+    // Preserve the user's question while tool output can still be shortened.
+    const candidates = fitted
+      .map((message, index) => ({message, index}))
+      .filter(({message}) => message.role !== "system" && Array.from(message.content).length > marker.length * 2)
+      .sort(
+        (a, b) =>
+          Number(b.message.role === "tool") - Number(a.message.role === "tool") ||
+          b.message.content.length - a.message.content.length,
+      )
+    const candidate = candidates[0]
+    if (!candidate) {
+      throw new Error(
+        `On-device prompt instructions and tool definitions exceed the context budget (${promptTokens} > ${budgetTokens} tokens)`,
+      )
+    }
+    // Keep both ends of long transcripts and questions, without splitting a
+    // Unicode code point. Each pass shrinks strictly, then remeasures the whole
+    // rendered prompt; character ratios are never used as a safety check.
+    const chars = Array.from(candidate.message.content)
+    const keep = Math.floor(chars.length / 4)
+    fitted[candidate.index] = {
+      ...candidate.message,
+      content: `${chars.slice(0, keep).join("")}${marker}${chars.slice(-keep).join("")}`,
+    }
   }
-  return fitted
 }
 
 /** Remove a leading `<think>…</think>` block (or an unterminated one). */

package/src/services/LocalLlmProvider.ts

@@ -22,9 +22,8 @@ import {BgTimer} from "../utils/timers"
 import localLlmModelManager, {type LocalLlmModelId} from "./LocalLlmModelManager"
 import {
   buildTemplateMessages,
-  estimateTokens,
   finalizeLocalResult,
-  fitToBudget,
+  fitPromptToContext,
   toTemplateTools,
   type LocalCompletion,
 } from "./LocalLlmPrompt"
@@ -33,6 +32,7 @@ import {SETTINGS, useSettingsStore} from "../stores/settings"
 type Rne = typeof import("react-native-executorch")
 // The runner API lives on the `llm` namespace; derive its types from there so
 // the import stays a type-only reference (no native module at load time).
+type Tokenizer = ReturnType<Rne["nlp"]["loadTokenizer"]>
 type LLMRunner = ReturnType<Rne["llm"]["createLLMRunner"]>
 type LLMGenerationStats = ReturnType<LLMRunner["generate"]>
 type LLMGenerationConfig = NonNullable<Parameters<LLMRunner["generate"]>[1]>
@@ -45,12 +45,14 @@ const IDLE_UNLOAD_MS = 5 * 60_000
 const GENERATION_DEADLINE_MS = 100_000
 const DEFAULT_MAX_NEW_TOKENS = 512
 const MAX_NEW_TOKENS_CEILING = 1_024
-/** Headroom between the estimated prompt and the context end; estimates are rough. */
+/** Reserve space for runner-added BOS/EOS tokens and context boundary bookkeeping. */
 const CONTEXT_MARGIN_TOKENS = 96
 
 interface LoadedModel {
   modelId: LocalLlmModelId
   runner: LLMRunner
+  tokenizer: Tokenizer
+  countTokens: (tokenizer: Tokenizer, prompt: string) => Promise<number>
   template: Template
   stopTokens: string[]
   maxSeqLen: number
@@ -85,6 +87,11 @@ function generateWorklet(
   return {text, stats}
 }
 
+function countTokensWorklet(tokenizer: Tokenizer, prompt: string): number {
+  "worklet"
+  return tokenizer.encode(prompt).length
+}
+
 class LocalLlmProvider {
   private static instance: LocalLlmProvider
 
@@ -150,6 +157,8 @@ class LocalLlmProvider {
       loaded.runner.dispose()
     } catch (error) {
       console.warn(`${LOG_TAG}: runner dispose failed:`, error)
+    } finally {
+      loaded.tokenizer.dispose()
     }
   }
 
@@ -159,21 +168,21 @@ class LocalLlmProvider {
       const model = await this.ensureLoaded()
       const jsonMode = request.jsonMode === true
       const tools = toTemplateTools(request.tools)
-      const toolTokens = tools ? estimateTokens(JSON.stringify(tools)) + 64 : 0
 
       const maxNewTokens = Math.max(
         16,
         Math.min(request.maxTokens ?? DEFAULT_MAX_NEW_TOKENS, MAX_NEW_TOKENS_CEILING, Math.floor(model.maxSeqLen / 4)),
       )
       const promptBudget = model.maxSeqLen - maxNewTokens - CONTEXT_MARGIN_TOKENS
-      const messages = fitToBudget(buildTemplateMessages({...request, jsonMode}), promptBudget, toolTokens)
-
-      const prompt = model.template.render({
-        messages,
-        tools,
-        add_generation_prompt: true,
-        enable_thinking: false,
-      })
+      const {prompt, promptTokens, truncated} = await fitPromptToContext(
+        buildTemplateMessages({...request, jsonMode}),
+        promptBudget,
+        (messages) => model.template.render({messages, tools, add_generation_prompt: true, enable_thinking: false}),
+        (text) => model.countTokens(model.tokenizer, text),
+      )
+      console.log(
+        `${LOG_TAG}: ${model.modelId} prompt=${promptTokens} budget=${promptBudget} context=${model.maxSeqLen} maxNewTokens=${maxNewTokens} truncated=${truncated}`,
+      )
 
       model.runner.reset(0)
       const deadline = BgTimer.setTimeout(() => {
@@ -189,6 +198,12 @@ class LocalLlmProvider {
           {maxNewTokens, temperature: request.temperature ?? 0.7},
           model.stopTokens,
         )
+      } catch (error) {
+        console.warn(
+          `${LOG_TAG}: generation failed: model=${model.modelId} prompt=${promptTokens} context=${model.maxSeqLen} maxNewTokens=${maxNewTokens}`,
+          error,
+        )
+        throw error
       } finally {
         BgTimer.clearTimeout(deadline)
       }
@@ -246,16 +261,28 @@ class LocalLlmProvider {
     const startedAt = Date.now()
     const createRunner = rne.wrapAsync(rne.llm.createLLMRunner)
     const runner = await createRunner(paths.model, paths.tokenizer, [])
-    const maxSeqLen = runner.getKVCacheState().maxSeqLen
-    console.log(`${LOG_TAG}: loaded ${modelId} (context ${maxSeqLen}) in ${Date.now() - startedAt} ms`)
-
-    return {
-      modelId,
-      runner,
-      template,
-      stopTokens: [...tokenizerConfig.stopTokens],
-      maxSeqLen,
-      generate: rne.wrapAsync(generateWorklet),
+    let tokenizer: Tokenizer | undefined
+    try {
+      const maxSeqLen = runner.getKVCacheState().maxSeqLen
+      if (!Number.isInteger(maxSeqLen) || maxSeqLen <= CONTEXT_MARGIN_TOKENS + 16) {
+        throw new Error(`On-device model ${modelId} reported an invalid context size: ${maxSeqLen}`)
+      }
+      tokenizer = await rne.wrapAsync(rne.nlp.loadTokenizer)(paths.tokenizer)
+      console.log(`${LOG_TAG}: loaded ${modelId} (context ${maxSeqLen}) in ${Date.now() - startedAt} ms`)
+      return {
+        modelId,
+        runner,
+        tokenizer,
+        countTokens: rne.wrapAsync(countTokensWorklet),
+        template,
+        stopTokens: [...tokenizerConfig.stopTokens],
+        maxSeqLen,
+        generate: rne.wrapAsync(generateWorklet),
+      }
+    } catch (error) {
+      tokenizer?.dispose()
+      runner.dispose()
+      throw error
     }
   }