docs: improve AI Assistant (#11208)

* initial implementation

* integrate new ai assistant in other projects + ux improvements

* fix chat window on mobile devices

* fixes to mobile

* allow pre

* change shortcut to i

* improved responsiveness

* align version in navbar

www/apps/api-reference/components/Tags/Operation/CodeSection/index.tsx

@@ -36,7 +36,7 @@ const TagOperationCodeSection = ({
       >
         <div className={clsx("flex w-[calc(100%-36px)] gap-1")}>
           <MethodLabel method={method} className="h-fit" />
-          <code className="text-medusa-fg-subtle =break-words break-all">
+          <code className="text-medusa-fg-base =break-words break-all">
             {endpointPath}
           </code>
         </div>

www/apps/api-reference/components/Tags/Operation/DescriptionSection/index.tsx

@@ -53,7 +53,7 @@ const TagsOperationDescriptionSection = ({
         {operation["x-featureFlag"] && (
           <FeatureFlagNotice
             featureFlag={operation["x-featureFlag"]}
-            tooltipTextClassName="font-normal text-medusa-fg-subtle"
+            tooltipTextClassName="font-normal text-medusa-fg-base"
             badgeClassName="ml-0.5"
           />
         )}

www/apps/api-reference/components/VersionSwitcher/index.tsx

@@ -16,7 +16,7 @@ const VersionSwitcher = () => {
           location.href = process.env.NEXT_PUBLIC_API_V1_URL + pathname
         }}
       />
-      <span className={clsx("text-medusa-fg-subtle")}>V2</span>
+      <span className={clsx("text-medusa-fg-base")}>V2</span>
     </div>
   )
 }

www/apps/api-reference/providers/index.tsx

@@ -1,4 +1,5 @@
 import {
+  AiAssistantProvider,
   AnalyticsProvider,
   PageLoadingProvider,
   ScrollControllerProvider,
@@ -21,7 +22,20 @@ const Providers = ({ children }: ProvidersProps) => {
           <ScrollControllerProvider scrollableSelector="#main">
             <SidebarProvider>
               <MainNavProvider>
-                <SearchProvider>{children}</SearchProvider>
+                <SearchProvider>
+                  <AiAssistantProvider
+                    apiUrl={process.env.NEXT_PUBLIC_AI_ASSISTANT_URL || "temp"}
+                    websiteId={process.env.NEXT_PUBLIC_AI_WEBSITE_ID || "temp"}
+                    recaptchaSiteKey={
+                      process.env
+                        .NEXT_PUBLIC_AI_API_ASSISTANT_RECAPTCHA_SITE_KEY ||
+                      "temp"
+                    }
+                    chatType="popover"
+                  >
+                    {children}
+                  </AiAssistantProvider>
+                </SearchProvider>
               </MainNavProvider>
             </SidebarProvider>
           </ScrollControllerProvider>

www/apps/api-reference/providers/search.tsx

@@ -4,8 +4,6 @@ import {
   usePageLoading,
   SearchProvider as UiSearchProvider,
   searchFilters,
-  AiAssistantIcon,
-  AiAssistantProvider,
 } from "docs-ui"
 import { config } from "../config"
 import basePathUrl from "../utils/base-path-url"
@@ -53,28 +51,6 @@ const SearchProvider = ({ children }: SearchProviderProps) => {
         ),
         filterOptions: searchFilters,
       }}
-      commands={[
-        {
-          name: "ai-assistant",
-          icon: <AiAssistantIcon />,
-          component: (
-            <AiAssistantProvider
-              apiUrl={process.env.NEXT_PUBLIC_AI_ASSISTANT_URL || "temp"}
-              websiteId={process.env.NEXT_PUBLIC_AI_WEBSITE_ID || "temp"}
-              recaptchaSiteKey={
-                process.env.NEXT_PUBLIC_AI_API_ASSISTANT_RECAPTCHA_SITE_KEY ||
-                "temp"
-              }
-            />
-          ),
-          title: "AI Assistant",
-          badge: {
-            variant: "blue",
-            badgeType: "shaded",
-            children: "Beta",
-          },
-        },
-      ]}
     >
       {children}
     </UiSearchProvider>

www/apps/book/app/page.tsx

@@ -1,5 +1,5 @@
 import clsx from "clsx"
-import { MainNav, RootProviders } from "docs-ui"
+import { AiAssistantChatWindow, MainNav, RootProviders } from "docs-ui"
 import HomepageTopSection from "../components/Homepage/TopSection"
 import Providers from "../providers"
 import HomepageLinksSection from "../components/Homepage/LinksSection"
@@ -12,12 +12,20 @@ const Homepage = () => {
     <body
       className={clsx(
         "bg-medusa-bg-subtle font-base text-medium w-full",
-        "text-medusa-fg-subtle px-0.25 pt-0.25",
+        "text-medusa-fg-base px-0.25 pt-0.25",
         "h-screen overflow-hidden"
       )}
     >
-      <RootProviders>
+      <RootProviders
-        <Providers>
+        layoutProviderProps={{
+          disableResizeObserver: true,
+        }}
+      >
+        <Providers
+          aiAssistantProps={{
+            chatType: "popover",
+          }}
+        >
           <div
             className={clsx(
               "rounded-t bg-medusa-bg-base",
@@ -40,6 +48,7 @@ const Homepage = () => {
             <HomepageModulesSection />
             <HomepageFooter />
           </div>
+          <AiAssistantChatWindow />
         </Providers>
       </RootProviders>
     </body>

www/apps/book/providers/index.tsx

@@ -1,6 +1,8 @@
 "use client"
 
 import {
+  AiAssistantProvider,
+  AiAssistantProviderProps,
   AnalyticsProvider,
   HooksLoader,
   LearningPathProvider,
@@ -16,9 +18,10 @@ import { MainNavProvider } from "./main-nav"
 
 type ProvidersProps = {
   children?: React.ReactNode
+  aiAssistantProps?: Partial<Omit<AiAssistantProviderProps, "children">>
 }
 
-const Providers = ({ children }: ProvidersProps) => {
+const Providers = ({ children, aiAssistantProps = {} }: ProvidersProps) => {
   return (
     <AnalyticsProvider writeKey={process.env.NEXT_PUBLIC_SEGMENT_API_KEY}>
       <SiteConfigProvider config={config}>
@@ -29,14 +32,29 @@ const Providers = ({ children }: ProvidersProps) => {
                 <PaginationProvider>
                   <MainNavProvider>
                     <SearchProvider>
-                      <HooksLoader
+                      <AiAssistantProvider
-                        options={{
+                        {...aiAssistantProps}
-                          pageScrollManager: true,
+                        apiUrl={
-                          currentLearningPath: false,
+                          process.env.NEXT_PUBLIC_AI_ASSISTANT_URL || "temp"
-                        }}
+                        }
+                        websiteId={
+                          process.env.NEXT_PUBLIC_AI_WEBSITE_ID || "temp"
+                        }
+                        recaptchaSiteKey={
+                          process.env
+                            .NEXT_PUBLIC_AI_API_ASSISTANT_RECAPTCHA_SITE_KEY ||
+                          "temp"
+                        }
                       >
-                        {children}
+                        <HooksLoader
-                      </HooksLoader>
+                          options={{
+                            pageScrollManager: true,
+                            currentLearningPath: false,
+                          }}
+                        >
+                          {children}
+                        </HooksLoader>
+                      </AiAssistantProvider>
                     </SearchProvider>
                   </MainNavProvider>
                 </PaginationProvider>

www/apps/book/providers/search.tsx

@@ -1,11 +1,6 @@
 "use client"
 
-import {
+import { SearchProvider as UiSearchProvider, searchFilters } from "docs-ui"
-  AiAssistantIcon,
-  AiAssistantProvider,
-  SearchProvider as UiSearchProvider,
-  searchFilters,
-} from "docs-ui"
 import { config } from "../config"
 
 type SearchProviderProps = {
@@ -52,28 +47,6 @@ const SearchProvider = ({ children }: SearchProviderProps) => {
         ),
         filterOptions: searchFilters,
       }}
-      commands={[
-        {
-          name: "ai-assistant",
-          icon: <AiAssistantIcon />,
-          component: (
-            <AiAssistantProvider
-              apiUrl={process.env.NEXT_PUBLIC_AI_ASSISTANT_URL || "temp"}
-              websiteId={process.env.NEXT_PUBLIC_AI_WEBSITE_ID || "temp"}
-              recaptchaSiteKey={
-                process.env.NEXT_PUBLIC_AI_API_ASSISTANT_RECAPTCHA_SITE_KEY ||
-                "temp"
-              }
-            />
-          ),
-          title: "AI Assistant",
-          badge: {
-            variant: "blue",
-            badgeType: "shaded",
-            children: "Beta",
-          },
-        },
-      ]}
       initialDefaultFilters={["guides"]}
     >
       {children}

www/apps/resources/providers/index.tsx

@@ -1,6 +1,7 @@
 "use client"
 
 import {
+  AiAssistantProvider,
   AnalyticsProvider,
   HooksLoader,
   LearningPathProvider,
@@ -31,14 +32,28 @@ const Providers = ({ children }: ProvidersProps) => {
                 <PaginationProvider>
                   <MainNavProvider>
                     <SearchProvider>
-                      <HooksLoader
+                      <AiAssistantProvider
-                        options={{
+                        apiUrl={
-                          pageScrollManager: true,
+                          process.env.NEXT_PUBLIC_AI_ASSISTANT_URL || "temp"
-                          currentLearningPath: false,
+                        }
-                        }}
+                        websiteId={
+                          process.env.NEXT_PUBLIC_AI_WEBSITE_ID || "temp"
+                        }
+                        recaptchaSiteKey={
+                          process.env
+                            .NEXT_PUBLIC_AI_API_ASSISTANT_RECAPTCHA_SITE_KEY ||
+                          "temp"
+                        }
                       >
-                        {children}
+                        <HooksLoader
-                      </HooksLoader>
+                          options={{
+                            pageScrollManager: true,
+                            currentLearningPath: false,
+                          }}
+                        >
+                          {children}
+                        </HooksLoader>
+                      </AiAssistantProvider>
                     </SearchProvider>
                   </MainNavProvider>
                 </PaginationProvider>

www/apps/resources/providers/search.tsx

@@ -1,11 +1,6 @@
 "use client"
 
-import {
+import { SearchProvider as UiSearchProvider, searchFilters } from "docs-ui"
-  AiAssistantIcon,
-  AiAssistantProvider,
-  SearchProvider as UiSearchProvider,
-  searchFilters,
-} from "docs-ui"
 import { config } from "../config"
 
 type SearchProviderProps = {
@@ -41,28 +36,6 @@ const SearchProvider = ({ children }: SearchProviderProps) => {
         checkInternalPattern: new RegExp(`^${config.baseUrl}/resources/.*`),
         filterOptions: searchFilters,
       }}
-      commands={[
-        {
-          name: "ai-assistant",
-          icon: <AiAssistantIcon />,
-          component: (
-            <AiAssistantProvider
-              apiUrl={process.env.NEXT_PUBLIC_AI_ASSISTANT_URL || "temp"}
-              websiteId={process.env.NEXT_PUBLIC_AI_WEBSITE_ID || "temp"}
-              recaptchaSiteKey={
-                process.env.NEXT_PUBLIC_AI_API_ASSISTANT_RECAPTCHA_SITE_KEY ||
-                "temp"
-              }
-            />
-          ),
-          title: "AI Assistant",
-          badge: {
-            variant: "blue",
-            badgeType: "shaded",
-            children: "Beta",
-          },
-        },
-      ]}
       initialDefaultFilters={["guides"]}
     >
       {children}

www/apps/ui/src/components/feedback.tsx

@@ -23,7 +23,7 @@ export const Feedback = ({ title, ...props }: FeedbackProps) => {
         section: title,
       }}
       {...props}
-      className={clsx("text-medusa-fg-subtle", props.className)}
+      className={clsx("text-medusa-fg-base", props.className)}
     />
   )
 }

www/apps/ui/src/providers/index.tsx

@@ -1,6 +1,7 @@
 "use client"
 
 import {
+  AiAssistantProvider,
   AnalyticsProvider,
   ScrollControllerProvider,
   SiteConfigProvider,
@@ -23,7 +24,16 @@ const Providers = ({ children }: ProvidersProps) => {
           <SidebarProvider>
             <MainNavProvider>
               <SearchProvider>
-                <TooltipProvider>{children}</TooltipProvider>
+                <AiAssistantProvider
+                  apiUrl={process.env.NEXT_PUBLIC_AI_ASSISTANT_URL || "temp"}
+                  websiteId={process.env.NEXT_PUBLIC_AI_WEBSITE_ID || "temp"}
+                  recaptchaSiteKey={
+                    process.env
+                      .NEXT_PUBLIC_AI_API_ASSISTANT_RECAPTCHA_SITE_KEY || "temp"
+                  }
+                >
+                  <TooltipProvider>{children}</TooltipProvider>
+                </AiAssistantProvider>
               </SearchProvider>
             </MainNavProvider>
           </SidebarProvider>

www/apps/ui/src/providers/search.tsx

@@ -1,11 +1,6 @@
 "use client"
 
-import {
+import { SearchProvider as UiSearchProvider, searchFilters } from "docs-ui"
-  AiAssistantIcon,
-  AiAssistantProvider,
-  SearchProvider as UiSearchProvider,
-  searchFilters,
-} from "docs-ui"
 import { absoluteUrl } from "../lib/absolute-url"
 
 type SearchProviderProps = {
@@ -37,29 +32,6 @@ const SearchProvider = ({ children }: SearchProviderProps) => {
         filterOptions: searchFilters,
       }}
       initialDefaultFilters={["ui"]}
-      commands={[
-        {
-          name: "ai-assistant",
-          icon: <AiAssistantIcon />,
-          component: (
-            <AiAssistantProvider
-              apiUrl={process.env.NEXT_PUBLIC_AI_ASSISTANT_URL || "temp"}
-              websiteId={process.env.NEXT_PUBLIC_AI_WEBSITE_ID || "temp"}
-              recaptchaSiteKey={
-                process.env.NEXT_PUBLIC_AI_API_ASSISTANT_RECAPTCHA_SITE_KEY ||
-                "temp"
-              }
-              version="v1"
-            />
-          ),
-          title: "AI Assistant",
-          badge: {
-            variant: "blue",
-            badgeType: "shaded",
-            children: "Beta",
-          },
-        },
-      ]}
     >
       {children}
     </UiSearchProvider>

www/apps/ui/src/styles/globals.css

@@ -28,7 +28,7 @@
   }
 
   body {
-    @apply text-ui-fg-subtle;
+    @apply text-ui-fg-base;
   }
 
   /* Hack to hide navbar / toc when some components like prompt are opened. */

www/apps/user-guide/providers/index.tsx

@@ -1,6 +1,7 @@
 "use client"
 
 import {
+  AiAssistantProvider,
   AnalyticsProvider,
   ColorModeProvider,
   HooksLoader,
@@ -35,13 +36,28 @@ const Providers = ({ children }: ProvidersProps) => {
                       <PaginationProvider>
                         <MainNavProvider>
                           <SearchProvider>
-                            <HooksLoader
+                            <AiAssistantProvider
-                              options={{
+                              apiUrl={
-                                pageScrollManager: true,
+                                process.env.NEXT_PUBLIC_AI_ASSISTANT_URL ||
-                              }}
+                                "temp"
+                              }
+                              websiteId={
+                                process.env.NEXT_PUBLIC_AI_WEBSITE_ID || "temp"
+                              }
+                              recaptchaSiteKey={
+                                // eslint-disable-next-line prettier/prettier, max-len
+                                process.env.NEXT_PUBLIC_AI_API_ASSISTANT_RECAPTCHA_SITE_KEY ||
+                                "temp"
+                              }
                             >
-                              {children}
+                              <HooksLoader
-                            </HooksLoader>
+                                options={{
+                                  pageScrollManager: true,
+                                }}
+                              >
+                                {children}
+                              </HooksLoader>
+                            </AiAssistantProvider>
                           </SearchProvider>
                         </MainNavProvider>
                       </PaginationProvider>

www/apps/user-guide/providers/search.tsx

@@ -1,11 +1,6 @@
 "use client"
 
-import {
+import { SearchProvider as UiSearchProvider, searchFilters } from "docs-ui"
-  SearchProvider as UiSearchProvider,
-  AiAssistantIcon,
-  AiAssistantProvider,
-  searchFilters,
-} from "docs-ui"
 import { config } from "../config"
 
 type SearchProviderProps = {
@@ -46,28 +41,6 @@ const SearchProvider = ({ children }: SearchProviderProps) => {
       }}
       // TODO change later when we have a user guide filter
       initialDefaultFilters={["guides"]}
-      commands={[
-        {
-          name: "ai-assistant",
-          icon: <AiAssistantIcon />,
-          component: (
-            <AiAssistantProvider
-              apiUrl={process.env.NEXT_PUBLIC_AI_ASSISTANT_URL || "temp"}
-              websiteId={process.env.NEXT_PUBLIC_AI_WEBSITE_ID || "temp"}
-              recaptchaSiteKey={
-                process.env.NEXT_PUBLIC_AI_API_ASSISTANT_RECAPTCHA_SITE_KEY ||
-                "temp"
-              }
-            />
-          ),
-          title: "AI Assistant",
-          badge: {
-            variant: "blue",
-            children: "Beta",
-            badgeType: "shaded",
-          },
-        },
-      ]}
     >
       {children}
     </UiSearchProvider>

www/packages/docs-ui/src/components/AiAssistant/ChatWindow/Footer/index.tsx

@@ -0,0 +1,29 @@
+import React from "react"
+import clsx from "clsx"
+import { Kbd } from "../../../Kbd"
+
+export const AiAssistantChatWindowFooter = () => {
+  return (
+    <div
+      className={clsx(
+        "bg-medusa-bg-component border-t border-medusa-border-base",
+        "flex items-center justify-end gap-docs_0.75 text-compact-x-small",
+        "py-docs_0.75 px-docs_1"
+      )}
+    >
+      <span className="text-medusa-fg-muted">Chat is cleared on refresh</span>
+      <span className="h-full w-px bg-medusa-border-base"></span>
+      <div className="flex items-center gap-docs_0.5">
+        <span className="text-medusa-fg-subtle">Line break</span>
+        <div className="flex items-center gap-[5px]">
+          <Kbd className="bg-medusa-bg-field-component border-medusa-border-strong w-[18px] h-[18px] inline-block">
+            ⇧
+          </Kbd>
+          <Kbd className="bg-medusa-bg-field-component border-medusa-border-strong w-[18px] h-[18px] inline-block">
+            ↵
+          </Kbd>
+        </div>
+      </div>
+    </div>
+  )
+}

www/packages/docs-ui/src/components/AiAssistant/ChatWindow/Header/index.tsx

@@ -0,0 +1,48 @@
+"use client"
+
+import clsx from "clsx"
+import React from "react"
+import { Tooltip } from "../../../Tooltip"
+import { Link } from "../../../Link"
+import { ShieldCheck, XMark } from "@medusajs/icons"
+import { Button } from "../../../Button"
+import { useAiAssistant } from "../../../../providers"
+
+export const AiAssistantChatWindowHeader = () => {
+  const { setChatOpened } = useAiAssistant()
+  return (
+    <div
+      className={clsx(
+        "flex gap-docs_0.5 items-center justify-between",
+        "w-full px-docs_1 py-docs_0.75 rounded-t-docs_sm",
+        "border-medusa-border-base border-b"
+      )}
+    >
+      <div className="flex gap-[6px] items-center">
+        <span className="text-h3 text-medusa-fg-base">Ask Anything</span>
+        <Tooltip
+          tooltipChildren={
+            <>
+              This site is protected by reCAPTCHA and the{" "}
+              <Link href="https://policies.google.com/privacy">
+                Google Privacy Policy
+              </Link>{" "}
+              and <Link href="https://policies.google.com/terms">ToS</Link>{" "}
+              apply
+            </>
+          }
+          clickable={true}
+        >
+          <ShieldCheck className="text-medusa-fg-muted" />
+        </Tooltip>
+      </div>
+      <Button
+        variant="transparent-clear"
+        className="!p-[6.5px] rounded-docs_sm"
+        onClick={() => setChatOpened(false)}
+      >
+        <XMark className="text-medusa-fg-muted" height={15} width={15} />
+      </Button>
+    </div>
+  )
+}

www/packages/docs-ui/src/components/AiAssistant/ChatWindow/Input/index.tsx

@@ -0,0 +1,106 @@
+import React, { useEffect, useRef } from "react"
+import clsx from "clsx"
+import { useAiAssistantChat } from "../../../../providers/AiAssistant/Chat"
+import { ArrowUpCircleSolid } from "@medusajs/icons"
+
+export const AiAssistantChatWindowInput = () => {
+  const {
+    inputRef,
+    question,
+    setQuestion,
+    handleSubmit: submitQuestion,
+    loading,
+    getThreadItems,
+  } = useAiAssistantChat()
+  const formRef = useRef<HTMLFormElement | null>(null)
+
+  const onSubmit = (e?: React.FormEvent<HTMLFormElement>) => {
+    e?.preventDefault()
+    submitQuestion()
+  }
+
+  const handleKeyboardDown = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
+    if (e.key === "ArrowUp" && !question) {
+      const lastQuestion = getThreadItems()
+        .reverse()
+        .find((item) => item.type === "question")
+      if (lastQuestion) {
+        setQuestion(lastQuestion.content)
+      }
+      return
+    }
+    if (e.key !== "Enter") {
+      return
+    }
+    if (e.shiftKey) {
+      setQuestion((prev) => `${prev}\n`)
+    } else {
+      onSubmit()
+    }
+  }
+
+  const adjustTextareaHeight = () => {
+    if (!inputRef.current) {
+      return
+    }
+    if (!question.length) {
+      inputRef.current.style.height = "auto"
+      return
+    }
+    inputRef.current.style.height = `${inputRef.current.scrollHeight}px`
+  }
+
+  useEffect(() => {
+    adjustTextareaHeight()
+  }, [question])
+
+  const handleTouch = (e: React.TouchEvent<HTMLTextAreaElement>) => {
+    e.preventDefault()
+    inputRef.current?.focus({
+      preventScroll: true,
+    })
+  }
+
+  return (
+    <div
+      className={clsx(
+        "px-docs_1 py-docs_0.75 border-t border-medusa-border-base"
+      )}
+    >
+      <form
+        className="flex flex-col gap-docs_0.5"
+        onSubmit={onSubmit}
+        ref={formRef}
+      >
+        <textarea
+          className={clsx(
+            "appearance-none text-base md:text-small placeholder:text-medusa-fg-muted",
+            "text-medusa-fg-base max-h-[210px] overflow-auto resize-none bg-transparent",
+            "focus:outline-none focus:ring-0 disabled:cursor-not-allowed max-h-[210px]",
+            "disabled:!bg-transparent disabled:text-medusa-fg-disabled"
+          )}
+          value={question}
+          onChange={(e) => setQuestion(e.target.value)}
+          onKeyDown={handleKeyboardDown}
+          onTouchStart={handleTouch}
+          onTouchMove={handleTouch}
+          onTouchEnd={handleTouch}
+          ref={inputRef as React.RefObject<HTMLTextAreaElement | null>}
+          placeholder="Ask me a question about Medusa..."
+          disabled={loading}
+        />
+        <div className="flex items-center justify-end">
+          <button
+            className={clsx(
+              "appearance-none p-0 text-medusa-fg-base disabled:text-medusa-fg-disabled",
+              "transition-colors"
+            )}
+            disabled={!question || loading}
+          >
+            <ArrowUpCircleSolid />
+          </button>
+        </div>
+      </form>
+    </div>
+  )
+}

www/packages/docs-ui/src/components/AiAssistant/ChatWindow/index.tsx

@@ -0,0 +1,196 @@
+"use client"
+
+import clsx from "clsx"
+import React, { useCallback, useEffect, useRef, useState } from "react"
+import { useAiAssistant, useIsBrowser } from "../../../providers"
+import { AiAssistantChatWindowHeader } from "./Header"
+import { useAiAssistantChat } from "../../../providers/AiAssistant/Chat"
+import { AiAssistantSuggestions } from "../Suggestions"
+import { AiAssistantThreadItem } from "../ThreadItem"
+import { AiAssistantChatWindowInput } from "./Input"
+import { useAiAssistantChatNavigation, useKeyboardShortcut } from "../../.."
+import { AiAssistantChatWindowFooter } from "./Footer"
+
+const DEFAULT_HEIGHT = "calc(100% - 8px)"
+
+export const AiAssistantChatWindow = () => {
+  const { chatOpened, setChatOpened, chatType: type } = useAiAssistant()
+  const [height, setHeight] = useState(DEFAULT_HEIGHT)
+  const [showFade, setShowFade] = useState(false)
+  const { isBrowser } = useIsBrowser()
+  const {
+    inputRef,
+    thread,
+    getThreadItems: getChatThreadItems,
+    answer,
+    loading,
+    contentRef,
+  } = useAiAssistantChat()
+  const chatWindowRef = useRef<HTMLDivElement | null>(null)
+
+  useEffect(() => {
+    if (chatOpened) {
+      inputRef.current?.focus({
+        preventScroll: true,
+      })
+    } else {
+      inputRef.current?.blur()
+    }
+  }, [chatOpened])
+
+  const getThreadItems = useCallback(() => {
+    const sortedThread = getChatThreadItems()
+
+    return sortedThread.map((item, index) => (
+      <AiAssistantThreadItem item={item} key={index} />
+    ))
+  }, [getChatThreadItems])
+
+  useAiAssistantChatNavigation({
+    getChatWindowElm: () => chatWindowRef.current as HTMLElement | null,
+    getInputElm: () => inputRef.current as HTMLTextAreaElement | null,
+    focusInput: () =>
+      inputRef.current?.focus({
+        preventScroll: true,
+      }),
+  })
+
+  useKeyboardShortcut({
+    metakey: false,
+    shortcutKeys: ["escape"],
+    checkEditing: false,
+    action: () => {
+      if (!chatWindowRef.current?.contains(document.activeElement)) {
+        return
+      }
+
+      setChatOpened(false)
+    },
+  })
+
+  const checkShowFade = () => {
+    const parentElm = contentRef.current?.parentElement
+    if (!parentElm) {
+      return
+    }
+    setShowFade(
+      !loading &&
+        parentElm.offsetHeight + parentElm.scrollTop < parentElm.scrollHeight
+    )
+  }
+
+  useEffect(() => {
+    if (!contentRef.current?.parentElement) {
+      return
+    }
+    contentRef.current.parentElement.addEventListener("scroll", checkShowFade)
+
+    return () => {
+      contentRef.current?.parentElement?.removeEventListener(
+        "scroll",
+        checkShowFade
+      )
+    }
+  }, [contentRef.current])
+
+  useEffect(() => {
+    if (loading) {
+      setShowFade(false)
+    } else {
+      checkShowFade()
+    }
+  }, [loading])
+
+  const changeHeightForViewport = () => {
+    if (!window.visualViewport?.height) {
+      setHeight(DEFAULT_HEIGHT)
+      return
+    }
+
+    setHeight(`${window.visualViewport.height - 8}px`)
+  }
+
+  useEffect(() => {
+    if (!isBrowser) {
+      return
+    }
+
+    window.visualViewport?.addEventListener("resize", changeHeightForViewport)
+
+    return () => {
+      window.visualViewport?.removeEventListener(
+        "resize",
+        changeHeightForViewport
+      )
+    }
+  }, [isBrowser])
+
+  useEffect(() => {
+    checkShowFade()
+  }, [height])
+
+  return (
+    <>
+      <div
+        className={clsx(
+          "fixed top-0 left-0 h-screen w-screen z-50 bg-medusa-bg-overlay",
+          !chatOpened && "hidden",
+          chatOpened && "block",
+          type === "default" && "xxl:hidden"
+        )}
+        onClick={() => setChatOpened(false)}
+      />
+      <div
+        className={clsx(
+          "flex z-50 w-[calc(100%-8px)] md:w-ai-assistant transition-[height,right]",
+          "absolute -right-[150%] sm:-right-full top-0",
+          type === "default" && [
+            "xxl:w-0 xxl:relative xxl:transition-[height,right,width]",
+            "xxl:shadow-elevation-card-rest xxl:dark:shadow-elevation-card-rest-dark",
+            chatOpened && "xxl:!w-ai-assistant",
+          ],
+          "shadow-elevation-modal dark:shadow-elevation-modal-dark",
+          "bg-medusa-bg-base rounded-docs_DEFAULT overflow-x-hidden",
+          "flex-col justify-between m-docs_0.25 max-w-ai-assistant",
+          chatOpened && ["!right-0"]
+        )}
+        style={{
+          height,
+        }}
+        ref={chatWindowRef}
+      >
+        <AiAssistantChatWindowHeader />
+        <div className="flex flex-auto overflow-auto relative">
+          <div
+            className={clsx(
+              "overflow-y-auto flex-auto px-docs_0.5 pt-docs_0.25 pb-docs_2"
+            )}
+          >
+            <div ref={contentRef}>
+              {!thread.length && <AiAssistantSuggestions />}
+              {getThreadItems()}
+              {(answer.length || loading) && (
+                <AiAssistantThreadItem
+                  item={{
+                    type: "answer",
+                    content: answer,
+                    order: 0,
+                  }}
+                />
+              )}
+            </div>
+          </div>
+          <span
+            className={clsx(
+              "bg-ai-assistant-bottom content-[''] absolute pointer-events-none",
+              "bottom-0 left-0 w-full h-docs_6 z-10 opacity-0 transition-opacity",
+              showFade && "opacity-100"
+            )}
+          ></span>
+        </div>
+        <AiAssistantChatWindowInput />
+        <AiAssistantChatWindowFooter />
+      </div>
+    </>
+  )
+}

www/packages/docs-ui/src/components/AiAssistant/SearchWindow/index.tsx

@@ -0,0 +1,186 @@
+"use client"
+
+import React, { useCallback } from "react"
+import { Badge, Button, InputText, Kbd, Tooltip, Link } from "@/components"
+import { useSearch } from "@/providers"
+import { ArrowUturnLeft } from "@medusajs/icons"
+import clsx from "clsx"
+import { AiAssistantThreadItem } from "../ThreadItem"
+import { AiAssistantSuggestions } from "../Suggestions"
+import { useAiAssistantChat } from "../../../providers/AiAssistant/Chat"
+import { useSearchNavigation } from "../../.."
+
+export const AiAssistantSearchWindow = () => {
+  const {
+    handleSubmit,
+    getThreadItems: getChatThreadItems,
+    question,
+    setQuestion,
+    inputRef,
+    contentRef,
+    loading,
+    thread,
+    answer,
+  } = useAiAssistantChat()
+  const { setCommand } = useSearch()
+
+  const getThreadItems = useCallback(() => {
+    const sortedThread = getChatThreadItems()
+
+    return sortedThread.map((item, index) => (
+      <AiAssistantThreadItem item={item} key={index} />
+    ))
+  }, [getChatThreadItems])
+
+  useSearchNavigation({
+    getInputElm: () => inputRef.current as HTMLInputElement | null,
+    focusInput: () => inputRef.current?.focus(),
+    handleSubmit,
+  })
+
+  return (
+    <div className="h-full">
+      <div className={clsx("px-docs_1 pt-docs_1")}>
+        <Tooltip
+          tooltipChildren={
+            <>
+              This site is protected by reCAPTCHA and the{" "}
+              <Link href="https://policies.google.com/privacy">
+                Google Privacy Policy
+              </Link>{" "}
+              and <Link href="https://policies.google.com/terms">ToS</Link>{" "}
+              apply
+            </>
+          }
+          clickable={true}
+        >
+          <Badge variant="neutral">AI Assistant</Badge>
+        </Tooltip>
+      </div>
+      <div
+        className={clsx(
+          "flex gap-docs_1 px-docs_1 py-docs_0.75",
+          "h-[57px] w-full md:rounded-t-docs_xl relative border-0 border-solid",
+          "border-b border-medusa-border-base relative"
+        )}
+      >
+        <Button
+          variant="transparent"
+          onClick={() => setCommand(null)}
+          className="text-medusa-fg-muted p-[6.5px]"
+        >
+          <ArrowUturnLeft />
+        </Button>
+        <InputText
+          value={question}
+          onChange={(e) => setQuestion(e.target.value)}
+          className={clsx(
+            "bg-transparent border-0 focus:outline-none hover:!bg-transparent",
+            "!shadow-none flex-1 text-medusa-fg-base",
+            "disabled:!bg-transparent disabled:cursor-not-allowed"
+          )}
+          placeholder="Ask me a question about Medusa..."
+          autoFocus={true}
+          passedRef={inputRef as React.RefObject<HTMLInputElement | null>}
+          disabled={loading}
+        />
+        <span
+          onClick={() => {
+            setQuestion("")
+            inputRef.current?.focus()
+          }}
+          className={clsx(
+            "text-medusa-fg-muted hover:text-medusa-fg-subtle",
+            "absolute top-docs_0.75 right-docs_1",
+            "cursor-pointer",
+            question.length === 0 && "hidden"
+          )}
+        >
+          Clear
+        </span>
+      </div>
+      <div className="h-[calc(100%-95px)] lg:max-h-[calc(100%-140px)] lg:min-h-[calc(100%-140px)] overflow-auto">
+        <div ref={contentRef}>
+          {!thread.length && <AiAssistantSuggestions className="mx-docs_0.5" />}
+          {getThreadItems()}
+          {(answer.length || loading) && (
+            <AiAssistantThreadItem
+              item={{
+                type: "answer",
+                content: answer,
+                order: 0,
+              }}
+            />
+          )}
+        </div>
+      </div>
+      <div
+        className={clsx(
+          "py-docs_0.75 hidden md:flex items-center justify-end px-docs_1",
+          "border-medusa-border-base border-t",
+          "bg-medusa-bg-field-component"
+        )}
+      >
+        <div className="flex items-center gap-docs_0.75">
+          <div className="flex items-center gap-docs_0.5">
+            {thread.length === 0 && (
+              <>
+                <span
+                  className={clsx(
+                    "text-medusa-fg-subtle",
+                    "text-compact-x-small"
+                  )}
+                >
+                  Navigate FAQ
+                </span>
+                <span className="gap-[5px] flex">
+                  <Kbd
+                    className={clsx(
+                      "!bg-medusa-bg-field-component !border-medusa-border-strong",
+                      "!text-medusa-fg-subtle h-[18px] w-[18px] p-0"
+                    )}
+                  >
+                    ↑
+                  </Kbd>
+                  <Kbd
+                    className={clsx(
+                      "!bg-medusa-bg-field-component !border-medusa-border-strong",
+                      "!text-medusa-fg-subtle h-[18px] w-[18px] p-0"
+                    )}
+                  >
+                    ↓
+                  </Kbd>
+                </span>
+              </>
+            )}
+            {thread.length > 0 && (
+              <span
+                className={clsx("text-medusa-fg-muted", "text-compact-x-small")}
+              >
+                Chat is cleared on exit
+              </span>
+            )}
+          </div>
+          <div
+            className={clsx("h-docs_0.75 w-px bg-medusa-border-strong")}
+          ></div>
+          <div className="flex items-center gap-docs_0.5">
+            <span
+              className={clsx("text-medusa-fg-subtle", "text-compact-x-small")}
+            >
+              Ask Question
+            </span>
+            <Kbd
+              className={clsx(
+                "!bg-medusa-bg-field-component !border-medusa-border-strong",
+                "!text-medusa-fg-subtle h-[18px] w-[18px] p-0"
+              )}
+            >
+              ↵
+            </Kbd>
+          </div>
+        </div>
+      </div>
+    </div>
+  )
+}

www/packages/docs-ui/src/components/AiAssistant/Suggestions/index.tsx

@@ -0,0 +1,76 @@
+"use client"
+
+import React, { useMemo } from "react"
+import { SearchSuggestionType } from "../../Search/Suggestions"
+import { useAiAssistant } from "../../../providers"
+import { SearchHitGroupName } from "../../Search/Hits/GroupName"
+import { SearchSuggestionItem } from "../../Search/Suggestions/Item"
+import { useAiAssistantChat } from "../../../providers/AiAssistant/Chat"
+
+type AiAssistantSuggestionsProps = React.AllHTMLAttributes<HTMLDivElement>
+
+export const AiAssistantSuggestions = (props: AiAssistantSuggestionsProps) => {
+  const { version } = useAiAssistant()
+  const { setQuestion, handleSubmit } = useAiAssistantChat()
+  const suggestions: SearchSuggestionType[] = useMemo(() => {
+    return version === "v2"
+      ? [
+          {
+            title: "FAQ",
+            items: [
+              "What is Medusa?",
+              "How can I create a module?",
+              "How can I create a data model?",
+              "How do I create a workflow?",
+              "How can I extend a data model in the Product Module?",
+            ],
+          },
+          {
+            title: "Recipes",
+            items: [
+              "How do I build a marketplace with Medusa?",
+              "How do I build digital products with Medusa?",
+              "How do I build subscription-based purchases with Medusa?",
+              "What other recipes are available in the Medusa documentation?",
+            ],
+          },
+        ]
+      : [
+          {
+            title: "FAQ",
+            items: [
+              "What is Medusa?",
+              "How can I create an ecommerce store with Medusa?",
+              "How can I build a marketplace with Medusa?",
+              "How can I build subscription-based purchases with Medusa?",
+              "How can I build digital products with Medusa?",
+              "What can I build with Medusa?",
+              "What is Medusa Admin?",
+              "How do I configure the database in Medusa?",
+            ],
+          },
+        ]
+  }, [version])
+
+  return (
+    <div {...props}>
+      {suggestions.map((suggestion, index) => (
+        <React.Fragment key={index}>
+          <SearchHitGroupName name={suggestion.title} />
+          {suggestion.items.map((item, itemIndex) => (
+            <SearchSuggestionItem
+              onClick={() => {
+                setQuestion(item)
+                handleSubmit(item)
+              }}
+              key={itemIndex}
+              tabIndex={itemIndex}
+            >
+              {item}
+            </SearchSuggestionItem>
+          ))}
+        </React.Fragment>
+      ))}
+    </div>
+  )
+}

www/packages/docs-ui/src/components/AiAssistant/ThreadItem/Actions/index.tsx

@@ -1,19 +1,17 @@
 import React, { useState } from "react"
-import { ThreadType } from "../.."
 import clsx from "clsx"
-import { Button, type ButtonProps } from "@/components"
+import { Badge, Button, Link, type ButtonProps } from "@/components"
-import { Check, SquareTwoStackMini, ThumbDown, ThumbUp } from "@medusajs/icons"
+import { ThumbDown, ThumbUp } from "@medusajs/icons"
-import { useCopy } from "@/hooks"
 import { AiAssistantFeedbackType, useAiAssistant } from "@/providers"
+import { AiAssistantThread } from "../../../../providers/AiAssistant/Chat"
 
 export type AiAssistantThreadItemActionsProps = {
-  item: ThreadType
+  item: AiAssistantThread
 }
 
 export const AiAssistantThreadItemActions = ({
   item,
 }: AiAssistantThreadItemActionsProps) => {
-  const { isCopied, handleCopy } = useCopy(item.content)
   const [feedback, setFeedback] = useState<AiAssistantFeedbackType | null>(null)
   const { sendFeedback } = useAiAssistant()
 
@@ -37,28 +35,36 @@ export const AiAssistantThreadItemActions = ({
   }
 
   return (
-    <div
+    <div className={clsx("flex gap-docs_0.75 justify-between items-center")}>
-      className={clsx("hidden md:flex gap-docs_0.25", "text-medusa-fg-muted")}
+      {item.sources !== undefined && item.sources.length > 0 && (
-    >
+        <div className="flex gap-[6px] items-center flex-wrap">
-      <ActionButton onClick={handleCopy}>
+          {item.sources.map((source) => (
-        {isCopied ? <Check /> : <SquareTwoStackMini />}
+            <Badge key={source.source_url} variant="neutral">
-      </ActionButton>
+              <Link href={source.source_url} className="!text-inherit">
-      {(feedback === null || feedback === "upvote") && (
+                {source.title}
-        <ActionButton
+              </Link>
-          onClick={async () => handleFeedback("upvote", item.question_id)}
+            </Badge>
-          className={clsx(feedback === "upvote" && "!text-medusa-fg-subtle")}
+          ))}
-        >
+        </div>
-          <ThumbUp />
-        </ActionButton>
-      )}
-      {(feedback === null || feedback === "downvote") && (
-        <ActionButton
-          onClick={async () => handleFeedback("downvote", item.question_id)}
-          className={clsx(feedback === "downvote" && "!text-medusa-fg-subtle")}
-        >
-          <ThumbDown />
-        </ActionButton>
       )}
+      <div className="flex gap-docs_0.25 items-center text-medusa-fg-muted">
+        {(feedback === null || feedback === "upvote") && (
+          <ActionButton
+            onClick={async () => handleFeedback("upvote", item.question_id)}
+            className={clsx(feedback === "upvote" && "!text-medusa-fg-muted")}
+          >
+            <ThumbUp />
+          </ActionButton>
+        )}
+        {(feedback === null || feedback === "downvote") && (
+          <ActionButton
+            onClick={async () => handleFeedback("downvote", item.question_id)}
+            className={clsx(feedback === "downvote" && "!text-medusa-fg-muted")}
+          >
+            <ThumbDown />
+          </ActionButton>
+        )}
+      </div>
     </div>
   )
 }
@@ -68,7 +74,7 @@ const ActionButton = ({ children, className, ...props }: ButtonProps) => {
     <Button
       variant="transparent"
       className={clsx(
-        "text-medusa-fg-muted hover:text-medusa-fg-subtle",
+        "text-medusa-fg-muted hover:text-medusa-fg-muted",
         "hover:bg-medusa-bg-subtle-hover",
         "!p-[4.5px] rounded-docs_sm",
         className

www/packages/docs-ui/src/components/AiAssistant/ThreadItem/index.tsx

@@ -1,11 +1,11 @@
 import clsx from "clsx"
 import React from "react"
-import { ThreadType } from ".."
 import { AiAssistantIcon, DotsLoading, MarkdownContent } from "@/components"
 import { AiAssistantThreadItemActions } from "./Actions"
+import { AiAssistantThread } from "../../../providers/AiAssistant/Chat"
 
 export type AiAssistantThreadItemProps = {
-  item: ThreadType
+  item: AiAssistantThread
 }
 
 export const AiAssistantThreadItem = ({ item }: AiAssistantThreadItemProps) => {
@@ -24,16 +24,24 @@ export const AiAssistantThreadItem = ({ item }: AiAssistantThreadItemProps) => {
       )}
       <div
         className={clsx(
-          "txt-small text-medusa-fg-subtle",
+          "txt-small text-medusa-fg-base",
           item.type === "question" && [
             "rounded-docs_xl bg-medusa-tag-neutral-bg",
             "px-docs_0.75 py-docs_0.5",
           ],
           item.type !== "question" && "flex-1",
-          item.type === "answer" && "text-pretty flex-1"
+          item.type === "answer" && "text-pretty flex-1 max-w-[calc(100%-20px)]"
         )}
       >
-        {item.type === "question" && <>{item.content}</>}
+        {item.type === "question" && (
+          <MarkdownContent
+            className="[&>*:last-child]:mb-0"
+            allowedElements={["br", "p", "code", "pre"]}
+            unwrapDisallowed={true}
+          >
+            {item.content}
+          </MarkdownContent>
+        )}
         {item.type === "answer" && (
           <div className="flex flex-col gap-docs_0.75">
             {!item.question_id && item.content.length === 0 && <DotsLoading />}

www/packages/docs-ui/src/components/AiAssistant/TriggerButton/index.tsx

@@ -0,0 +1,71 @@
+"use client"
+
+import React, { useMemo, useState } from "react"
+import { Button } from "../../Button"
+import { Tooltip } from "../../Tooltip"
+import { Kbd } from "../../Kbd"
+import { getOsShortcut } from "../../../utils"
+import { useAiAssistant, useSearch, useSiteConfig } from "../../../providers"
+import { useKeyboardShortcut } from "../../../hooks"
+import Image from "next/image"
+
+const AI_ASSISTANT_ICON = "/images/ai-assistent-luminosity.png"
+const AI_ASSISTANT_ICON_ACTIVE = "/images/ai-assistent.png"
+
+export const AiAssistantTriggerButton = () => {
+  const [hovered, setHovered] = useState(false)
+  const { config } = useSiteConfig()
+  const { chatOpened, setChatOpened } = useAiAssistant()
+  const { setIsOpen } = useSearch()
+  const isActive = useMemo(() => {
+    return hovered || chatOpened
+  }, [hovered, chatOpened])
+  const osShortcut = getOsShortcut()
+
+  useKeyboardShortcut({
+    metakey: true,
+    shortcutKeys: ["i"],
+    action: () => {
+      setChatOpened((prev) => !prev)
+      setIsOpen(false)
+    },
+    checkEditing: false,
+  })
+
+  return (
+    <Tooltip
+      render={() => (
+        <span className="flex gap-[6px] items-center">
+          <span className="text-compact-x-small-plus text-medusa-fg-base">
+            Ask AI
+          </span>
+          <span className="flex gap-[5px] items-center">
+            <Kbd className="bg-medusa-bg-field-component border-medusa-border-strong w-[18px] h-[18px] inline-block">
+              {osShortcut}
+            </Kbd>
+            <Kbd className="bg-medusa-bg-field-component border-medusa-border-strong w-[18px] h-[18px] inline-block">
+              i
+            </Kbd>
+          </span>
+        </span>
+      )}
+    >
+      <Button
+        variant="transparent-clear"
+        className="!p-[6.5px]"
+        onMouseOver={() => setHovered(true)}
+        onMouseOut={() => setHovered(false)}
+        onTouchStart={() => setHovered(true)}
+        onTouchEnd={() => setHovered(false)}
+        onClick={() => setChatOpened((prev) => !prev)}
+      >
+        <Image
+          src={`${config.basePath}${isActive ? AI_ASSISTANT_ICON_ACTIVE : AI_ASSISTANT_ICON}`}
+          width={15}
+          height={15}
+          alt="AI Assistant"
+        />
+      </Button>
+    </Tooltip>
+  )
+}

www/packages/docs-ui/src/components/AiAssistant/index.tsx

@@ -1,481 +0,0 @@
-"use client"
-
-import React, { useState, useEffect, useCallback, useMemo, useRef } from "react"
-import {
-  Badge,
-  Button,
-  InputText,
-  Kbd,
-  SearchSuggestionItem,
-  SearchSuggestionType,
-  SearchHitGroupName,
-  Tooltip,
-  Link,
-} from "@/components"
-import { useAiAssistant, useSearch } from "@/providers"
-import { ArrowUturnLeft } from "@medusajs/icons"
-import clsx from "clsx"
-import { useSearchNavigation } from "@/hooks"
-import { AiAssistantThreadItem } from "./ThreadItem"
-import useResizeObserver from "@react-hook/resize-observer"
-
-export type ChunkType = {
-  stream_end: boolean
-} & (
-  | {
-      type: "relevant_sources"
-      content: {
-        relevant_sources: RelevantSourcesType[]
-      }
-    }
-  | {
-      type: "partial_answer"
-      content: PartialAnswerType
-    }
-  | {
-      type: "identifiers"
-      content: IdentifierType
-    }
-  | {
-      type: "error"
-      content: ErrorType
-    }
-)
-
-export type RelevantSourcesType = {
-  source_url: string
-}
-
-export type PartialAnswerType = {
-  text: string
-}
-
-export type IdentifierType = {
-  thread_id: string
-  question_answer_id: string
-}
-
-export type ErrorType = {
-  reason: string
-}
-
-export type ThreadType = {
-  type: "question" | "answer" | "error"
-  content: string
-  question_id?: string
-  // for some reason, items in the array get reordered
-  // sometimes, so this is one way to avoid it
-  order: number
-}
-
-export const AiAssistant = () => {
-  const [question, setQuestion] = useState("")
-  // this helps set the `order` field of the threadtype
-  const [messagesCount, setMessagesCount] = useState(0)
-  const [thread, setThread] = useState<ThreadType[]>([])
-  const [answer, setAnswer] = useState("")
-  const [identifiers, setIdentifiers] = useState<IdentifierType | null>(null)
-  const [loading, setLoading] = useState(false)
-  const { getAnswer, version } = useAiAssistant()
-  const { setCommand } = useSearch()
-  const inputRef = useRef<HTMLInputElement>(null)
-  const contentRef = useRef<HTMLDivElement>(null)
-
-  const suggestions: SearchSuggestionType[] = useMemo(() => {
-    return version === "v2"
-      ? [
-          {
-            title: "FAQ",
-            items: [
-              "What is Medusa?",
-              "How can I create a module?",
-              "How can I create a data model?",
-              "How do I create a workflow?",
-              "How can I extend a data model in the Product Module?",
-            ],
-          },
-          {
-            title: "Recipes",
-            items: [
-              "How do I build a marketplace with Medusa?",
-              "How do I build digital products with Medusa?",
-              "How do I build subscription-based purchases with Medusa?",
-              "What other recipes are available in the Medusa documentation?",
-            ],
-          },
-        ]
-      : [
-          {
-            title: "FAQ",
-            items: [
-              "What is Medusa?",
-              "How can I create an ecommerce store with Medusa?",
-              "How can I build a marketplace with Medusa?",
-              "How can I build subscription-based purchases with Medusa?",
-              "How can I build digital products with Medusa?",
-              "What can I build with Medusa?",
-              "What is Medusa Admin?",
-              "How do I configure the database in Medusa?",
-            ],
-          },
-        ]
-  }, [version])
-
-  const handleSubmit = (selectedQuestion?: string) => {
-    if (!selectedQuestion?.length && !question.length) {
-      return
-    }
-    setLoading(true)
-    setAnswer("")
-    setThread((prevThread) => [
-      ...prevThread,
-      {
-        type: "question",
-        content: selectedQuestion || question,
-        order: getNewOrder(prevThread),
-      },
-    ])
-    setMessagesCount((prev) => prev + 1)
-  }
-
-  useSearchNavigation({
-    getInputElm: () => inputRef.current,
-    focusInput: () => inputRef.current?.focus(),
-    handleSubmit,
-  })
-
-  const sortThread = (threadArr: ThreadType[]) => {
-    const sortedThread = [...threadArr]
-    sortedThread.sort((itemA, itemB) => {
-      if (itemA.order < itemB.order) {
-        return -1
-      }
-
-      return itemA.order < itemB.order ? 1 : 0
-    })
-    return sortedThread
-  }
-
-  const getNewOrder = (prevThread: ThreadType[]) => {
-    const sortedThread = sortThread(prevThread)
-
-    return sortedThread.length === 0
-      ? messagesCount + 1
-      : sortedThread[prevThread.length - 1].order + 1
-  }
-
-  const setError = (logMessage?: string) => {
-    if (logMessage?.length) {
-      console.error(`[AI ERROR]: ${logMessage}`)
-    }
-    setThread((prevThread) => [
-      ...prevThread,
-      {
-        type: "error",
-        content:
-          "I'm sorry, but I'm having trouble connecting to my knowledge base. Please try again. If the issue keeps persisting, please consider reporting an issue.",
-        order: getNewOrder(prevThread),
-      },
-    ])
-    setMessagesCount((prev) => prev + 1)
-    setLoading(false)
-    setQuestion("")
-    setAnswer("")
-    inputRef.current?.focus()
-  }
-
-  const scrollToBottom = () => {
-    const parent = contentRef.current?.parentElement as HTMLElement
-
-    parent.scrollTop = parent.scrollHeight
-  }
-
-  const lastAnswerIndex = useMemo(() => {
-    const index = thread.reverse().findIndex((item) => item.type === "answer")
-    return index !== -1 ? index : 0
-  }, [thread])
-
-  const process_stream = useCallback(async (response: Response) => {
-    const reader = response.body?.getReader()
-    if (!reader) {
-      return
-    }
-    const decoder = new TextDecoder("utf-8")
-    const delimiter = "\u241E"
-    const delimiterBytes = new TextEncoder().encode(delimiter)
-    let buffer = new Uint8Array()
-
-    const findDelimiterIndex = (arr: Uint8Array) => {
-      for (let i = 0; i < arr.length - delimiterBytes.length + 1; i++) {
-        let found = true
-        for (let j = 0; j < delimiterBytes.length; j++) {
-          if (arr[i + j] !== delimiterBytes[j]) {
-            found = false
-            break
-          }
-        }
-        if (found) {
-          return i
-        }
-      }
-      return -1
-    }
-
-    let result
-    let loop = true
-    while (loop) {
-      result = await reader.read()
-      if (result.done) {
-        loop = false
-        continue
-      }
-      buffer = new Uint8Array([...buffer, ...result.value])
-      let delimiterIndex
-      while ((delimiterIndex = findDelimiterIndex(buffer)) !== -1) {
-        const chunkBytes = buffer.slice(0, delimiterIndex)
-        const chunkText = decoder.decode(chunkBytes)
-        buffer = buffer.slice(delimiterIndex + delimiterBytes.length)
-        const chunk = JSON.parse(chunkText).chunk as ChunkType
-
-        if (chunk.type === "partial_answer") {
-          setAnswer((prevAnswer) => {
-            return prevAnswer + chunk.content.text
-          })
-        } else if (chunk.type === "identifiers") {
-          setIdentifiers(chunk.content)
-        } else if (chunk.type === "error") {
-          setError(chunk.content.reason)
-          loop = false
-          return
-        }
-      }
-    }
-
-    setLoading(false)
-    setQuestion("")
-  }, [])
-
-  const fetchAnswer = useCallback(async () => {
-    try {
-      const response = await getAnswer(question, identifiers?.thread_id)
-
-      if (response.status === 200) {
-        await process_stream(response)
-      } else {
-        const message = await response.text()
-        setError(message)
-      }
-    } catch (error: any) {
-      setError(JSON.stringify(error))
-    }
-  }, [question, identifiers, process_stream])
-
-  useEffect(() => {
-    if (loading && !answer) {
-      void fetchAnswer()
-    }
-  }, [loading, fetchAnswer])
-
-  useEffect(() => {
-    if (
-      !loading &&
-      answer.length &&
-      thread[lastAnswerIndex]?.content !== answer
-    ) {
-      setThread((prevThread) => [
-        ...prevThread,
-        {
-          type: "answer",
-          content: answer,
-          question_id: identifiers?.question_answer_id,
-          order: getNewOrder(prevThread),
-        },
-      ])
-      setAnswer("")
-      setMessagesCount((prev) => prev + 1)
-      inputRef.current?.focus()
-    }
-  }, [loading, answer, thread, lastAnswerIndex, inputRef.current])
-
-  useResizeObserver(contentRef as React.RefObject<HTMLDivElement>, () => {
-    if (!loading) {
-      return
-    }
-
-    scrollToBottom()
-  })
-
-  const getThreadItems = useCallback(() => {
-    const sortedThread = sortThread(thread)
-
-    return sortedThread.map((item, index) => (
-      <AiAssistantThreadItem item={item} key={index} />
-    ))
-  }, [thread])
-
-  return (
-    <div className="h-full">
-      <div className={clsx("px-docs_1 pt-docs_1")}>
-        <Tooltip
-          tooltipChildren={
-            <>
-              This site is protected by reCAPTCHA and the{" "}
-              <Link href="https://policies.google.com/privacy">
-                Google Privacy Policy
-              </Link>{" "}
-              and <Link href="https://policies.google.com/terms">ToS</Link>{" "}
-              apply
-            </>
-          }
-          clickable={true}
-        >
-          <Badge variant="neutral">AI Assistant</Badge>
-        </Tooltip>
-      </div>
-      <div
-        className={clsx(
-          "flex gap-docs_1 px-docs_1 py-docs_0.75",
-          "h-[57px] w-full md:rounded-t-docs_xl relative border-0 border-solid",
-          "border-b border-medusa-border-base relative"
-        )}
-      >
-        <Button
-          variant="transparent"
-          onClick={() => setCommand(null)}
-          className="text-medusa-fg-muted p-[6.5px]"
-        >
-          <ArrowUturnLeft />
-        </Button>
-        <InputText
-          value={question}
-          onChange={(e) => setQuestion(e.target.value)}
-          className={clsx(
-            "bg-transparent border-0 focus:outline-none hover:!bg-transparent",
-            "!shadow-none flex-1 text-medusa-fg-base",
-            "disabled:!bg-transparent disabled:cursor-not-allowed"
-          )}
-          placeholder="Ask me a question about Medusa..."
-          autoFocus={true}
-          passedRef={inputRef}
-          disabled={loading}
-        />
-        <span
-          onClick={() => {
-            setQuestion("")
-            inputRef.current?.focus()
-          }}
-          className={clsx(
-            "text-medusa-fg-muted hover:text-medusa-fg-subtle",
-            "absolute top-docs_0.75 right-docs_1",
-            "cursor-pointer",
-            question.length === 0 && "hidden"
-          )}
-        >
-          Clear
-        </span>
-      </div>
-      <div className="h-[calc(100%-95px)] lg:max-h-[calc(100%-140px)] lg:min-h-[calc(100%-140px)] overflow-auto">
-        <div ref={contentRef}>
-          {!thread.length && (
-            <div className="mx-docs_0.5">
-              {suggestions.map((suggestion, index) => (
-                <React.Fragment key={index}>
-                  <SearchHitGroupName name={suggestion.title} />
-                  {suggestion.items.map((item, itemIndex) => (
-                    <SearchSuggestionItem
-                      onClick={() => {
-                        setQuestion(item)
-                        handleSubmit(item)
-                      }}
-                      key={itemIndex}
-                      tabIndex={itemIndex}
-                    >
-                      {item}
-                    </SearchSuggestionItem>
-                  ))}
-                </React.Fragment>
-              ))}
-            </div>
-          )}
-          {getThreadItems()}
-          {(answer.length || loading) && (
-            <AiAssistantThreadItem
-              item={{
-                type: "answer",
-                content: answer,
-                order: 0,
-              }}
-            />
-          )}
-        </div>
-      </div>
-      <div
-        className={clsx(
-          "py-docs_0.75 hidden md:flex items-center justify-end px-docs_1",
-          "border-medusa-border-base border-t",
-          "bg-medusa-bg-field-component"
-        )}
-      >
-        <div className="flex items-center gap-docs_0.75">
-          <div className="flex items-center gap-docs_0.5">
-            {thread.length === 0 && (
-              <>
-                <span
-                  className={clsx(
-                    "text-medusa-fg-subtle",
-                    "text-compact-x-small"
-                  )}
-                >
-                  Navigate FAQ
-                </span>
-                <span className="gap-[5px] flex">
-                  <Kbd
-                    className={clsx(
-                      "!bg-medusa-bg-field-component !border-medusa-border-strong",
-                      "!text-medusa-fg-subtle h-[18px] w-[18px] p-0"
-                    )}
-                  >
-                    ↑
-                  </Kbd>
-                  <Kbd
-                    className={clsx(
-                      "!bg-medusa-bg-field-component !border-medusa-border-strong",
-                      "!text-medusa-fg-subtle h-[18px] w-[18px] p-0"
-                    )}
-                  >
-                    ↓
-                  </Kbd>
-                </span>
-              </>
-            )}
-            {thread.length > 0 && (
-              <span
-                className={clsx("text-medusa-fg-muted", "text-compact-x-small")}
-              >
-                Chat is cleared on exit
-              </span>
-            )}
-          </div>
-          <div
-            className={clsx("h-docs_0.75 w-px bg-medusa-border-strong")}
-          ></div>
-          <div className="flex items-center gap-docs_0.5">
-            <span
-              className={clsx("text-medusa-fg-subtle", "text-compact-x-small")}
-            >
-              Ask Question
-            </span>
-            <Kbd
-              className={clsx(
-                "!bg-medusa-bg-field-component !border-medusa-border-strong",
-                "!text-medusa-fg-subtle h-[18px] w-[18px] p-0"
-              )}
-            >
-              ↵
-            </Kbd>
-          </div>
-        </div>
-      </div>
-    </div>
-  )
-}

www/packages/docs-ui/src/components/ApiRunner/ParamInputs/index.tsx

@@ -16,7 +16,7 @@ export const ApiRunnerParamInputs = ({
 }: ApiRunnerParamInputsProps) => {
   return (
     <div className="flex flex-col gap-docs_0.25 w-full">
-      <span className="txt-small-plus text-medusa-fg-subtle">{title}</span>
+      <span className="txt-small-plus text-medusa-fg-base">{title}</span>
       <div className="flex flex-col gap-docs_0.5">
         {Object.keys(data).map((pathParam, index) => (
           <ApiRunnerParamInput

www/packages/docs-ui/src/components/Feedback/index.tsx

@@ -163,7 +163,7 @@ export const Feedback = ({
                 )}
                 ref={inlineFeedbackRef}
               >
-                <Label className={"text-compact-small text-medusa-fg-subtle"}>
+                <Label className={"text-compact-small text-medusa-fg-base"}>
                   {question}
                 </Label>
                 <div

www/packages/docs-ui/src/components/Kbd/index.tsx

@@ -16,9 +16,9 @@ export const Kbd = ({
       className={clsx(
         "rounded-docs_xs border-solid border border-medusa-border-base",
         "inline-flex items-center justify-center",
-        "py-0 px-docs_0.25",
+        "p-0",
         "bg-medusa-bg-field",
-        "text-medusa-fg-subtle",
+        "text-medusa-fg-base",
         "font-base shadow-none",
         variant === "small"
           ? "text-compact-x-small"

www/packages/docs-ui/src/components/MDXComponents/index.tsx

@@ -38,7 +38,7 @@ export const MDXComponents: MDXComponentsType = {
     return (
       <p
         className={clsx(
-          "text-medusa-fg-subtle [&:not(:last-child)]:mb-docs_1.5 last:!mb-0",
+          "text-medusa-fg-base [&:not(:last-child)]:mb-docs_1.5 last:!mb-0",
           className
         )}
         {...props}
@@ -87,7 +87,7 @@ export const MDXComponents: MDXComponentsType = {
     return (
       <li
         className={clsx(
-          "text-medusa-fg-subtle [&:not(:last-child)]:mb-docs_0.5",
+          "text-medusa-fg-base [&:not(:last-child)]:mb-docs_0.5",
           "[&_ol]:mt-docs_0.5 [&_ul]:mt-docs_0.5",
           className
         )}

www/packages/docs-ui/src/components/MainNav/Version/index.tsx

@@ -47,7 +47,9 @@ export const MainNavVersion = () => {
             className="relative text-compact-small-plus"
             onMouseOut={afterHover}
           >
-            <span>v{version.number}</span>
+            <span className="flex justify-center items-center">
+              v{version.number}
+            </span>
             {showNewBadge && (
               <span
                 className={clsx(

www/packages/docs-ui/src/components/MainNav/index.tsx

@@ -7,6 +7,7 @@ import {
   Button,
   LinkButton,
   SearchModalOpener,
+  useLayout,
   useMainNav,
   useSidebar,
   useSiteConfig,
@@ -18,6 +19,7 @@ import { SidebarLeftIcon } from "../Icons/SidebarLeft"
 import { MainNavMobileMenu } from "./MobileMenu"
 import Link from "next/link"
 import { MainNavVersion } from "./Version"
+import { AiAssistantTriggerButton } from "../AiAssistant/TriggerButton"
 
 type MainNavProps = {
   className?: string
@@ -28,17 +30,18 @@ export const MainNav = ({ className, itemsClassName }: MainNavProps) => {
   const { editDate } = useMainNav()
   const { setMobileSidebarOpen, isSidebarShown } = useSidebar()
   const { config } = useSiteConfig()
+  const { showCollapsedNavbar } = useLayout()
 
   return (
     <div
-      className={clsx(
+      className={clsx("w-full z-20 sticky top-0 bg-medusa-bg-base", className)}
-        "flex justify-between items-center",
-        "px-docs_1 w-full z-20",
-        "sticky top-0 bg-medusa-bg-base",
-        className
-      )}
     >
-      <div className="flex items-center gap-docs_1">
+      <div
+        className={clsx(
+          "flex justify-between items-center px-docs_1 w-full gap-docs_1",
+          showCollapsedNavbar && "border-b border-medusa-border-base"
+        )}
+      >
         <div className="flex items-center gap-[10px]">
           {isSidebarShown && (
             <Button
@@ -53,32 +56,46 @@ export const MainNav = ({ className, itemsClassName }: MainNavProps) => {
             <BorderedIcon
               icon={config.logo}
               iconWrapperClassName="my-[14px]"
+              wrapperClassName="w-[20px] h-[20px]"
               iconWidth={20}
               iconHeight={20}
             />
           </Link>
         </div>
-        <MainNavItems className={itemsClassName} />
+        {!showCollapsedNavbar && (
-      </div>
+          <MainNavItems className={clsx("flex-grow", itemsClassName)} />
-      <div className="flex items-center gap-docs_0.75 my-docs_0.75">
+        )}
-        <div className="lg:flex items-center gap-docs_0.5 text-medusa-fg-subtle hidden">
+        <div
-          <MainNavVersion />
+          className={clsx(
-          {editDate && <MainNavEditDate date={editDate} />}
+            "flex items-center gap-docs_0.75 my-docs_0.75",
-          <LinkButton
+            showCollapsedNavbar && "flex-grow justify-between"
-            href={config.reportIssueLink || ""}
+          )}
-            variant="subtle"
+        >
-            target="_blank"
+          <div className="lg:flex items-center gap-docs_0.5 text-medusa-fg-subtle hidden">
-            className="text-compact-small-plus"
+            <MainNavVersion />
-          >
+            {editDate && <MainNavEditDate date={editDate} />}
-            Report Issue
+            <LinkButton
-          </LinkButton>
+              href={config.reportIssueLink || ""}
-        </div>
+              variant="subtle"
-        <div className="flex items-center gap-docs_0.25">
+              target="_blank"
-          <SearchModalOpener />
+              className="text-compact-small-plus"
-          <MainNavDesktopMenu />
+            >
-          <MainNavMobileMenu />
+              Report Issue
+            </LinkButton>
+          </div>
+          <div className="flex items-center gap-docs_0.25">
+            <AiAssistantTriggerButton />
+            <SearchModalOpener />
+            <MainNavDesktopMenu />
+            <MainNavMobileMenu />
+          </div>
         </div>
       </div>
+      {showCollapsedNavbar && (
+        <div className={clsx("border-b border-medusa-border-base px-docs_1")}>
+          <MainNavItems className={clsx("flex-wrap", itemsClassName)} />
+        </div>
+      )}
     </div>
   )
 }

www/packages/docs-ui/src/components/RootProviders/index.tsx

@@ -2,21 +2,29 @@ import React from "react"
 import {
   BrowserProvider,
   ColorModeProvider,
+  LayoutProvider,
+  LayoutProviderProps,
   MobileProvider,
   ModalProvider,
 } from "../../providers"
 
 type RootProvidersProps = {
   children: React.ReactNode
+  layoutProviderProps?: Omit<LayoutProviderProps, "children">
 }
 
-export const RootProviders = ({ children }: RootProvidersProps) => {
+export const RootProviders = ({
+  children,
+  layoutProviderProps = {},
+}: RootProvidersProps) => {
   return (
     <BrowserProvider>
       <MobileProvider>
-        <ColorModeProvider>
+        <LayoutProvider {...layoutProviderProps}>
-          <ModalProvider>{children}</ModalProvider>
+          <ColorModeProvider>
-        </ColorModeProvider>
+            <ModalProvider>{children}</ModalProvider>
+          </ColorModeProvider>
+        </LayoutProvider>
       </MobileProvider>
     </BrowserProvider>
   )

www/packages/docs-ui/src/components/index.ts

@@ -1,4 +1,5 @@
-export * from "./AiAssistant"
+export * from "./AiAssistant/SearchWindow"
+export * from "./AiAssistant/ChatWindow"
 export * from "./ApiRunner"
 export * from "./Badge"
 export * from "./BetaBadge"

www/packages/docs-ui/src/hooks/index.ts

@@ -1,4 +1,5 @@
 export * from "./use-active-on-scroll"
+export * from "./use-ai-assistant-chat-navigation"
 export * from "./use-child-docs"
 export * from "./use-click-outside"
 export * from "./use-collapsible"

www/packages/docs-ui/src/hooks/use-ai-assistant-chat-navigation/index.ts

@@ -0,0 +1,141 @@
+"use client"
+
+import { useCallback, useEffect, useMemo } from "react"
+import { useAiAssistant } from "@/providers"
+import { findNextSibling, findPrevSibling } from "@/utils"
+import {
+  useKeyboardShortcut,
+  type useKeyboardShortcutOptions,
+} from "../use-keyboard-shortcut"
+
+export type UseAiAssistantChatNavigationProps = {
+  getChatWindowElm: () => HTMLElement | null
+  getInputElm: () => HTMLTextAreaElement | null
+  focusInput: () => void
+  keyboardProps?: Partial<useKeyboardShortcutOptions>
+}
+
+export const useAiAssistantChatNavigation = ({
+  getInputElm,
+  focusInput,
+  keyboardProps,
+  getChatWindowElm,
+}: UseAiAssistantChatNavigationProps) => {
+  const shortcutKeys = useMemo(() => ["ArrowUp", "ArrowDown", "Enter"], [])
+  const { chatOpened } = useAiAssistant()
+
+  const handleKeyAction = (e: KeyboardEvent) => {
+    const chatElm = getChatWindowElm()
+    if (
+      !chatOpened ||
+      e.metaKey ||
+      e.ctrlKey ||
+      !chatElm?.contains(document.activeElement)
+    ) {
+      return
+    }
+    e.preventDefault()
+    const focusedItem = chatElm?.querySelector(":focus") as HTMLElement
+    if (!focusedItem) {
+      // focus the first data-hit
+      const nextItem = chatElm?.querySelector("[data-hit]") as HTMLElement
+      nextItem?.focus()
+      return
+    }
+
+    const isHit = focusedItem.hasAttribute("data-hit")
+    const isInput = focusedItem.tagName.toLowerCase() === "textarea"
+
+    if (!isHit && !isInput) {
+      // ignore if focused items aren't input/data-hit
+      return
+    }
+
+    const lowerPressedKey = e.key.toLowerCase()
+
+    if (lowerPressedKey === "enter") {
+      if (isHit) {
+        // trigger click event of the focused element
+        focusedItem.click()
+      }
+      return
+    }
+
+    if (lowerPressedKey === "arrowdown") {
+      // only hit items has action on arrow down
+      if (isHit) {
+        // find if there's a data-hit item before this one
+        const beforeItem = findNextSibling(focusedItem, "[data-hit]")
+        if (!beforeItem) {
+          // focus the input
+          focusInput()
+        } else {
+          // focus the previous item
+          beforeItem.focus()
+        }
+      }
+    } else if (lowerPressedKey === "arrowup") {
+      // check if item is input or hit
+      if (isInput) {
+        // go to the first data-hit item
+        const nextItem = chatElm?.querySelector(
+          "[data-hit]:last-child"
+        ) as HTMLElement
+        nextItem?.focus()
+      } else {
+        // handle go down for hit items
+        // find if there's a data-hit item after this one
+        const afterItem = findPrevSibling(focusedItem, "[data-hit]")
+        if (afterItem) {
+          // focus the next item
+          afterItem.focus()
+        }
+      }
+    }
+  }
+
+  /** Handles starting to type which focuses the input */
+  const handleKeyDown = useCallback(
+    (e: KeyboardEvent) => {
+      if (!chatOpened) {
+        return
+      }
+      // check if shortcut keys were pressed
+      const lowerPressedKey = e.key.toLowerCase()
+      const pressedShortcut = [...shortcutKeys, "Escape"].some(
+        (s) => s.toLowerCase() === lowerPressedKey
+      )
+      if (pressedShortcut) {
+        return
+      }
+
+      const chatElm = getChatWindowElm()
+      if (!chatElm?.contains(document.activeElement)) {
+        return
+      }
+      const focusedItem = chatElm?.querySelector(":focus") as HTMLElement
+      const inputElm = getInputElm()
+      if (inputElm && focusedItem !== inputElm) {
+        focusInput()
+      }
+    },
+    [shortcutKeys, chatOpened, shortcutKeys, getInputElm, focusInput]
+  )
+
+  useEffect(() => {
+    window.addEventListener("keydown", handleKeyDown)
+
+    return () => {
+      window.removeEventListener("keydown", handleKeyDown)
+    }
+  }, [handleKeyDown])
+
+  useKeyboardShortcut({
+    metakey: false,
+    shortcutKeys: shortcutKeys,
+    checkEditing: false,
+    isLoading: false,
+    action: handleKeyAction,
+    ...keyboardProps,
+  })
+}

www/packages/docs-ui/src/hooks/use-child-docs/index.tsx

@@ -379,10 +379,10 @@ export const useChildDocs = ({
         {!searchResult.length && (
           <div className="flex flex-col justify-center items-center gap-docs_0.75">
             <ExclamationCircle className="text-medusa-fg-subtle" />
-            <span className="text-compact-small-plus text-medusa-fg-subtle text-center">
+            <span className="text-compact-small-plus text-medusa-fg-base text-center">
               No results found matching your query.
             </span>
-            <span className="text-compact-small text-medusa-fg-muted text-center">
+            <span className="text-compact-small text-medusa-fg-subtle text-center">
               Try searching with another term or clearing the search.
             </span>
           </div>

www/packages/docs-ui/src/hooks/use-keyboard-shortcut/index.tsx

@@ -40,7 +40,7 @@ export const useKeyboardShortcut = ({
     [shortcutKeys]
   )
 
-  const sidebarShortcut = useCallback(
+  const onKeyDown = useCallback(
     (e: KeyboardEvent) => {
       // the event is triggered when an input
       // autocompletes, and in that case
@@ -63,10 +63,10 @@ export const useKeyboardShortcut = ({
   )
 
   useEffect(() => {
-    window.addEventListener("keydown", sidebarShortcut)
+    window.addEventListener("keydown", onKeyDown)
 
     return () => {
-      window.removeEventListener("keydown", sidebarShortcut)
+      window.removeEventListener("keydown", onKeyDown)
     }
-  }, [sidebarShortcut])
+  }, [onKeyDown])
 }

www/packages/docs-ui/src/layouts/main-content.tsx

@@ -3,7 +3,7 @@
 import React, { useEffect } from "react"
 import { useSidebar } from "../providers/Sidebar"
 import clsx from "clsx"
-import { MainNav, useIsBrowser } from ".."
+import { MainNav, useIsBrowser, useLayout } from ".."
 
 export type MainContentLayoutProps = {
   mainWrapperClasses?: string
@@ -18,6 +18,7 @@ export const MainContentLayout = ({
 }: MainContentLayoutProps) => {
   const { isBrowser } = useIsBrowser()
   const { desktopSidebarOpen } = useSidebar()
+  const { mainContentRef } = useLayout()
 
   useEffect(() => {
     if (!isBrowser) {
@@ -52,6 +53,7 @@ export const MainContentLayout = ({
           mainWrapperClasses
         )}
         id="main"
+        ref={mainContentRef}
       >
         <MainNav />
         <div

www/packages/docs-ui/src/layouts/root.tsx

@@ -3,6 +3,7 @@ import clsx from "clsx"
 import { RootProviders, Sidebar, SidebarProps } from "@/components"
 import { Toc } from "../components/Toc"
 import { MainContentLayout, MainContentLayoutProps } from "./main-content"
+import { AiAssistantChatWindow } from "../components/AiAssistant/ChatWindow"
 
 export type RootLayoutProps = {
   bodyClassName?: string
@@ -26,7 +27,7 @@ export const RootLayout = ({
     <body
       className={clsx(
         "bg-medusa-bg-subtle font-base text-medium w-full",
-        "text-medusa-fg-subtle",
+        "text-medusa-fg-base",
         "h-screen overflow-hidden",
         "grid grid-cols-1 lg:mx-auto lg:grid-cols-[221px_1fr]",
         bodyClassName
@@ -38,6 +39,7 @@ export const RootLayout = ({
           <div className={clsx("relative", "h-screen", "flex")}>
             <MainContentLayout {...mainProps} />
             {showToc && <Toc />}
+            <AiAssistantChatWindow />
           </div>
         </ProvidersComponent>
       </RootProviders>

www/packages/docs-ui/src/providers/AiAssistant/Chat/index.tsx

@@ -0,0 +1,345 @@
+"use client"
+
+import React, {
+  createContext,
+  useCallback,
+  useContext,
+  useEffect,
+  useMemo,
+  useRef,
+  useState,
+} from "react"
+import { useAiAssistant } from ".."
+import useResizeObserver from "@react-hook/resize-observer"
+
+export type AiAssistantChatContextType = {
+  handleSubmit: (selectedQuestion?: string) => void
+  getThreadItems: () => AiAssistantThread[]
+  question: string
+  setQuestion: React.Dispatch<React.SetStateAction<string>>
+  inputRef: React.RefObject<HTMLInputElement | HTMLTextAreaElement | null>
+  contentRef: React.RefObject<HTMLDivElement | null>
+  loading: boolean
+  thread: AiAssistantThread[]
+  answer: string
+}
+
+const AiAssistantChatContext = createContext<AiAssistantChatContextType | null>(
+  null
+)
+
+export type AiAssistantChunk = {
+  stream_end: boolean
+} & (
+  | {
+      type: "relevant_sources"
+      content: {
+        relevant_sources: AiAssistantRelevantSources[]
+      }
+    }
+  | {
+      type: "partial_answer"
+      content: AiAssistantPartialAnswer
+    }
+  | {
+      type: "identifiers"
+      content: AiAssistantIdentifier
+    }
+  | {
+      type: "error"
+      content: AiAssistantError
+    }
+)
+
+export type AiAssistantRelevantSources = {
+  title: string
+  source_url: string
+}
+
+export type AiAssistantPartialAnswer = {
+  text: string
+}
+
+export type AiAssistantIdentifier = {
+  thread_id: string
+  question_answer_id: string
+}
+
+export type AiAssistantError = {
+  reason: string
+}
+
+export type AiAssistantThread = {
+  type: "question" | "answer" | "error"
+  content: string
+  question_id?: string
+  sources?: AiAssistantRelevantSources[]
+  // for some reason, items in the array get reordered
+  // sometimes, so this is one way to avoid it
+  order: number
+}
+
+type AiAssistantChatProvider = {
+  children: React.ReactNode
+}
+
+export const AiAssistantChatProvider = ({
+  children,
+}: AiAssistantChatProvider) => {
+  const [question, setQuestion] = useState("")
+  // this helps set the `order` field of the threadtype
+  const [messagesCount, setMessagesCount] = useState(0)
+  const [thread, setThread] = useState<AiAssistantThread[]>([])
+  const [answer, setAnswer] = useState("")
+  const [answerSources, setAnswerSources] = useState<
+    AiAssistantRelevantSources[]
+  >([])
+  const [identifiers, setIdentifiers] = useState<AiAssistantIdentifier | null>(
+    null
+  )
+  const [loading, setLoading] = useState(false)
+  const { getAnswer } = useAiAssistant()
+  const inputRef = useRef<HTMLInputElement | HTMLTextAreaElement>(null)
+  const contentRef = useRef<HTMLDivElement>(null)
+
+  const handleSubmit = (selectedQuestion?: string) => {
+    if (!selectedQuestion?.length && !question.length) {
+      return
+    }
+    setLoading(true)
+    setAnswer("")
+    setThread((prevThread) => [
+      ...prevThread,
+      {
+        type: "question",
+        content: selectedQuestion || question,
+        order: getNewOrder(prevThread),
+      },
+    ])
+    setMessagesCount((prev) => prev + 1)
+  }
+
+  const sortThread = (threadArr: AiAssistantThread[]) => {
+    const sortedThread = [...threadArr]
+    sortedThread.sort((itemA, itemB) => {
+      if (itemA.order < itemB.order) {
+        return -1
+      }
+
+      return itemA.order < itemB.order ? 1 : 0
+    })
+    return sortedThread
+  }
+
+  const getNewOrder = (prevThread: AiAssistantThread[]) => {
+    const sortedThread = sortThread(prevThread)
+
+    return sortedThread.length === 0
+      ? messagesCount + 1
+      : sortedThread[prevThread.length - 1].order + 1
+  }
+
+  const setError = (logMessage?: string) => {
+    if (logMessage?.length) {
+      console.error(`[AI ERROR]: ${logMessage}`)
+    }
+    setThread((prevThread) => [
+      ...prevThread,
+      {
+        type: "error",
+        content:
+          "I'm sorry, but I'm having trouble connecting to my knowledge base. Please try again. If the issue keeps persisting, please consider reporting an issue.",
+        order: getNewOrder(prevThread),
+      },
+    ])
+    setMessagesCount((prev) => prev + 1)
+    setLoading(false)
+    setQuestion("")
+    setAnswer("")
+    inputRef.current?.focus()
+  }
+
+  const scrollToBottom = () => {
+    const parent = contentRef.current?.parentElement as HTMLElement
+
+    parent.scrollTop = parent.scrollHeight
+  }
+
+  const lastAnswerIndex = useMemo(() => {
+    const index = thread.reverse().findIndex((item) => item.type === "answer")
+    return index !== -1 ? index : 0
+  }, [thread])
+
+  const process_stream = useCallback(async (response: Response) => {
+    const reader = response.body?.getReader()
+    if (!reader) {
+      return
+    }
+    const decoder = new TextDecoder("utf-8")
+    const delimiter = "\u241E"
+    const delimiterBytes = new TextEncoder().encode(delimiter)
+    let buffer = new Uint8Array()
+
+    const findDelimiterIndex = (arr: Uint8Array) => {
+      for (let i = 0; i < arr.length - delimiterBytes.length + 1; i++) {
+        let found = true
+        for (let j = 0; j < delimiterBytes.length; j++) {
+          if (arr[i + j] !== delimiterBytes[j]) {
+            found = false
+            break
+          }
+        }
+        if (found) {
+          return i
+        }
+      }
+      return -1
+    }
+
+    let result
+    let loop = true
+    while (loop) {
+      result = await reader.read()
+      if (result.done) {
+        loop = false
+        continue
+      }
+      buffer = new Uint8Array([...buffer, ...result.value])
+      let delimiterIndex
+      while ((delimiterIndex = findDelimiterIndex(buffer)) !== -1) {
+        const chunkBytes = buffer.slice(0, delimiterIndex)
+        const chunkText = decoder.decode(chunkBytes)
+        buffer = buffer.slice(delimiterIndex + delimiterBytes.length)
+        const chunk = JSON.parse(chunkText).chunk as AiAssistantChunk
+
+        switch (chunk.type) {
+          case "partial_answer":
+            setAnswer((prevAnswer) => prevAnswer + chunk.content.text)
+            break
+          case "identifiers":
+            setIdentifiers(chunk.content)
+            break
+          case "error":
+            setError(chunk.content.reason)
+            loop = false
+            return
+          case "relevant_sources":
+            setAnswerSources((prev) => [
+              ...prev,
+              ...chunk.content.relevant_sources,
+            ])
+            break
+        }
+      }
+    }
+
+    setLoading(false)
+    setQuestion("")
+  }, [])
+
+  const fetchAnswer = useCallback(async () => {
+    try {
+      const response = await getAnswer(question, identifiers?.thread_id)
+
+      if (response.status === 200) {
+        await process_stream(response)
+      } else {
+        const message = await response.text()
+        setError(message)
+      }
+    } catch (error: any) {
+      setError(JSON.stringify(error))
+    }
+  }, [question, identifiers, process_stream])
+
+  useEffect(() => {
+    if (loading && !answer) {
+      void fetchAnswer()
+    }
+  }, [loading, fetchAnswer])
+
+  useEffect(() => {
+    if (
+      !loading &&
+      answer.length &&
+      thread[lastAnswerIndex]?.content !== answer
+    ) {
+      const uniqueAnswerSources = answerSources
+        .filter(
+          (source, index) =>
+            answerSources.findIndex(
+              (s) => s.source_url === source.source_url
+            ) === index
+        )
+        .map((source) => {
+          const separatorIndex = source.title.indexOf("|")
+          return {
+            ...source,
+            title:
+              separatorIndex !== -1
+                ? source.title.slice(0, separatorIndex)
+                : source.title,
+          }
+        })
+      setThread((prevThread) => [
+        ...prevThread,
+        {
+          type: "answer",
+          content: answer,
+          question_id: identifiers?.question_answer_id,
+          order: getNewOrder(prevThread),
+          sources:
+            uniqueAnswerSources.length > 3
+              ? uniqueAnswerSources.slice(0, 3)
+              : uniqueAnswerSources,
+        },
+      ])
+      setAnswer("")
+      setAnswerSources([])
+      setMessagesCount((prev) => prev + 1)
+      inputRef.current?.focus()
+    }
+  }, [loading, answer, thread, lastAnswerIndex, inputRef.current])
+
+  useResizeObserver(contentRef as React.RefObject<HTMLDivElement>, () => {
+    if (!loading) {
+      return
+    }
+
+    scrollToBottom()
+  })
+
+  const getThreadItems = useCallback(() => {
+    return sortThread(thread)
+  }, [thread])
+
+  return (
+    <AiAssistantChatContext.Provider
+      value={{
+        handleSubmit,
+        getThreadItems,
+        question,
+        setQuestion,
+        inputRef,
+        contentRef,
+        loading,
+        thread,
+        answer,
+      }}
+    >
+      {children}
+    </AiAssistantChatContext.Provider>
+  )
+}
+
+export const useAiAssistantChat = () => {
+  const context = useContext(AiAssistantChatContext)
+
+  if (!context) {
+    throw new Error(
+      "useAiAssistantChat must be used within a AiAssistantChatContext"
+    )
+  }
+
+  return context
+}

www/packages/docs-ui/src/providers/AiAssistant/index.tsx

@@ -1,12 +1,15 @@
 "use client"
 
-import React, { createContext, useContext } from "react"
+import React, { createContext, useContext, useEffect, useState } from "react"
-import { useAnalytics } from "@/providers"
+import { useAnalytics, useSearch } from "@/providers"
-import { AiAssistant } from "@/components"
+import { AiAssistantIcon, AiAssistantSearchWindow } from "@/components"
 import { RecaptchaAction, useRecaptcha } from "../../hooks/use-recaptcha"
+import { AiAssistantChatProvider } from "./Chat"
 
 export type AiAssistantFeedbackType = "upvote" | "downvote"
 
+export type AiAssistantChatType = "default" | "popover"
+
 export type AiAssistantContextType = {
   getAnswer: (question: string, thread_id?: string) => Promise<Response>
   sendFeedback: (
@@ -14,6 +17,9 @@ export type AiAssistantContextType = {
     reaction: AiAssistantFeedbackType
   ) => Promise<Response>
   version: "v1" | "v2"
+  chatOpened: boolean
+  setChatOpened: React.Dispatch<React.SetStateAction<boolean>>
+  chatType: AiAssistantChatType
 }
 
 const AiAssistantContext = createContext<AiAssistantContextType | null>(null)
@@ -24,6 +30,8 @@ export type AiAssistantProviderProps = {
   recaptchaSiteKey: string
   websiteId: string
   version?: "v1" | "v2"
+  type?: "search" | "chat"
+  chatType?: AiAssistantChatType
 }
 
 export const AiAssistantProvider = ({
@@ -32,7 +40,11 @@ export const AiAssistantProvider = ({
   websiteId,
   version = "v2",
   children,
+  type = "chat",
+  chatType = "default",
 }: AiAssistantProviderProps) => {
+  const [chatOpened, setChatOpened] = useState(false)
+  const { setCommands, setIsOpen, setCommand } = useSearch()
   const { analytics } = useAnalytics()
   const { execute: getReCaptchaToken } = useRecaptcha({
     siteKey: recaptchaSiteKey,
@@ -85,16 +97,47 @@ export const AiAssistantProvider = ({
     )
   }
 
+  useEffect(() => {
+    setCommands((prevCommands) => {
+      const newCommands = [...prevCommands]
+
+      if (!newCommands.find((c) => c.name === "ai-assistant")) {
+        newCommands.push({
+          name: "ai-assistant",
+          icon: <AiAssistantIcon />,
+          title: "AI Assistant",
+          badge: {
+            variant: "blue",
+            badgeType: "shaded",
+            children: "Beta",
+          },
+          action: () => {
+            setIsOpen(false)
+            setChatOpened(true)
+            setCommand(null)
+          },
+        })
+      }
+
+      return newCommands
+    })
+  }, [])
+
   return (
     <AiAssistantContext.Provider
       value={{
         getAnswer,
         sendFeedback,
         version,
+        chatOpened,
+        setChatOpened,
+        chatType,
       }}
     >
-      {children}
+      <AiAssistantChatProvider>
-      <AiAssistant />
+        {children}
+        {type === "search" && <AiAssistantSearchWindow />}
+      </AiAssistantChatProvider>
     </AiAssistantContext.Provider>
   )
 }

www/packages/docs-ui/src/providers/Layout/index.tsx

@@ -0,0 +1,57 @@
+"use client"
+
+import useResizeObserver from "@react-hook/resize-observer"
+import React, { createContext, createRef, useContext, useState } from "react"
+
+export type LayoutProviderContextType = {
+  mainContentRef: React.RefObject<HTMLDivElement | null>
+  showCollapsedNavbar: boolean
+}
+
+export const LayoutProviderContext =
+  createContext<LayoutProviderContextType | null>(null)
+
+export type LayoutProviderProps = {
+  children: React.ReactNode
+  disableResizeObserver?: boolean
+}
+
+export const LayoutProvider = ({
+  children,
+  disableResizeObserver = false,
+}: LayoutProviderProps) => {
+  const mainContentRef = createRef<HTMLDivElement>()
+  const [showCollapsedNavbar, setShowCollapsedNavbar] = useState(false)
+
+  useResizeObserver(mainContentRef as React.RefObject<HTMLDivElement>, () => {
+    if (disableResizeObserver) {
+      setShowCollapsedNavbar(false)
+      return
+    }
+    if (window.innerWidth < 992) {
+      setShowCollapsedNavbar(false)
+      return
+    }
+    if (mainContentRef.current) {
+      setShowCollapsedNavbar(mainContentRef.current.clientWidth < 992)
+    }
+  })
+
+  return (
+    <LayoutProviderContext.Provider
+      value={{ mainContentRef, showCollapsedNavbar }}
+    >
+      {children}
+    </LayoutProviderContext.Provider>
+  )
+}
+
+export const useLayout = (): LayoutProviderContextType => {
+  const context = useContext(LayoutProviderContext)
+
+  if (!context) {
+    throw new Error("useLayout must be used inside a LayoutProvider")
+  }
+
+  return context
+}

www/packages/docs-ui/src/providers/Search/index.tsx

@@ -23,7 +23,8 @@ import { CSSTransition, SwitchTransition } from "react-transition-group"
 
 export type SearchCommand = {
   name: string
-  component: React.ReactNode
+  component?: React.ReactNode
+  action?: () => void
   icon?: React.ReactNode
   title: string
   badge?: BadgeProps
@@ -38,6 +39,7 @@ export type SearchContextType = {
   commands: SearchCommand[]
   command: SearchCommand | null
   setCommand: React.Dispatch<React.SetStateAction<SearchCommand | null>>
+  setCommands: React.Dispatch<React.SetStateAction<SearchCommand[]>>
   modalRef: React.MutableRefObject<HTMLDialogElement | null>
 }
 
@@ -64,13 +66,14 @@ export const SearchProvider = ({
   initialDefaultFilters = [],
   searchProps,
   algolia,
-  commands = [],
+  commands: initialCommands = [],
   modalClassName,
 }: SearchProviderProps) => {
   const [isOpen, setIsOpen] = useState(false)
   const [defaultFilters, setDefaultFilters] = useState<string[]>(
     initialDefaultFilters
   )
+  const [commands, setCommands] = useState<SearchCommand[]>(initialCommands)
   const [command, setCommand] = useState<SearchCommand | null>(null)
 
   const modalRef = useRef<HTMLDialogElement | null>(null)
@@ -211,6 +214,10 @@ export const SearchProvider = ({
 
   const componentWrapperRef = useRef(null)
 
+  useEffect(() => {
+    command?.action?.()
+  }, [command])
+
   return (
     <SearchContext.Provider
       value={{
@@ -223,6 +230,7 @@ export const SearchProvider = ({
         command,
         setCommand,
         modalRef,
+        setCommands,
       }}
     >
       {children}
@@ -241,20 +249,20 @@ export const SearchProvider = ({
           <CSSTransition
             classNames={{
               enter:
-                command === null
+                command === null || !command.component
                   ? "animate-fadeInLeft animate-fast"
                   : "animate-fadeInRight animate-fast",
               exit:
-                command === null
+                command === null || !command.component
                   ? "animate-fadeOutLeft animate-fast"
                   : "animate-fadeOutRight animate-fast",
             }}
             timeout={250}
-            key={command?.name || "search"}
+            key={command?.component ? command.name : "search"}
             nodeRef={componentWrapperRef}
           >
             <div ref={componentWrapperRef} className="h-full">
-              {command === null && (
+              {!command?.component && (
                 <Search {...searchProps} algolia={algolia} />
               )}
               {command?.component}

www/packages/docs-ui/src/providers/index.ts

@@ -2,6 +2,7 @@ export * from "./AiAssistant"
 export * from "./Analytics"
 export * from "./BrowserProvider"
 export * from "./ColorMode"
+export * from "./Layout"
 export * from "./LearningPath"
 export * from "./MainNav"
 export * from "./Mobile"

www/packages/tailwind/base.tailwind.config.js

@@ -275,6 +275,8 @@ module.exports = {
         "subtle-code-fade-right-to-left-dark": `linear-gradient(90deg, #30303380, #303033)`,
         "border-dotted":
           "linear-gradient(90deg,var(--docs-border-strong) 1px,transparent 1px)",
+        "ai-assistant-bottom":
+          "linear-gradient(180deg, rgba(255, 255, 255, 0.00) 0%, var(--docs-bg-base) 100%)",
       },
       screens: {
         xs: "568px",
@@ -283,7 +285,8 @@ module.exports = {
         lg: "1024px",
         xl: "1280px",
         xxl: "1536px",
-        xxxl: "3840px",
+        xxxl: "1800px",
+        xxxxl: "3840px",
       },
       transitionTimingFunction: {
         ease: "ease",
@@ -291,6 +294,7 @@ module.exports = {
       width: {
         toc: "221px",
         "sidebar-xs": "calc(100% - 20px)",
+        "ai-assistant": "500px"
       },
       maxWidth: {
         // sidebar
@@ -326,6 +330,8 @@ module.exports = {
         "modal-sm": "624px",
         "modal-md": "752px",
         "modal-lg": "640px",
+        // ai-assistant
+        "ai-assistant": "500px"
       },
       minWidth: {
         xl: "1419px",

www/packages/tailwind/theme-presets.js

@@ -22,8 +22,8 @@ const light = {
 
   "--docs-fg-base": "rgba(24, 24, 27, 1)",
   "--docs-fg-subtle": "rgba(82, 82, 91, 1)",
-  "--docs-fg-muted": "rgba(161, 161, 170, 1)",
+  "--docs-fg-muted": "rgba(113, 113, 122, 1)",
-  "--docs-fg-disabled": "rgba(212, 212, 216, 1)",
+  "--docs-fg-disabled": "rgba(161, 161, 170, 1)",
   "--docs-fg-on-color": "rgba(255, 255, 255, 1)",
   "--docs-fg-on-inverted": "rgba(255, 255, 255, 1)",
   "--docs-fg-interactive": "rgba(59, 130, 246, 1)",