diff --git a/editors/demo/src/Editoria/config/config.js b/editors/demo/src/Editoria/config/config.js
index e7f95b75be28699f9c8ddf4ece965f364b1eb98e..5aef8c600517acc61ee3c445abaa53cccb824af5 100644
--- a/editors/demo/src/Editoria/config/config.js
+++ b/editors/demo/src/Editoria/config/config.js
@@ -41,10 +41,6 @@ import invisibles, {
 
 import CharactersList from './CharactersList';
 
-// const updateTitle = title => {
-//   console.log(title);
-// };
-
 async function DummyPromise(userInput) {
   return new Promise((resolve, reject) => {
     setTimeout(() => {
@@ -52,9 +48,15 @@ async function DummyPromise(userInput) {
       if (userInput === 'reject') {
         reject('Your request could not be processed for now');
       } else {
-        resolve(
-          'He made significant contributions to theoretical physics, including achievements in quantum mechanics',
-        );
+        // JSON response test
+        const json = JSON.stringify({
+          content:
+            'From the documents you have this is the most accurate answer',
+          citations: ['citation 1', 'citation 2', 'citation 3'],
+          links: ['https://coko.foundation/', 'https://waxjs.net/about/'],
+        });
+
+        resolve(json);
       }
     }, 3150);
   });
@@ -313,6 +315,9 @@ export default {
   AskAiContentService: {
     AskAiContentTransformation: DummyPromise,
     AiOn: true,
+    AskKb: false,
+    GenerateImages: false,
+    CustomPromptsOn: true,
     FreeTextPromptsOn: true,
     CustomPrompts: [],
   },
diff --git a/wax-prosemirror-core/src/components/icons/icons.js b/wax-prosemirror-core/src/components/icons/icons.js
index 8f4bee2e2e61cdecaefbad8f94c97ef4b3c09717..e79ca05f4d9554966aff3aae9bbb73bc441aa7f8 100644
--- a/wax-prosemirror-core/src/components/icons/icons.js
+++ b/wax-prosemirror-core/src/components/icons/icons.js
@@ -737,4 +737,178 @@ export default {
       <path d="m83.44,72.27l-4.73-2.06c-1.14-.5-2.05-1.41-2.55-2.55l-2.06-4.73c-.16-.37-.68-.37-.84,0l-2.06,4.73c-.5,1.14-1.41,2.05-2.55,2.55l-4.73,2.07c-.37.16-.37.68,0,.84l4.82,2.13c1.14.5,2.05,1.42,2.54,2.57l2,4.64c.16.37.68.37.84,0l2.06-4.72c.5-1.14,1.41-2.05,2.55-2.55l4.73-2.06c.37-.16.37-.68,0-.84Z" />
     </Svg>
   ),
+  settings: ({ className }) => (
+    <svg
+      className={className}
+      height="28px"
+      preserveAspectRatio="none"
+      version="1.1"
+      viewBox="0 0 28 28"
+      width="28px"
+      x="0px"
+      xmlns="http://www.w3.org/2000/svg"
+      y="0px"
+    >
+      <defs>
+        <g id="icon-Layer0_0_FILL">
+          <path
+            d="
+M 21.3 8.8
+Q 21.75 8.35 21.25 7.85
+L 20.2 6.8
+Q 19.7 6.3 19.25 6.8
+L 18.05 7.95
+Q 16.85 7.15 15.4 6.9
+L 15.4 5.15
+Q 15.4 4.5 14.75 4.5
+L 13.25 4.5
+Q 12.55 4.5 12.55 5.15
+L 12.55 6.9
+Q 11.1 7.15 9.9 8
+L 8.65 6.75
+Q 8.15 6.3 7.7 6.75
+L 6.65 7.8
+Q 6.15 8.3 6.6 8.8
+L 7.9 10.1
+Q 7.1 11.25 6.85 12.65
+L 5.05 12.65
+Q 4.35 12.65 4.35 13.35
+L 4.35 14.8
+Q 4.35 15.5 5.05 15.5
+L 6.85 15.5
+Q 7.1 16.95 7.9 18.1
+L 6.65 19.35
+Q 6.2 19.85 6.65 20.3
+L 7.7 21.35
+Q 8.25 21.9 8.7 21.4
+L 10 20.15
+Q 11.15 20.95 12.55 21.2
+L 12.55 22.95
+Q 12.55 23.6 13.25 23.6
+L 14.75 23.6
+Q 15.4 23.6 15.4 22.95
+L 15.4 21.2
+Q 16.8 20.95 18 20.2
+L 19.2 21.35
+Q 19.65 21.85 20.15 21.35
+L 21.2 20.3
+Q 21.7 19.8 21.2 19.35
+L 20.05 18.15
+Q 20.9 16.95 21.15 15.5
+L 22.8 15.5
+Q 23.5 15.5 23.5 14.8
+L 23.5 13.35
+Q 23.5 12.65 22.8 12.65
+L 21.2 12.65
+Q 20.9 11.25 20.1 10.05
+L 21.3 8.8
+M 10.45 10.45
+Q 11.95 9 14 9 16.05 9 17.5 10.45 19 11.95 19 14 19 16.05 17.5 17.5 16.05 19 14 19 11.95 19 10.45 17.5 9 16.05 9 14 9 11.95 10.45 10.45 Z"
+            fill="var(--color-blue)"
+            stroke="var(--color-fill-stroke)"
+          />
+        </g>
+
+        <path
+          d="
+M 13.144646453857423 2.0080963134765626
+L 15.023285675048829 2.0080963134765626
+Q 15.83736267089844 2.0080963134765626 15.83736267089844 2.8253768920898445
+L 15.83736267089844 5.025747680664063
+Q 17.6533805847168 5.340086364746094 19.156291961669922 6.345970153808594
+L 20.659203338623048 4.900012207031249
+Q 21.22279510498047 4.271334838867187 21.84900817871094 4.900012207031249
+L 23.164055633544923 6.22023468017578
+Q 23.79026870727539 6.848912048339843 23.22667694091797 7.414721679687501
+L 21.723765563964847 8.986415100097657
+Q 22.725706481933592 10.495240783691406 23.101434326171876 12.255537414550782
+L 25.105316162109375 12.255537414550782
+Q 25.982014465332032 12.255537414550782 25.982014465332032 13.135685729980468
+L 25.982014465332032 14.958850097656251
+Q 25.982014465332032 15.838998413085937 25.105316162109375 15.838998413085937
+L 23.038813018798827 15.838998413085937
+Q 22.725706481933592 17.66216278076172 21.661144256591797 19.170988464355467
+L 23.101434326171876 20.679814147949223
+Q 23.727647399902345 21.245623779296878 23.101434326171876 21.87430114746094
+L 21.78638687133789 23.194523620605473
+Q 21.16017379760742 23.823200988769536 20.59658203125 23.194523620605473
+L 19.093670654296876 21.748565673828125
+Q 17.59075927734375 22.69158172607422 15.83736267089844 23.00592041015625
+L 15.83736267089844 25.20629119873047
+Q 15.83736267089844 26.023571777343754 15.023285675048829 26.023571777343754
+L 13.144646453857423 26.023571777343754
+Q 12.267948150634766 26.023571777343754 12.267948150634766 25.20629119873047
+L 12.267948150634766 23.00592041015625
+Q 10.514551544189452 22.69158172607422 9.074261474609376 21.685697937011717
+L 7.446107482910155 23.257391357421874
+Q 6.882515716552734 23.886068725585936 6.193681335449219 23.194523620605473
+L 4.878633880615234 21.87430114746094
+Q 4.315042114257812 21.308491516113286 4.878633880615234 20.679814147949223
+L 6.444166564941406 19.108120727539067
+Q 5.442225646972656 17.66216278076172 5.129119110107422 15.838998413085937
+L 2.8747520446777344 15.838998413085937
+Q 1.9980537414550774 15.838998413085937 1.9980537414550774 14.958850097656251
+L 1.9980537414550774 13.135685729980468
+Q 1.9980537414550774 12.255537414550782 2.8747520446777344 12.255537414550782
+L 5.129119110107422 12.255537414550782
+Q 5.442225646972656 10.495240783691406 6.444166564941406 9.049282836914061
+L 4.816012573242188 7.414721679687501
+Q 4.252420806884766 6.786044311523439 4.878633880615234 6.157366943359374
+L 6.193681335449219 4.837144470214843
+Q 6.75727310180664 4.271334838867187 7.383486175537109 4.837144470214843
+L 8.94901885986328 6.408837890625
+Q 10.451930236816406 5.340086364746094 12.267948150634766 5.025747680664063
+L 12.267948150634766 2.8253768920898445
+Q 12.267948150634766 2.0080963134765626 13.144646453857423 2.0080963134765626 Z
+M 20.346096801757813 13.95296630859375
+Q 20.346096801757813 16.53054351806641 18.467457580566407 18.35370788574219 16.651439666748047 20.239739990234376 14.083966064453126 20.239739990234376 11.516492462158201 20.239739990234376 9.637853240966795 18.35370788574219 7.821835327148437 16.53054351806641 7.821835327148437 13.95296630859375 7.821835327148437 11.375389099121092 9.637853240966795 9.489356994628904 11.516492462158201 7.666192626953125 14.083966064453126 7.666192626953125 16.651439666748047 7.666192626953125 18.467457580566407 9.489356994628904 20.346096801757813 11.375389099121092 20.346096801757813 13.95296630859375 Z"
+          fill="none"
+          id="iconLayer0_0_1_STROKES"
+          stroke="var(--color-stroke)"
+          strokeLinecap="square"
+          strokeLinejoin="miter"
+          strokeMiterlimit="3"
+          strokeWidth="1"
+        />
+      </defs>
+
+      <g transform="matrix( 1.2524261474609375, 0, 0, 1.257354736328125, -3.45,-3.65) ">
+        <use xlinkHref="#icon-Layer0_0_FILL" />
+
+        <use
+          transform="matrix( 0.7984466552734375, 0, 0, 0.795318603515625, 2.75,2.9) "
+          xlinkHref="#iconLayer0_0_1_STROKES"
+        />
+      </g>
+    </svg>
+  ),
+  send: ({ className }) => (
+    <svg
+      className={className}
+      height="18px"
+      preserveAspectRatio="none"
+      version="1.1"
+      viewBox="0 0 18 18"
+      width="18px"
+      x="0px"
+      xmlns="http://www.w3.org/2000/svg"
+      y="0px"
+    >
+      <defs>
+        <g id="sendicon-Layer0_0_FILL_2">
+          <path
+            d="
+M 6.65 9.9
+L 5.65 14.6 15.4 9.4 5.65 4.2 6.65 8.9 9.25 9.4 6.65 9.9 Z"
+            fill="var(--ai-tool-icon-color, #777)"
+            stroke="none"
+          />
+        </g>
+      </defs>
+
+      <g transform="matrix( 1, 0, 0, 1, 0,0) ">
+        <use xlinkHref="#sendicon-Layer0_0_FILL_2" />
+      </g>
+    </svg>
+  ),
 };
diff --git a/wax-prosemirror-services/src/AiService/ReplaceSelectedText.js b/wax-prosemirror-services/src/AiService/ReplaceSelectedText.js
index 9aee495b486bd99796690dec0fd44dce29a22358..3fd6459968d914b96b047ddd5133c35a2340bd69 100644
--- a/wax-prosemirror-services/src/AiService/ReplaceSelectedText.js
+++ b/wax-prosemirror-services/src/AiService/ReplaceSelectedText.js
@@ -8,6 +8,7 @@ const elementFromString = string => {
 };
 
 const replaceSelectedText = (view, responseText, replace = false) => {
+  if (!responseText) return;
   let { state } = view;
   let { tr } = state;
   const { from, to } = tr.selection;
diff --git a/wax-prosemirror-services/src/AiService/components/AiSettingsMenu.js b/wax-prosemirror-services/src/AiService/components/AiSettingsMenu.js
new file mode 100644
index 0000000000000000000000000000000000000000..b53e50665d8a19b55c371ce8aa42933775141212
--- /dev/null
+++ b/wax-prosemirror-services/src/AiService/components/AiSettingsMenu.js
@@ -0,0 +1,198 @@
+import React, { useContext, useEffect } from 'react';
+import PropTypes from 'prop-types';
+import styled from 'styled-components';
+import { th } from '@pubsweet/ui-toolkit';
+import { WaxContext } from 'wax-prosemirror-core';
+import { keys } from 'lodash';
+import SwitchComponent from './Switch';
+
+// #region CONSTANTS ---------------------------------------------
+const LABEL_POS = 'left';
+
+const SETTINGS = [
+  {
+    key: 'AskKb',
+    label: 'Ask knowledge base',
+  },
+  {
+    key: 'CustomPromptsOn',
+    label: 'Custom prompts',
+  },
+  {
+    key: 'GenerateImages',
+    label: 'Generate image',
+  },
+];
+// #endregion CONSTANTS ---------------------------------------------
+
+// #region STYLED COMPONENTS ---------------------------------------------
+
+export const ToggleInput = styled(SwitchComponent)`
+  --width: 20px;
+  --divisor: 2.1;
+  --height: calc(var(--width) / var(--divisor));
+  --padding: 1px;
+  --border-width: 0px;
+  --label-font-size: 12px;
+
+  align-items: center;
+  display: flex;
+  gap: 10px;
+  padding: 0.5rem 0.2rem;
+
+  button {
+    align-items: center;
+    border-radius: calc(
+      var(--height) + (var(--padding) * 2) + var(--border-width)
+    );
+    display: flex;
+    height: fit-content;
+    padding: var(--padding);
+    transition: all 0.3s;
+    width: var(--width);
+
+    &:focus {
+      box-shadow: none;
+      outline: 1px solid #d5f1fd;
+      outline-offset: 1px;
+    }
+  }
+
+  label {
+    color: var(--ai-tool-switch-label-color, #777);
+    font-size: var(--label-font-size);
+    margin: 0;
+  }
+
+  .rc-switch {
+    border-width: var(--border-width);
+
+    &::after {
+      height: calc(var(--height) - (var(--padding) * 2));
+      position: unset;
+      transition: margin 0.2s cubic-bezier(0.35, 0, 0.25, 1);
+      width: calc(var(--height) - (var(--padding) * 2));
+    }
+  }
+
+  .rc-switch-checked {
+    border: var(--border-width) solid ${th('colorPrimary')};
+
+    &::after {
+      border: none;
+      margin-left: calc(var(--width) - var(--height));
+    }
+  }
+`;
+
+const SettingsDropdown = styled.ul`
+  --item-height: var(--ai-settings-menu-height, 30px);
+  --items-length: ${p => p.$listlng};
+  --max-height: calc(var(--item-height) * var(--items-length));
+
+  background-color: #f8f8f8;
+  border: ${p => (p.$show ? 'var(--ai-tool-border)' : '0 solid #0000')};
+  border-radius: 0 0 0.3rem 0.3rem;
+  border-top: none;
+  display: flex;
+  flex-direction: column;
+  list-style: none;
+  margin: 0;
+  max-height: ${p => (p.$show ? 'var(--max-height)' : 0)};
+  overflow: hidden;
+  padding: 0 0.5rem;
+  position: absolute;
+  right: calc(-1 * var(--ai-tool-border-width));
+  top: calc(100% + var(--ai-tool-border-width));
+  transition: all ${p => `${0.15 * p.$listlng}s`};
+  z-index: 15;
+
+  > li {
+    align-items: center;
+    display: flex;
+    height: var(--item-height);
+    justify-content: ${LABEL_POS === 'right' ? 'flex-start' : 'flex-end'};
+    padding: 0;
+    user-select: none;
+  }
+
+  > li:not(:first-child) {
+    border-top: 1px solid #0001;
+  }
+`;
+// #endregion STYLED COMPONENTS ---------------------------------------------
+
+const AiSettingsMenu = ({ show, aiService, optionsState, setOptionsState }) => {
+  const ctx = useContext(WaxContext);
+
+  const onAiService = ({ key }) => keys(aiService).includes(key);
+
+  const enabledSettings = SETTINGS.filter(onAiService).map(setting => ({
+    ...setting,
+    type: typeof aiService[setting.key],
+  }));
+
+  useEffect(() => {
+    enabledSettings.forEach(({ key }) => setOption(key, aiService[key]));
+  }, []);
+
+  const setOption = (key, state) => {
+    ctx.setOption({ [key]: state });
+    setOptionsState(prev => ({
+      ...prev,
+      [key]: state,
+    }));
+  };
+
+  const handlers = (key, type, value) => {
+    const typeBasedHandlers = {
+      boolean: () => {
+        setOption(key, !value);
+      },
+    };
+    return typeBasedHandlers[type] ?? (() => {});
+  };
+
+  const renderSettings = ({ key, label, type }) => {
+    const value = optionsState[key];
+    const typeBasedSettingUI = {
+      boolean: (
+        <li key={key}>
+          <ToggleInput
+            checked={optionsState[key]}
+            label={label}
+            labelPosition={LABEL_POS}
+            onChange={handlers(key, type, value)}
+          />
+        </li>
+      ),
+    };
+
+    return typeBasedSettingUI[type] ?? typeBasedSettingUI.boolean;
+  };
+
+  return (
+    <SettingsDropdown
+      $listlng={enabledSettings.length}
+      $show={show}
+      onClick={e => e.stopPropagation()} // to prevent !showSettings
+    >
+      {enabledSettings?.map(renderSettings)}
+    </SettingsDropdown>
+  );
+};
+
+AiSettingsMenu.propTypes = {
+  show: PropTypes.bool,
+  aiService: PropTypes.shape({}),
+  optionsState: PropTypes.shape({}),
+  setOptionsState: PropTypes.func,
+};
+AiSettingsMenu.defaultProps = {
+  show: false,
+  aiService: {},
+  optionsState: {},
+  setOptionsState: () => {},
+};
+
+export default AiSettingsMenu;
diff --git a/wax-prosemirror-services/src/AiService/components/AskAIOverlay.js b/wax-prosemirror-services/src/AiService/components/AskAIOverlay.js
index 4607e7834235573b4b4abc0c3734b884ec8c89b0..7c9e4d13467c4136f1b1462ef6c8cbb371f0a68e 100644
--- a/wax-prosemirror-services/src/AiService/components/AskAIOverlay.js
+++ b/wax-prosemirror-services/src/AiService/components/AskAIOverlay.js
@@ -1,62 +1,103 @@
-/* eslint-disable react/prop-types */
+/* eslint-disable react/jsx-props-no-spreading */
 import React, { useRef, useLayoutEffect, useContext, useState } from 'react';
 import styled from 'styled-components';
-import { isEmpty } from 'lodash';
+import { capitalize, isEmpty } from 'lodash';
 import { useTranslation } from 'react-i18next';
 import { WaxContext, icons } from 'wax-prosemirror-core';
+import { PropTypes } from 'prop-types';
 import replaceSelectedText from '../ReplaceSelectedText';
+import AiSettingsMenu from './AiSettingsMenu';
+import { safeParse, resultsToHtml, getUpdatedPosition } from '../helpers';
 
-const Wrapper = styled.div`
+const AI_TOOL_ID = 'ai-overlay';
+const DEFAULT_KEY = 'content';
+
+// #region STYLED COMPONENTS ------------------------------------------------
+
+// #region MAIN & MISC-------------------
+
+const Root = styled.div`
+  --ai-tool-border-width: 1px;
+  --ai-tool-border: var(--ai-tool-border-width) solid #0001;
+  align-items: flex-end;
+  background: #fff;
+  border: var(--ai-tool-border);
   display: flex;
+  filter: drop-shadow(0 0 1px #0002);
   flex-direction: column;
-`;
 
-const AskAIForm = styled.div`
-  background: #fafafa;
-  border: 0.5px #dcdcdc solid;
-  border-top-left-radius: 4px;
-  border-top-right-radius: 4px;
-  box-shadow: 0px 4px 8px rgba(0, 0, 0, 0.08);
-  display: ${prop => (prop.hidden ? 'none' : 'inline-flex;')};
-  justify-content: space-between;
-  padding: 8px 12px;
-  width: 458px;
-`;
+  div {
+    ::-webkit-scrollbar {
+      height: 5px;
+      width: 5px;
+    }
+
+    ::-webkit-scrollbar-thumb {
+      background: var(--scrollbar-color, #0004);
+      border-radius: 5px;
+      width: 5px;
+    }
 
-const ActionButton = styled.button`
+    ::-webkit-scrollbar-track {
+      background: #fff0;
+      padding: 5px;
+    }
+  }
+`;
+const Heading = styled.header`
+  --gradient-direction: to bottom;
   align-items: center;
-  align-self: stretch;
-  background: white;
-  border: 0.5px #f0f0f0 solid;
-  cursor: pointer;
-  display: inline-flex;
-  gap: 8px;
-  justify-content: flex-start;
-  padding: 8px 12px;
+  background: linear-gradient(
+    var(--gradient-direction, to bottom),
+    #dadada 0%,
+    #f5f5f5 3%
+  );
+  /* border-block: var(--ai-tool-border); */
+  display: flex;
+  height: 23px;
+  margin: 0;
+  max-height: 23px;
+  min-height: 23px;
+  padding-left: 3px;
+  width: 100%;
 `;
-
-const ActionSection = styled.div`
-  align-items: flex-start;
-  box-shadow: 0px 4px 8px rgba(0, 0, 0, 0.04);
+const FlexRow = styled.div`
   display: flex;
-  flex-direction: column;
-  justify-content: flex-start;
-  width: 250px;
+  flex-direction: row;
+  width: 100%;
+`;
+const ButtonBase = styled.button.attrs({ type: 'button' })`
+  align-items: center;
+  background: none;
+  border: none;
+  cursor: ${p => (p.disabled ? 'not-allowed' : 'pointer')};
+  display: flex;
+  justify-content: center;
+  margin: 0;
+  opacity: ${p => (p.disabled ? '0.4' : '1')};
+  outline: none;
+  user-select: none;
 `;
 
-const ActionText = styled.div`
-  color: ${props => props.color || '#434343'};
-  font-family: 'Helvetica Neue', sans-serif;
-  font-size: 14px;
-  font-weight: 400;
-  line-height: 22px;
-  word-wrap: break-word;
+// #endregion MAIN & MISC----------------
+
+// #region FORM -------------------------
+
+const AskAIForm = styled.form`
+  background: #fafafa;
+  border-radius: 0.3rem 0.3rem 0 0;
+  display: ${p => (p.$show ? 'flex' : 'none')};
+  flex-direction: column;
+  justify-content: space-between;
+  padding: 8px 12px;
+  position: relative;
+  width: 458px;
 `;
 
-const AskAIFormInput = styled.input`
+const PromptInput = styled.input`
   background: transparent;
   border: none;
-  color: #000;
+  color: #555;
   font-family: 'Helvetica Neue', sans-serif;
   font-size: 14px;
   font-weight: 400;
@@ -65,238 +106,466 @@ const AskAIFormInput = styled.input`
   width: 100%;
 `;
 
-const ResultDiv = styled.div`
+const SendButton = styled(ButtonBase)`
+  border-right: 1px solid #0001;
+  margin-bottom: 3px;
+  padding: 0 8px;
+`;
+
+const SettingsButton = styled(ButtonBase)`
+  padding: 0 0 0 8px;
+
+  > svg {
+    fill: #777;
+    height: 16px;
+    width: 16px;
+  }
+`;
+
+// #endregion FORM --------------------------
+
+// #region RESULT -----------------------
+
+const ResultContainer = styled.div`
   align-items: center;
   background: white;
-  border: 0.5px #dcdcdc solid;
-  border-top-left-radius: 4px;
-  border-top-right-radius: 4px;
-  box-shadow: 0px 4px 8px rgba(0, 0, 0, 0.08);
-  display: inline-flex;
+  display: flex;
   flex-direction: column;
-  gap: 10px;
+  gap: 0;
+  height: 150px;
   justify-content: flex-start;
-  max-height: 200px;
+  max-height: ${p => (p.$show ? '150px' : '0')};
   overflow-y: auto;
-  padding: 8px 12px;
+  padding: 0;
+  position: relative;
+  transition: all 0.3s;
   width: 458px;
 `;
 
-const ResultText = styled.div`
-  color: black;
-  flex: 1 1 0;
+const ResultHeading = styled(Heading)`
+  align-items: flex-end;
+`;
+
+const ResultTab = styled(ButtonBase)`
+  background: ${p => (p.$selected ? '#fff' : '#fff4')};
+  border: var(--ai-tool-border);
+  border-bottom-color: ${p => (p.$selected ? '#fff' : '#0001')};
+  border-radius: 5px 5px 0 0;
+  color: #777;
+  font-size: 11px;
+  font-style: italic;
+  margin-bottom: -1px;
+  padding: 3px 0.5rem;
+  text-decoration: underline;
+  text-decoration-color: #bbb0;
+  text-underline-offset: 5px;
+  transform: scale(${p => (p.$selected ? '1' : '0.93')});
+  transform-origin: bottom;
+  transition: all 0.2s;
+  z-index: 9;
+
+  &:focus {
+    text-decoration-color: #bbb;
+    text-underline-offset: 2px;
+  }
+`;
+
+const ResultActions = styled.div`
+  --separation: 7px;
+  align-self: flex-end;
+  background-color: #f8f8f8;
+  border: var(--ai-tool-border);
+  bottom: var(--separation);
+  display: flex;
+  max-height: ${p => (p.$show ? '150px' : '0')};
+  opacity: ${p => (p.$show ? '1' : '0')};
+  overflow: hidden;
+  padding: 0;
+  position: absolute;
+  right: var(--separation);
+  transition: all 0.3s;
+  width: fit-content;
+
+  > button:not(:first-child) {
+    border-left: 1px solid #0001;
+  }
+`;
+
+const ResultActionButton = styled(ButtonBase)`
+  background-color: #fff0;
+  outline: none;
+  overflow: hidden;
+  padding: 3.2px;
+  transition: background-color 0.3s;
+
+  &:hover {
+    background-color: #f2f2f2;
+  }
+`;
+
+const ResultContent = styled.div`
+  border-top: var(--ai-tool-border);
+  color: #555;
   font-family: Roboto, sans-serif;
   font-size: 14px;
   font-weight: 400;
   line-height: 19px;
+  overflow-y: scroll;
+  padding: 14px 20px 32px;
   white-space: pre-line;
+  width: 100%;
   word-wrap: break-word;
-`;
 
-const SubmitButton = styled.button`
-  background: none;
-  border: none;
-  cursor: ${props => (props.disabled ? 'not-allowed' : 'pointer')};
-  outline: none;
-  padding: 0 8px; /* Adjust padding as needed */
+  * {
+    margin: 0;
+  }
+
+  h4 {
+    margin: 0 0 10px 0;
+  }
 `;
 
-const PromptDiv = styled.div`
-  background: white;
-  border: 0.5px #dcdcdc solid;
-  box-shadow: 0px 4px 8px rgba(0, 0, 0, 0.08);
-  display: inline-flex;
-  justify-content: space-between;
-  padding: 8px 12px;
+// #endregion RESULT --------------------
+
+// #region CUSTOM PROMPTS ---------------
+
+const CustomPromptContainer = styled.div`
+  border-top: ${p => (p.$show ? 'var(--ai-tool-border)' : '0 solid #0000')};
+  display: flex;
+  flex-direction: column;
+  font-style: italic;
+  max-height: ${p => (p.$show ? '205px' : '0')};
+  overflow: hidden;
+  padding: 0;
+  transition: all 0.3s linear;
   width: 458px;
 `;
 
+const CustomPromptsHeading = styled(Heading)`
+  --gradient-direction: to top;
+  color: #777;
+  font-size: 11px;
+  font-style: italic;
+  line-height: 1;
+  padding: 0.4rem;
+`;
+
+const CustomPromptsList = styled.div`
+  background: #f2f2f2;
+  display: flex;
+  flex-direction: column;
+  overflow-y: scroll;
+  padding: 0;
+  transition: all 0.3s linear;
+
+  > *:not(:first-child) {
+    border-top: var(--ai-tool-border);
+  }
+`;
+
+const CustomPromptButton = styled(ButtonBase)`
+  align-items: center;
+  padding: 0;
+
+  p {
+    background-color: ${p => (p.$selected ? '#fffa' : '#fff6')};
+    border: var(--ai-tool-border);
+    border-left: 4px solid ${p => (p.$selected ? '#ddd' : '#aaa')};
+    color: #777;
+    margin: 0;
+    padding: 0.5rem 1rem;
+    text-align: left;
+    transform: ${p => (p.$selected ? 'scale(1)' : 'scale(0.95)')};
+    transition: all 0.3s;
+    width: 100%;
+  }
+
+  &:focus {
+    p {
+      background-color: #fffa;
+      border-color: #aaa;
+    }
+  }
+`;
+
+const NoCustomPrompts = styled(FlexRow)`
+  color: #555;
+  justify-content: center;
+  text-align: center;
+`;
+// #endregion CUSTOM PROMPTS ------------
+
+// #endregion STYLED COMPONENTS ---------------------------------------------
+
 const AskAIOverlay = ({ setPosition, position, config }) => {
+  // #region HOOKS & INIT ------------------------
   const { t, i18n } = useTranslation();
+  const ctx = useContext(WaxContext);
   const {
     app,
     pmViews: { main },
     options,
-  } = useContext(WaxContext);
-  const [result, setResult] = useState('');
+  } = ctx;
+  const inputRef = useRef(null);
+  const [userPrompt, setUserPrompt] = useState('');
+  const [optionsState, setOptionsState] = useState({ ...options });
+
   const [isSubmitted, setIsSubmitted] = useState(false);
   const [isLoading, setIsLoading] = useState(false);
-  const [promptIdx, setPromptIdx] = useState(-1);
-  const { AskAiContentTransformation } = config;
-  const inputRef = useRef(null);
+
+  const [result, setResult] = useState({ [DEFAULT_KEY]: '' });
+  const [resultKey, setResultKey] = useState(DEFAULT_KEY);
+  const [showSettings, setShowSettings] = useState(false);
 
   const aiService = app.config.get('config.AskAiContentService');
-  const prompts = aiService.CustomPrompts;
+  const { AskAiContentTransformation } = config;
+  const { CustomPrompts, FreeTextPromptsOn } = aiService;
 
   useLayoutEffect(() => {
-    const WaxSurface = main.dom.getBoundingClientRect();
-    const { selection } = main.state;
-    const { to } = selection;
-    // const start = main.coordsAtPos(from);
-    const end = main.coordsAtPos(to - 1);
-    const overLayComponent = document.getElementById('ai-overlay');
-    if (!overLayComponent) return;
-    const overLayComponentCoords = overLayComponent.getBoundingClientRect();
-    const top = end.top - WaxSurface.top + 20;
-    let left = end.left - WaxSurface.left - overLayComponentCoords.width / 2;
-
-    if (end.left - overLayComponentCoords.width / 2 < WaxSurface.left) {
-      left += WaxSurface.left - (end.left - overLayComponentCoords.width / 2);
-    }
+    const aiOverlay = document.getElementById(AI_TOOL_ID);
+    if (!aiOverlay) return;
+    const coords = {
+      surface: main.dom.getBoundingClientRect(),
+      end: main.coordsAtPos(main.state.selection.to - 1),
+      overlay: aiOverlay.getBoundingClientRect(),
+    };
 
-    // Don't get out of right boundary of the surface
-    if (end.left + overLayComponentCoords.width / 2 > WaxSurface.right) {
-      left -= end.left + overLayComponentCoords.width / 2 - WaxSurface.right;
-    }
+    const { left, top } = getUpdatedPosition(coords);
 
     setPosition({ ...position, left, top });
   }, [position.left, options.AiOn]);
 
-  const tryAgain = () => {
-    // Reset the state to initial values
-    setIsSubmitted(false);
-    setResult('');
+  // #endregion HOOKS & INIT --------------------
+
+  // #region HELPERS -----------------------------
 
-    // Call the handleSubmit function again
-    handleSubmit(new Event('submit'));
-    // add underline
+  const fillAndFocusInput = customPrompt => {
+    if (userPrompt === customPrompt) return;
+    setUserPrompt(customPrompt);
+    inputRef.current.focus();
   };
 
-  const handleInsertTextBelow = () => {
-    replaceSelectedText(main, result);
+  const safeTranslation = (translation, fallback) => {
+    return !isEmpty(i18n) && i18n.exists(translation)
+      ? t(translation)
+      : fallback;
+  };
+
+  // #endregion HELPERS --------------------------
+
+  // #region HANDLERS ----------------------------
+
+  const handleInputChange = e => {
+    setUserPrompt(e.target.value);
+  };
+
+  const handleKeyDown = e => {
+    if (e.key === 'Enter') {
+      handleSubmit();
+    }
+  };
+
+  const handleShowSettings = e => {
+    e.stopPropagation(); // to prevent !showSettings
+    setShowSettings(!showSettings);
   };
 
   const handleSubmit = async () => {
-    const inputValue = inputRef.current.value;
-    if (inputValue === '') {
+    if (!enabled.send) {
       inputRef.current.focus();
       return;
     }
     setIsSubmitted(false);
     setIsLoading(true);
 
-    // Get the highlighted text from the editor
     const { from, to } = main.state.selection;
     const highlightedText = main.state.doc.textBetween(from, to);
 
-    // Combine the user's input and the highlighted text
-    const combinedInput = `${inputValue}\n\n${highlightedText}`;
+    // Updated to the new input format from gpt4o, We can pass an array of base64 images under image_url prop
+    const input = { text: [userPrompt, highlightedText] };
 
     try {
-      const response = await AskAiContentTransformation(combinedInput);
-      setResult(response);
-      setIsSubmitted(true);
+      const response = await AskAiContentTransformation(input);
+      const processedRes = safeParse(response, DEFAULT_KEY);
+      setResult(processedRes);
     } catch (error) {
-      setResult(error);
-      setIsSubmitted(true);
+      setResult({ [DEFAULT_KEY]: error });
     } finally {
+      setResultKey(DEFAULT_KEY);
+      setIsSubmitted(true);
       setIsLoading(false);
-      setPromptIdx(-1);
     }
   };
 
-  const handleReplaceText = () => {
-    replaceSelectedText(main, result, true);
-  };
+  // #endregion HANDLERS -------------------------
 
-  const discardResults = () => {
-    // Clear the input field
-    inputRef.current.value = '';
+  // #region UI ----------------------------------
 
-    // Reset the state variables
-    setResult('');
-    setIsSubmitted(false);
-  };
+  const submitIcon = isLoading ? <icons.loaderIco /> : <icons.send />;
 
-  const handleKeyDown = e => {
-    if (e.key === 'Enter') {
-      handleSubmit();
-    }
-  };
+  const resultKeys = Object.keys(result);
+
+  // To ensure we pass the response string in the correct format to the editor
+  const resultString =
+    typeof result[resultKey] === 'string'
+      ? result[resultKey]
+      : result[resultKey]?.join('\n\n') ?? '';
 
-  const handleSubmitPrompt = idx => {
-    setPromptIdx(idx);
+  const enabled = {
+    component: !!options?.AiOn,
+    input: !!FreeTextPromptsOn,
+    results: isSubmitted && !!result[resultKey],
+    customprompts: !!options?.CustomPromptsOn,
+    send: userPrompt.length > 1,
+  };
 
-    inputRef.current.value = prompts[idx];
-    handleSubmit(new Event('submit'));
+  const resultActions = {
+    replace: {
+      Icon: icons.replaceIco,
+      title: safeTranslation(
+        `Wax.AI.Replace selected text`,
+        'Replace selected text',
+      ),
+      onClick: () => {
+        replaceSelectedText(main, resultString, true);
+      },
+      tabIndex: result.enabled ? 0 : -1,
+    },
+    insert: {
+      Icon: icons.insertIco,
+      title: safeTranslation(`Wax.AI.Insert`, 'Insert'),
+      onClick: () => {
+        replaceSelectedText(main, resultString);
+      },
+      tabIndex: result.enabled ? 0 : -1,
+    },
+    tryAgain: {
+      Icon: icons.tryAgain,
+      title: safeTranslation(`Wax.AI. Try again`, 'Try again'),
+      onClick: () => {
+        setIsSubmitted(false);
+        setResult({});
+        handleSubmit();
+      },
+      tabIndex: result.enabled ? 0 : -1,
+    },
+    discard: {
+      Icon: icons.deleteIco,
+      title: safeTranslation(`Wax.AI.Discard`, 'Discard'),
+      onClick: () => {
+        setUserPrompt('');
+        setResult({});
+        setIsSubmitted(false);
+      },
+      tabIndex: result.enabled ? 0 : -1,
+    },
   };
 
-  return options?.AiOn ? (
-    <Wrapper id="ai-overlay">
-      <AskAIForm hidden={!aiService.FreeTextPromptsOn}>
-        <AskAIFormInput
-          id="askAiInput"
-          onKeyPress={handleKeyDown}
-          placeholder={
-            !isEmpty(i18n) && i18n.exists(`Wax.AI.Placeholder`)
-              ? t(`Wax.AI.Placeholder`)
-              : 'How can I help you? Type your prompt here.'
-          }
-          ref={inputRef}
-          type="text"
+  // #endregion UI -------------------------------
+
+  return enabled.component ? (
+    <Root id={AI_TOOL_ID} onClick={() => setShowSettings(false)}>
+      <AskAIForm $show>
+        <FlexRow>
+          <PromptInput
+            disabled={!enabled.input}
+            id="askAiInput"
+            onChange={handleInputChange}
+            onKeyDown={handleKeyDown}
+            placeholder={safeTranslation(
+              `Wax.AI.Placeholder`,
+              'How can I help you? Type your prompt here.',
+            )}
+            ref={inputRef}
+            type="text"
+            value={userPrompt}
+          />
+          <SendButton disabled={!enabled.send} onClick={handleSubmit}>
+            {submitIcon}
+          </SendButton>
+          <SettingsButton onClick={handleShowSettings}>
+            <icons.settings />
+          </SettingsButton>
+        </FlexRow>
+        <AiSettingsMenu
+          aiService={aiService}
+          optionsState={optionsState}
+          setOptionsState={setOptionsState}
+          show={showSettings}
         />
-        <SubmitButton onClick={handleSubmit}>
-          {isLoading ? <icons.loaderIco /> : <icons.submitIco />}
-        </SubmitButton>
       </AskAIForm>
-      {prompts.map((prompt, idx) => (
-        <PromptDiv key={prompt}>
-          <ResultText>{prompt}</ResultText>
-          <SubmitButton
-            disabled={isLoading}
-            onClick={() => handleSubmitPrompt(idx)}
-          >
-            {isLoading && idx === promptIdx ? (
-              <icons.loaderIco />
-            ) : (
-              <icons.submitIco />
-            )}
-          </SubmitButton>
-        </PromptDiv>
-      ))}
-      {isSubmitted && (
-        <>
-          <ResultDiv>
-            <ResultText dangerouslySetInnerHTML={{ __html: result }} />
-          </ResultDiv>
-          <ActionSection>
-            <ActionButton onClick={handleReplaceText}>
-              <ActionText>
-                <icons.replaceIco />{' '}
-                {!isEmpty(i18n) && i18n.exists(`Wax.AI.Replace selected text`)
-                  ? t(`Wax.AI.Replace selected text`)
-                  : 'Replace selected text'}
-              </ActionText>
-            </ActionButton>
-            <ActionButton onClick={handleInsertTextBelow}>
-              <ActionText>
-                <icons.insertIco />{' '}
-                {!isEmpty(i18n) && i18n.exists(`Wax.AI.Insert`)
-                  ? t(`Wax.AI.Insert`)
-                  : 'Insert'}
-              </ActionText>
-            </ActionButton>
-            <ActionButton onClick={tryAgain}>
-              <ActionText>
-                <icons.tryAgain />{' '}
-                {!isEmpty(i18n) && i18n.exists(`Wax.AI.Try again`)
-                  ? t(`Wax.AI. Try again`)
-                  : 'Try again'}
-              </ActionText>
-            </ActionButton>
-            <ActionButton onClick={discardResults}>
-              <ActionText color="#FF4E4E">
-                <icons.deleteIco />{' '}
-                {!isEmpty(i18n) && i18n.exists(`Wax.AI.Discard`)
-                  ? t(`Wax.AI.Discard`)
-                  : 'Discard'}
-              </ActionText>
-            </ActionButton>
-          </ActionSection>
-        </>
-      )}
-    </Wrapper>
+
+      <ResultContainer $show={enabled.results}>
+        <ResultHeading>
+          {resultKeys?.map(k => (
+            <ResultTab
+              $selected={resultKey === k}
+              key={k}
+              onClick={() => setResultKey(k)}
+            >
+              {capitalize(k)}
+            </ResultTab>
+          ))}
+        </ResultHeading>
+        <ResultContent
+          dangerouslySetInnerHTML={{
+            __html: `<h4>${capitalize(resultKey)}</h4>${resultsToHtml(
+              resultKey,
+              result[resultKey],
+            )}</br>`,
+          }}
+        />
+        <ResultActions $show={enabled.results}>
+          {Object.values(resultActions).map(({ title, Icon, ...rest }) => (
+            <ResultActionButton key={title} {...rest}>
+              <Icon />
+            </ResultActionButton>
+          ))}
+        </ResultActions>
+      </ResultContainer>
+
+      <CustomPromptContainer $show={enabled.customprompts}>
+        <CustomPromptsHeading>Custom Prompts</CustomPromptsHeading>
+        <CustomPromptsList>
+          {CustomPrompts.length > 0 ? (
+            CustomPrompts?.map(prompt => (
+              <CustomPromptButton
+                $selected={userPrompt === prompt}
+                key={prompt}
+                onClick={() => fillAndFocusInput(prompt)}
+              >
+                <p>{`"${prompt}"`}</p>
+              </CustomPromptButton>
+            ))
+          ) : (
+            <NoCustomPrompts>
+              <p>
+                --
+                {safeTranslation(
+                  `Wax.AI.No custom prompts`,
+                  `You don't have any custom prompts`,
+                )}{' '}
+                --
+              </p>
+            </NoCustomPrompts>
+          )}
+        </CustomPromptsList>
+      </CustomPromptContainer>
+    </Root>
   ) : null;
 };
 
+AskAIOverlay.propTypes = {
+  position: PropTypes.shape({
+    left: PropTypes.number,
+  }),
+  setPosition: PropTypes.func,
+  config: PropTypes.shape({ AskAiContentTransformation: PropTypes.shape({}) }),
+};
+AskAIOverlay.defaultProps = {
+  position: {},
+  setPosition: () => {},
+  config: {},
+};
+
 export default AskAIOverlay;
diff --git a/wax-prosemirror-services/src/AiService/components/Switch.js b/wax-prosemirror-services/src/AiService/components/Switch.js
index 8ee857968a999e989d55bf3c8f8c34ac05960738..841ea41c330f3810dd3a28c8b8ae7de21d16eb41 100644
--- a/wax-prosemirror-services/src/AiService/components/Switch.js
+++ b/wax-prosemirror-services/src/AiService/components/Switch.js
@@ -15,12 +15,13 @@ const Wrapper = styled.span`
   }
 
   .rc-switch-checked {
-    border: 1px solid ${th('colorPrimary')};
     background-color: ${th('colorPrimary')};
+    border: 1px solid ${th('colorPrimary')};
 
     .rc-switch-inner {
       left: 6px;
     }
+
     :after {
       left: 42px;
     }
@@ -28,6 +29,7 @@ const Wrapper = styled.span`
 `;
 
 const Label = styled.label`
+  cursor: pointer;
   ${props =>
     props.labelPosition === 'left' &&
     css`
@@ -43,7 +45,6 @@ const Label = styled.label`
       font-size: 14px;
       margin-left: ${grid(2)};
     `}
-    cursor: pointer;
 `;
 
 const SwitchComponent = props => {
diff --git a/wax-prosemirror-services/src/AiService/helpers/index.js b/wax-prosemirror-services/src/AiService/helpers/index.js
new file mode 100644
index 0000000000000000000000000000000000000000..0c74fb22e65738c89a7f4108d9213c6d91f91d42
--- /dev/null
+++ b/wax-prosemirror-services/src/AiService/helpers/index.js
@@ -0,0 +1,41 @@
+export const safeParse = (str, fallbackKey) => {
+  try {
+    const parsed = JSON.parse(str);
+    return parsed;
+  } catch (e) {
+    return { [fallbackKey]: str };
+  }
+};
+
+export const resultsToHtml = (key, value) => {
+  const variations = {
+    content: value,
+    citations:
+      typeof value === 'string'
+        ? value
+        : value?.map(item => `<blockquote>${item}</blockquote>`).join(''),
+    links:
+      typeof value === 'string'
+        ? value
+        : value
+            ?.map(item => `<a href="${item}" target="_blank">${item}</a></br>`)
+            .join(''),
+  };
+  return variations[key] ?? value;
+};
+
+export const getUpdatedPosition = ({ surface, end, overlay }) => {
+  const top = end.top - surface.top + 20;
+  let left = end.left - surface.left - overlay.width / 2;
+
+  if (end.left - overlay.width / 2 < surface.left) {
+    left += surface.left - (end.left - overlay.width / 2);
+  }
+
+  // Don't get out of right boundary of the surface
+  if (end.left + overlay.width / 2 > surface.right) {
+    left -= end.left + overlay.width / 2 - surface.right;
+  }
+
+  return { left, top };
+};