diff --git a/webview_panels/.storybook/preview-head.html b/webview_panels/.storybook/preview-head.html
index 639dacbeb..2fa90bea2 100644
--- a/webview_panels/.storybook/preview-head.html
+++ b/webview_panels/.storybook/preview-head.html
@@ -4,6 +4,12 @@
 /> -->
 
 <style>
+  body {
+    --vscode-panel-background: rgb(248, 248, 248);
+  }
+  body.vscode-dark {
+    --vscode-panel-background: rgb(30, 30, 30);
+  }
   perspective-viewer {
     height: 100vh !important;
   }
diff --git a/webview_panels/src/modules/documentationEditor/DocumentationEditor.tsx b/webview_panels/src/modules/documentationEditor/DocumentationEditor.tsx
index 05fa636c9..5ccc6617f 100644
--- a/webview_panels/src/modules/documentationEditor/DocumentationEditor.tsx
+++ b/webview_panels/src/modules/documentationEditor/DocumentationEditor.tsx
@@ -22,6 +22,7 @@ import CoachAiIfModified from "./components/docGenerator/CoachAiIfModified";
 import Citations from "./components/docGenerator/Citations";
 import { Citation } from "@lib";
 import SearchColumnsInput from "./components/search/SearchColumnsInput";
+import BulkGenerateButton from "./components/docGenerator/BulkGenerateButton";
 
 const DocumentationEditor = (): JSX.Element => {
   const {
@@ -112,10 +113,11 @@ const DocumentationEditor = (): JSX.Element => {
   return (
     <div className={`${classes.documentationWrapper} ${classes.limitWidth}`}>
       <Stack className="mb-2 justify-content-between">
+        <SearchColumnsInput />
         <Stack>
-          <SearchColumnsInput />
+          <BulkGenerateButton />
+          <CommonActionButtons />
         </Stack>
-        <CommonActionButtons />
       </Stack>
       <div className={classes.docGenerator}>
         <Stack className={classes.bodyWrap}>
diff --git a/webview_panels/src/modules/documentationEditor/components/docGenerator/BulkGenerateButton.tsx b/webview_panels/src/modules/documentationEditor/components/docGenerator/BulkGenerateButton.tsx
index 4b439c15f..0299dd1cb 100644
--- a/webview_panels/src/modules/documentationEditor/components/docGenerator/BulkGenerateButton.tsx
+++ b/webview_panels/src/modules/documentationEditor/components/docGenerator/BulkGenerateButton.tsx
@@ -7,18 +7,18 @@ import { updateColumnsInCurrentDocsData } from "@modules/documentationEditor/sta
 import { DBTDocumentationColumn } from "@modules/documentationEditor/state/types";
 import useDocumentationContext from "@modules/documentationEditor/state/useDocumentationContext";
 import { panelLogger } from "@modules/logger";
-import { Button, DropdownButton, List, Popover, PopoverBody } from "@uicore";
+import { Button, DropdownButton, List, PopoverWithButton } from "@uicore";
 import { useRef, useState } from "react";
 import classes from "../../styles.module.scss";
 import { mergeCurrentAndIncomingDocumentationColumns } from "@modules/documentationEditor/utils";
 import DocGenSelectedColumns from "./DocGenSelectedColumns";
+import { noop } from "antd/es/_util/warning";
 
 enum SidePanelState {
   DOCUMENTATION_SELECTED = "documentationSelected",
 }
 
 const BulkGenerateButton = (): JSX.Element => {
-  const [openPopover, setOpenPopover] = useState(false);
   const [sidePanelState, setSidePanelState] = useState<
     SidePanelState | undefined
   >();
@@ -28,10 +28,6 @@ const BulkGenerateButton = (): JSX.Element => {
     dispatch,
   } = useDocumentationContext();
 
-  const onToggleClick = () => {
-    setOpenPopover((prev) => !prev);
-  };
-
   const resetSidepanelState = () => {
     setSidePanelState(undefined);
   };
@@ -117,7 +113,6 @@ const BulkGenerateButton = (): JSX.Element => {
   };
 
   const onOptionSelect = async (value: string) => {
-    setOpenPopover(false);
     const startTime = Date.now();
     sendTelemetryEvent(value);
 
@@ -151,41 +146,50 @@ const BulkGenerateButton = (): JSX.Element => {
   return (
     <>
       <div ref={ref}>
-        <DropdownButton onToggleClick={onToggleClick} onClick={onToggleClick}>
-          <ShinesIcon /> Bulk generate
-        </DropdownButton>
-        <Popover
-          isOpen={openPopover}
-          target={ref}
-          placement="bottom"
-          hideArrow
-          className={classes.popover}
+        <PopoverWithButton
+          width="auto"
+          button={
+            <DropdownButton onToggleClick={noop} onClick={noop}>
+              <ShinesIcon /> Bulk generate
+            </DropdownButton>
+          }
+          popoverProps={{
+            placement: "bottom",
+            hideArrow: true,
+          }}
         >
-          <PopoverBody>
-            <List>
-              {Object.entries(options).map(([key, actions]) => (
-                <>
-                  <li>{key}</li>
-                  {actions.map((option) => (
-                    <li key={option.label}>
-                      <Button
-                        color="link"
-                        onClick={() => onOptionSelect(option.value)}
-                      >
-                        {option.label}
-                      </Button>
-                    </li>
+          {({ styles, close }) => (
+            <div className={classes.popover}>
+              <div className={styles.popoverActions}>
+                <List>
+                  {Object.entries(options).map(([key, actions]) => (
+                    <>
+                      <li className={classes.sectionTitle}>{key}</li>
+                      {actions.map((option) => (
+                        <li key={option.label}>
+                          <Button
+                            color="secondary"
+                            onClick={() => {
+                              close();
+                              void onOptionSelect(option.value);
+                            }}
+                          >
+                            {option.label}
+                          </Button>
+                        </li>
+                      ))}
+                    </>
                   ))}
-                </>
-              ))}
-            </List>
-          </PopoverBody>
-        </Popover>
+                </List>
+              </div>
+            </div>
+          )}
+        </PopoverWithButton>
       </div>
       {sidePanelState === SidePanelState.DOCUMENTATION_SELECTED ? (
         <DocGenSelectedColumns
           onClose={resetSidepanelState}
-          generateForColumns={(c) => bulkGenerateDocs(c, false)}
+          generateForColumns={bulkGenerateDocs}
         />
       ) : null}
     </>
diff --git a/webview_panels/src/modules/documentationEditor/components/docGenerator/DocGenSelectedColumns.tsx b/webview_panels/src/modules/documentationEditor/components/docGenerator/DocGenSelectedColumns.tsx
index 4ffb300b3..789caee29 100644
--- a/webview_panels/src/modules/documentationEditor/components/docGenerator/DocGenSelectedColumns.tsx
+++ b/webview_panels/src/modules/documentationEditor/components/docGenerator/DocGenSelectedColumns.tsx
@@ -1,3 +1,4 @@
+import { DocsIcon, SearchIcon } from "@assets/icons";
 import { DBTDocumentationColumn } from "@modules/documentationEditor/state/types";
 import useDocumentationContext from "@modules/documentationEditor/state/useDocumentationContext";
 import {
@@ -7,19 +8,27 @@ import {
   Input,
   List,
   ListGroupItem,
+  LoadingButton,
   Stack,
+  Tooltip,
 } from "@uicore";
 import { ChangeEvent, useEffect, useMemo, useRef, useState } from "react";
+import classes from "../../styles.module.scss";
+import { panelLogger } from "@modules/logger";
 
 interface Props {
   onClose: () => void;
-  generateForColumns: (columns: DBTDocumentationColumn[]) => void;
+  generateForColumns: (
+    columns: DBTDocumentationColumn[],
+    isAll: boolean,
+  ) => Promise<DBTDocumentationColumn[]>;
 }
 const DocGenSelectedColumns = ({
   onClose,
   generateForColumns,
 }: Props): JSX.Element => {
   const drawerRef = useRef<DrawerRef | null>(null);
+  const [isGenerating, setIsGenerating] = useState(false);
   const [searchQuery, setsearchQuery] = useState("");
   const [selectedColumns, setSelectedColumns] = useState<string[]>([]);
   const {
@@ -38,12 +47,21 @@ const DocGenSelectedColumns = ({
     );
   }, [searchQuery, currentDocsData?.columns]);
 
-  const handleGenerate = () => {
+  const handleGenerate = async () => {
     const columns = currentDocsData?.columns.filter((column) =>
       selectedColumns.includes(column.name),
     );
     if (!columns) return;
-    generateForColumns(columns);
+    setIsGenerating(true);
+    try {
+      await generateForColumns(columns, false);
+      drawerRef.current?.close();
+      onClose();
+    } catch (error) {
+      panelLogger.error("error while generating documentation", error);
+    } finally {
+      setIsGenerating(false);
+    }
   };
 
   const handleFilterChange = (e: ChangeEvent<HTMLInputElement>) => {
@@ -74,23 +92,42 @@ const DocGenSelectedColumns = ({
     }
   };
 
+  const handleSelectedColumns = (e: ChangeEvent<HTMLInputElement>) => {
+    const { value } = e.target;
+    if (selectedColumns.includes(value)) {
+      setSelectedColumns(selectedColumns.filter((column) => column !== value));
+    } else {
+      setSelectedColumns([...selectedColumns, value]);
+    }
+  };
+
   return (
     <Drawer ref={drawerRef} onClose={onClose}>
       <Stack direction="column" className="h-100">
         <Stack className="justify-content-between">
           Generate documentation{" "}
-          <Button onClick={handleGenerate}>Generate</Button>
+          <LoadingButton
+            disabled={selectedColumns.length === 0}
+            loading={isGenerating}
+            color="primary"
+            onClick={handleGenerate}
+          >
+            Generate{" "}
+            {selectedColumns.length ? `(${selectedColumns.length})` : ""}
+          </LoadingButton>
         </Stack>
-        <div>
+        <Stack className={classes.search}>
+          <SearchIcon />
           <Input placeholder="Search columns" onChange={handleFilterChange} />
-        </div>
+        </Stack>
+        <div></div>
         <Stack>
           <Button onClick={() => handleSelectColumns("all")}>Select All</Button>
           <Button
             color="success"
             onClick={() => handleSelectColumns("documented")}
           >
-            Documented
+            <DocsIcon className={classes.docsIconInBtn} /> Documented
           </Button>
           <Button
             color="warning"
@@ -100,18 +137,27 @@ const DocGenSelectedColumns = ({
           </Button>
         </Stack>
         <Stack className="overflow-auto">
-          <List>
+          <List className="m-0 p-0">
             {filteredColumns.map((column) => (
               <ListGroupItem
                 key={`${column.name}-${column.type}`}
                 tag={"label"}
+                className="mb-1"
               >
-                <Input
-                  type="checkbox"
-                  value={column.name}
-                  checked={selectedColumns.includes(column.name)}
-                />
-                {column.name}
+                <Stack>
+                  <Input
+                    type="checkbox"
+                    value={column.name}
+                    checked={selectedColumns.includes(column.name)}
+                    onChange={handleSelectedColumns}
+                  />
+                  {column.name}
+                  {column.description ? (
+                    <Tooltip title="Documented" placement="top">
+                      <DocsIcon className={classes.docsIcon} />
+                    </Tooltip>
+                  ) : null}
+                </Stack>
               </ListGroupItem>
             ))}
           </List>
diff --git a/webview_panels/src/modules/documentationEditor/components/docGenerator/DocGeneratorColumnsList.tsx b/webview_panels/src/modules/documentationEditor/components/docGenerator/DocGeneratorColumnsList.tsx
index 07cd4c4ad..302de73d5 100644
--- a/webview_panels/src/modules/documentationEditor/components/docGenerator/DocGeneratorColumnsList.tsx
+++ b/webview_panels/src/modules/documentationEditor/components/docGenerator/DocGeneratorColumnsList.tsx
@@ -2,7 +2,6 @@ import useDocumentationContext from "@modules/documentationEditor/state/useDocum
 import { Alert, Stack } from "@uicore";
 import DocGeneratorColumn from "./DocGeneratorColumn";
 import classes from "../../styles.module.scss";
-import BulkGenerateButton from "./BulkGenerateButton";
 import SyncWithDatabase from "./SyncWithDatabase";
 import { useMemo } from "react";
 import { DBTModelTest } from "@modules/documentationEditor/state/types";
@@ -49,7 +48,6 @@ const DocGeneratorColumnsList = (): JSX.Element => {
             <h3>Columns</h3>
             <SyncWithDatabase />
           </Stack>
-          <BulkGenerateButton />
         </Stack>
         <Alert color="warning">
           Note: If you don’t want to override existing documentation, please
diff --git a/webview_panels/src/modules/documentationEditor/components/search/SearchColumnsInput.tsx b/webview_panels/src/modules/documentationEditor/components/search/SearchColumnsInput.tsx
index b96c99296..c9191423c 100644
--- a/webview_panels/src/modules/documentationEditor/components/search/SearchColumnsInput.tsx
+++ b/webview_panels/src/modules/documentationEditor/components/search/SearchColumnsInput.tsx
@@ -1,6 +1,8 @@
+import { SearchIcon } from "@assets/icons";
 import { setSearchQuery } from "@modules/documentationEditor/state/documentationSlice";
 import useDocumentationContext from "@modules/documentationEditor/state/useDocumentationContext";
-import { Input } from "@uicore";
+import { Input, Stack } from "@uicore";
+import styles from "../../styles.module.scss";
 
 const SearchColumnsInput = (): JSX.Element => {
   const { dispatch } = useDocumentationContext();
@@ -9,7 +11,18 @@ const SearchColumnsInput = (): JSX.Element => {
     dispatch(setSearchQuery(e.target.value));
   };
 
-  return <Input onChange={handleChange} placeholder="Search by column name" />;
+  return (
+    <Stack className={styles.search}>
+      <SearchIcon />
+      <Input
+        aria-label="Search by column name"
+        role="search"
+        type="search"
+        onChange={handleChange}
+        placeholder="Search by column name"
+      />
+    </Stack>
+  );
 };
 
 export default SearchColumnsInput;
diff --git a/webview_panels/src/modules/documentationEditor/styles.module.scss b/webview_panels/src/modules/documentationEditor/styles.module.scss
index 41b022d4e..30196d1eb 100644
--- a/webview_panels/src/modules/documentationEditor/styles.module.scss
+++ b/webview_panels/src/modules/documentationEditor/styles.module.scss
@@ -292,28 +292,23 @@
 }
 
 .popover {
-  :global .popover {
-    margin-top: -18px;
-    :global .popover-body {
-      padding: 0;
-
-      & ul {
-        padding: 0.5rem 0;
-        margin: 0;
-        list-style: none;
-        li {
-          :global .btn-link {
-            text-decoration: none;
-            color: #1e1e1e;
-            width: 100%;
-            text-align: left;
-            &:hover,
-            &.active {
-              color: var(--primary-color);
-              background-color: #eee;
-            }
-          }
-        }
+  margin-top: -1rem;
+  
+  ul {
+    padding: 0.5rem 0;
+    margin: 0;
+    list-style: none;
+    .section_title{
+      background-color: var(--background--base);
+      padding: 0.5rem;
+      border-radius: 6px;
+    }
+    li {
+      :global .btn {
+        text-decoration: none;
+        width: 100%;
+        text-align: left;
+        background-color: transparent;
       }
     }
   }
@@ -336,3 +331,35 @@
     max-width: 1440px;
   }
 }
+
+.docsIcon {
+  width: 13px;
+  height: 13px;
+
+  path {
+    fill: #198754;
+  }
+}
+
+.docsIconInBtn {
+  height: 13px;
+  width: 13px;
+  vertical-align: text-bottom;
+
+  path {
+    fill: #fff;
+  }
+}
+
+.search{
+  position: relative;
+  :global .codicon{
+    position: absolute;
+    top: 8px;
+    left: 5px;
+  }
+  input {
+    padding-left: 25px;
+    min-width: 200px;
+  }
+}
diff --git a/webview_panels/src/uiCore/components/dropdownButton/styles.module.scss b/webview_panels/src/uiCore/components/dropdownButton/styles.module.scss
index 00eb4089f..fad4344a4 100644
--- a/webview_panels/src/uiCore/components/dropdownButton/styles.module.scss
+++ b/webview_panels/src/uiCore/components/dropdownButton/styles.module.scss
@@ -1,5 +1,5 @@
 .dropdownButton {
-  gap: 0;
+  gap: 0 !important;
   & button:first-child {
     border-radius: 2px 0 0 2px;
   }