fix: improve thought segmentation syntax extraction

Author: giladgd

src/chatWrappers/QwenChatWrapper.ts

@@ -84,8 +84,8 @@ export class QwenChatWrapper extends ChatWrapper {
             segments: {
                 reiterateStackAfterFunctionCalls: true,
                 thought: {
-                    prefix: LlamaText(new SpecialTokensText("<think>")),
-                    suffix: LlamaText(new SpecialTokensText("</think>"))
+                    prefix: LlamaText(new SpecialTokensText("<think>\n")),
+                    suffix: LlamaText(new SpecialTokensText("\n</think>"))
                 }
             }
         };

src/chatWrappers/generic/utils/extractSegmentSettingsFromTokenizerAndChatTemplate.ts

@@ -8,6 +8,42 @@ export function extractSegmentSettingsFromTokenizerAndChatTemplate(
     function tryMatchPrefixSuffixPair(tryMatchGroups: [prefix: string, suffix: string][]) {
         if (chatTemplate != null) {
             for (const [prefix, suffix] of tryMatchGroups) {
+                if (
+                    (
+                        hasAll(chatTemplate.replaceAll(prefix + "\\n\\n" + suffix, ""), [
+                            prefix + "\\n\\n",
+                            "\\n\\n" + suffix
+                        ])
+                    ) || (
+                        hasAll(chatTemplate.replaceAll(prefix + "\n\n" + suffix, ""), [
+                            prefix + "\n\n",
+                            "\n\n" + suffix
+                        ])
+                    )
+                )
+                    return {
+                        prefix: LlamaText(new SpecialTokensText(prefix + "\n\n")),
+                        suffix: LlamaText(new SpecialTokensText("\n\n" + suffix))
+                    };
+
+                if (
+                    (
+                        hasAll(chatTemplate.replaceAll(prefix + "\\n" + suffix, ""), [
+                            prefix + "\\n",
+                            "\\n" + suffix
+                        ])
+                    ) || (
+                        hasAll(chatTemplate.replaceAll(prefix + "\n" + suffix, ""), [
+                            prefix + "\n",
+                            "\n" + suffix
+                        ])
+                    )
+                )
+                    return {
+                        prefix: LlamaText(new SpecialTokensText(prefix + "\n")),
+                        suffix: LlamaText(new SpecialTokensText("\n" + suffix))
+                    };
+
                 if (chatTemplate.includes(prefix) && chatTemplate.includes(suffix))
                     return {
                         prefix: LlamaText(new SpecialTokensText(prefix)),
@@ -46,3 +82,7 @@ export function extractSegmentSettingsFromTokenizerAndChatTemplate(
         ])
     });
 }
+
+function hasAll(text: string, matches: string[]) {
+    return matches.every((match) => text.includes(match));
+}

test/modelDependent/qwen3-0.6b/functions.test.ts

@@ -41,9 +41,12 @@ describe("qwen3 0.6b", () => {
                 }
             } as const;
 
-            const res = await chatSession.prompt("What is the second word?", {
+            const res = await chatSession.prompt("What is the second word? No yapping, no formatting", {
                 ...promptOptions,
-                maxTokens: 200
+                maxTokens: 250,
+                budgets: {
+                    thoughtTokens: 100
+                }
             });
 
             expect(res.trim()).to.be.eq('The second word is "secret".');
@@ -96,12 +99,15 @@ describe("qwen3 0.6b", () => {
                 }
             } as const;
 
-            const res = await chatSession.prompt("What is the second word?", {
+            const res = await chatSession.prompt("What is the second word? No yapping, no formatting", {
                 ...promptOptions,
-                maxTokens: 205
+                maxTokens: 250,
+                budgets: {
+                    thoughtTokens: 100
+                }
             });
 
-            expect(res.trim()).to.be.eq('The second word is **"secret"**.');
+            expect(res.trim()).to.be.eq('The second word is "secret".');
 
             const res2 = await chatSession.prompt("Explain what this word means", {
                 ...promptOptions,