From a4b8883473b91f852d37afc890988a9294dc0cdd Mon Sep 17 00:00:00 2001
From: Pabloader <pabloader@pabloader.ru>
Date: Sat, 2 Nov 2024 13:10:31 +0000
Subject: [PATCH] Prompts editor

---
 src/games/ai/components/autoTextarea.tsx      |  34 +++++
 .../ai/components/header/header.module.css    |  16 +-
 src/games/ai/components/header/header.tsx     |  66 ++++++--
 src/games/ai/components/input.tsx             |   9 +-
 src/games/ai/components/message/message.tsx   |   7 +-
 src/games/ai/components/minichat/minichat.tsx |   9 +-
 src/games/ai/const.ts                         | 141 +-----------------
 src/games/ai/contexts/llm.tsx                 |  66 +++++---
 src/games/ai/contexts/state.tsx               |  40 ++++-
 src/games/ai/dom.ts                           |  12 --
 10 files changed, 193 insertions(+), 207 deletions(-)
 create mode 100644 src/games/ai/components/autoTextarea.tsx

diff --git a/src/games/ai/components/autoTextarea.tsx b/src/games/ai/components/autoTextarea.tsx
new file mode 100644
index 0000000..b56dcfc
--- /dev/null
+++ b/src/games/ai/components/autoTextarea.tsx
@@ -0,0 +1,34 @@
+import { useEffect, useRef, useState } from "preact/hooks";
+import type { JSX } from "preact/jsx-runtime"
+
+export const AutoTextarea = (props: JSX.HTMLAttributes<HTMLTextAreaElement>) => {
+    const { value } = props;
+    const ref = useRef<HTMLTextAreaElement>(null);
+    const [isVisible, setVisible] = useState(false);
+
+    useEffect(() => {
+        if (ref.current) {
+            const area = ref.current;
+
+            area.style.height = '0'; // reset
+            area.style.height = `${area.scrollHeight}px`;
+        }
+    }, [value, isVisible]);
+
+    useEffect(() => {
+        if (ref.current) {
+            const observer = new IntersectionObserver(([entry]) => {
+                setVisible(entry.isIntersecting);
+                if (entry.isIntersecting) {
+                    observer.disconnect();
+                }
+            });
+
+            observer.observe(ref.current);
+
+            return () => observer.disconnect();
+        }
+    }, [ref.current]);
+
+    return <textarea {...props} ref={ref} />
+};
diff --git a/src/games/ai/components/header/header.module.css b/src/games/ai/components/header/header.module.css
index 865fd77..2336e06 100644
--- a/src/games/ai/components/header/header.module.css
+++ b/src/games/ai/components/header/header.module.css
@@ -25,5 +25,19 @@
 }
 
 .modalTitle {
-    margin-top: 0;   
+    margin: 0;   
+}
+
+.scrollPane {
+    overflow-x: hidden;
+    overflow-y: auto;
+    margin: 8px 0;
+
+    textarea {
+        overflow: hidden;
+    }
+}
+
+.template {
+    font-family: 'Ubuntu Mono', 'Courier New', Courier, monospace;
 }
\ No newline at end of file
diff --git a/src/games/ai/components/header/header.tsx b/src/games/ai/components/header/header.tsx
index 567eded..5091ce6 100644
--- a/src/games/ai/components/header/header.tsx
+++ b/src/games/ai/components/header/header.tsx
@@ -1,4 +1,4 @@
-import { useCallback, useContext, useEffect, useRef, useState } from "preact/hooks";
+import { useCallback, useContext, useEffect, useMemo, useRef, useState } from "preact/hooks";
 import { useBool } from "@common/hooks/useBool";
 import { Modal } from "@common/components/modal/modal";
 
@@ -8,18 +8,23 @@ import { MiniChat } from "../minichat/minichat";
 
 import styles from './header.module.css';
 import { DOMTools } from "../../dom";
+import { AutoTextarea } from "../autoTextarea";
 
 export const Header = () => {
-    const llm = useContext(LLMContext);
-    const { messages, connectionUrl, lore, setConnectionUrl, setLore, addSwipe } = useContext(StateContext);
+    const { getContextLength } = useContext(LLMContext);
+    const {
+        messages, connectionUrl, systemPrompt, lore, userPrompt, bannedWords,
+        setConnectionUrl, setSystemPrompt, setLore, setUserPrompt, addSwipe, setBannedWords,
+    } = useContext(StateContext);
     const [urlValid, setUrlValid] = useState(false);
     const [urlEditing, setUrlEditing] = useState(false);
 
-    const loreAreaRef = useRef<HTMLTextAreaElement>(null);
-
     const loreOpen = useBool();
+    const promptsOpen = useBool();
     const assistantOpen = useBool();
 
+    const bannedWordsInput = useMemo(() => bannedWords.join('\n'), [bannedWords]);
+
     const handleFocusUrl = useCallback(() => setUrlEditing(true), []);
 
     const handleBlurUrl = useCallback(() => {
@@ -33,22 +38,32 @@ export const Header = () => {
 
     useEffect(() => {
         if (!urlEditing) {
-            llm.getContextLength().then(length => {
+            getContextLength().then(length => {
                 setUrlValid(length > 0);
             });
         }
     }, [connectionUrl, urlEditing]);
 
-    useEffect(() => {
-        DOMTools.fixHeight(loreAreaRef.current);
-    }, [lore, loreOpen.value]);
-
     const handleAssistantAddSwipe = useCallback((answer: string) => {
         const index = messages.findLastIndex(m => m.role === 'assistant');
         addSwipe(index, answer);
         assistantOpen.setFalse();
     }, [addSwipe, messages]);
 
+    const handleSetBannedWords = useCallback((e: Event) => {
+        if (e.target instanceof HTMLTextAreaElement) {
+            const words = e.target.value.split('\n');
+            setBannedWords(words);
+        }
+    }, [setBannedWords]);
+
+    const handleBlurBannedWords = useCallback((e: Event) => {
+        if (e.target instanceof HTMLTextAreaElement) {
+            const words = e.target.value.split('\n').sort();
+            setBannedWords(words);
+        }
+    }, [setBannedWords]);
+
     return (
         <div class={styles.header}>
             <input value={connectionUrl}
@@ -57,16 +72,39 @@ export const Header = () => {
                 onBlur={handleBlurUrl}
                 class={urlEditing ? '' : urlValid ? styles.valid : styles.invalid} />
             <div class={styles.buttons}>
-                <button class='icon' onClick={assistantOpen.setTrue} title='Ask assistant'>
-                    ❓
-                </button>
                 <button class='icon color' title='Edit lore' onClick={loreOpen.setTrue}>
                     🌍
                 </button>
+                <button class='icon color' title='Edit prompts' onClick={promptsOpen.setTrue}>
+                    📃
+                </button>
+            </div>
+            <div class={styles.buttons}>
+                <button class='icon' onClick={assistantOpen.setTrue} title='Ask assistant'>
+                    ❓
+                </button>
             </div>
             <Modal open={loreOpen.value} onClose={loreOpen.setFalse}>
                 <h3 class={styles.modalTitle}>Lore Editor</h3>
-                <textarea value={lore} onInput={setLore} ref={loreAreaRef} />
+                <AutoTextarea value={lore} onInput={setLore} />
+            </Modal>
+            <Modal open={promptsOpen.value} onClose={promptsOpen.setFalse}>
+                <h3 class={styles.modalTitle}>Prompts Editor</h3>
+                <div className={styles.scrollPane}>
+                    <h4 class={styles.modalTitle}>System prompt</h4>
+                    <AutoTextarea value={systemPrompt} onInput={setSystemPrompt} />
+                    <hr />
+                    <h4 class={styles.modalTitle}>User prompt template</h4>
+                    <AutoTextarea value={userPrompt} onInput={setUserPrompt} class={styles.template} />
+                    <hr />
+                    <h4 class={styles.modalTitle}>Banned phrases</h4>
+                    <AutoTextarea
+                        value={bannedWordsInput}
+                        onInput={handleSetBannedWords}
+                        onBlur={handleBlurBannedWords}
+                        class={styles.template}
+                    />
+                </div>
             </Modal>
             <MiniChat
                 history={messages}
diff --git a/src/games/ai/components/input.tsx b/src/games/ai/components/input.tsx
index fe499b8..df02a3c 100644
--- a/src/games/ai/components/input.tsx
+++ b/src/games/ai/components/input.tsx
@@ -1,17 +1,12 @@
 import { useCallback, useContext } from "preact/hooks";
-import { DOMTools } from "../dom";
 import { StateContext } from "../contexts/state";
 import { LLMContext } from "../contexts/llm";
+import { AutoTextarea } from "./autoTextarea";
 
 export const Input = () => {
     const { input, setInput, addMessage, continueMessage } = useContext(StateContext);
     const { generating } = useContext(LLMContext);
 
-    const handleChange = useCallback((e: Event) => {
-        setInput(e);
-        DOMTools.fixHeight(e.target);
-    }, [setInput]);
-
     const handleSend = useCallback(async () => {
         if (!generating) {
             const newInput = input.trim();
@@ -33,7 +28,7 @@ export const Input = () => {
 
     return (
         <div class="chat-input">
-            <textarea onInput={handleChange} onKeyDown={handleKeyDown} value={input} />
+            <AutoTextarea onInput={setInput} onKeyDown={handleKeyDown} value={input} />
             <button onClick={handleSend} class={`${generating ? 'disabled' : ''}`}>{input ? 'Send' : 'Continue'}</button>
         </div>
     );
diff --git a/src/games/ai/components/message/message.tsx b/src/games/ai/components/message/message.tsx
index 4a3b45a..493882d 100644
--- a/src/games/ai/components/message/message.tsx
+++ b/src/games/ai/components/message/message.tsx
@@ -4,6 +4,7 @@ import { StateContext } from "../../contexts/state";
 import { DOMTools } from "../../dom";
 
 import styles from './message.module.css';
+import { AutoTextarea } from "../autoTextarea";
 
 interface IProps {
     message: IMessage;
@@ -17,7 +18,6 @@ export const Message = ({ message, index, isLastUser, isLastAssistant }: IProps)
     const [editing, setEditing] = useState(false);
     const [savedMessage, setSavedMessage] = useState('');
     const textRef = useRef<HTMLDivElement>(null);
-    const textareaRef = useRef<HTMLTextAreaElement>(null);
 
     const content = useMemo(() => MessageTools.getSwipe(message)?.content, [message]);
     const htmlContent = useMemo(() => MessageTools.format(content ?? ''), [content]);
@@ -46,7 +46,6 @@ export const Message = ({ message, index, isLastUser, isLastAssistant }: IProps)
             const newContent = e.target.value;
             editMessage(index, newContent);
         }
-        DOMTools.fixHeight(e.target);
     }, [editMessage, index]);
 
     const handleSwipeLeft = useCallback(() => {
@@ -59,13 +58,11 @@ export const Message = ({ message, index, isLastUser, isLastAssistant }: IProps)
         DOMTools.animate(textRef.current, 'swipe-from-right');
     }, [setCurrentSwipe, index, message]);
 
-    useEffect(() => DOMTools.fixHeight(textareaRef.current), [editing]);
-
     return (
         <div class={`${styles.message} ${styles[message.role]} ${isLastUser ? styles.lastUser : ''}`}>
             <div class={styles.content}>
                 {editing
-                    ? <textarea onInput={handleEdit} value={content} ref={textareaRef} />
+                    ? <AutoTextarea onInput={handleEdit} value={content} />
                     : <div class={styles.text} dangerouslySetInnerHTML={{ __html: htmlContent }} ref={textRef} />
                 }
                 {(isLastUser || message.role === 'assistant') &&
diff --git a/src/games/ai/components/minichat/minichat.tsx b/src/games/ai/components/minichat/minichat.tsx
index 0b38411..152347d 100644
--- a/src/games/ai/components/minichat/minichat.tsx
+++ b/src/games/ai/components/minichat/minichat.tsx
@@ -6,6 +6,7 @@ import { DOMTools } from "../../dom";
 import styles from './minichat.module.css';
 import { LLMContext } from "../../contexts/llm";
 import { FormattedMessage } from "../message/formattedMessage";
+import { AutoTextarea } from "../autoTextarea";
 
 interface IProps {
     open: boolean;
@@ -24,10 +25,6 @@ export const MiniChat = ({ history = [], buttons = {}, open, onClose }: IProps)
         [messages]
     );
 
-    const fitTextareas = useCallback(() => {
-        ref.current?.querySelectorAll('textarea').forEach(DOMTools.fixHeight);
-    }, []);
-
     const handleInit = useCallback((force = false) => {
         if (force || confirm('Clear chat?')) {
             setMessages([MessageTools.create('', 'user', true)]);
@@ -35,7 +32,6 @@ export const MiniChat = ({ history = [], buttons = {}, open, onClose }: IProps)
     }, []);
 
     useEffect(() => {
-        fitTextareas();
         DOMTools.scrollDown(ref.current, false);
     }, [generating, open]);
 
@@ -84,7 +80,6 @@ export const MiniChat = ({ history = [], buttons = {}, open, onClose }: IProps)
         if (e.target instanceof HTMLTextAreaElement) {
             setMessages(MessageTools.updateContent(messages, i, e.target.value));
         }
-        DOMTools.fixHeight(e.target);
     }, [messages]);
 
     if (!open) {
@@ -99,7 +94,7 @@ export const MiniChat = ({ history = [], buttons = {}, open, onClose }: IProps)
                             ? <FormattedMessage key={i} class={`${styles[m.role]} ${styles.message}`}>
                                 {MessageTools.getSwipe(m)?.content ?? ''}
                             </FormattedMessage>
-                            : <textarea
+                            : <AutoTextarea
                                 key={i}
                                 class={styles[m.role]}
                                 value={MessageTools.getSwipe(m)?.content ?? ''}
diff --git a/src/games/ai/const.ts b/src/games/ai/const.ts
index 792a645..172d434 100644
--- a/src/games/ai/const.ts
+++ b/src/games/ai/const.ts
@@ -1,140 +1 @@
-export const p = (strings: TemplateStringsArray, ...args: any[]) =>
-    String.raw(strings, ...args).trim().replace(/^ +| +$/img, '')
-
-export const SYSTEM_PROMPT = p`
-    You are creative writer.
-    Write a story based on the world description below.
-    Make sure you're following the provided lore exactly and not making up impossible things. 
-`.replace(/[ \r\n]+/ig, ' ');
-
-export const START_PROMPT = p`
-    Write a novel using information above as a reference. Make sure to follow the lore exactly and avoid cliffhangers.
-`;
-
-export const CONTINUE_PROPMT = (prompt?: string, isStart = false) =>
-    prompt?.trim()
-        ? p`
-            ${isStart ? 'Start' : 'Continue'} this story, taking information into account: ${prompt.trim()}
-
-            Remember that this story should be infinite and go forever. Avoid cliffhangers and pauses, be creative.
-        `
-        : `Continue the story forward. Avoid cliffhangers and pauses.`;
-
-export const LLAMA_TEMPLATE = `{% for message in messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}`;
-
-export const BANNED_WORDS = [
-    " await",
-    " lie ahead",
-    " lies ahead",
-    " lay ahead",
-    " lays ahead",
-    "a testament to",
-    "anticipat",
-    "journey",
-    "voyage ",
-    "heart race",
-    "heart racing",
-    "mind race",
-    "mind racing",
-    "minds racing",
-    "minds race",
-    "Would you like that",
-    "What do you say",
-    "possibilities",
-    "predator",
-    " explor",
-    " intimate",
-    " bond",
-    " desire",
-    ".desire",
-    " biting",
-    " bites her plump lower lip",
-    " bites her lower lip",
-    " bites her lip",
-    " bit her lower lip",
-    " bit her lip",
-    " bites his plump lower lip",
-    " bites his lower lip",
-    " bites his lip",
-    " bit his lip",
-    "barely above a whisper",
-    "barely audible",
-    "barely a whisper",
-    " vulnerab",
-    " shiver",
-    " chill run",
-    " chill down",
-    " chill up",
-    "sparkling with mischief",
-    "tracing patterns",
-    "traced patterns",
-    "idly traces patterns",
-    " air is thick ",
-    " air was thick ",
-    " air thick ",
-    "aldric",
-    "elara",
-    "aedric",
-    "zephyr",
-    "lyra",
-    "timmy",
-    " indulg",
-    " embrac",
-    " hot and bothered",
-    "sex",
-    "manhood",
-    "boyhood",
-    "taboo",
-    "spill the beans",
-    "with need",
-    "slick folds",
-    "sensitive folds",
-    "glistening folds",
-    "wet folds",
-    "swollen folds",
-    "her folds",
-    "my folds",
-    "slippery folds",
-    "claim me as your",
-    "screaming your name",
-    "scream your name",
-    "awaits your next command",
-    "awaits your next move",
-    "waits for your next move",
-    "plump bottom lip",
-    "plump lip",
-    "claim",
-    " carnal",
-    " primal",
-    "...",
-    "…",
-    "throbbing",
-    " a mix of ",
-    " a blend of ",
-    "camaraderie",
-    "What do you think",
-    "boundaries",
-    " tapestr",
-    "dynamic",
-    "I see,",
-    "I see.",
-    " mournful ",
-    "voluptuous",
-    "eerie",
-    " ye ",
-    " ya ",
-].sort();
-
-export const GENERATION_SETTINGS = {
-    temperature: 0.8,
-    min_p: 0.1,
-    rep_pen: 1.08,
-    rep_pen_range: -1,
-    rep_pen_slope: 0.7,
-    top_k: 100,
-    top_p: 0.92,
-    banned_tokens: BANNED_WORDS,
-    max_length: 512,
-    trim_stop: true,
-    stop_sequence: ['[INST]', '[/INST]', '</s>', '<|']
-}
\ No newline at end of file
+export const LLAMA_TEMPLATE = `{% for message in messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}`;
\ No newline at end of file
diff --git a/src/games/ai/contexts/llm.tsx b/src/games/ai/contexts/llm.tsx
index cc246b0..f31d3c9 100644
--- a/src/games/ai/contexts/llm.tsx
+++ b/src/games/ai/contexts/llm.tsx
@@ -1,6 +1,6 @@
 import Lock from "@common/lock";
 import SSE from "@common/sse";
-import { CONTINUE_PROPMT, GENERATION_SETTINGS, LLAMA_TEMPLATE, START_PROMPT, SYSTEM_PROMPT } from "../const";
+import { LLAMA_TEMPLATE } from "../const";
 import { createContext } from "preact";
 import { useContext, useEffect, useMemo } from "preact/hooks";
 import { MessageTools, type IMessage } from "../messages";
@@ -28,10 +28,26 @@ interface IContext {
     generating: boolean;
 }
 
+const DEFAULT_GENERATION_SETTINGS = {
+    temperature: 0.8,
+    min_p: 0.1,
+    rep_pen: 1.08,
+    rep_pen_range: -1,
+    rep_pen_slope: 0.7,
+    top_k: 100,
+    top_p: 0.92,
+    banned_tokens: [],
+    max_length: 512,
+    trim_stop: true,
+    stop_sequence: ['[INST]', '[/INST]', '</s>', '<|']
+}
+
+type IGenerationSettings = Partial<typeof DEFAULT_GENERATION_SETTINGS>;
+
 interface IActions {
     applyChatTemplate: (messages: ITemplateMessage[], templateString: string, eosToken?: string) => string;
     compilePrompt: (messages: IMessage[], args?: ICompileArgs) => Promise<ICompiledPrompt>;
-    generate: (prompt: string, extraSettings?: Partial<typeof GENERATION_SETTINGS>) => AsyncGenerator<string>;
+    generate: (prompt: string, extraSettings?: IGenerationSettings) => AsyncGenerator<string>;
     countTokens(prompt: string): Promise<number>;
     getContextLength(): Promise<number>;
 
@@ -41,9 +57,14 @@ export type ILLMContext = IContext & IActions;
 export const LLMContext = createContext<ILLMContext>({} as ILLMContext);
 
 export const LLMContextProvider = ({ children }: { children?: any }) => {
-    const { connectionUrl, messages, triggerNext, lore, setTriggerNext, addMessage, editMessage } = useContext(StateContext);
+    const {
+        connectionUrl, messages, triggerNext, lore, userPrompt, systemPrompt, bannedWords,
+        setTriggerNext, addMessage, editMessage,
+    } = useContext(StateContext);
     const generating = useBool(false);
 
+    const userPromptTemplate = useMemo(() => new Template(userPrompt), [userPrompt]);
+
     const actions: IActions = useMemo(() => ({
         applyChatTemplate: (messages: ITemplateMessage[], templateString: string, eosToken = '</s>') => {
             const template = new Template(templateString);
@@ -64,11 +85,15 @@ export const LLMContextProvider = ({ children }: { children?: any }) => {
             const isRegen = isAssistantLast && !MessageTools.getSwipe(lastMessage)?.content;
             const isContinue = isAssistantLast && !isRegen;
 
+            const userMessages = promptMessages.filter(m => m.role === 'user');
+            const lastUserMessage = userMessages.at(-1);
+            const firstUserMessage = userMessages.at(0);
+
             if (isContinue) {
-                promptMessages.push(MessageTools.create(CONTINUE_PROPMT()));
+                promptMessages.push(MessageTools.create(userPromptTemplate.render({})));
             }
 
-            const system = `${SYSTEM_PROMPT}\n\n${lore}`.trim();
+            const system = `${systemPrompt}\n\n${lore}`.trim();
 
             const templateMessages: ITemplateMessage[] = [
                 { role: 'system', content: system },
@@ -87,7 +112,10 @@ export const LLMContextProvider = ({ children }: { children?: any }) => {
                         wasStory = true;
                         templateMessages.at(-1).content += '\n\n' + content;
                     } else if (role === 'user' && !message.technical) {
-                        templateMessages.push({ role: message.role, content: CONTINUE_PROPMT(content, !wasStory) });
+                        templateMessages.push({
+                            role: message.role,
+                            content: userPromptTemplate.render({ prompt: content, isStart: !wasStory }),
+                        });
                     } else {
                         if (role === 'assistant') {
                             wasStory = true;
@@ -97,23 +125,26 @@ export const LLMContextProvider = ({ children }: { children?: any }) => {
                 }
 
                 if (templateMessages[1]?.role !== 'user') {
-                    templateMessages.splice(1, 0, { role: 'user', content: START_PROMPT });
+                    const prompt = MessageTools.getSwipe(firstUserMessage)?.content;
+
+                    templateMessages.splice(1, 0, {
+                        role: 'user',
+                        content: userPromptTemplate.render({ prompt, isStart: true }),
+                    });
                 }
             } else {
                 const story = promptMessages.filter(m => m.role === 'assistant')
                     .map(m => MessageTools.getSwipe(m)?.content.trim()).join('\n\n');
 
-                const userMessages = promptMessages.filter(m => m.role === 'user');
-                const lastUserMessage = userMessages.at(-1);
+                if (story.length > 0) {
+                    const prompt = MessageTools.getSwipe(firstUserMessage)?.content;
+                    templateMessages.push({ role: 'user', content: userPromptTemplate.render({ prompt, isStart: true }) });
+                    templateMessages.push({ role: 'assistant', content: story });
+                }
 
                 let userPrompt = MessageTools.getSwipe(lastUserMessage)?.content;
                 if (!lastUserMessage?.technical && !isContinue && userPrompt) {
-                    userPrompt = CONTINUE_PROPMT(userPrompt, story.length === 0);
-                }
-
-                if (story.length > 0) {
-                    templateMessages.push({ role: 'user', content: START_PROMPT });
-                    templateMessages.push({ role: 'assistant', content: story });
+                    userPrompt = userPromptTemplate.render({ prompt: userPrompt, isStart: story.length === 0 });
                 }
 
                 if (userPrompt) {
@@ -139,7 +170,8 @@ export const LLMContextProvider = ({ children }: { children?: any }) => {
 
                 const sse = new SSE(`${connectionUrl}/api/extra/generate/stream`, {
                     payload: JSON.stringify({
-                        ...GENERATION_SETTINGS,
+                        ...DEFAULT_GENERATION_SETTINGS,
+                        banned_tokens: bannedWords.filter(w => w.trim()),
                         ...extraSettings,
                         prompt,
                     }),
@@ -229,7 +261,7 @@ export const LLMContextProvider = ({ children }: { children?: any }) => {
 
             return 0;
         },
-    }), [connectionUrl]);
+    }), [connectionUrl, lore, userPromptTemplate, systemPrompt, bannedWords]);
 
     useEffect(() => void (async () => {
         if (triggerNext && !generating.value) {
diff --git a/src/games/ai/contexts/state.tsx b/src/games/ai/contexts/state.tsx
index 7310742..adb7b50 100644
--- a/src/games/ai/contexts/state.tsx
+++ b/src/games/ai/contexts/state.tsx
@@ -7,6 +7,9 @@ interface IContext {
     connectionUrl: string;
     input: string;
     lore: string;
+    systemPrompt: string;
+    userPrompt: string;
+    bannedWords: string[];
     messages: IMessage[];
     triggerNext: boolean;
 }
@@ -15,6 +18,9 @@ interface IActions {
     setConnectionUrl: (url: string | Event) => void;
     setInput: (url: string | Event) => void;
     setLore: (lore: string | Event) => void;
+    setSystemPrompt: (prompt: string | Event) => void;
+    setUserPrompt: (prompt: string | Event) => void;
+    setBannedWords: (words: string[]) => void;
     setTriggerNext: (triggerNext: boolean) => void;
 
     setMessages: (messages: IMessage[]) => void;
@@ -30,17 +36,33 @@ interface IActions {
 const SAVE_KEY = 'ai_game_save_state';
 
 export const saveContext = (context: IContext) => {
-    localStorage.setItem(SAVE_KEY, JSON.stringify({
-        ...context,
-        triggerNext: false,
-    }));
+    const contextToSave: Partial<IContext> = { ...context };
+    delete contextToSave.triggerNext;
+
+    localStorage.setItem(SAVE_KEY, JSON.stringify(contextToSave));
 }
 
 export const loadContext = (): IContext => {
     const defaultContext: IContext = {
         connectionUrl: 'http://192.168.10.102:5001',
         input: '',
+        systemPrompt: 'You are creative writer. Write a story based on the world description below.',
         lore: '',
+        userPrompt: `
+{%- if prompt -%}
+    {%- if isStart -%}
+        Start
+    {%- else -%}
+        Continue
+    {%- endif %} this story, taking information into account: {{ prompt | trim }}
+Remember that this story should be infinite and go forever. Avoid cliffhangers and pauses, be creative.
+{%- elif isStart -%}
+    Write a novel using information above as a reference. Make sure to follow the lore exactly and avoid cliffhangers.
+{%- else -%}
+    Continue the story forward. Avoid cliffhangers and pauses.
+{%- endif -%}
+        `.trim(),
+        bannedWords: [],
         messages: [],
         triggerNext: false,
     };
@@ -66,14 +88,21 @@ export const StateContextProvider = ({ children }: { children?: any }) => {
     const [connectionUrl, setConnectionUrl] = useInputState(loadedContext.connectionUrl);
     const [input, setInput] = useInputState(loadedContext.input);
     const [lore, setLore] = useInputState(loadedContext.lore);
+    const [systemPrompt, setSystemPrompt] = useInputState(loadedContext.systemPrompt);
+    const [userPrompt, setUserPrompt] = useInputState(loadedContext.userPrompt);
+    const [bannedWords, setBannedWords] = useState<string[]>(loadedContext.bannedWords);
     const [messages, setMessages] = useState(loadedContext.messages);
+
     const [triggerNext, setTriggerNext] = useState(false);
 
     const actions: IActions = useMemo(() => ({
         setConnectionUrl,
         setInput,
+        setSystemPrompt,
+        setUserPrompt,
         setLore,
         setTriggerNext,
+        setBannedWords: (words) => setBannedWords([...words]),
 
         setMessages: (newMessages) => setMessages(newMessages.slice()),
         addMessage: (content, role, triggerNext = false) => {
@@ -145,7 +174,10 @@ export const StateContextProvider = ({ children }: { children?: any }) => {
     const rawContext: IContext = {
         connectionUrl,
         input,
+        systemPrompt,
         lore,
+        userPrompt,
+        bannedWords,
         messages,
         triggerNext,
     };
diff --git a/src/games/ai/dom.ts b/src/games/ai/dom.ts
index 6419366..42fa6cf 100644
--- a/src/games/ai/dom.ts
+++ b/src/games/ai/dom.ts
@@ -1,16 +1,4 @@
 export namespace DOMTools {
-    export const fixHeight = (e: unknown, onlyGrow?: any) => {
-        if (e instanceof Event) {
-            e = e.target;
-        }
-        if (e instanceof HTMLElement) {
-            if (onlyGrow !== true) {
-                e.style.height = '0'; // reset
-            }
-            e.style.height = `${e.scrollHeight + 10}px`;
-        }
-    }
-
     export const animate = (e: unknown, animationName: string) => {
         if (e instanceof Event) {
             e = e.target;