From eb778e43951286ea5da142146cbe1c48e03156b1 Mon Sep 17 00:00:00 2001 From: Gabor Berenyi <205867+ber4444@users.noreply.github.com> Date: Fri, 17 Jul 2026 10:26:03 -0700 Subject: [PATCH] fix(coach): strip CoT blocks from LiteRT-LM output before emitting MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Qwen3 (the default .litertlm model) emits ... chain-of-thought before its answer. The Move Coach contract wants only the delivered answer — leaking internal deliberation to the user is a UX bug, and downstream validation (MoveCoachResponseValidator) should judge only the answer, not the reasoning. Adds LitertLmTextGenerator.stripThinkBlocks (companion-level, unit-tested): strips closed ... blocks and an unterminated trailing . Applied in generate() before the Token is emitted. Mirrors the Python-side stripping in ferryman's score_litert_outputs.py (defense in depth — either side stripping is sufficient, but both makes a model swap or generation hiccup safe). --- .../litertlm/LitertLmTextGenerator.kt | 31 ++++++++++- .../LitertLmTextGeneratorStripThinkTest.kt | 53 +++++++++++++++++++ 2 files changed, 83 insertions(+), 1 deletion(-) create mode 100644 onDeviceAi/src/desktopTest/kotlin/com/example/ondeviceai/litertlm/LitertLmTextGeneratorStripThinkTest.kt diff --git a/onDeviceAi/src/desktopMain/kotlin/com/example/ondeviceai/litertlm/LitertLmTextGenerator.kt b/onDeviceAi/src/desktopMain/kotlin/com/example/ondeviceai/litertlm/LitertLmTextGenerator.kt index 1026818..4090386 100644 --- a/onDeviceAi/src/desktopMain/kotlin/com/example/ondeviceai/litertlm/LitertLmTextGenerator.kt +++ b/onDeviceAi/src/desktopMain/kotlin/com/example/ondeviceai/litertlm/LitertLmTextGenerator.kt @@ -104,7 +104,16 @@ class LitertLmTextGenerator( } if (output.isNotEmpty()) { - emit(AiTokenOrFinal.Token(output.toString())) + // Qwen3 (the default .litertlm model) emits … + // chain-of-thought blocks before its answer. Strip them so the + // Move Coach never shows internal deliberation to the user — and + // so downstream validation (MoveCoachResponseValidator) judges + // only the delivered answer, not the reasoning. Matches both a + // closed … and an unterminated trailing …. + val cleaned = stripThinkBlocks(output.toString()) + if (cleaned.isNotEmpty()) { + emit(AiTokenOrFinal.Token(cleaned)) + } } // getTokenCount() is @ExperimentalApi in litertlm-jvm and BenchmarkInfo is // gated too; we avoid both and derive a rough token count from output length @@ -174,5 +183,25 @@ class LitertLmTextGenerator( const val DEFAULT_CONTEXT_SIZE = 1024 const val DEFAULT_TOP_K = 40 const val DEFAULT_TOP_P = 1.0 + + // Matches … (case-insensitive, DOTALL so it spans newlines). + // LiteRT-LM's Qwen3 model emits chain-of-thought before its answer; the + // Move Coach contract wants only the delivered answer, so [generate] + // strips these blocks before emitting. The unterminated variant catches + // a the model opened but never closed (generation cut off). + private val THINK_BLOCK = Regex("(?is).*?") + private val THINK_UNTERMINATED = Regex("(?is).*") + + /** + * Remove `…` chain-of-thought blocks from [text]. Handles + * both closed blocks and an unterminated trailing ``. Returns the + * cleaned text, trimmed. If no `` is present, returns the text + * trimmed. Public on the companion so it can be unit-tested directly. + */ + fun stripThinkBlocks(text: String): String { + var cleaned = THINK_BLOCK.replace(text, "") + cleaned = THINK_UNTERMINATED.replace(cleaned, "") + return cleaned.trim() + } } } diff --git a/onDeviceAi/src/desktopTest/kotlin/com/example/ondeviceai/litertlm/LitertLmTextGeneratorStripThinkTest.kt b/onDeviceAi/src/desktopTest/kotlin/com/example/ondeviceai/litertlm/LitertLmTextGeneratorStripThinkTest.kt new file mode 100644 index 0000000..47e41b6 --- /dev/null +++ b/onDeviceAi/src/desktopTest/kotlin/com/example/ondeviceai/litertlm/LitertLmTextGeneratorStripThinkTest.kt @@ -0,0 +1,53 @@ +package com.example.ondeviceai.litertlm + +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue +import com.example.ondeviceai.litertlm.LitertLmTextGenerator.Companion.stripThinkBlocks + +/** + * Tests for [LitertLmTextGenerator.stripThinkBlocks] — the CoT-stripping that + * keeps the Move Coach from leaking Qwen3's `` deliberation to the user. + * + * Desktop-only because LitertLmTextGenerator is desktop-only. The cases mirror + * the real outputs observed from the LiteRT-LM driver: a leading `…` + * block followed by the final answer, and the unterminated edge case. + */ +class LitertLmTextGeneratorStripThinkTest { + + @Test + fun `strips a closed think block and keeps the answer`() { + val raw = "Let me consider this move. Does it give check? No.\nNf3 develops the knight." + assertEquals("Nf3 develops the knight.", stripThinkBlocks(raw)) + } + + @Test + fun `strips think block with newlines inside`() { + val raw = "\nOkay, the move is g1 to h3.\nThis develops the knight.\n\n\nNh3 develops the knight toward the edge." + assertEquals("Nh3 develops the knight toward the edge.", stripThinkBlocks(raw)) + } + + @Test + fun `strips an unterminated trailing think block`() { + // A generation cut off mid-reasoning: opened but never closed. + val raw = "Nh3 develops the knight. Wait, let me reconsider whether" + assertEquals("Nh3 develops the knight.", stripThinkBlocks(raw)) + } + + @Test + fun `returns_text_unchanged_when_no_think_block`() { + val raw = "Nf3 develops the knight toward the center." + assertEquals("Nf3 develops the knight toward the center.", stripThinkBlocks(raw)) + } + + @Test + fun `handles_case_insensitive_think_tag`() { + val raw = "reasoningAnswer." + assertEquals("Answer.", stripThinkBlocks(raw)) + } + + @Test + fun `returns_empty_string_for_input_that_was_only_think`() { + assertEquals("", stripThinkBlocks("all reasoning, no answer")) + } +}