feat(EE): AI vector search (#1691)

* WIP

* AI module - init

* WIP

* sync

* WIP

* refactor naming

* new columns

* sync

* sync

* fix search bug

* stream response

* WIP

* feat embeddings sync

* refine

* Add workspaceId to page events

* refine

* WIP

* add translation string

* sync

* reset ai answer on query change

* hide AI search in cloud

* capture streaming error

* sync

apps/client/src/features/search/components/search-spotlight-filters.tsx

@@ -9,6 +9,7 @@ import {
   ScrollArea,
   Avatar,
   Group,
+  Switch,
   getDefaultZIndex,
 } from "@mantine/core";
 import {
@@ -17,6 +18,7 @@ import {
   IconFileDescription,
   IconSearch,
   IconCheck,
+  IconSparkles,
 } from "@tabler/icons-react";
 import { useTranslation } from "react-i18next";
 import { useDebouncedValue } from "@mantine/hooks";
@@ -24,15 +26,21 @@ import { useGetSpacesQuery } from "@/features/space/queries/space-query";
 import { useLicense } from "@/ee/hooks/use-license";
 import classes from "./search-spotlight-filters.module.css";
 import { isCloud } from "@/lib/config.ts";
+import { useAtom } from "jotai/index";
+import { workspaceAtom } from "@/features/user/atoms/current-user-atom.ts";
 
 interface SearchSpotlightFiltersProps {
   onFiltersChange?: (filters: any) => void;
+  onAskClick?: () => void;
   spaceId?: string;
+  isAiMode?: boolean;
 }
 
 export function SearchSpotlightFilters({
   onFiltersChange,
+  onAskClick,
   spaceId,
+  isAiMode = false,
 }: SearchSpotlightFiltersProps) {
   const { t } = useTranslation();
   const { hasLicenseKey } = useLicense();
@@ -42,6 +50,7 @@ export function SearchSpotlightFilters({
   const [spaceSearchQuery, setSpaceSearchQuery] = useState("");
   const [debouncedSpaceQuery] = useDebouncedValue(spaceSearchQuery, 300);
   const [contentType, setContentType] = useState<string | null>("page");
+  const [workspace] = useAtom(workspaceAtom);
 
   const { data: spacesData } = useGetSpacesQuery({
     page: 1,
@@ -120,6 +129,31 @@ export function SearchSpotlightFilters({
 
   return (
     <div className={classes.filtersContainer}>
+      {workspace?.settings?.ai?.search === true && (
+        <div
+          style={{
+            display: "flex",
+            alignItems: "center",
+            height: "32px",
+            paddingLeft: "8px",
+            paddingRight: "8px",
+          }}
+        >
+          <Switch
+            checked={isAiMode}
+            onChange={(event) => onAskClick()}
+            label={t("Ask AI")}
+            size="sm"
+            color="blue"
+            labelPosition="left"
+            styles={{
+              root: { display: "flex", alignItems: "center" },
+              label: { paddingRight: "8px", fontSize: "13px", fontWeight: 500 },
+            }}
+          />
+        </div>
+      )}
+
       <Menu
         shadow="md"
         width={250}
@@ -231,7 +265,7 @@ export function SearchSpotlightFilters({
                 contentType !== option.value &&
                 handleFilterChange("contentType", option.value)
               }
-              disabled={option.disabled}
+              disabled={option.disabled || (isAiMode && option.value === "attachment")}
             >
               <Group flex="1" gap="xs">
                 <div>
@@ -241,6 +275,11 @@ export function SearchSpotlightFilters({
                       {t("Enterprise")}
                     </Badge>
                   )}
+                  {!option.disabled && isAiMode && option.value === "attachment" && (
+                    <Text size="xs" mt={4}>
+                      {t("Ask AI not available for attachments")}
+                    </Text>
+                  )}
                 </div>
                 {contentType === option.value && <IconCheck size={20} />}
               </Group>

apps/client/src/features/search/components/search-spotlight.tsx

@@ -1,12 +1,16 @@
 import { Spotlight } from "@mantine/spotlight";
-import { IconSearch } from "@tabler/icons-react";
-import React, { useState, useMemo } from "react";
+import { IconSearch, IconSparkles } from "@tabler/icons-react";
+import { Group, Button } from "@mantine/core";
+import React, { useState, useMemo, useEffect } from "react";
 import { useDebouncedValue } from "@mantine/hooks";
 import { useTranslation } from "react-i18next";
+import { notifications } from "@mantine/notifications";
 import { searchSpotlightStore } from "../constants.ts";
 import { SearchSpotlightFilters } from "./search-spotlight-filters.tsx";
 import { useUnifiedSearch } from "../hooks/use-unified-search.ts";
+import { useAiSearch } from "../../../ee/ai/hooks/use-ai-search.ts";
 import { SearchResultItem } from "./search-result-item.tsx";
+import { AiSearchResult } from "../../../ee/ai/components/ai-search-result.tsx";
 import { useLicense } from "@/ee/hooks/use-license.tsx";
 import { isCloud } from "@/lib/config.ts";
 
@@ -24,6 +28,7 @@ export function SearchSpotlight({ spaceId }: SearchSpotlightProps) {
   }>({
     contentType: "page",
   });
+  const [isAiMode, setIsAiMode] = useState(false);
 
   // Build unified search params
   const searchParams = useMemo(() => {
@@ -40,7 +45,42 @@ export function SearchSpotlight({ spaceId }: SearchSpotlightProps) {
     return params;
   }, [debouncedSearchQuery, filters]);
 
-  const { data: searchResults, isLoading } = useUnifiedSearch(searchParams);
+  const { data: searchResults, isLoading } = useUnifiedSearch(
+    searchParams,
+    !isAiMode // Disable regular search when in AI mode
+  );
+  const {
+    //@ts-ignore
+    data: aiSearchResult,
+    //@ts-ignore
+    isPending: isAiLoading,
+    //@ts-ignore
+    mutate: triggerAiSearchMutation,
+    //@ts-ignore
+    reset: resetAiMutation,
+    //@ts-ignore
+    error: aiSearchError,
+    streamingAnswer,
+    streamingSources,
+    clearStreaming,
+  } = useAiSearch();
+
+  // Clear streaming state and mutation data when query changes (user is typing a new query)
+  useEffect(() => {
+    clearStreaming();
+    resetAiMutation();
+  }, [query, clearStreaming, resetAiMutation]);
+
+  // Show error notification when AI search fails
+  useEffect(() => {
+    if (aiSearchError) {
+      notifications.show({
+        message: aiSearchError.message || t("AI search failed. Please try again."),
+        color: "red",
+        position: "top-center"
+      });
+    }
+  }, [aiSearchError, t]);
 
   // Determine result type for rendering
   const isAttachmentSearch =
@@ -59,6 +99,16 @@ export function SearchSpotlight({ spaceId }: SearchSpotlightProps) {
     setFilters(newFilters);
   };
 
+  const handleAskClick = () => {
+    setIsAiMode(!isAiMode);
+  };
+
+  const handleAiSearchTrigger = () => {
+    if (query.trim() && isAiMode) {
+      triggerAiSearchMutation(searchParams);
+    }
+  };
+
   return (
     <>
       <Spotlight.Root
@@ -72,10 +122,30 @@ export function SearchSpotlight({ spaceId }: SearchSpotlightProps) {
           backgroundOpacity: 0.55,
         }}
       >
-        <Spotlight.Search
-          placeholder={t("Search...")}
-          leftSection={<IconSearch size={20} stroke={1.5} />}
-        />
+        <Group gap="xs" px="sm" pt="sm" pb="xs">
+          <Spotlight.Search
+            placeholder={isAiMode ? t("Ask a question...") : t("Search...")}
+            leftSection={<IconSearch size={20} stroke={1.5} />}
+            style={{ flex: 1 }}
+            onKeyDown={(e) => {
+              if (e.key === "Enter" && isAiMode && query.trim() && !isAiLoading) {
+                e.preventDefault();
+                handleAiSearchTrigger();
+              }
+            }}
+          />
+          {isAiMode && hasLicenseKey && (
+            <Button
+              size="xs"
+              leftSection={<IconSparkles size={16} />}
+              onClick={handleAiSearchTrigger}
+              disabled={!query.trim()}
+              loading={isAiLoading}
+            >
+              Ask
+            </Button>
+          )}
+        </Group>
 
         <div
           style={{
@@ -84,20 +154,43 @@ export function SearchSpotlight({ spaceId }: SearchSpotlightProps) {
         >
           <SearchSpotlightFilters
             onFiltersChange={handleFiltersChange}
+            onAskClick={handleAskClick}
             spaceId={spaceId}
+            isAiMode={isAiMode}
           />
         </div>
 
         <Spotlight.ActionsList>
-          {query.length === 0 && resultItems.length === 0 && (
-            <Spotlight.Empty>{t("Start typing to search...")}</Spotlight.Empty>
-          )}
+          {isAiMode ? (
+            <>
+              {query.length === 0 && (
+                <Spotlight.Empty>{t("Ask a question...")}</Spotlight.Empty>
+              )}
+              {query.length > 0 && (isAiLoading || aiSearchResult || streamingAnswer) && (
+                <AiSearchResult
+                  result={aiSearchResult}
+                  isLoading={isAiLoading}
+                  streamingAnswer={streamingAnswer}
+                  streamingSources={streamingSources}
+                />
+              )}
+              {query.length > 0 && !isAiLoading && !aiSearchResult && (
+                <Spotlight.Empty>{t("No answer available")}</Spotlight.Empty>
+              )}
+            </>
+          ) : (
+            <>
+              {query.length === 0 && resultItems.length === 0 && (
+                <Spotlight.Empty>{t("Start typing to search...")}</Spotlight.Empty>
+              )}
 
-          {query.length > 0 && !isLoading && resultItems.length === 0 && (
-            <Spotlight.Empty>{t("No results found...")}</Spotlight.Empty>
-          )}
+              {query.length > 0 && !isLoading && resultItems.length === 0 && (
+                <Spotlight.Empty>{t("No results found...")}</Spotlight.Empty>
+              )}
 
-          {resultItems.length > 0 && <>{resultItems}</>}
+              {resultItems.length > 0 && <>{resultItems}</>}
+            </>
+          )}
         </Spotlight.ActionsList>
       </Spotlight.Root>
     </>

apps/client/src/features/search/hooks/use-unified-search.ts

@@ -19,6 +19,7 @@ export interface UseUnifiedSearchParams extends IPageSearchParams {
 
 export function useUnifiedSearch(
   params: UseUnifiedSearchParams,
+  enabled: boolean = true,
 ): UseQueryResult<UnifiedSearchResult[], Error> {
   const { hasLicenseKey } = useLicense();
 
@@ -38,6 +39,6 @@ export function useUnifiedSearch(
         return await searchPage(backendParams);
       }
     },
-    enabled: !!params.query,
+    enabled: !!params.query && enabled,
   });
 }

apps/client/src/features/workspace/services/workspace-service.ts

@@ -42,7 +42,7 @@ export async function deleteWorkspaceMember(data: {
   await api.post("/workspace/members/delete", data);
 }
 
-export async function updateWorkspace(data: Partial<IWorkspace>) {
+export async function updateWorkspace(data: Partial<IWorkspace> & { aiSearch?: boolean }) {
   const req = await api.post<IWorkspace>("/workspace/update", data);
   return req.data;
 }

apps/client/src/features/workspace/types/workspace.types.ts

@@ -9,7 +9,7 @@ export interface IWorkspace {
   defaultSpaceId: string;
   customDomain: string;
   enableInvite: boolean;
-  settings: any;
+  settings: IWorkspaceSettings;
   status: string;
   enforceSso: boolean;
   stripeCustomerId: string;
@@ -24,6 +24,14 @@ export interface IWorkspace {
   enforceMfa?: boolean;
 }
 
+export interface IWorkspaceSettings {
+  ai?: IWorkspaceAiSettings;
+}
+
+export interface IWorkspaceAiSettings {
+  search?: boolean;
+}
+
 export interface ICreateInvite {
   role: string;
   emails: string[];