feat: implement thinking budget configuration for Gemma-4 think mode

[lokhor]

Problem

Gemma-4 E4B natively supports a "thinking" mode where it generates internal reasoning tokens (<think>...</think>) before producing its visible response. Our app already captures these tokens:

// LiteRtInferenceEngine.kt, lines 180-183
val thought = message.channels["thought"]
if (thought != null) {
    // Routes as GenerationResult.Thinking
}

However, we have zero control over this behavior:

1. No thinking budget — we can't tell the model to use more or fewer thinking tokens. It uses whatever its default is.
2. No on/off toggle — users can't disable thinking mode for simple queries (wastes tokens and adds latency) or enable it for complex reasoning tasks.
3. No UI configuration — ModelSettingsEntity has no thinking-related fields. The settings screen has no thinking controls.
4. Unknown API surface — we need to investigate what LiteRT-LM actually exposes for thinking budget configuration.

Impact

• TTFT (Time to First Token): Thinking adds significant latency because the model generates reasoning tokens before the visible response. For simple queries like "turn on torch" or "what time is it", this is wasted time.
• Token budget: Thinking tokens consume context window space. On a 4096-token window, even 200 thinking tokens is meaningful.
• Quality: For complex queries (multi-step reasoning, math, planning), thinking dramatically improves output quality. We should enable users to leverage this.

Research Needed

LiteRT-LM Thinking Configuration API

Before implementation, investigate:

1. Does LiteRT-LM support thinkingConfig / thinkingBudget?

   • Check the LiteRT-LM Kotlin SDK for configuration options when creating/initializing the engine
   • Look for ThinkingConfig, thinkingBudget, enableThinking, or similar parameters
   • Reference: https://github.com/AIDevelopersMonster/LiteRT-LM (check Kotlin API)

2. How does Gemma-4 E4B handle thinking natively?

   • The model likely uses <think> tags in its tokenizer vocabulary
   • Can we suppress thinking by including "Do not think" in the system prompt? (crude but may work)
   • Is there a generation config parameter that controls thinking token budget?

3. What does Edge Gallery do?

   • Check if AI Edge Gallery configures thinking budget in its LlmChatModelHelper.initialize()
   • Look for any thinkingConfig in their model initialization code

Implementation Plan

Phase A: Research + Basic Toggle

1. Investigate LiteRT-LM API — document what's available for thinking configuration:

   • Check LlmModelConfig, SamplerConfig, or engine initialization params
   • File findings in a comment on this issue

2. Add thinkingEnabled: Boolean to ModelSettingsEntity:

   @Entity(tableName = "model_settings")
   data class ModelSettingsEntity(
       @PrimaryKey val modelName: String,
       val contextWindowSize: Int = 4096,
       val temperature: Float = 0.7f,
       val topP: Float = 0.95f,
       val topK: Int = 40,
       val thinkingEnabled: Boolean = true  // NEW
   )

   • Room migration required

3. Add thinking toggle to ModelSettingsScreen.kt:

   • Switch: "Enable thinking mode"
   • Description: "Model reasons internally before responding. Improves quality but adds latency."
   • Default: ON (for Gemma-4 models that support it)

4. Apply toggle in LiteRtInferenceEngine:

   • If thinkingEnabled == false:
     • Option A (API): Pass thinkingBudget = 0 to LiteRT-LM (if supported)
     • Option B (prompt): Prepend "Do not use internal reasoning" to system prompt
     • Option C (filter): Suppress channels["thought"] output (model still thinks, we just ignore it — least useful)

Phase B: Budget Control (if API supports it)

5. Add thinkingBudget: Int to ModelSettingsEntity (nullable, null = auto):

   val thinkingBudget: Int? = null  // null = model default, 0 = off, 100-2000 = custom

6. Add thinking budget slider to settings UI:

   • Only visible when thinking toggle is ON
   • Range: 0-2000 tokens (step 50)
   • Special values: "Auto" (null/default), "Off" (0)
   • Label: "Thinking budget (tokens)"

7. Pass budget to LiteRT-LM during inference:

   // Pseudocode — actual API TBD from research
   val thinkingConfig = ThinkingConfig(
       enabled = settings.thinkingEnabled,
       budget = settings.thinkingBudget ?: -1  // -1 = auto
   )
   engine.configure(thinkingConfig = thinkingConfig)

Phase C: Smart Defaults

8. Auto-detect thinking capability — not all models support thinking:

   • If model doesn't emit channels["thought"] tokens, hide thinking settings
   • Show thinking settings only for compatible models (Gemma-4 E4B, future thinking models)

9. Task-aware thinking (advanced, optional):

   • Disable thinking for simple tool calls (alarm, torch, timer)
   • Enable thinking for complex queries (math, reasoning, planning)
   • Implementation: lightweight classifier before inference, or just use low thinking budget for tool-call turns
   • Decision needed: worth the complexity? May defer.

UI Mockup

Model Settings
├── Context Window: [====|========] 4096
├── Temperature:    [===|=========] 0.7
├── Top-P:          [========|====] 0.95
├── Top-K:          [====|========] 40
├── ── Presets: [Precise] [Balanced] [Creative]
├── 
├── Thinking Mode
│   ├── [✓] Enable thinking mode
│   ├── Budget: [Auto ▼] / [=====|=====] 500 tokens
│   └── ℹ️ "Model reasons internally before responding"

Files to Modify

File	Changes
ModelSettingsEntity.kt	Add thinkingEnabled, thinkingBudget fields
ModelSettingsDao.kt	Room migration
ModelSettingsScreen.kt	Add thinking toggle + budget slider
ChatViewModel.kt	Pass thinking config to engine
LiteRtInferenceEngine.kt	Apply thinking configuration to LiteRT-LM
ModelConfig.kt	Add thinking defaults

Success Criteria

• LiteRT-LM thinking API documented (or workaround identified)
• Thinking toggle works — OFF disables thinking tokens, ON enables them
• Budget slider controls thinking depth (if API supports it)
• TTFT measurably faster with thinking OFF on simple queries
• Quality measurably better with thinking ON for complex reasoning
• Settings hidden for models that don't support thinking

Related

• research: review skill loading and system prompt architecture vs AI Edge Gallery #340 (parent research issue)
• feat: implement lazy skill loading (load_skill pattern) to reduce system prompt bloat #341 (skill loading — shorter prompt + thinking = cleaner architecture)
• feat: expand model settings — topK, temperature presets, and per-task optimization #342 (model settings — thinking is part of the settings expansion)
• Phase 4 roadmap (advanced inference features)
• LiteRT-LM SDK: https://github.com/AIDevelopersMonster/LiteRT-LM

[lokhor]

🔬 Deep Research: LiteRT-LM Thinking API & Implementation Feasibility

Full analysis of the LiteRT-LM 0.10.0 SDK, thinking token handling, and implementation options.

Current Thinking Token Support — What Works

✅ Thinking tokens ARE being generated and captured:

// LiteRtInferenceEngine.kt, onMessage() callback
val thinkingText = message.channels["thought"]
if (!thinkingText.isNullOrEmpty()) {
    trySend(GenerationResult.Thinking(thinkingText))
}

The Gemma-4 model natively generates thinking tokens, and LiteRT-LM routes them through message.channels["thought"]. The app already captures and displays these via GenerationResult.Thinking.

LiteRT-LM 0.10.0 API Surface — What Is Available

From the full SDK imports and usage analysis:

// Available configuration objects:
EngineConfig(modelPath, backend, maxNumTokens, cacheDir)
SamplerConfig(topK, topP, temperature)
ConversationConfig(samplerConfig, systemInstruction)

// Message callback:
message.channels["thought"]  // ← thinking tokens
message.toString()            // ← visible output tokens

❌ NO thinking configuration parameters found in the 0.10.0 API:

• No ThinkingConfig class
• No thinkingBudget parameter on ConversationConfig
• No enableThinking flag
• No maxThinkingTokens on EngineConfig

The thinking behavior is entirely model-controlled — when Gemma-4 decides to think, it thinks. The SDK just routes the tokens.

Three Implementation Options

Option A: SDK Upgrade (Best, but blocked)

If a future LiteRT-LM version (0.11.0+) exposes ThinkingConfig, this becomes trivial:

ConversationConfig(
    samplerConfig = samplerConfig,
    systemInstruction = systemInstruction,
    thinkingConfig = ThinkingConfig(budget = settings.thinkingBudget)  // hypothetical
)

Status: ⏸️ Blocked — need to monitor LiteRT-LM releases. Current version is 0.10.0 per libs.versions.toml.

Option B: System Prompt Directive (Practical, untested)

Add a thinking control directive to the system prompt:

// thinkingEnabled = false:
"Respond directly without internal reasoning. Do not use extended thinking."

// thinkingBudget = 200:
"When reasoning internally, limit your thinking to 200 tokens maximum."

Pros: No SDK changes, works today
Cons: Model may ignore the directive (prompt-following varies by quantisation), still generates thinking tokens internally (just shorter)
Recommendation: Worth testing empirically on device — if Gemma-4 respects "do not think", this is the quickest win

Option C: Application-Level Token Counting (Reliable, more work)

Track cumulative thinking tokens during generation and force-stop when budget exceeded:

var thinkingTokenCount = 0
val thinkingBudget = settings.thinkingBudget ?: Int.MAX_VALUE

override fun onMessage(message: Message) {
    val thought = message.channels["thought"]
    if (!thought.isNullOrEmpty()) {
        thinkingTokenCount += estimateTokens(thought)
        if (thinkingTokenCount > thinkingBudget) {
            // Option: stop generation and restart without thinking
            // Option: just suppress further thinking display
        }
        if (thinkingTokenCount <= thinkingBudget) {
            trySend(GenerationResult.Thinking(thought))
        }
    }
}

Pros: Works regardless of SDK support, precise control
Cons: Stopping generation mid-think may produce garbage output; suppressing display doesn't save compute
Recommendation: Best paired with Option B — prompt directive reduces thinking, app-level counting provides a hard ceiling

Implementation Plan (Recommended Phased Approach)

Phase 1 — Toggle + Prompt Directive (Option B):

1. Add thinkingEnabled: Boolean = true to ModelSettingsEntity
2. Add toggle to ModelSettingsScreen ("Enable thinking mode")
3. When disabled, prepend to system prompt: "Respond directly. Do not use extended internal reasoning."
4. Test on device — measure TTFT delta and quality impact
5. Bundle with feat: expand model settings — topK, temperature presets, and per-task optimization #342 migration → single MIGRATION_9_10

Phase 2 — Budget Slider + App-Level Counting (Option C):

1. Add thinkingBudget: Int? = null to entity (null = unlimited)
2. Add slider (0-2000, step 50) — visible only when toggle ON
3. Implement token counting in LiteRtInferenceEngine.onMessage()
4. When budget exceeded, suppress further Thinking emissions (model finishes internally but UI stops showing reasoning)

Phase 3 — SDK Upgrade (Option A, when available):

1. Monitor LiteRT-LM releases for thinking budget API
2. When available, replace prompt directive + app counting with native SDK call
3. Remove system prompt workaround

Schema Changes (Coordinate with #342)

-- Single MIGRATION_9_10 for both #342 and #343:
ALTER TABLE model_settings ADD COLUMN topK INTEGER NOT NULL DEFAULT 40;
ALTER TABLE model_settings ADD COLUMN thinkingEnabled INTEGER NOT NULL DEFAULT 1;
ALTER TABLE model_settings ADD COLUMN thinkingBudget INTEGER;  -- nullable, null = unlimited

Token Budget Impact

Thinking tokens consume context window space. On constrained devices:

Tier	Context Window	Available After System Prompt	Thinking Budget Recommendation
FLAGSHIP	8000	~5100 (with history)	500-1000 (auto)
MID_RANGE	2000	~0 (no history room!)	100-200 or OFF
LOW_POWER	1000	Negative (!)	OFF by default

⚠️ On MID_RANGE and LOW_POWER, thinking should probably default to OFF — the token budget is too tight.

Cross-References

• feat: expand model settings — topK, temperature presets, and per-task optimization #342 — share MIGRATION_9_10, temperature/topK settings
• Consider token strategy for Jandal identity + NZ Knowledge (Research) #367 — token budget strategy; thinking competes for same window
• feat: implement lazy skill loading (load_skill pattern) to reduce system prompt bloat #341 (merged) — shorter prompt freed ~230 tokens, partially offsets thinking cost
• model can no longer do math #371 — LaTeX/math quality may improve with higher thinking budget