├── manifest.json
├── README.md
├── style.css
├── rewrite_settings.html
└── index.js


/manifest.json:
--------------------------------------------------------------------------------
 1 | {
 2 |     "display_name": "Rewrite Extension",
 3 |     "loading_order": 39,
 4 |     "requires": [],
 5 |     "optional": [],
 6 |     "js": "index.js",
 7 |     "css": "style.css",
 8 |     "author": "SplitClover",
 9 |     "version": "1.3.0",
10 |     "homePage": ""
11 | }
12 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # Rewrite Extension for SillyTavern
 2 | 
 3 | > [!IMPORTANT]
 4 | > **2025.09.02** - This is extension is no longer maintained for the foreseeable future.
 5 | 
 6 | ## Overview
 7 | 
 8 | The Rewrite Extension enhances the chat experience in SillyTavern by allowing users to dynamically rewrite, shorten, or expand selected text within messages. Works for chat completion, text completion and NovelAI.
 9 | 
10 | ## Features
11 | 
12 | - Custom {{rewrite}} macro that contains the selected text
13 | - Custom {{targetmessage}} macro that contains the full targeted message
14 | - Custom {{rewritecount}} macro that returns a numeric (39) count of words selected
15 | - Rewrite, shorten, or expand selected text in chat messages, with an added delete button
16 | - Convenient undo button
17 | - Real-time streaming of rewritten text
18 | - Temporary highlighting of modified text for easy identification
19 | - Ability to abort ongoing rewrites
20 | 
21 | ## Installation
22 | 
23 | Use SillyTavern's built-in extension installer:
24 | `https://github.com/splitclover/rewrite-extension`
25 | 
26 | ## Usage
27 | 
28 | To use the Rewrite Extension:
29 | 
30 | 1. Configure the extension (see below)
31 | 2. Select text within a single(!) chat message
32 | 3. A context menu will appear with options to Rewrite, Shorten, or Expand
33 | 4. Choose the desired option
34 | 5. The selected text will be replaced with the AI-generated modification
35 | 
36 | ## Configuration
37 | 
38 | 1. Open the Extension tab -> Rewrite Extension
39 | 2. Set presets for Rewrite, Shorten, and Expand operations
40 | 3. Adjust the highlight duration for modified text
41 | 
42 | ## Contributing
43 | 
44 | Contributions to improve the Rewrite Extension are welcome. Please fork the repository and submit a pull request with your changes.


--------------------------------------------------------------------------------
/style.css:
--------------------------------------------------------------------------------
  1 | @keyframes vibrantWavyBackground {
  2 |     0% {
  3 |         background-position: 0% 50%;
  4 |         box-shadow: 0 0 10px rgba(26, 35, 126, 0.7);
  5 |     }
  6 |     25% {
  7 |         background-position: 50% 100%;
  8 |         box-shadow: 0 0 15px rgba(74, 20, 140, 0.7);
  9 |     }
 10 |     50% {
 11 |         background-position: 100% 50%;
 12 |         box-shadow: 0 0 20px rgba(26, 35, 126, 0.7);
 13 |     }
 14 |     75% {
 15 |         background-position: 50% 0%;
 16 |         box-shadow: 0 0 15px rgba(74, 20, 140, 0.7);
 17 |     }
 18 |     100% {
 19 |         background-position: 0% 50%;
 20 |         box-shadow: 0 0 10px rgba(26, 35, 126, 0.7);
 21 |     }
 22 | }
 23 | 
 24 | .animated-highlight {
 25 |     background: linear-gradient(45deg, #1a237e, #4a148c, #7b1fa2, #4a148c, #1a237e);
 26 |     background-size: 300% 300%;
 27 |     animation: vibrantWavyBackground 6s ease infinite;
 28 |     border-radius: 8px;
 29 |     padding: 2px 6px;
 30 |     color: white;
 31 |     display: inline-block;
 32 |     box-decoration-break: clone;
 33 |     -webkit-box-decoration-break: clone;
 34 |     transition: all 0.3s ease;
 35 | }
 36 | 
 37 | .rewrite-extension_block textarea.text_pole {
 38 |     width: 100%;
 39 |     min-height: 60px;
 40 |     resize: vertical;
 41 |     padding: 5px;
 42 |     box-sizing: border-box;
 43 | }
 44 | 
 45 | .rewrite-extension-settings .flex-container {
 46 |     display: flex;
 47 |     align-items: center;
 48 |     margin-bottom: 10px;
 49 |     position: relative;
 50 | }
 51 | 
 52 | .rewrite-extension-settings label {
 53 |     flex-grow: 1;
 54 |     margin-right: 10px;
 55 | }
 56 | 
 57 | .rewrite-extension-settings .tooltip-icon {
 58 |     font-size: var(--mainFontSize);
 59 |     color: var(--SmartThemeBodyColor);
 60 |     cursor: pointer;
 61 |     margin-right: 10px;
 62 | }
 63 | 
 64 | .rewrite-extension-settings .tooltip-icon:hover::after {
 65 |     content: attr(data-tooltip);
 66 |     position: absolute;
 67 |     right: 30px; /* Adjust this value as needed */
 68 |     top: 50%;
 69 |     transform: translateY(-50%);
 70 |     background-color: var(--SmartThemeBlurTintColor);
 71 |     color: var(--SmartThemeBodyColor);
 72 |     padding: 8px 12px;
 73 |     border-radius: 5px;
 74 |     font-size: calc(var(--mainFontSize) * 0.9);
 75 |     font-weight: normal;
 76 |     font-family: var(--mainFontFamily);
 77 |     white-space: normal;
 78 |     max-width: 250px;
 79 |     z-index: 1000;
 80 |     text-shadow: 0px 0px calc(var(--shadowWidth) * 1px) var(--SmartThemeShadowColor);
 81 |     border: 1px solid var(--SmartThemeBorderColor);
 82 |     box-shadow: 0 2px 10px rgba(0, 0, 0, 0.1);
 83 | }
 84 | 
 85 | .rewrite-extension-settings .inline-drawer-content > div {
 86 |     margin-bottom: 15px;
 87 | }
 88 | 
 89 | .rewrite-extension-settings .checkbox-container {
 90 |     justify-content: space-between;
 91 | }
 92 | 
 93 | .rewrite-extension-settings .checkbox-label-group {
 94 |     display: flex;
 95 |     align-items: center;
 96 | }
 97 | 
 98 | .rewrite-extension-settings .checkbox-label-group label {
 99 |     margin-left: 5px;
100 |     order: 1;
101 | }
102 | 
103 | .rewrite-extension-settings .checkbox-label-group input[type="checkbox"] {
104 |     order: 0;
105 |     margin-right: 0;
106 | }
107 | 
108 | .rewrite-extension-settings .checkbox-container .tooltip-icon {
109 |     margin-left: 10px;
110 | }
111 | 
112 | .rewrite-extension-settings .checkbox-container .tooltip-icon:hover::after {
113 |     right: 25px;
114 | }
115 | 
116 | .rewrite-extension_block .flex-container > div {
117 |     display: flex;
118 |     align-items: center;
119 |     justify-content: space-between;
120 | }
121 | 
122 | .rewrite-extension_block .flex-container > div .tooltip-icon {
123 |     margin-left: 5px;
124 | }
125 | 
126 | .rewrite-extension-settings .inline-drawer {
127 |     width: 100%;
128 | }
129 | 
130 | .rewrite-extension-settings .inline-drawer-content {
131 |     width: 100%;
132 | }
133 | 
134 | .rewrite-extension-settings .inline-drawer-toggle {
135 |     width: 100%;
136 |     display: flex;
137 |     justify-content: space-between;
138 |     align-items: center;
139 |     cursor: pointer;
140 |     padding: 5px 0;
141 | }
142 | 
143 | .rewrite-extension-settings .inline-drawer-content .checkbox-container {
144 |     display: flex;
145 |     justify-content: flex-start;
146 |     align-items: center;
147 |     margin-bottom: 5px;
148 | }
149 | 
150 | .rewrite-extension-settings .inline-drawer-content .checkbox-container input[type="checkbox"] {
151 |     margin-right: 5px;
152 | }
153 | 
154 | .rewrite-extension-settings .inline-drawer-content .checkbox-container label {
155 |     margin-right: 0;
156 | }
157 | 
158 | /* Ensure the tooltip doesn't go off-screen on smaller screens */
159 | @media (max-width: 768px) {
160 |     .rewrite-extension-settings .tooltip-icon:hover::after {
161 |         right: 0;
162 |         left: 50%;
163 |         transform: translateX(-50%);
164 |         top: 100%;
165 |         margin-top: 5px;
166 |     }
167 | }
168 | 


--------------------------------------------------------------------------------
/rewrite_settings.html:
--------------------------------------------------------------------------------
  1 | <div class="rewrite-extension-settings">
  2 |     <div class="inline-drawer">
  3 |         <div class="inline-drawer-toggle inline-drawer-header">
  4 |             <b>Rewrite Extension</b>
  5 |             <div class="inline-drawer-icon fa-solid fa-circle-chevron-down down"></div>
  6 |         </div>
  7 |         <div class="inline-drawer-content">
  8 |             <div class="rewrite-extension_block flex-container">
  9 |                 <label for="rewrite_extension_model_select">Configure Model:</label>
 10 |                 <span class="tooltip-icon fa-solid fa-circle-question" data-tooltip="Choose which Completion type to configure. For Chat Completion this affects if presets or templates are used."></span>
 11 |                 <select id="rewrite_extension_model_select" class="text_pole">
 12 |                     <option value="chat_completion">Chat Completion</option>
 13 |                     <option value="textgenerationwebui">Text Completion/NovelAI</option>
 14 |                 </select>
 15 |             </div>
 16 | 
 17 |             <div id="chat_completion_settings" class="model-specific-settings">
 18 |                 <div class="rewrite-extension_block flex-container">
 19 |                     <label for="rewrite_preset">Rewrite Preset:</label>
 20 |                     <span class="tooltip-icon fa-solid fa-circle-question" data-tooltip="To use a simple 'System: prompt', switch to Text Completion. This will always use the preset's max tokens setting."></span>
 21 |                     <select id="rewrite_preset" class="text_pole" data-preset-manager-for="rewrite">
 22 |                         <!-- Options will be populated dynamically -->
 23 |                     </select>
 24 |                 </div>
 25 | 
 26 |                 <div class="rewrite-extension_block flex-container">
 27 |                     <label for="shorten_preset">Shorten Preset:</label>
 28 |                     <select id="shorten_preset" class="text_pole" data-preset-manager-for="shorten">
 29 |                         <!-- Options will be populated dynamically -->
 30 |                     </select>
 31 |                 </div>
 32 | 
 33 |                 <div class="rewrite-extension_block flex-container">
 34 |                     <label for="expand_preset">Expand Preset:</label>
 35 |                     <select id="expand_preset" class="text_pole" data-preset-manager-for="expand">
 36 |                         <!-- Options will be populated dynamically -->
 37 |                     </select>
 38 |                 </div>
 39 | 
 40 |                 <div class="rewrite-extension_block flex-container">
 41 |                     <label for="custom_preset">Custom Preset:</label>
 42 |                     <select id="custom_preset" class="text_pole" data-preset-manager-for="custom">
 43 |                         <!-- Options will be populated dynamically -->
 44 |                     </select>
 45 |                 </div>
 46 | 
 47 |                 <div class="rewrite-extension_block flex-container checkbox-container">
 48 |                     <div class="checkbox-label-group">
 49 |                         <input type="checkbox" id="override_max_tokens" class="checkbox" checked>
 50 |                         <label for="override_max_tokens">Override Preset Max Tokens</label>
 51 |                     </div>
 52 |                     <span class="tooltip-icon fa-solid fa-circle-question" data-tooltip="If checked, the extension will override the max tokens set in the chat completion preset."></span>
 53 |                 </div>
 54 |             </div>
 55 | 
 56 |             <div id="text_based_settings" class="model-specific-settings" style="display: none;">
 57 |                 <div class="rewrite-extension_block flex-container">
 58 |                     <label for="text_rewrite_prompt">Rewrite Prompt:</label>
 59 |                     <textarea id="text_rewrite_prompt" class="text_pole" rows="3" placeholder="Enter rewrite prompt"></textarea>
 60 |                 </div>
 61 | 
 62 |                 <div class="rewrite-extension_block flex-container">
 63 |                     <label for="text_shorten_prompt">Shorten Prompt:</label>
 64 |                     <textarea id="text_shorten_prompt" class="text_pole" rows="3" placeholder="Enter shorten prompt"></textarea>
 65 |                 </div>
 66 | 
 67 |                 <div class="rewrite-extension_block flex-container">
 68 |                     <label for="text_expand_prompt">Expand Prompt:</label>
 69 |                     <textarea id="text_expand_prompt" class="text_pole" rows="3" placeholder="Enter expand prompt"></textarea>
 70 |                 </div>
 71 | 
 72 |                 <div class="rewrite-extension_block flex-container">
 73 |                     <label for="text_custom_prompt">Custom Prompt:</label>
 74 |                     <textarea id="text_custom_prompt" class="text_pole" rows="3" placeholder="Enter custom prompt. Use {{custom_instructions}} macro."></textarea>
 75 |                 </div>
 76 |             </div>
 77 | 
 78 |             <hr class="sysHR" />
 79 | 
 80 |             <!-- Universal settings -->
 81 |             <div class="rewrite-extension_block flex-container checkbox-container">
 82 |                 <div class="checkbox-label-group">
 83 |                     <input type="checkbox" id="use_streaming" class="checkbox">
 84 |                     <label for="use_streaming">Use Streaming</label>
 85 |                 </div>
 86 |                 <span class="tooltip-icon fa-solid fa-circle-question" data-tooltip="Enable to see the rewritten text appear gradually. Might not work for all APIs."></span>
 87 |             </div>
 88 | 
 89 |             <div class="rewrite-extension_block flex-container checkbox-container">
 90 |                 <div class="checkbox-label-group">
 91 |                     <input type="checkbox" id="apply_regex_on_rewrite" class="checkbox">
 92 |                     <label for="apply_regex_on_rewrite">Apply Regex on Rewrite</label>
 93 |                 </div>
 94 |                 <span class="tooltip-icon fa-solid fa-circle-question" data-tooltip="If checked, applies configured AI Output regex scripts to the rewritten text."></span>
 95 |             </div>
 96 | 
 97 |             <div class="rewrite-extension_block flex-container checkbox-container">
 98 |                 <div class="checkbox-label-group">
 99 |                     <input type="checkbox" id="use_dynamic_tokens" class="checkbox">
100 |                     <label for="use_dynamic_tokens">Use Dynamic Token Calculation</label>
101 |                 </div>
102 |                 <span class="tooltip-icon fa-solid fa-circle-question" data-tooltip="Adjust token count based on the length of the selected text."></span>
103 |             </div>
104 | 
105 |             <!-- Token settings -->
106 |             <div id="static_token_settings">
107 |                 <div class="rewrite-extension_block flex-container">
108 |                     <label for="rewrite_tokens">Rewrite Tokens:</label>
109 |                     <input type="number" id="rewrite_tokens" class="text_pole" min="0" value="100">
110 |                 </div>
111 |                 <div class="rewrite-extension_block flex-container">
112 |                     <label for="shorten_tokens">Shorten Tokens:</label>
113 |                     <input type="number" id="shorten_tokens" class="text_pole" min="0" value="50">
114 |                 </div>
115 |                 <div class="rewrite-extension_block flex-container">
116 |                     <label for="expand_tokens">Expand Tokens:</label>
117 |                     <input type="number" id="expand_tokens" class="text_pole" min="0" value="150">
118 |                 </div>
119 |                 <div class="rewrite-extension_block flex-container">
120 |                     <label for="custom_tokens">Custom Tokens:</label>
121 |                     <input type="number" id="custom_tokens" class="text_pole" min="0" value="100">
122 |                 </div>
123 |             </div>
124 | 
125 |             <div id="dynamic_token_settings" style="display: none;">
126 |                 <div class="rewrite-extension_block flex-container">
127 |                     <label for="dynamic_token_mode">Dynamic Token Mode:</label>
128 |                     <span class="tooltip-icon fa-solid fa-circle-question" data-tooltip="Multiplicative: Adjust tokens by percentage (e.g., 1.2 for 20% more). Additive: Adjust by adding/subtracting tokens (e.g., +50 or -20)."></span>
129 |                     <select id="dynamic_token_mode" class="text_pole">
130 |                         <option value="multiplicative">Multiplicative</option>
131 |                         <option value="additive">Additive</option>
132 |                     </select>
133 |                 </div>
134 |                 <div id="additive_settings">
135 |                     <div class="rewrite-extension_block flex-container">
136 |                         <label for="rewrite_tokens_add">Rewrite Tokens Difference:</label>
137 |                         <input type="number" id="rewrite_tokens_add" class="text_pole" value="0">
138 |                     </div>
139 |                     <div class="rewrite-extension_block flex-container">
140 |                         <label for="shorten_tokens_add">Shorten Tokens Difference:</label>
141 |                         <input type="number" id="shorten_tokens_add" class="text_pole" value="-20">
142 |                     </div>
143 |                     <div class="rewrite-extension_block flex-container">
144 |                         <label for="expand_tokens_add">Expand Tokens Difference:</label>
145 |                         <input type="number" id="expand_tokens_add" class="text_pole" value="20">
146 |                     </div>
147 |                     <div class="rewrite-extension_block flex-container">
148 |                         <label for="custom_tokens_add">Custom Tokens Difference:</label>
149 |                         <input type="number" id="custom_tokens_add" class="text_pole" value="0">
150 |                     </div>
151 |                 </div>
152 |                 <div id="multiplicative_settings">
153 |                     <div class="rewrite-extension_block flex-container">
154 |                         <label for="rewrite_tokens_mult">Rewrite Tokens Multiplier:</label>
155 |                         <input type="number" id="rewrite_tokens_mult" class="text_pole" value="1.0" step="0.1">
156 |                     </div>
157 |                     <div class="rewrite-extension_block flex-container">
158 |                         <label for="shorten_tokens_mult">Shorten Tokens Multiplier:</label>
159 |                         <input type="number" id="shorten_tokens_mult" class="text_pole" value="0.8" step="0.1">
160 |                     </div>
161 |                     <div class="rewrite-extension_block flex-container">
162 |                         <label for="expand_tokens_mult">Expand Tokens Multiplier:</label>
163 |                         <input type="number" id="expand_tokens_mult" class="text_pole" value="1.2" step="0.1">
164 |                     </div>
165 |                     <div class="rewrite-extension_block flex-container">
166 |                         <label for="custom_tokens_mult">Custom Tokens Multiplier:</label>
167 |                         <input type="number" id="custom_tokens_mult" class="text_pole" value="1.0" step="0.1">
168 |                     </div>
169 |                 </div>
170 |             </div>
171 | 
172 |             <div class="rewrite-extension_block" style="margin-top: 15px;">
173 |                 <div class="flex-container">
174 |                     <div style="width: 49%;">
175 |                         <label for="remove_prefix">Remove Prefix:</label>
176 |                         <span class="tooltip-icon fa-solid fa-circle-question" data-tooltip="Text to remove from the beginning of the AI's response."></span>
177 |                     </div>
178 |                     <div style="width: 49%;">
179 |                         <label for="remove_suffix">Remove Suffix:</label>
180 |                         <span class="tooltip-icon fa-solid fa-circle-question" data-tooltip="Text to remove from the end of the AI's response."></span>
181 |                     </div>
182 |                 </div>
183 |                 <div class="flex-container" style="margin-top: 5px;">
184 |                     <input type="text" id="remove_prefix" class="text_pole" style="width: 49%;" placeholder="Prefix to remove">
185 |                     <input type="text" id="remove_suffix" class="text_pole" style="width: 49%;" placeholder="Suffix to remove">
186 |                 </div>
187 |             </div>
188 | 
189 |             <div class="rewrite-extension_block flex-container">
190 |                 <label for="highlight_duration">Highlight Duration (ms):</label>
191 |                 <span class="tooltip-icon fa-solid fa-circle-question" data-tooltip="Duration in milliseconds for which the rewritten text will be highlighted."></span>
192 |                 <input type="number" id="highlight_duration" class="text_pole" min="0" step="100" value="3000">
193 |             </div>
194 | 
195 |             <hr class="sysHR" />
196 |             <div class="rewrite-extension_block flex-container">
197 |                 <div class="inline-drawer">
198 |                     <div class="inline-drawer-toggle inline-drawer-header">
199 |                         <b>Button visibility</b>
200 |                         <div class="inline-drawer-icon fa-solid fa-circle-chevron-up up"></div>
201 |                     </div>
202 |                     <div class="inline-drawer-content" style="display: none;">
203 |                         <div class="checkbox-container">
204 |                             <input type="checkbox" id="show_rewrite" class="checkbox" checked>
205 |                             <label for="show_rewrite">Show Rewrite</label>
206 |                         </div>
207 |                         <div class="checkbox-container">
208 |                             <input type="checkbox" id="show_shorten" class="checkbox" checked>
209 |                             <label for="show_shorten">Show Shorten</label>
210 |                         </div>
211 |                         <div class="checkbox-container">
212 |                             <input type="checkbox" id="show_expand" class="checkbox" checked>
213 |                             <label for="show_expand">Show Expand</label>
214 |                         </div>
215 |                         <div class="checkbox-container">
216 |                             <input type="checkbox" id="show_custom" class="checkbox" checked>
217 |                             <label for="show_custom">Show Custom</label>
218 |                         </div>
219 |                         <div class="checkbox-container">
220 |                             <input type="checkbox" id="show_delete" class="checkbox" checked>
221 |                             <label for="show_delete">Show Delete</label>
222 |                         </div>
223 |                     </div>
224 |                 </div>
225 |             </div>
226 |         </div>
227 |     </div>
228 | </div>
229 | 


--------------------------------------------------------------------------------
/index.js:
--------------------------------------------------------------------------------
   1 | import { sendOpenAIRequest, oai_settings } from "../../../openai.js";
   2 | import { extractAllWords } from "../../../utils.js";
   3 | import { getTokenCount } from "../../../tokenizers.js";
   4 | import { getNovelGenerationData, generateNovelWithStreaming, nai_settings } from "../../../nai-settings.js";
   5 | import { generateHorde, MIN_LENGTH } from "../../../horde.js";
   6 | import { getTextGenGenerationData, generateTextGenWithStreaming } from "../../../textgen-settings.js";
   7 | import {
   8 |     main_api,
   9 |     novelai_settings,
  10 |     novelai_setting_names,
  11 |     eventSource,
  12 |     event_types,
  13 |     saveSettingsDebounced,
  14 |     messageFormatting,
  15 |     addCopyToCodeBlocks,
  16 |     getRequestHeaders,
  17 |     generateRaw,
  18 | } from "../../../../script.js";
  19 | import { extension_settings, getContext } from "../../../extensions.js";
  20 | import { getRegexedString, regex_placement } from '../../regex/engine.js'; // Import from built-in regex extension
  21 | 
  22 | const extensionName = "rewrite-extension";
  23 | const extensionFolderPath = `scripts/extensions/third-party/${extensionName}`;
  24 | 
  25 | const undo_steps = 15;
  26 | 
  27 | // Default settings
  28 | const defaultSettings = {
  29 |     rewritePreset: "",
  30 |     shortenPreset: "",
  31 |     expandPreset: "",
  32 |     customPreset: "", 
  33 |     highlightDuration: 3000,
  34 |     selectedModel: "chat_completion",
  35 |     textRewritePrompt: `[INST]Rewrite this section of text: """{{rewrite}}""" while keeping the same content, general style and length. Do not list alternatives and only print the result without prefix or suffix.[/INST]
  36 | 
  37 | Sure, here is only the rewritten text without any comments: `,
  38 |     textShortenPrompt: `[INST]Rewrite this section of text: """{{rewrite}}""" while keeping the same content, general style. Do not list alternatives and only print the result without prefix or suffix. Shorten it by roughly 20%.[/INST]
  39 | 
  40 | Sure, here is only the rewritten text without any comments: `,
  41 |     textExpandPrompt: `[INST]Rewrite this section of text: """{{rewrite}}""" while keeping the same content, general style. Do not list alternatives and only print the result without prefix or suffix. Lengthen it by roughly 20%.[/INST]
  42 | 
  43 | Sure, here is only the rewritten text without any comments: `,
  44 |     textCustomPrompt: `[INST]Rewrite this section of text: """{{rewrite}}""" according to the following instructions: "{{custom_instructions}}". Keep the general style. Do not list alternatives and only print the result without prefix or suffix.[/INST]
  45 | 
  46 | Sure, here is only the rewritten text without any comments: `, 
  47 |     useStreaming: true,
  48 |     useDynamicTokens: true,
  49 |     dynamicTokenMode: 'multiplicative',
  50 |     rewriteTokens: 100,
  51 |     shortenTokens: 50,
  52 |     expandTokens: 150,
  53 |     customTokens: 100, 
  54 |     rewriteTokensAdd: 0,
  55 |     shortenTokensAdd: -50,
  56 |     expandTokensAdd: 50,
  57 |     customTokensAdd: 0, 
  58 |     rewriteTokensMult: 1.05,
  59 |     shortenTokensMult: 0.8,
  60 |     expandTokensMult: 1.5,
  61 |     customTokensMult: 1.0, 
  62 |     removePrefix: `"`,
  63 |     removeSuffix: `"`,
  64 |     overrideMaxTokens: true,
  65 |     showRewrite: true,
  66 |     showShorten: true,
  67 |     showExpand: true,
  68 |     showCustom: true, 
  69 |     showDelete: true,
  70 |     applyRegexOnRewrite: true, // New setting to control regex application
  71 | };
  72 | 
  73 | let rewriteMenu = null;
  74 | let lastSelection = null;
  75 | let abortController;
  76 | 
  77 | let changeHistory = [];
  78 | 
  79 | // Load settings
  80 | function loadSettings() {
  81 |     extension_settings[extensionName] = extension_settings[extensionName] || {};
  82 | 
  83 |     // Helper function to get a setting with a default value
  84 |     const getSetting = (key, defaultValue) => {
  85 |         return extension_settings[extensionName][key] !== undefined
  86 |             ? extension_settings[extensionName][key]
  87 |             : defaultValue;
  88 |     };
  89 | 
  90 |     // Load settings, using defaults if not set
  91 |     $("#rewrite_preset").val(getSetting('rewritePreset', defaultSettings.rewritePreset));
  92 |     $("#shorten_preset").val(getSetting('shortenPreset', defaultSettings.shortenPreset));
  93 |     $("#expand_preset").val(getSetting('expandPreset', defaultSettings.expandPreset));
  94 |     $("#custom_preset").val(getSetting('customPreset', defaultSettings.customPreset)); 
  95 |     $("#highlight_duration").val(getSetting('highlightDuration', defaultSettings.highlightDuration));
  96 |     $("#rewrite_extension_model_select").val(getSetting('selectedModel', defaultSettings.selectedModel));
  97 |     $("#text_rewrite_prompt").val(getSetting('textRewritePrompt', defaultSettings.textRewritePrompt));
  98 |     $("#text_shorten_prompt").val(getSetting('textShortenPrompt', defaultSettings.textShortenPrompt));
  99 |     $("#text_expand_prompt").val(getSetting('textExpandPrompt', defaultSettings.textExpandPrompt));
 100 |     $("#text_custom_prompt").val(getSetting('textCustomPrompt', defaultSettings.textCustomPrompt)); 
 101 |     $("#use_streaming").prop('checked', getSetting('useStreaming', defaultSettings.useStreaming));
 102 |     $("#use_dynamic_tokens").prop('checked', getSetting('useDynamicTokens', defaultSettings.useDynamicTokens));
 103 |     $("#dynamic_token_mode").val(getSetting('dynamicTokenMode', defaultSettings.dynamicTokenMode));
 104 |     $("#rewrite_tokens").val(getSetting('rewriteTokens', defaultSettings.rewriteTokens));
 105 |     $("#shorten_tokens").val(getSetting('shortenTokens', defaultSettings.shortenTokens));
 106 |     $("#expand_tokens").val(getSetting('expandTokens', defaultSettings.expandTokens));
 107 |     $("#custom_tokens").val(getSetting('customTokens', defaultSettings.customTokens)); 
 108 |     $("#rewrite_tokens_add").val(getSetting('rewriteTokensAdd', defaultSettings.rewriteTokensAdd));
 109 |     $("#shorten_tokens_add").val(getSetting('shortenTokensAdd', defaultSettings.shortenTokensAdd));
 110 |     $("#expand_tokens_add").val(getSetting('expandTokensAdd', defaultSettings.expandTokensAdd));
 111 |     $("#custom_tokens_add").val(getSetting('customTokensAdd', defaultSettings.customTokensAdd)); 
 112 |     $("#rewrite_tokens_mult").val(getSetting('rewriteTokensMult', defaultSettings.rewriteTokensMult));
 113 |     $("#shorten_tokens_mult").val(getSetting('shortenTokensMult', defaultSettings.shortenTokensMult));
 114 |     $("#expand_tokens_mult").val(getSetting('expandTokensMult', defaultSettings.expandTokensMult));
 115 |     $("#custom_tokens_mult").val(getSetting('customTokensMult', defaultSettings.customTokensMult)); 
 116 |     $("#remove_prefix").val(getSetting('removePrefix', defaultSettings.removePrefix));
 117 |     $("#remove_suffix").val(getSetting('removeSuffix', defaultSettings.removeSuffix));
 118 |     $("#override_max_tokens").prop('checked', getSetting('overrideMaxTokens', defaultSettings.overrideMaxTokens));
 119 |     $("#show_rewrite").prop('checked', getSetting('showRewrite', defaultSettings.showRewrite));
 120 |     $("#show_shorten").prop('checked', getSetting('showShorten', defaultSettings.showShorten));
 121 |     $("#show_expand").prop('checked', getSetting('showExpand', defaultSettings.showExpand));
 122 |     $("#show_custom").prop('checked', getSetting('showCustom', defaultSettings.showCustom)); 
 123 |     $("#show_delete").prop('checked', getSetting('showDelete', defaultSettings.showDelete));
 124 |     $("#apply_regex_on_rewrite").prop('checked', getSetting('applyRegexOnRewrite', defaultSettings.applyRegexOnRewrite)); // Load new setting
 125 | 
 126 |     // Update the UI based on loaded settings
 127 |     updateModelSettings();
 128 |     updateTokenSettings();
 129 | }
 130 | 
 131 | function saveSettings() {
 132 |     extension_settings[extensionName] = {
 133 |         rewritePreset: $("#rewrite_preset").val(),
 134 |         shortenPreset: $("#shorten_preset").val(),
 135 |         expandPreset: $("#expand_preset").val(),
 136 |         customPreset: $("#custom_preset").val(), 
 137 |         highlightDuration: parseInt($("#highlight_duration").val()),
 138 |         selectedModel: $("#rewrite_extension_model_select").val(),
 139 |         textRewritePrompt: $("#text_rewrite_prompt").val(),
 140 |         textShortenPrompt: $("#text_shorten_prompt").val(),
 141 |         textExpandPrompt: $("#text_expand_prompt").val(),
 142 |         textCustomPrompt: $("#text_custom_prompt").val(), 
 143 |         useStreaming: $("#use_streaming").is(':checked'),
 144 |         useDynamicTokens: $("#use_dynamic_tokens").is(':checked'),
 145 |         dynamicTokenMode: $("#dynamic_token_mode").val(),
 146 |         rewriteTokens: parseInt($("#rewrite_tokens").val()),
 147 |         shortenTokens: parseInt($("#shorten_tokens").val()),
 148 |         expandTokens: parseInt($("#expand_tokens").val()),
 149 |         customTokens: parseInt($("#custom_tokens").val()), 
 150 |         rewriteTokensAdd: parseInt($("#rewrite_tokens_add").val()),
 151 |         shortenTokensAdd: parseInt($("#shorten_tokens_add").val()),
 152 |         expandTokensAdd: parseInt($("#expand_tokens_add").val()),
 153 |         customTokensAdd: parseInt($("#custom_tokens_add").val()), 
 154 |         rewriteTokensMult: parseFloat($("#rewrite_tokens_mult").val()),
 155 |         shortenTokensMult: parseFloat($("#shorten_tokens_mult").val()),
 156 |         expandTokensMult: parseFloat($("#expand_tokens_mult").val()),
 157 |         customTokensMult: parseFloat($("#custom_tokens_mult").val()), 
 158 |         removePrefix: $("#remove_prefix").val(),
 159 |         removeSuffix: $("#remove_suffix").val(),
 160 |         overrideMaxTokens: $("#override_max_tokens").is(':checked'),
 161 |         showRewrite: $("#show_rewrite").is(':checked'),
 162 |         showShorten: $("#show_shorten").is(':checked'),
 163 |         showExpand: $("#show_expand").is(':checked'),
 164 |         showCustom: $("#show_custom").is(':checked'), 
 165 |         showDelete: $("#show_delete").is(':checked'),
 166 |         applyRegexOnRewrite: $("#apply_regex_on_rewrite").is(':checked'), // Save new setting
 167 |     };
 168 | 
 169 |     // Ensure all settings have a value, using defaults if necessary
 170 |     for (const [key, value] of Object.entries(defaultSettings)) {
 171 |         if (extension_settings[extensionName][key] === undefined) {
 172 |             extension_settings[extensionName][key] = value;
 173 |         }
 174 |     }
 175 | 
 176 |     saveSettingsDebounced();
 177 | }
 178 | 
 179 | // Populate dropdowns
 180 | async function populateDropdowns() {
 181 |     const result = await fetch('/api/settings/get', {
 182 |         method: 'POST',
 183 |         headers: getContext().getRequestHeaders(),
 184 |         body: JSON.stringify({}),
 185 |     });
 186 | 
 187 |     if (result.ok) {
 188 |         const data = await result.json();
 189 |         const presets = data.openai_setting_names;
 190 |         const dropdowns = ['rewrite_preset', 'shorten_preset', 'expand_preset', 'custom_preset']; // Added custom_preset
 191 |         dropdowns.forEach(dropdown => {
 192 |             const select = $(`#${dropdown}`);
 193 |             select.empty();
 194 |             presets.forEach(preset => {
 195 |                 select.append($('<option>', {
 196 |                     value: preset,
 197 |                     text: preset
 198 |                 }));
 199 |             });
 200 |         });
 201 | 
 202 |         // Set the selected values after populating
 203 |         loadSettings();
 204 |     }
 205 | }
 206 | 
 207 | function updateModelSettings() {
 208 |     const modelSelect = document.getElementById('rewrite_extension_model_select');
 209 |     const chatCompletionSettings = document.getElementById('chat_completion_settings');
 210 |     const textBasedSettings = document.getElementById('text_based_settings');
 211 | 
 212 |     if (modelSelect.value === 'chat_completion') {
 213 |         chatCompletionSettings.style.display = 'block';
 214 |         textBasedSettings.style.display = 'none';
 215 |     } else {
 216 |         chatCompletionSettings.style.display = 'none';
 217 |         textBasedSettings.style.display = 'block';
 218 |     }
 219 | }
 220 | 
 221 | function updateTokenSettings() {
 222 |     const useDynamicTokens = $("#use_dynamic_tokens").is(':checked');
 223 |     const dynamicTokenMode = $("#dynamic_token_mode").val();
 224 |     $("#static_token_settings").toggle(!useDynamicTokens);
 225 |     $("#dynamic_token_settings").toggle(useDynamicTokens);
 226 |     $("#additive_settings").toggle(dynamicTokenMode === 'additive');
 227 |     $("#multiplicative_settings").toggle(dynamicTokenMode === 'multiplicative');
 228 | }
 229 | 
 230 | // Initialize
 231 | jQuery(async () => {
 232 |     const settingsHtml = await $.get(`${extensionFolderPath}/rewrite_settings.html`);
 233 |     $("#extensions_settings2").append(settingsHtml);
 234 | 
 235 |     // Populate dropdowns
 236 |     await populateDropdowns();
 237 | 
 238 |     // Add event listeners
 239 |     $(".rewrite-extension-settings select, #highlight_duration, #text_rewrite_prompt, #text_shorten_prompt, #text_expand_prompt, #text_custom_prompt").on("change", saveSettings); // Added #text_custom_prompt
 240 |     $("#use_streaming").on("change", saveSettings);
 241 |     $("#use_dynamic_tokens, #dynamic_token_mode").on("change", () => {
 242 |         updateTokenSettings();
 243 |         saveSettings();
 244 |     });
 245 |     $("#rewrite_tokens, #shorten_tokens, #expand_tokens, #custom_tokens, #rewrite_tokens_add, #shorten_tokens_add, #expand_tokens_add, #custom_tokens_add, #rewrite_tokens_mult, #shorten_tokens_mult, #expand_tokens_mult, #custom_tokens_mult").on("input", saveSettings); // Added custom token inputs
 246 |     $("#remove_prefix, #remove_suffix").on("change", saveSettings);
 247 |     $("#override_max_tokens").on("change", saveSettings);
 248 |     $("#show_rewrite, #show_shorten, #show_expand, #show_custom, #show_delete").on("change", saveSettings); // Added #show_custom
 249 |     $("#apply_regex_on_rewrite").on("change", saveSettings); // Add listener for new checkbox
 250 | 
 251 |     $("#rewrite_extension_model_select").on("change", () => {
 252 |         updateModelSettings();
 253 |         saveSettings();
 254 |     });
 255 | 
 256 |     // Load settings
 257 |     loadSettings();
 258 | 
 259 |     // Add event listener for SETTINGS_UPDATED
 260 |     eventSource.on(event_types.SETTINGS_UPDATED, () => {
 261 |         populateDropdowns();
 262 |     });
 263 | 
 264 |     eventSource.on(event_types.CHAT_CHANGED, () => {
 265 |         changeHistory = [];
 266 |         updateUndoButtons();
 267 |     });
 268 | 
 269 |     eventSource.on(event_types.MESSAGE_EDITED, (editedMesId) => {
 270 |         removeUndoButton(editedMesId);
 271 |     });
 272 | 
 273 |     updateModelSettings();
 274 | });
 275 | 
 276 | // Initialize the rewrite menu functionality
 277 | initRewriteMenu();
 278 | 
 279 | function initRewriteMenu() {
 280 |     // document.addEventListener('mouseup', handleSelectionEnd);
 281 |     // document.addEventListener('touchend', handleSelectionEnd);
 282 |     document.addEventListener('selectionchange', handleSelectionChange);
 283 |     document.addEventListener('mousedown', hideMenuOnOutsideClick);
 284 |     document.addEventListener('touchstart', hideMenuOnOutsideClick);
 285 | 
 286 |     let chatContainer = document.getElementById('chat');
 287 |     chatContainer.addEventListener('scroll', positionMenu);
 288 | 
 289 |     $('#mes_stop').on('click', handleStopRewrite);
 290 | }
 291 | 
 292 | 
 293 | function handleStopRewrite() {
 294 |     if (abortController) {
 295 |         const { mesDiv, mesId, swipeId, highlightDuration } = abortController.signal;
 296 |         abortController.abort();
 297 |         // Restore the original settings
 298 |         if (abortController.signal.prev_oai_settings) {
 299 |             Object.assign(oai_settings, abortController.signal.prev_oai_settings);
 300 |         }
 301 | 
 302 |         getContext().activateSendButtons();
 303 | 
 304 |         // Call removeHighlight with the stored arguments
 305 |         setTimeout(() => removeHighlight(mesDiv, mesId, swipeId), highlightDuration);
 306 |     }
 307 | }
 308 | 
 309 | // function handleSelectionEnd(e) {
 310 | //     if (e.target && e.target.closest('.ctx-menu')) return;
 311 | //     removeRewriteMenu();
 312 | //     setTimeout(processSelection, 50);
 313 | // }
 314 | 
 315 | function handleSelectionChange() {
 316 |     // Use a small timeout to ensure the selection has been updated
 317 |     setTimeout(processSelection, 50);
 318 | }
 319 | 
 320 | function processSelection() {
 321 |     // First, check if getContext().chatId is defined
 322 |     if (getContext().chatId === undefined) {
 323 |         return; // Exit the function if chatId is undefined
 324 |     }
 325 | 
 326 |     let selection = window.getSelection();
 327 |     let selectedText = selection.toString().trim();
 328 | 
 329 |     // Always remove the existing menu first
 330 |     removeRewriteMenu();
 331 | 
 332 |     if (selectedText.length > 0) {
 333 |         let range = selection.getRangeAt(0);
 334 | 
 335 |         // Find the mes_text elements for both start and end of the selection
 336 |         let startMesText = range.startContainer.nodeType === Node.ELEMENT_NODE
 337 |             ? range.startContainer.closest('.mes_text')
 338 |             : range.startContainer.parentElement.closest('.mes_text');
 339 | 
 340 |         let endMesText = range.endContainer.nodeType === Node.ELEMENT_NODE
 341 |             ? range.endContainer.closest('.mes_text')
 342 |             : range.endContainer.parentElement.closest('.mes_text');
 343 | 
 344 |         // Check if both start and end are within the same mes_text element
 345 |         if (startMesText && endMesText && startMesText === endMesText) {
 346 |             createRewriteMenu();
 347 |         }
 348 |     }
 349 | 
 350 |     lastSelection = selectedText.length > 0 ? selectedText : null;
 351 | }
 352 | 
 353 | async function getCustomInstructionsFromPopup() {
 354 |     const { callPopup } = getContext();
 355 |     try {
 356 |         const instructions = await callPopup('Enter custom rewrite instructions:', 'input');
 357 | 
 358 |         // Introduce a zero-delay setTimeout to yield to the event loop
 359 |         await new Promise(resolve => setTimeout(resolve, 0));
 360 | 
 361 |         return instructions;
 362 |     } catch (error) {
 363 |         console.error("[Rewrite Extension] Error during custom instruction popup:", error);
 364 |         return null;
 365 |     } finally {
 366 |     }
 367 | }
 368 | 
 369 | async function handleMenuItemClick(e) {
 370 |     e.preventDefault();
 371 |     e.stopPropagation();
 372 | 
 373 |     const option = e.target.dataset.option;
 374 |     const selection = window.getSelection();
 375 | 
 376 |     // Ensure there's a selection and a range
 377 |     if (!selection || selection.rangeCount === 0) {
 378 |         removeRewriteMenu();
 379 |         return;
 380 |     }
 381 | 
 382 |     // Capture the range *before* any awaits or potential selection changes
 383 |     const initialRange = selection.getRangeAt(0).cloneRange();
 384 |     const selectedText = initialRange.toString().trim();
 385 | 
 386 |     if (selectedText) {
 387 |         const mesTextElement = findClosestMesText(selection.anchorNode);
 388 |         if (mesTextElement) {
 389 |             const messageDiv = findMessageDiv(mesTextElement);
 390 |             if (messageDiv) {
 391 |                 const mesId = messageDiv.getAttribute('mesid');
 392 |                 const swipeId = messageDiv.getAttribute('swipeid');
 393 | 
 394 |                 if (option === 'Delete') {
 395 |                     // Pass the initially captured range to handleDeleteSelection
 396 |                     await handleDeleteSelection(mesId, swipeId, initialRange);
 397 |                 } else if (option === 'Custom') {
 398 |                     const customInstructions = await getCustomInstructionsFromPopup();
 399 |                     if (customInstructions !== null && customInstructions.trim() !== '') { // Proceed only if user entered text and didn't cancel
 400 |                         // Get selectionInfo *after* await and *before* handleRewrite
 401 |                         // Pass the initially captured range
 402 |                         const selectionInfo = getSelectedTextInfo(mesId, mesTextElement, initialRange);
 403 |                         if (!selectionInfo) {
 404 |                              console.error("[Rewrite Extension] Failed to get selectionInfo for Custom rewrite!");
 405 |                              return; // Prevent calling with undefined
 406 |                         }
 407 |                         await handleRewrite(mesId, swipeId, option, customInstructions, selectionInfo); // Use the locally scoped selectionInfo
 408 |                     } else {
 409 |                         // User cancelled or entered empty instructions
 410 |                     }
 411 |                 } else {
 412 |                     // For other rewrite options, get selectionInfo right before the call
 413 |                     // Pass the initially captured range
 414 |                     const selectionInfo = getSelectedTextInfo(mesId, mesTextElement, initialRange); // Get selectionInfo here
 415 |                     if (!selectionInfo) {
 416 |                          console.error(`[Rewrite Extension] Failed to get selectionInfo for ${option} rewrite!`);
 417 |                          return; // Prevent calling with undefined
 418 |                     }
 419 |                     await handleRewrite(mesId, swipeId, option, null, selectionInfo); // Use the locally scoped selectionInfo
 420 |                 }
 421 |             }
 422 |         }
 423 |     }
 424 | 
 425 |     removeRewriteMenu();
 426 |     window.getSelection().removeAllRanges();
 427 | }
 428 | 
 429 | // Modify signature to accept the captured range
 430 | async function handleDeleteSelection(mesId, swipeId, range) {
 431 |     const mesDiv = document.querySelector(`[mesid="${mesId}"] .mes_text`);
 432 |     // Use the passed-in range to get selection info
 433 |     const { fullMessage, selectedRawText, rawStartOffset, rawEndOffset } = getSelectedTextInfo(mesId, mesDiv, range);
 434 | 
 435 |     // Create the new message with the deleted section removed
 436 |     const newMessage = fullMessage.slice(0, rawStartOffset) + fullMessage.slice(rawEndOffset);
 437 | 
 438 |     // Save the change to the history (this also calls updateUndoButtons)
 439 |     saveLastChange(mesId, swipeId, fullMessage, newMessage);
 440 | 
 441 |     // Update the message in the chat context
 442 |     getContext().chat[mesId].mes = newMessage;
 443 |     if (swipeId !== undefined && getContext().chat[mesId].swipes) {
 444 |         getContext().chat[mesId].swipes[swipeId] = newMessage;
 445 |     }
 446 | 
 447 |     // Update the UI
 448 |     mesDiv.innerHTML = messageFormatting(newMessage, getContext().name2, getContext().chat[mesId].isSystem, getContext().chat[mesId].isUser, mesId);
 449 |     addCopyToCodeBlocks(mesDiv);
 450 | 
 451 |     // Save the chat
 452 |     await getContext().saveChat();
 453 | }
 454 | 
 455 | function hideMenuOnOutsideClick(e) {
 456 |     if (rewriteMenu && !rewriteMenu.contains(e.target)) {
 457 |         removeRewriteMenu();
 458 |     }
 459 | }
 460 | 
 461 | function createRewriteMenu() {
 462 |     removeRewriteMenu();
 463 | 
 464 |     rewriteMenu = document.createElement('ul');
 465 |     rewriteMenu.className = 'list-group ctx-menu';
 466 |     rewriteMenu.style.position = 'absolute';
 467 |     rewriteMenu.style.zIndex = '1000';
 468 |     rewriteMenu.style.position = 'fixed';
 469 | 
 470 |     const options = [
 471 |         { name: 'Rewrite', show: extension_settings[extensionName].showRewrite },
 472 |         { name: 'Shorten', show: extension_settings[extensionName].showShorten },
 473 |         { name: 'Expand', show: extension_settings[extensionName].showExpand },
 474 |         { name: 'Custom', show: extension_settings[extensionName].showCustom }, 
 475 |         { name: 'Delete', show: extension_settings[extensionName].showDelete }
 476 |     ];
 477 |     options.forEach(option => {
 478 |         if (option.show) {
 479 |             let li = document.createElement('li');
 480 |             li.className = 'list-group-item ctx-item';
 481 |             li.textContent = option.name;
 482 |             li.addEventListener('mousedown', handleMenuItemClick);
 483 |             li.addEventListener('touchstart', handleMenuItemClick);
 484 |             li.dataset.option = option.name;
 485 |             rewriteMenu.appendChild(li);
 486 |         }
 487 |     });
 488 | 
 489 |     document.body.appendChild(rewriteMenu);
 490 |     positionMenu();
 491 | }
 492 | 
 493 | function positionMenu() {
 494 |     if (!rewriteMenu) return;
 495 | 
 496 |     let selection = window.getSelection();
 497 |     let range = selection.getRangeAt(0);
 498 |     let rect = range.getBoundingClientRect();
 499 | 
 500 |     // Calculate the menu's position
 501 |     let left = rect.left + window.pageXOffset;
 502 |     let top = rect.bottom + window.pageYOffset + 5;
 503 | 
 504 |     // Get the viewport dimensions
 505 |     let viewportWidth = window.innerWidth;
 506 |     let viewportHeight = window.innerHeight;
 507 | 
 508 |     // Get the menu's dimensions
 509 |     let menuWidth = rewriteMenu.offsetWidth;
 510 |     let menuHeight = rewriteMenu.offsetHeight;
 511 | 
 512 |     // Adjust the position if the menu overflows the viewport
 513 |     if (left + menuWidth > viewportWidth) {
 514 |         left = viewportWidth - menuWidth;
 515 |     }
 516 |     if (top + menuHeight > viewportHeight) {
 517 |         top = rect.top + window.pageYOffset - menuHeight - 5;
 518 |     }
 519 | 
 520 |     rewriteMenu.style.left = `${left}px`;
 521 |     rewriteMenu.style.top = `${top}px`;
 522 | }
 523 | 
 524 | function removeRewriteMenu() {
 525 |     if (rewriteMenu) {
 526 |         rewriteMenu.remove();
 527 |         rewriteMenu = null;
 528 |     }
 529 | }
 530 | 
 531 | function addUndoButton(mesId) {
 532 |     const messageDiv = document.querySelector(`[mesid="${mesId}"]`);
 533 |     if (messageDiv) {
 534 |         const mesButtons = messageDiv.querySelector('.mes_buttons');
 535 |         if (mesButtons) {
 536 |             const undoButton = document.createElement('div');
 537 |             undoButton.className = 'mes_button mes_undo_rewrite fa-solid fa-undo interactable';
 538 |             undoButton.title = 'Undo rewrite';
 539 |             undoButton.dataset.mesId = mesId;
 540 |             undoButton.addEventListener('click', handleUndo);
 541 | 
 542 |             if (mesButtons.children.length >= 1) {
 543 |                 mesButtons.insertBefore(undoButton, mesButtons.children[1]);
 544 |             } else {
 545 |                 mesButtons.appendChild(undoButton);
 546 |             }
 547 |         }
 548 |     }
 549 | }
 550 | 
 551 | function removeUndoButton(editedMesId) {
 552 |     // Remove all changes for this message from the changeHistory
 553 |     changeHistory = changeHistory.filter(change => change.mesId !== editedMesId);
 554 | 
 555 |     // Update undo buttons for other messages
 556 |     updateUndoButtons();
 557 | }
 558 | 
 559 | async function removeHighlight(mesDiv, mesId, swipeId) {
 560 |     const highlightSpan = mesDiv.querySelector('.animated-highlight');
 561 |     if (highlightSpan) {
 562 |         const textNode = document.createTextNode(highlightSpan.textContent);
 563 |         highlightSpan.parentNode.replaceChild(textNode, highlightSpan);
 564 |     }
 565 | 
 566 |     const context = getContext();
 567 |     const messageData = context.chat[mesId];
 568 | 
 569 |     if (messageData) {
 570 |         let messageContent;
 571 |         if (swipeId !== undefined && messageData.swipes && messageData.swipes[swipeId]) {
 572 |             messageContent = messageData.swipes[swipeId];
 573 |         } else {
 574 |             messageContent = messageData.mes;
 575 |         }
 576 | 
 577 |         // Format the message into HTML
 578 |         const formattedMessage = messageFormatting(
 579 |             messageContent,
 580 |             context.name2,
 581 |             messageData.isSystem,
 582 |             messageData.isUser,
 583 |             mesId
 584 |         );
 585 | 
 586 |         // Create a temporary div to hold the formatted message
 587 |         const tempDiv = document.createElement('div');
 588 |         tempDiv.innerHTML = formattedMessage;
 589 | 
 590 |         // Apply addCopyToCodeBlocks to the temporary div
 591 |         addCopyToCodeBlocks(tempDiv);
 592 | 
 593 |         // Find the mes_text element within the message div
 594 |         const mesTextElement = mesDiv.closest('.mes').querySelector('.mes_text');
 595 |         if (mesTextElement) {
 596 |             // Replace the content of mes_text with the new formatted content
 597 |             mesTextElement.innerHTML = tempDiv.innerHTML;
 598 |         }
 599 |     }
 600 | }
 601 | 
 602 | function findClosestMesText(element) {
 603 |     while (element && element.nodeType !== 1) {
 604 |         element = element.parentElement;
 605 |     }
 606 |     while (element) {
 607 |         if (element.classList && element.classList.contains('mes_text')) {
 608 |             return element;
 609 |         }
 610 |         element = element.parentElement;
 611 |     }
 612 |     return null;
 613 | }
 614 | 
 615 | function findMessageDiv(element) {
 616 |     while (element) {
 617 |         if (element.hasAttribute('mesid') && element.hasAttribute('swipeid')) {
 618 |             return element;
 619 |         }
 620 |         element = element.parentElement;
 621 |     }
 622 |     return null;
 623 | }
 624 | 
 625 | function createTextMapping(rawText, formattedHtml) {
 626 |     const formattedText = stripHtml(formattedHtml);
 627 |     const mapping = [];
 628 |     let rawIndex = 0;
 629 |     let formattedIndex = 0;
 630 | 
 631 |     while (rawIndex < rawText.length && formattedIndex < formattedText.length) {
 632 |         if (rawText[rawIndex] === formattedText[formattedIndex]) {
 633 |             mapping.push([rawIndex, formattedIndex]);
 634 |             rawIndex++;
 635 |             formattedIndex++;
 636 |         } else if (rawText.substr(rawIndex, 3) === '...' && formattedText[formattedIndex] === '…') {
 637 |             // Handle ellipsis
 638 |             mapping.push([rawIndex, formattedIndex]);
 639 |             mapping.push([rawIndex + 1, formattedIndex]);
 640 |             mapping.push([rawIndex + 2, formattedIndex]);
 641 |             rawIndex += 3;
 642 |             formattedIndex++;
 643 |         } else if (formattedText[formattedIndex] === ' ' || formattedText[formattedIndex] === '\n') {
 644 |             // Skip extra whitespace in formatted text
 645 |             formattedIndex++;
 646 |         } else {
 647 |             // Skip characters in raw text that don't appear in formatted text
 648 |             rawIndex++;
 649 |         }
 650 |     }
 651 | 
 652 |     return {
 653 |         formattedToRaw: (formattedOffset) => {
 654 |             let low = 0;
 655 |             let high = mapping.length - 1;
 656 | 
 657 |             while (low <= high) {
 658 |                 let mid = Math.floor((low + high) / 2);
 659 |                 if (mapping[mid][1] === formattedOffset) {
 660 |                     return mapping[mid][0];
 661 |                 } else if (mapping[mid][1] < formattedOffset) {
 662 |                     low = mid + 1;
 663 |                 } else {
 664 |                     high = mid - 1;
 665 |                 }
 666 |             }
 667 | 
 668 |             // If we didn't find an exact match, return the closest one
 669 |             if (low > 0) low--;
 670 |             return mapping[low][0] + (formattedOffset - mapping[low][1]);
 671 |         }
 672 |     };
 673 | }
 674 | 
 675 | function stripHtml(html) {
 676 |     const tmp = document.createElement('DIV');
 677 |     tmp.innerHTML = html;
 678 |     return tmp.textContent || tmp.innerText || '';
 679 | }
 680 | 
 681 | function getTextOffset(parent, node) {
 682 |     const treeWalker = document.createTreeWalker(
 683 |         parent,
 684 |         NodeFilter.SHOW_TEXT,
 685 |         null,
 686 |         false
 687 |     );
 688 | 
 689 |     let offset = 0;
 690 |     while (treeWalker.nextNode() !== node) {
 691 |         offset += treeWalker.currentNode.length;
 692 |     }
 693 | 
 694 |     return offset;
 695 | }
 696 | 
 697 | // Modify signature to accept the captured range
 698 | function getSelectedTextInfo(mesId, mesDiv, range) {
 699 |     // Removed: const selection = window.getSelection();
 700 |     // Removed: const range = selection.getRangeAt(0); - Use the passed-in range directly
 701 | 
 702 |     // Get the full message content
 703 |     const fullMessage = getContext().chat[mesId].mes;
 704 | 
 705 |     // Get the formatted message
 706 |     const formattedMessage = messageFormatting(fullMessage, undefined, getContext().chat[mesId].isSystem, getContext().chat[mesId].isUser, mesId);
 707 | 
 708 |     // Create a mapping between raw and formatted text
 709 |     const mapping = createTextMapping(fullMessage, formattedMessage);
 710 | 
 711 |     // Calculate the start and end offsets relative to the formatted text content
 712 |     const startOffset = getTextOffset(mesDiv, range.startContainer) + range.startOffset;
 713 |     const endOffset = getTextOffset(mesDiv, range.endContainer) + range.endOffset;
 714 | 
 715 |     // Map these offsets back to the raw message
 716 |     let rawStartOffset = mapping.formattedToRaw(startOffset);
 717 |     let rawEndOffset = mapping.formattedToRaw(endOffset);
 718 | 
 719 |     // Heuristic: Adjust offsets to include surrounding markdown if selection seems to abut it
 720 |     // Check for italics (*)
 721 |     if (rawStartOffset > 0 && rawEndOffset < fullMessage.length &&
 722 |         fullMessage[rawStartOffset - 1] === '*' && fullMessage[rawEndOffset] === '*') {
 723 |         // Avoid expanding if it looks like bold/bold-italics boundary
 724 |         const prevChar = rawStartOffset > 1 ? fullMessage[rawStartOffset - 2] : null;
 725 |         const nextChar = rawEndOffset + 1 < fullMessage.length ? fullMessage[rawEndOffset + 1] : null;
 726 |         if (prevChar !== '*' && nextChar !== '*') {
 727 |             rawStartOffset--;
 728 |             rawEndOffset++;
 729 |         }
 730 |     }
 731 |     // Check for bold (**) - ensure we don't double-adjust if italics check already expanded
 732 |     else if (rawStartOffset > 1 && rawEndOffset < fullMessage.length - 1 &&
 733 |              fullMessage.substring(rawStartOffset - 2, rawStartOffset) === '**' &&
 734 |              fullMessage.substring(rawEndOffset, rawEndOffset + 2) === '**') {
 735 |         // Avoid expanding if it looks like bold-italics boundary
 736 |         const prevChar = rawStartOffset > 2 ? fullMessage[rawStartOffset - 3] : null;
 737 |         const nextChar = rawEndOffset + 2 < fullMessage.length ? fullMessage[rawEndOffset + 2] : null;
 738 |         if (prevChar !== '*' && nextChar !== '*') {
 739 |             rawStartOffset -= 2;
 740 |             rawEndOffset += 2;
 741 |         }
 742 |     }
 743 |     // Note: This doesn't handle ***bold italics*** or nested cases perfectly, but covers common scenarios.
 744 | 
 745 |     // Get the selected raw text using potentially adjusted offsets
 746 |     const selectedRawText = fullMessage.substring(rawStartOffset, rawEndOffset);
 747 | 
 748 |     return {
 749 |         fullMessage,
 750 |         selectedRawText,
 751 |         rawStartOffset,
 752 |         rawEndOffset,
 753 |         range
 754 |     };
 755 | }
 756 | 
 757 | function saveLastChange(mesId, swipeId, originalContent, newContent) {
 758 |     changeHistory.push({
 759 |         mesId,
 760 |         swipeId,
 761 |         originalContent,
 762 |         newContent,
 763 |         timestamp: Date.now()
 764 |     });
 765 | 
 766 |     // Limit history to last n changes
 767 |     if (changeHistory.length > undo_steps) {
 768 |         changeHistory.shift();
 769 |     }
 770 | 
 771 |     updateUndoButtons();
 772 | }
 773 | 
 774 | function updateUndoButtons() {
 775 |     // Remove all existing undo buttons
 776 |     document.querySelectorAll('.mes_undo_rewrite').forEach(button => button.remove());
 777 | 
 778 |     // Add undo buttons for all messages with changes
 779 |     const changedMessageIds = [...new Set(changeHistory.map(change => change.mesId))];
 780 |     changedMessageIds.forEach(mesId => addUndoButton(mesId));
 781 | }
 782 | 
 783 | // Updated handleRewrite signature to accept selectionInfo
 784 | async function handleRewrite(mesId, swipeId, option, customInstructions = null, selectionInfo) {
 785 |     if (!selectionInfo) {
 786 |         console.error("[Rewrite Extension] handleRewrite called without selectionInfo!");
 787 |         return; // Cannot proceed without selection info
 788 |     }
 789 | 
 790 |     if (main_api === 'openai') {
 791 |         const selectedModel = extension_settings[extensionName].selectedModel;
 792 |         if (selectedModel === 'chat_completion') {
 793 |             return handleChatCompletionRewrite(mesId, swipeId, option, customInstructions, selectionInfo); // Pass selectionInfo
 794 |         } else {
 795 |             return handleSimplifiedChatCompletionRewrite(mesId, swipeId, option, customInstructions, selectionInfo); // Pass selectionInfo
 796 |         }
 797 |     } else {
 798 |         return handleTextBasedRewrite(mesId, swipeId, option, customInstructions, selectionInfo); // Pass selectionInfo
 799 |     }
 800 | }
 801 | 
 802 | // Updated signature to accept selectionInfo
 803 | async function handleChatCompletionRewrite(mesId, swipeId, option, customInstructions, selectionInfo) {
 804 |     // Use pre-captured selection info
 805 |     const { fullMessage, selectedRawText, rawStartOffset, rawEndOffset, range } = selectionInfo;
 806 |     const mesDiv = document.querySelector(`[mesid="${mesId}"] .mes_text`); // Keep getting mesDiv for highlight/DOM ops
 807 |     if (!mesDiv) { // Add check for mesDiv existence
 808 |         console.error("[Rewrite Extension] Could not find mesDiv in handleChatCompletionRewrite.");
 809 |         return;
 810 |     }
 811 | 
 812 |     // Get the selected preset based on the option
 813 |     let selectedPreset;
 814 |     switch (option) {
 815 |         case 'Rewrite':
 816 |             selectedPreset = extension_settings[extensionName].rewritePreset;
 817 |             break;
 818 |         case 'Shorten':
 819 |             selectedPreset = extension_settings[extensionName].shortenPreset;
 820 |             break;
 821 |         case 'Expand':
 822 |             selectedPreset = extension_settings[extensionName].expandPreset;
 823 |             break;
 824 |         case 'Custom': // New case
 825 |             selectedPreset = extension_settings[extensionName].customPreset;
 826 |             break;
 827 |         default:
 828 |             console.error("Unknown rewrite option:", option);
 829 |             return; // Exit if the option is not recognized
 830 |     }
 831 | 
 832 |     // Fetch the settings
 833 |     const result = await fetch('/api/settings/get', {
 834 |         method: 'POST',
 835 |         headers: getContext().getRequestHeaders(),
 836 |         body: JSON.stringify({}),
 837 |     });
 838 | 
 839 |     if (!result.ok) {
 840 |         console.error('Failed to fetch settings');
 841 |         return;
 842 |     }
 843 | 
 844 |     const data = await result.json();
 845 |     const presetIndex = data.openai_setting_names.indexOf(selectedPreset);
 846 |     if (presetIndex === -1) {
 847 |         console.error('Selected preset not found');
 848 |         return;
 849 |     }
 850 | 
 851 |     // Save the current settings
 852 |     const prev_oai_settings = Object.assign({}, oai_settings);
 853 | 
 854 |     // Parse the selected preset settings
 855 |     let selectedPresetSettings;
 856 |     try {
 857 |         selectedPresetSettings = JSON.parse(data.openai_settings[presetIndex]);
 858 |     } catch (error) {
 859 |         console.error('Error parsing preset settings:', error);
 860 |         return;
 861 |     }
 862 | 
 863 |     // Extension streaming overrides preset streaming
 864 |     selectedPresetSettings.stream_openai = extension_settings[extensionName].useStreaming;
 865 | 
 866 |     if (extension_settings[extensionName].overrideMaxTokens) {
 867 |         selectedPresetSettings.openai_max_tokens = calculateTargetTokenCount(selectedRawText, option);
 868 |     }
 869 | 
 870 |     // Override oai_settings with the selected preset
 871 |     Object.assign(oai_settings, selectedPresetSettings);
 872 | 
 873 |     // Always generate the base prompt using the selected preset
 874 |     const promptReadyPromise = new Promise(resolve => {
 875 |         eventSource.once(event_types.CHAT_COMPLETION_PROMPT_READY, resolve);
 876 |     });
 877 |     getContext().generate('normal', {}, true); // Trigger prompt generation
 878 |     const promptData = await promptReadyPromise; // Wait for the generated prompt
 879 |     let chatToSend = promptData.chat; // Start with the generated chat array
 880 | 
 881 |     // Inject custom instructions if applicable
 882 |     if (option === 'Custom' && customInstructions) {
 883 |         // Find the last user message to append to
 884 |         let targetMessageIndex = -1;
 885 |         for (let i = chatToSend.length - 1; i >= 0; i--) {
 886 |             if (chatToSend[i].role === 'user') {
 887 |                 targetMessageIndex = i;
 888 |                 break;
 889 |             }
 890 |         }
 891 | 
 892 |         if (targetMessageIndex !== -1) {
 893 |             const targetMessage = chatToSend[targetMessageIndex];
 894 |             const instructionText = `\n\nAdditional Instructions:\n${customInstructions}`;
 895 | 
 896 |             if (Array.isArray(targetMessage.content)) {
 897 |                 // Find the last text part or add a new one
 898 |                 let lastTextPartIndex = -1;
 899 |                 for (let j = targetMessage.content.length - 1; j >= 0; j--) {
 900 |                     if (targetMessage.content[j].type === 'text') {
 901 |                         lastTextPartIndex = j;
 902 |                         break;
 903 |                     }
 904 |                 }
 905 |                 if (lastTextPartIndex !== -1) {
 906 |                     targetMessage.content[lastTextPartIndex].text += instructionText;
 907 |                 } else {
 908 |                     // Should not happen with standard prompts, but handle just in case
 909 |                     targetMessage.content.push({ type: 'text', text: instructionText });
 910 |                 }
 911 |             } else if (typeof targetMessage.content === 'string') {
 912 |                 targetMessage.content += instructionText;
 913 |             }
 914 |         } else {
 915 |             console.warn('[Rewrite Extension] Could not find a user message in the generated prompt to inject custom instructions into.');
 916 |             // Optionally, could append a new user message, but might break formatting
 917 |             // chatToSend.push({ role: "user", content: `Additional Instructions:\n${customInstructions}` });
 918 |         }
 919 |     }
 920 | 
 921 |     // Substitute standard macros AFTER potential custom instruction injection
 922 |     const wordCount = extractAllWords(selectedRawText).length;
 923 |     chatToSend = chatToSend.map(message => {
 924 |         if (Array.isArray(message.content)) {
 925 |             message.content = message.content.map(item => {
 926 |                 if (item.type === 'text') {
 927 |                     item.text = item.text.replace(/{{rewrite}}/gi, selectedRawText);
 928 |                     item.text = item.text.replace(/{{targetmessage}}/gi, fullMessage);
 929 |                     item.text = item.text.replace(/{{rewritecount}}/gi, wordCount);
 930 |                 }
 931 |                 return item;
 932 |             });
 933 |         } else if (typeof message.content === 'string') {
 934 |             message.content = message.content.replace(/{{rewrite}}/gi, selectedRawText);
 935 |             message.content = message.content.replace(/{{targetmessage}}/gi, fullMessage);
 936 |             message.content = message.content.replace(/{{rewritecount}}/gi, wordCount);
 937 |         }
 938 |         return message;
 939 |     });
 940 | 
 941 |     // Create a new AbortController
 942 |     abortController = new AbortController();
 943 | 
 944 |     // Store the necessary data in the signal
 945 |     abortController.signal.prev_oai_settings = prev_oai_settings;
 946 |     abortController.signal.mesDiv = mesDiv;
 947 |     abortController.signal.mesId = mesId;
 948 |     abortController.signal.swipeId = swipeId;
 949 |     abortController.signal.highlightDuration = extension_settings[extensionName].highlightDuration;
 950 | 
 951 |     // Show the stop button
 952 |     getContext().deactivateSendButtons();
 953 | 
 954 |     let res;
 955 |     try {
 956 | 
 957 |         // Send the request with the prepared chat
 958 |         res = await sendOpenAIRequest('normal', chatToSend, abortController.signal);
 959 |     } catch (error) {
 960 |         console.error('[Rewrite Extension] Error during sendOpenAIRequest:', error);
 961 |         toastr.error("Rewrite failed. Check browser console (F12) for details.", "Rewrite Error");
 962 |         // Ensure cleanup happens even on error
 963 |     } finally {
 964 |         window.getSelection().removeAllRanges();
 965 |         // Restore the original settings (moved to finally)
 966 |         Object.assign(oai_settings, prev_oai_settings);
 967 |         getContext().activateSendButtons();
 968 |     }
 969 | 
 970 |     // If the request failed, res will be undefined, stop further processing
 971 |     if (res === undefined) {
 972 |         // Remove highlight immediately if the request failed before starting streaming/display
 973 |         removeHighlight(mesDiv, mesId, swipeId);
 974 |         return;
 975 |     }
 976 | 
 977 |     let newText = '';
 978 |     try {
 979 |         if (typeof res === 'function') {
 980 |             // Streaming case
 981 |             const streamingSpan = document.createElement('span');
 982 |             streamingSpan.className = 'animated-highlight';
 983 | 
 984 |             // Replace the selected text with the streaming span
 985 |             range.deleteContents();
 986 |             range.insertNode(streamingSpan);
 987 | 
 988 |             for await (const chunk of res()) {
 989 |                 newText = chunk.text;
 990 |                 streamingSpan.textContent = newText;
 991 |             }
 992 |         } else {
 993 |             // Non-streaming case
 994 |             newText = res?.choices?.[0]?.message?.content ?? res?.choices?.[0]?.text ?? res?.text ?? '';
 995 |             const highlightedNewText = document.createElement('span');
 996 |             highlightedNewText.className = 'animated-highlight';
 997 |             highlightedNewText.textContent = newText;
 998 | 
 999 |             range.deleteContents();
1000 |             range.insertNode(highlightedNewText);
1001 |         }
1002 | 
1003 |         // Remove highlight after x seconds when processing is complete
1004 |         const highlightDuration = extension_settings[extensionName].highlightDuration;
1005 |         setTimeout(() => removeHighlight(mesDiv, mesId, swipeId), highlightDuration);
1006 | 
1007 |         await saveRewrittenText(mesId, swipeId, fullMessage, rawStartOffset, rawEndOffset, newText);
1008 | 
1009 |     } catch (error) {
1010 |         console.error('[Rewrite Extension] Error processing API response:', error);
1011 |         toastr.error("Failed to process rewrite response. Check console.", "Processing Error");
1012 |         // Ensure highlight is removed if processing fails
1013 |         removeHighlight(mesDiv, mesId, swipeId);
1014 |     }
1015 |     // activateSendButtons is now handled in the finally block above
1016 | }
1017 | 
1018 | // Updated signature to accept selectionInfo
1019 | async function handleSimplifiedChatCompletionRewrite(mesId, swipeId, option, customInstructions, selectionInfo) {
1020 |     // Use pre-captured selection info
1021 |     const { fullMessage, selectedRawText, rawStartOffset, rawEndOffset, range } = selectionInfo;
1022 |     const mesDiv = document.querySelector(`[mesid="${mesId}"] .mes_text`); // Keep getting mesDiv for highlight/DOM ops
1023 |     if (!mesDiv) { // Add check for mesDiv existence
1024 |         console.error("[Rewrite Extension] Could not find mesDiv in handleSimplifiedChatCompletionRewrite.");
1025 |         return;
1026 |     }
1027 |     // Get the text completion prompt based on the option
1028 |     let promptTemplate;
1029 |     switch (option) {
1030 |         case 'Rewrite':
1031 |             promptTemplate = extension_settings[extensionName].textRewritePrompt;
1032 |             break;
1033 |         case 'Shorten':
1034 |             promptTemplate = extension_settings[extensionName].textShortenPrompt;
1035 |             break;
1036 |         case 'Expand':
1037 |             promptTemplate = extension_settings[extensionName].textExpandPrompt;
1038 |             break;
1039 |         case 'Custom': // New case
1040 |             promptTemplate = extension_settings[extensionName].textCustomPrompt;
1041 |             break;
1042 |         default:
1043 |             console.error("Unknown rewrite option:", option);
1044 |             return; // Exit if the option is not recognized
1045 |     }
1046 | 
1047 |     // Get amount of words
1048 |     const wordCount = extractAllWords(selectedRawText).length;
1049 | 
1050 |     // Replace macros in the prompt template
1051 |     let prompt = getContext().substituteParams(promptTemplate);
1052 | 
1053 |     prompt = prompt
1054 |         .replace(/{{rewrite}}/gi, selectedRawText)
1055 |         .replace(/{{targetmessage}}/gi, fullMessage)
1056 |         .replace(/{{rewritecount}}/gi, wordCount);
1057 | 
1058 |     // Inject custom instructions if applicable
1059 |     if (option === 'Custom') {
1060 |         if (prompt.includes('{{custom_instructions}}')) {
1061 |             prompt = prompt.replace(/{{custom_instructions}}/gi, customInstructions);
1062 |         } else {
1063 |             // Append if macro is missing (basic fallback)
1064 |             prompt += `\n\nInstructions: ${customInstructions}`;
1065 |         }
1066 |     }
1067 | 
1068 |     // Create a simplified chat format
1069 |     const simplifiedChat = [
1070 |         {
1071 |             role: "system",
1072 |             content: prompt
1073 |         }
1074 |     ];
1075 | 
1076 |     // Create a new AbortController
1077 |     abortController = new AbortController();
1078 | 
1079 |     // Store the necessary data in the signal
1080 |     abortController.signal.mesDiv = mesDiv;
1081 |     abortController.signal.mesId = mesId;
1082 |     abortController.signal.swipeId = swipeId;
1083 |     abortController.signal.highlightDuration = extension_settings[extensionName].highlightDuration;
1084 | 
1085 |     // Show the stop button
1086 |     getContext().deactivateSendButtons();
1087 | 
1088 |     const res = await sendOpenAIRequest('normal', simplifiedChat, abortController.signal);
1089 |     window.getSelection().removeAllRanges();
1090 | 
1091 |     let newText = '';
1092 | 
1093 |     if (typeof res === 'function') {
1094 |         // Streaming case
1095 |         const streamingSpan = document.createElement('span');
1096 |         streamingSpan.className = 'animated-highlight';
1097 | 
1098 |         // Replace the selected text with the streaming span
1099 |         range.deleteContents();
1100 |         range.insertNode(streamingSpan);
1101 | 
1102 |         for await (const chunk of res()) {
1103 |             newText = chunk.text;
1104 |             streamingSpan.textContent = newText;
1105 |         }
1106 |     } else {
1107 |         // Non-streaming case
1108 |         newText = res?.choices?.[0]?.message?.content ?? '';
1109 |         const highlightedNewText = document.createElement('span');
1110 |         highlightedNewText.className = 'animated-highlight';
1111 |         highlightedNewText.textContent = newText;
1112 | 
1113 |         range.deleteContents();
1114 |         range.insertNode(highlightedNewText);
1115 |     }
1116 | 
1117 |     // Remove highlight after x seconds when streaming is complete
1118 |     const highlightDuration = extension_settings[extensionName].highlightDuration;
1119 |     setTimeout(() => removeHighlight(mesDiv, mesId, swipeId), highlightDuration);
1120 | 
1121 |     await saveRewrittenText(mesId, swipeId, fullMessage, rawStartOffset, rawEndOffset, newText);
1122 |     getContext().activateSendButtons();
1123 | }
1124 | 
1125 | // Updated signature to accept selectionInfo
1126 | async function handleTextBasedRewrite(mesId, swipeId, option, customInstructions, selectionInfo) {
1127 |     // Use pre-captured selection info
1128 |     const { fullMessage, selectedRawText, rawStartOffset, rawEndOffset, range } = selectionInfo;
1129 |     const mesDiv = document.querySelector(`[mesid="${mesId}"] .mes_text`); // Keep getting mesDiv for highlight/DOM ops
1130 |     if (!mesDiv) { // Add check for mesDiv existence
1131 |         console.error("[Rewrite Extension] Could not find mesDiv in handleTextBasedRewrite.");
1132 |         return;
1133 |     }
1134 |     // Get the selected model and option-specific prompt
1135 |     const selectedModel = extension_settings[extensionName].selectedModel;
1136 |     let promptTemplate;
1137 |     switch (option) {
1138 |         case 'Rewrite':
1139 |             promptTemplate = extension_settings[extensionName].textRewritePrompt;
1140 |             break;
1141 |         case 'Shorten':
1142 |             promptTemplate = extension_settings[extensionName].textShortenPrompt;
1143 |             break;
1144 |         case 'Expand':
1145 |             promptTemplate = extension_settings[extensionName].textExpandPrompt;
1146 |             break;
1147 |         case 'Custom': // New case
1148 |             promptTemplate = extension_settings[extensionName].textCustomPrompt;
1149 |             break;
1150 |         default:
1151 |             console.error('Unknown rewrite option:', option);
1152 |             return;
1153 |     }
1154 | 
1155 |     // Get amount of words
1156 |     const wordCount = extractAllWords(selectedRawText).length;
1157 | 
1158 |     // Replace macros in the prompt template
1159 |     let prompt = getContext().substituteParams(promptTemplate);
1160 | 
1161 |     prompt = prompt
1162 |         .replace(/{{rewrite}}/gi, selectedRawText)
1163 |         .replace(/{{targetmessage}}/gi, fullMessage)
1164 |         .replace(/{{rewritecount}}/gi, wordCount);
1165 | 
1166 |     // Inject custom instructions if applicable
1167 |     if (option === 'Custom') {
1168 |         if (prompt.includes('{{custom_instructions}}')) {
1169 |             prompt = prompt.replace(/{{custom_instructions}}/gi, customInstructions);
1170 |         } else {
1171 |             // Append if macro is missing (basic fallback)
1172 |             prompt += `\n\nInstructions: ${customInstructions}`;
1173 |         }
1174 |     }
1175 | 
1176 |     let generateData;
1177 |     let amount_gen;
1178 | 
1179 |     if (extension_settings[extensionName].useDynamicTokens) {
1180 |         amount_gen = calculateTargetTokenCount(selectedRawText, option);
1181 |     } else {
1182 |         switch (option) {
1183 |             case 'Rewrite':
1184 |                 amount_gen = extension_settings[extensionName].rewriteTokens;
1185 |                 break;
1186 |             case 'Shorten':
1187 |                 amount_gen = extension_settings[extensionName].shortenTokens;
1188 |                 break;
1189 |             case 'Expand':
1190 |                 amount_gen = extension_settings[extensionName].expandTokens;
1191 |                 break;
1192 |             case 'Custom': // New case
1193 |                 amount_gen = extension_settings[extensionName].customTokens;
1194 |                 break;
1195 |         }
1196 |     }
1197 | 
1198 |     // Prepare generation data based on the selected model
1199 |     switch (main_api) {
1200 |         case 'novel':
1201 |             const novelSettings = novelai_settings[novelai_setting_names[nai_settings.preset_settings_novel]];
1202 |             generateData = getNovelGenerationData(prompt, novelSettings, amount_gen, false, false, null, 'quiet');
1203 |             break;
1204 |         case 'textgenerationwebui':
1205 |             generateData = getTextGenGenerationData(prompt, amount_gen, false, false, null, 'quiet');
1206 |             break;
1207 |         case 'koboldhorde':
1208 |             if (option === 'Custom') {
1209 |                 // For Custom Horde, use the manually constructed prompt directly
1210 |                 // We need a basic structure for generateHorde, mimicking what getContext().generate would provide
1211 |                 generateData = {
1212 |                     prompt: prompt, // Use the manually constructed prompt
1213 |                     max_length: Math.max(amount_gen, MIN_LENGTH),
1214 |                     // Include other necessary default parameters if generateHorde requires them
1215 |                     // Based on generateHorde usage, 'quiet' and potentially others might be needed.
1216 |                     quiet: true, // Often used in background generation
1217 |                 };
1218 |             } else {
1219 |                 // Existing logic for non-custom Horde rewrites
1220 |                 const promptReadyPromise = new Promise(resolve => {
1221 |                     eventSource.once(event_types.GENERATE_AFTER_DATA, resolve);
1222 |                 });
1223 |                 getContext().generate('normal', {}, true); // Trigger standard prompt generation
1224 |                 generateData = await promptReadyPromise; // Wait for the generated data
1225 |                 generateData.max_length = Math.max(amount_gen, MIN_LENGTH);
1226 |             }
1227 |             break;
1228 |         // Add more cases for other text-based models as needed
1229 |         default:
1230 |             toastr.error('Unsupported model:', main_api);
1231 |             return;
1232 |     }
1233 | 
1234 |     // Create a new AbortController
1235 |     abortController = new AbortController();
1236 | 
1237 |     // Store the necessary data in the signal
1238 |     abortController.signal.mesDiv = mesDiv;
1239 |     abortController.signal.mesId = mesId;
1240 |     abortController.signal.swipeId = swipeId;
1241 |     abortController.signal.highlightDuration = extension_settings[extensionName].highlightDuration;
1242 | 
1243 |     // Show the stop button
1244 |     getContext().deactivateSendButtons();
1245 |     let res;
1246 |     if (extension_settings[extensionName].useStreaming) {
1247 |         switch (main_api) {
1248 |             case 'textgenerationwebui':
1249 |                 res = await generateTextGenWithStreaming(generateData, abortController.signal);
1250 |                 break;
1251 |             case 'novel':
1252 |                 res = await generateNovelWithStreaming(generateData, abortController.signal);
1253 |                 break;
1254 |             case 'koboldhorde':
1255 |                 toastr.warning('Rewrite streaming not supported for Kobold. Turn off in rewrite settings.');
1256 |             default:
1257 |                 throw new Error('Streaming is enabled, but the current API does not support streaming.');
1258 |         }
1259 |     } else {
1260 |         if (main_api === 'koboldhorde') {
1261 |             res = await generateHorde(prompt, generateData, abortController.signal, true);
1262 |         } else {
1263 |             const response = await generateRaw(prompt, null, false, false, null, generateData.max_length);
1264 |             res = {text: response};
1265 |             // Shamelessly copied from script.js
1266 |             /*function getGenerateUrl(api) {
1267 |                 switch (api) {
1268 |                     case 'textgenerationwebui':
1269 |                         return '/api/backends/text-completions/generate';
1270 |                     case 'novel':
1271 |                         return '/api/novelai/generate';
1272 |                     default:
1273 |                         throw new Error(`Unknown API: ${api}`);
1274 |                 }
1275 |             }
1276 | 
1277 |             const response = await fetch(getGenerateUrl(main_api), {
1278 |                 method: 'POST',
1279 |                 headers: getRequestHeaders(),
1280 |                 cache: 'no-cache',
1281 |                 body: JSON.stringify(generateData),
1282 |                 signal: abortController.signal,
1283 |             });
1284 | 
1285 |             if (!response.ok) {
1286 |                 const error = await response.json();
1287 |                 throw error;
1288 |             }
1289 | 
1290 |             res = await response.json();*/
1291 |         }
1292 |     }
1293 | 
1294 |     window.getSelection().removeAllRanges();
1295 | 
1296 |     let newText = '';
1297 | 
1298 |     if (typeof res === 'function') {
1299 |         // Streaming case
1300 | 
1301 |         const streamingSpan = document.createElement('span');
1302 |         streamingSpan.className = 'animated-highlight';
1303 | 
1304 |         // Replace the selected text with the streaming span
1305 |         range.deleteContents();
1306 |         range.insertNode(streamingSpan);
1307 | 
1308 |         for await (const chunk of res()) {
1309 |             newText = chunk.text;
1310 |             streamingSpan.textContent = newText;
1311 |         }
1312 |     } else {
1313 |         // Non-streaming case
1314 |         newText = res?.choices?.[0]?.message?.content ?? res?.choices?.[0]?.text ?? res?.text ?? '';
1315 |         if (main_api === 'novel') newText = res.output;
1316 |         const highlightedNewText = document.createElement('span');
1317 |         highlightedNewText.className = 'animated-highlight';
1318 |         highlightedNewText.textContent = newText;
1319 | 
1320 |         range.deleteContents();
1321 |         range.insertNode(highlightedNewText);
1322 |     }
1323 | 
1324 |     // Remove highlight after x seconds when streaming is complete
1325 |     const highlightDuration = extension_settings[extensionName].highlightDuration;
1326 |     setTimeout(() => removeHighlight(mesDiv, mesId, swipeId), highlightDuration);
1327 | 
1328 |     await saveRewrittenText(mesId, swipeId, fullMessage, rawStartOffset, rawEndOffset, newText);
1329 |     getContext().activateSendButtons();
1330 | }
1331 | 
1332 | function calculateTargetTokenCount(selectedText, option) {
1333 |     const baseTokenCount = getTokenCount(selectedText);
1334 |     const useDynamicTokens = extension_settings[extensionName].useDynamicTokens;
1335 |     const dynamicTokenMode = extension_settings[extensionName].dynamicTokenMode;
1336 |     let result;
1337 | 
1338 |     if (useDynamicTokens) {
1339 |         if (dynamicTokenMode === 'additive') {
1340 |             let modifier;
1341 |             switch (option) {
1342 |                 case 'Rewrite':
1343 |                     modifier = extension_settings[extensionName].rewriteTokensAdd;
1344 |                     break;
1345 |                 case 'Shorten':
1346 |                     modifier = extension_settings[extensionName].shortenTokensAdd;
1347 |                     break;
1348 |                 case 'Expand':
1349 |                     modifier = extension_settings[extensionName].expandTokensAdd;
1350 |                     break;
1351 |             }
1352 |             result = baseTokenCount + modifier;
1353 |         } else { // multiplicative
1354 |             let multiplier;
1355 |             switch (option) {
1356 |                 case 'Rewrite':
1357 |                     multiplier = extension_settings[extensionName].rewriteTokensMult;
1358 |                     break;
1359 |                 case 'Shorten':
1360 |                     multiplier = extension_settings[extensionName].shortenTokensMult;
1361 |                     break;
1362 |                 case 'Expand':
1363 |                     multiplier = extension_settings[extensionName].expandTokensMult;
1364 |                     break;
1365 |             }
1366 |             result = baseTokenCount * multiplier;
1367 |         }
1368 |     } else {
1369 |         switch (option) {
1370 |             case 'Rewrite':
1371 |                 result = extension_settings[extensionName].rewriteTokens;
1372 |                 break;
1373 |             case 'Shorten':
1374 |                 result = extension_settings[extensionName].shortenTokens;
1375 |                 break;
1376 |             case 'Expand':
1377 |                 result = extension_settings[extensionName].expandTokens;
1378 |                 break;
1379 |         }
1380 |     }
1381 | 
1382 |     return Math.max(1, Math.round(result)); // Ensure at least 1 token and round to nearest integer
1383 | }
1384 | 
1385 | async function handleUndo(event) {
1386 |     const mesId = event.target.dataset.mesId;
1387 |     const change = changeHistory.findLast(change => change.mesId === mesId);
1388 | 
1389 |     if (change) {
1390 |         const context = getContext();
1391 |         const messageDiv = document.querySelector(`[mesid="${mesId}"]`);
1392 | 
1393 |         if (!messageDiv || !context.chat[mesId]) {
1394 |             console.error('Message not found for undo operation');
1395 |             return;
1396 |         }
1397 | 
1398 |         // Update the chat context
1399 |         context.chat[mesId].mes = change.originalContent;
1400 | 
1401 |         // Only update swipes if they exist
1402 |         if (change.swipeId !== undefined && context.chat[mesId].swipes) {
1403 |             context.chat[mesId].swipes[change.swipeId] = change.originalContent;
1404 |         }
1405 | 
1406 |         // Update the UI
1407 |         const mesTextElement = messageDiv.querySelector('.mes_text');
1408 |         if (mesTextElement) {
1409 |             mesTextElement.innerHTML = messageFormatting(
1410 |                 change.originalContent,
1411 |                 context.name2,
1412 |                 context.chat[mesId].isSystem,
1413 |                 context.chat[mesId].isUser,
1414 |                 mesId
1415 |             );
1416 |             addCopyToCodeBlocks(mesTextElement);
1417 |         }
1418 | 
1419 |         // Save the chat
1420 |         await context.saveChat();
1421 | 
1422 |         // Remove this change from history
1423 |         changeHistory = changeHistory.filter(c => c !== change);
1424 | 
1425 |         // Update undo buttons
1426 |         updateUndoButtons();
1427 |     }
1428 | }
1429 | 
1430 | async function saveRewrittenText(mesId, swipeId, fullMessage, startOffset, endOffset, newText) {
1431 |     const context = getContext();
1432 | 
1433 |     // Get the prefix and suffix to remove from the settings
1434 |     const removePrefix = extension_settings[extensionName].removePrefix || '';
1435 |     const removeSuffix = extension_settings[extensionName].removeSuffix || '';
1436 | 
1437 |     // Remove prefix if present
1438 |     if (removePrefix && newText.startsWith(removePrefix)) {
1439 |         newText = newText.slice(removePrefix.length);
1440 |     }
1441 | 
1442 |     // Remove suffix if present
1443 |     if (removeSuffix && newText.endsWith(removeSuffix)) {
1444 |         newText = newText.slice(0, -removeSuffix.length);
1445 |     }
1446 | 
1447 |     // Apply AI Output regex scripts if setting is enabled
1448 |     let processedText = newText; // Default to original newText
1449 |     if (extension_settings[extensionName].applyRegexOnRewrite) {
1450 |         processedText = getRegexedString(newText, regex_placement.AI_OUTPUT);
1451 |     }
1452 | 
1453 |     // Create the new message with the rewritten and potentially processed section
1454 |     const newMessage =
1455 |         fullMessage.substring(0, startOffset) +
1456 |         processedText + // Use the processed text here
1457 |         fullMessage.substring(endOffset);
1458 | 
1459 |     // Save the change to the history
1460 |     saveLastChange(mesId, swipeId, fullMessage, newMessage);
1461 | 
1462 |     // Update the main message
1463 |     context.chat[mesId].mes = newMessage;
1464 | 
1465 |     // Update the swipe if it exists
1466 |     if (swipeId !== undefined && context.chat[mesId].swipes && context.chat[mesId].swipes[swipeId]) {
1467 |         context.chat[mesId].swipes[swipeId] = newMessage;
1468 |     }
1469 | 
1470 |     // Save and reload the chat
1471 |     await context.saveChat();
1472 | }
1473 | 


--------------------------------------------------------------------------------