diff --git a/src/frontend/apps/impress/src/features/docs/doc-editor/api/useAIRewrite.tsx b/src/frontend/apps/impress/src/features/docs/doc-editor/api/useAIRewrite.tsx
index c3aa94fa..9543ecc9 100644
--- a/src/frontend/apps/impress/src/features/docs/doc-editor/api/useAIRewrite.tsx
+++ b/src/frontend/apps/impress/src/features/docs/doc-editor/api/useAIRewrite.tsx
@@ -8,9 +8,7 @@ export type DocAIParams = {
   action: 'rephrase' | 'summarize';
 };
 
-export type DocAIResponse = {
-  textAI: string;
-};
+export type DocAIResponse = string;
 
 export const DocAI = async ({
   docId,
diff --git a/src/frontend/apps/impress/src/features/docs/doc-editor/components/AIButton.tsx b/src/frontend/apps/impress/src/features/docs/doc-editor/components/AIButton.tsx
index 5bf93671..061b50b4 100644
--- a/src/frontend/apps/impress/src/features/docs/doc-editor/components/AIButton.tsx
+++ b/src/frontend/apps/impress/src/features/docs/doc-editor/components/AIButton.tsx
@@ -7,12 +7,17 @@ import {
 } from '@blocknote/react';
 import { useMemo } from 'react';
 
+import { Doc } from '../../doc-management';
 import { useAIRewrite } from '../api/useAIRewrite';
 
+interface AIButtonProps {
+  doc: Doc;
+}
+
 /**
  * Custom Formatting Toolbar Button to convert markdown to json.
  */
-export function AIButton() {
+export function AIButton({ doc }: AIButtonProps) {
   const editor = useBlockNoteEditor();
   const Components = useComponentsContext();
   const selectedBlocks = useSelectedBlocks(editor);
@@ -22,49 +27,24 @@ export function AIButton() {
   useEditorContentOrSelectionChange(() => {
     console.log('Content or selection changed');
   }, editor);
-  const { mutateAsync: requestAI, data } = useAIRewrite();
-
-  console.log('Data', data);
+  const { mutateAsync: requestAI } = useAIRewrite();
 
   const handleRephraseAI = async () => {
-    console.log('Rephrase with AI', editor.getSelection());
-    console.log('Rephrase with BLockkk', selectedBlocks);
-
-    await requestAI({
-      docId: '901ac9b3-97a4-4f37-adcf-1941746b3c61',
-      text: 'toto',
+    const textCursorPosition = editor.getSelectedText();
+    const newText = await requestAI({
+      docId: doc.id,
+      text: textCursorPosition,
       action: 'rephrase',
     });
 
-    // Call backend endpoint to rephrase the selected block
-
     editor.replaceBlocks(
       [selectedBlocks[0]],
       [
         {
-          content:
-            'This block was replaced at ' + new Date().toLocaleTimeString(),
+          content: newText,
         },
       ],
     );
-
-    // forEach(blocks, async (block) => {
-    //   if (!isBlock(block as unknown as Block)) {
-    //     return;
-    //   }
-
-    //   try {
-    //     const fullContent = recursiveContent(
-    //       block.content as unknown as Block[],
-    //     );
-
-    //     const blockMarkdown =
-    //       await editor.tryParseMarkdownToBlocks(fullContent);
-    //     editor.replaceBlocks([block.id], blockMarkdown);
-    //   } catch (error) {
-    //     console.error('Error parsing Markdown:', error);
-    //   }
-    // });
   };
 
   const show = useMemo(() => {
diff --git a/src/frontend/apps/impress/src/features/docs/doc-editor/components/BlockNoteEditor.tsx b/src/frontend/apps/impress/src/features/docs/doc-editor/components/BlockNoteEditor.tsx
index 3d228a77..21a5fa0e 100644
--- a/src/frontend/apps/impress/src/features/docs/doc-editor/components/BlockNoteEditor.tsx
+++ b/src/frontend/apps/impress/src/features/docs/doc-editor/components/BlockNoteEditor.tsx
@@ -130,7 +130,7 @@ export const BlockNoteContent = ({
         editable={doc.abilities.partial_update && !isVersion}
         theme="light"
       >
-        <BlockNoteToolbar />
+        <BlockNoteToolbar doc={doc} />
       </BlockNoteView>
     </Box>
   );
diff --git a/src/frontend/apps/impress/src/features/docs/doc-editor/components/BlockNoteToolbar.tsx b/src/frontend/apps/impress/src/features/docs/doc-editor/components/BlockNoteToolbar.tsx
index ebed1ad2..79a03d3f 100644
--- a/src/frontend/apps/impress/src/features/docs/doc-editor/components/BlockNoteToolbar.tsx
+++ b/src/frontend/apps/impress/src/features/docs/doc-editor/components/BlockNoteToolbar.tsx
@@ -16,9 +16,15 @@ import {
 import { forEach, isArray } from 'lodash';
 import React, { useMemo } from 'react';
 
+import { Doc } from '../../doc-management';
+
 import { AIButton } from './AIButton';
 
-export const BlockNoteToolbar = () => {
+interface BlockNoteToolbarProps {
+  doc: Doc;
+}
+
+export const BlockNoteToolbar = ({ doc }: BlockNoteToolbarProps) => {
   return (
     <FormattingToolbarController
       formattingToolbar={() => (
@@ -26,7 +32,7 @@ export const BlockNoteToolbar = () => {
           <BlockTypeSelect key="blockTypeSelect" />
 
           {/* Extra button to convert from markdown to json */}
-          <AIButton key="AIButton" />
+          <AIButton key="AIButton" doc={doc} />
 
           {/* Extra button to convert from markdown to json */}
           <MarkdownButton key="customButton" />