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
2026-05-21 01:04:39 +08:00 · 2025-12-01 11:50:25 +00:00
parent c3b350d943
commit 9fb16bc842
46 changed files with 1608 additions and 68 deletions
@@ -0,0 +1,113 @@
+import React, { useMemo } from "react";
+import { Paper, Text, Group, Stack, Loader, Box } from "@mantine/core";
+import { IconSparkles, IconFileText } from "@tabler/icons-react";
+import { Link } from "react-router-dom";
+import { IAiSearchResponse } from "../services/ai-search-service.ts";
+import { buildPageUrl } from "@/features/page/page.utils.ts";
+import { markdownToHtml } from "@docmost/editor-ext";
+import DOMPurify from "dompurify";
+import { useTranslation } from "react-i18next";
+
+interface AiSearchResultProps {
+  result?: IAiSearchResponse;
+  isLoading?: boolean;
+  streamingAnswer?: string;
+  streamingSources?: any[];
+}
+
+export function AiSearchResult({
+  result,
+  isLoading,
+  streamingAnswer = "",
+  streamingSources = [],
+}: AiSearchResultProps) {
+  const { t } = useTranslation();
+
+  // Use streaming data if available, otherwise fall back to result
+  const answer = streamingAnswer || result?.answer || "";
+  const sources =
+    streamingSources.length > 0 ? streamingSources : result?.sources || [];
+
+  // Deduplicate sources by pageId, keeping the one with highest similarity
+  const deduplicatedSources = useMemo(() => {
+    if (!sources || sources.length === 0) return [];
+
+    const pageMap = new Map();
+    sources.forEach((source) => {
+      const existing = pageMap.get(source.pageId);
+      if (!existing || source.similarity > existing.similarity) {
+        pageMap.set(source.pageId, source);
+      }
+    });
+
+    return Array.from(pageMap.values());
+  }, [sources]);
+
+  if (isLoading && !answer) {
+    return (
+      <Paper p="md" radius="md" withBorder>
+        <Group>
+          <Loader size="sm" />
+          <Text size="sm">{t("AI is thinking...")}</Text>
+        </Group>
+      </Paper>
+    );
+  }
+
+  if (!answer && !isLoading) {
+    return null;
+  }
+
+  return (
+    <Stack gap="md" p="md">
+      <Paper p="md" radius="md" withBorder>
+        <Group gap="xs" mb="sm">
+          <IconSparkles size={20} color="var(--mantine-color-blue-6)" />
+          <Text fw={600} size="sm">
+            {t("AI Answer")}
+          </Text>
+          {isLoading && <Loader size="xs" />}
+        </Group>
+        <div
+          dangerouslySetInnerHTML={{
+            __html: DOMPurify.sanitize(markdownToHtml(answer) as string),
+          }}
+        />
+      </Paper>
+
+      {deduplicatedSources.length > 0 && (
+        <Stack gap="xs">
+          <Text size="xs" fw={600} c="dimmed">
+            {t("Sources")}
+          </Text>
+          {deduplicatedSources.map((source) => (
+            <Box
+              key={source.pageId}
+              component={Link}
+              to={buildPageUrl(source.spaceSlug, source.slugId, source.title)}
+              style={{
+                textDecoration: "none",
+                color: "inherit",
+                display: "block",
+              }}
+            >
+              <Paper
+                p="xs"
+                radius="sm"
+                withBorder
+                style={{ cursor: "pointer" }}
+              >
+                <Group gap="xs">
+                  <IconFileText size={16} />
+                  <Text size="sm" truncate>
+                    {source.title}
+                  </Text>
+                </Group>
+              </Paper>
+            </Box>
+          ))}
+        </Stack>
+      )}
+    </Stack>
+  );
+}
@@ -0,0 +1,69 @@
+import { Group, Text, Switch, MantineSize, Title } from "@mantine/core";
+import { useAtom } from "jotai";
+import { workspaceAtom } from "@/features/user/atoms/current-user-atom.ts";
+import React, { useState } from "react";
+import { useTranslation } from "react-i18next";
+import { updateWorkspace } from "@/features/workspace/services/workspace-service.ts";
+import { notifications } from "@mantine/notifications";
+import { isCloud } from "@/lib/config.ts";
+import useLicense from "@/ee/hooks/use-license.tsx";
+
+export default function EnableAiSearch() {
+  const { t } = useTranslation();
+
+  return (
+    <>
+      <Group justify="space-between" wrap="nowrap" gap="xl">
+        <div>
+          <Text size="md">{t("AI-powered search (Ask AI)")}</Text>
+          <Text size="sm" c="dimmed">
+            {t(
+              "AI search uses vector embeddings to provide semantic search capabilities across your workspace content.",
+            )}
+          </Text>
+        </div>
+
+        <AiSearchToggle />
+      </Group>
+    </>
+  );
+}
+
+interface AiSearchToggleProps {
+  size?: MantineSize;
+  label?: string;
+}
+export function AiSearchToggle({ size, label }: AiSearchToggleProps) {
+  const { t } = useTranslation();
+  const [workspace, setWorkspace] = useAtom(workspaceAtom);
+  const [checked, setChecked] = useState(workspace?.settings?.ai?.search);
+  const { hasLicenseKey } = useLicense();
+
+  const hasAccess = isCloud() || (!isCloud() && hasLicenseKey);
+
+  const handleChange = async (event: React.ChangeEvent<HTMLInputElement>) => {
+    const value = event.currentTarget.checked;
+    try {
+      const updatedWorkspace = await updateWorkspace({ aiSearch: value });
+      setChecked(value);
+      setWorkspace(updatedWorkspace);
+    } catch (err) {
+      notifications.show({
+        message: err?.response?.data?.message,
+        color: "red",
+      });
+    }
+  };
+
+  return (
+    <Switch
+      size={size}
+      label={label}
+      labelPosition="left"
+      defaultChecked={checked}
+      onChange={handleChange}
+      disabled={!hasAccess}
+      aria-label={t("Toggle AI search")}
+    />
+  );
+}