diff --git a/editors/demo/src/Editoria/config/config.js b/editors/demo/src/Editoria/config/config.js
index e8214f1daad42efb9d44179cbeca4835af58e62b..3c828db6d88f2261ab5daefe9f24aba7e2f465b5 100644
--- a/editors/demo/src/Editoria/config/config.js
+++ b/editors/demo/src/Editoria/config/config.js
@@ -106,7 +106,7 @@ async function DummyPromise(userInput) {
       if (userInput === 'reject') {
         reject('Your request could not be processed for now');
       } else {
-        resolve('hello how are you?');
+        resolve(text);
       }
     }, 150);
   });
diff --git a/wax-prosemirror-services/src/AiService/ReplaceSelectedText.js b/wax-prosemirror-services/src/AiService/ReplaceSelectedText.js
index 34f04ad175ea589eaba102549c861a5698e6ccb6..25b8c000f925a2fcbf6baf52e9f604e4b2828f90 100644
--- a/wax-prosemirror-services/src/AiService/ReplaceSelectedText.js
+++ b/wax-prosemirror-services/src/AiService/ReplaceSelectedText.js
@@ -15,7 +15,6 @@ const replaceSelectedText = (view, transformedText, replace = false) => {
   const parser = DOMParser.fromSchema(state.config.schema);
 
   if (from > state.doc.content.size || to > state.doc.content.size) {
-    console.error('Position out of range');
     return;
   }
 
@@ -49,7 +48,8 @@ const replaceSelectedText = (view, transformedText, replace = false) => {
   state = view.state;
 
   // Update the selection to the end of the new text
-  const newTo = from + transformedText.length;
+  const newFrom = replace ? from : to;
+  const newTo = newFrom + transformedText.length;
   const cursorPosition = paragraphNodes.length !== 0 ? newTo + 2 : newTo;
   const newSelection = TextSelection.create(
     state.doc,
diff --git a/wax-prosemirror-services/src/AiService/components/AskAIOverlay.js b/wax-prosemirror-services/src/AiService/components/AskAIOverlay.js
index cc697b9b63575325ef022d33eb1ce319a3a3bd78..0c33e7f0f6b4b73c8a19ad7381a3b0df4b668c6c 100644
--- a/wax-prosemirror-services/src/AiService/components/AskAIOverlay.js
+++ b/wax-prosemirror-services/src/AiService/components/AskAIOverlay.js
@@ -110,8 +110,8 @@ const AskAIOverlay = ({ setPosition, position, config }) => {
   useLayoutEffect(() => {
     const WaxSurface = activeView.dom.getBoundingClientRect();
     const { selection } = activeView.state;
-    const { from, to } = selection;
-    const start = activeView.coordsAtPos(from);
+    const { to } = selection;
+    // const start = activeView.coordsAtPos(from);
     const end = activeView.coordsAtPos(to - 1);
     const overLayComponent = document.getElementById('ai-overlay');
 
@@ -146,8 +146,12 @@ const AskAIOverlay = ({ setPosition, position, config }) => {
   };
 
   const handleSubmit = async () => {
-    setIsLoading(true);
     const inputValue = inputRef.current.value;
+    if (inputValue === '') {
+      inputRef.current.focus();
+      return;
+    }
+    setIsLoading(true);
 
     // Get the highlighted text from the editor
     const { from, to } = activeView.state.selection;