diff --git a/main.js b/main.js index ce36ae8..f1c197f 100644 --- a/main.js +++ b/main.js @@ -9,6 +9,7 @@ import { eventSource, event_types, streamingProcessor, saveSettingsDebounced } f import { Handlebars, hljs } from "../../../../lib.js"; import { NAME } from "./common.js"; import { loadSettings } from "./settings.js"; +import { processStreamBuffer } from "./qa_validator.js"; // @ts-ignore: Hack to suppress IDE errors due to SillyTavern's @@ -122,7 +123,7 @@ function runScript(script) { } function installStreamHook() { - if (hookToBeInstalled) { + if (hookToBeInstalled || settings.qaProfile !== "disabled") { const markerRegex = new RegExp(settings.markerRegex, "g"); const partialMarkerRegex = new RegExp(settings.partialMarkerRegex); @@ -130,7 +131,9 @@ function installStreamHook() { const originalOnProgressStreaming = streamingProcessor.onProgressStreaming.bind(streamingProcessor); - function onProgressStreaming(messageId, text, isFinal) { + async function onProgressStreaming(messageId, text, isFinal) { + text = await processStreamBuffer(messageId, text, isFinal, settings); + for (const match of text.matchAll(markerRegex)) { if (!processedIndices.has(match.index)) { // `match[0]` is the whole match, `match[1]` is the first capturing group, i.e. the script ID. @@ -218,6 +221,11 @@ $(async () => { saveSettingsDebounced(); }); + $("#sorcery-qa-profile").val(settings.qaProfile).on("change", function () { + settings.qaProfile = String($(this).val()); + saveSettingsDebounced(); + }); + $("#sorcery-flash-icon").prop("checked", settings.flashIcon).change(function () { settings.flashIcon = this.checked; saveSettingsDebounced(); diff --git a/qa_validator.js b/qa_validator.js new file mode 100644 index 0000000..5ac73a4 --- /dev/null +++ b/qa_validator.js @@ -0,0 +1,144 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later +// Copyright (C) 2025 Sorcery Stream QA contributors + +import { chat, eventSource, event_types, Generate, setSendButtonState, stopGeneration } from "../../../../script.js"; +import { textgenerationwebui_settings } from "../../../textgen-settings.js"; +import { oai_settings } from "../../../openai.js"; +import { kai_settings } from "../../../kai-settings.js"; +import { getContext } from "../../../extensions.js"; + +let rulesCache = null; +/** @type {object | null} Active profile waiting for GENERATION_STOPPED trampoline */ +let pendingCorrectionProfile = null; +/** @type {{ target: ReturnType, temp: number, top_k: number } | null} */ +let originalParams = null; + +function getSamplerTarget() { + const api = getContext().mainApi; + if (api === "textgenerationwebui") { + return { settings: textgenerationwebui_settings, tempKey: "temp", topKKey: "top_k" }; + } + if (api === "openai") { + return { settings: oai_settings, tempKey: "temp_openai", topKKey: "top_k_openai" }; + } + if (api === "kobold") { + return { settings: kai_settings, tempKey: "temp", topKKey: "top_k" }; + } + return null; +} + +function compileRegexRules(regexRules) { + const compiled = {}; + for (const [level, categories] of Object.entries(regexRules)) { + compiled[level] = {}; + for (const [category, patterns] of Object.entries(categories)) { + compiled[level][category] = patterns.map((pattern) => new RegExp(pattern, "gi")); + } + } + return compiled; +} + +async function loadRules() { + if (rulesCache) return rulesCache; + try { + const response = await fetch("/scripts/extensions/third-party/sorcery/rules.json"); + const data = await response.json(); + rulesCache = { profiles: data.profiles, compiled: compileRegexRules(data.regex_rules) }; + return rulesCache; + } catch (e) { + console.error("[Sorcery QA] Failed to load rules.json", e); + return null; + } +} + +function scrambleSamplers(profile) { + if (profile.temp_min == null || profile.top_k_min == null) return; + const target = getSamplerTarget(); + if (!target) return; + const { settings, tempKey, topKKey } = target; + if (!originalParams) { + originalParams = { + target, + temp: settings[tempKey], + top_k: settings[topKKey], + }; + } + const randomTemp = Math.random() * (profile.temp_max - profile.temp_min) + profile.temp_min; + const randomTopK = Math.floor(Math.random() * (profile.top_k_max - profile.top_k_min + 1)) + profile.top_k_min; + settings[tempKey] = parseFloat(randomTemp.toFixed(2)); + settings[topKKey] = randomTopK; +} + +function restoreSamplers() { + if (!originalParams) return; + const { target, temp, top_k } = originalParams; + target.settings[target.tempKey] = temp; + target.settings[target.topKKey] = top_k; + originalParams = null; +} + +function resetQaState() { + pendingCorrectionProfile = null; + restoreSamplers(); +} + +function scheduleContinueAfterStop() { + queueMicrotask(() => { + setSendButtonState(true); + void Generate("continue"); + }); +} + +eventSource.on(event_types.GENERATION_STOPPED, () => { + if (pendingCorrectionProfile) { + const profile = pendingCorrectionProfile; + pendingCorrectionProfile = null; + scrambleSamplers(profile); + scheduleContinueAfterStop(); + } else { + resetQaState(); + } +}); + +eventSource.on(event_types.GENERATION_ENDED, () => { + if (!pendingCorrectionProfile) { + resetQaState(); + } +}); + +export async function processStreamBuffer(messageId, text, isFinal, settings) { + if (!settings.qaProfile || settings.qaProfile === "disabled") return text; + + const rules = await loadRules(); + if (!rules) return text; + + const profile = rules.profiles[settings.qaProfile]; + const strictness = profile && rules.compiled[profile.regex_strictness]; + if (!profile || !strictness) return text; + + let cleanedText = text; + let hitDetected = false; + + for (const regex of strictness.looping ?? []) { + const next = cleanedText.replace(regex, "$1"); + if (next !== cleanedText) { + cleanedText = next; + hitDetected = true; + } + } + for (const regex of strictness.slop ?? []) { + const next = cleanedText.replace(regex, ""); + if (next !== cleanedText) { + cleanedText = next; + hitDetected = true; + } + } + + if (hitDetected && !isFinal && !pendingCorrectionProfile) { + pendingCorrectionProfile = profile; + if (chat[messageId]) chat[messageId].mes = cleanedText.trim(); + stopGeneration(); + } + + return cleanedText; +} diff --git a/rules.json b/rules.json new file mode 100644 index 0000000..2068b57 --- /dev/null +++ b/rules.json @@ -0,0 +1,24 @@ +{ + "profiles": { + "me-me": { "temp_min": 0.70, "temp_max": 1.35, "top_k_min": 20, "top_k_max": 120, "regex_strictness": "high" }, + "mild": { "temp_min": 0.10, "temp_max": 0.85, "top_k_min": 30, "top_k_max": 60, "regex_strictness": "low" }, + "med": { "temp_min": 0.50, "temp_max": 1.10, "top_k_min": 40, "top_k_max": 90, "regex_strictness": "medium" }, + "spicy-mild": { "temp_min": 0.85, "temp_max": 1.40, "top_k_min": 15, "top_k_max": 150, "regex_strictness": "medium" }, + "spicy-spicy": { "temp_min": 1.20, "temp_max": 3.50, "top_k_min": 5, "top_k_max": 300, "regex_strictness": "high" }, + "static": { "temp_min": null, "temp_max": null, "top_k_min": null, "top_k_max": null, "regex_strictness": "medium" } + }, + "regex_rules": { + "high": { + "slop": ["(?:shivers down|testament to|tapestry of|dance of|smirked|it's a cold morning|not a luxury but|beacon of|delve into).+?", "\\b(delve|tapestry|beacon|testament|nestled|commendable|furthermore)\\b"], + "looping": ["(.{15,})\\1+"] + }, + "medium": { + "slop": ["(?:testament to|tapestry of|delve into).+?", "\\b(delve|tapestry|testament)\\b"], + "looping": ["(.{30,})\\1+"] + }, + "low": { + "slop": ["(?:testament to|tapestry of).+?"], + "looping": ["(.{50,})\\1+"] + } + } +} diff --git a/settings.html b/settings.html index f2ef0a1..5c3ce53 100644 --- a/settings.html +++ b/settings.html @@ -38,6 +38,20 @@

+
+ + + Stream quality auto-correction. Does not affect marker scripts. +
+
diff --git a/settings.js b/settings.js index 7196322..dea843f 100644 --- a/settings.js +++ b/settings.js @@ -36,6 +36,7 @@ Character: "Great idea! Let's do that." *Her eyes sparkle with excitement.* const DEFAULT_SETTINGS = { enabled: true, + qaProfile: "disabled", flashIcon: true, instructions: INSTRUCTIONS, markerRegex: "%\\[(\\d+)\\]",