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-20 00:14:10 +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,46 @@
+import { useMutation, UseMutationResult } from "@tanstack/react-query";
+import { useState, useCallback } from "react";
+import { askAi, IAiSearchResponse } from "@/ee/ai/services/ai-search-service.ts";
+import { IPageSearchParams } from "@/features/search/types/search.types.ts";
+
+// @ts-ignore
+interface UseAiSearchResult extends UseMutationResult<IAiSearchResponse, Error, IPageSearchParams> {
+  streamingAnswer: string;
+  streamingSources: any[];
+  clearStreaming: () => void;
+}
+
+export function useAiSearch(): UseAiSearchResult {
+  const [streamingAnswer, setStreamingAnswer] = useState("");
+  const [streamingSources, setStreamingSources] = useState<any[]>([]);
+
+  const clearStreaming = useCallback(() => {
+    setStreamingAnswer("");
+    setStreamingSources([]);
+  }, []);
+
+  const mutation = useMutation({
+    mutationFn: async (params: IPageSearchParams & { contentType?: string }) => {
+      setStreamingAnswer("");
+      setStreamingSources([]);
+
+      const { contentType, ...apiParams } = params;
+
+      return await askAi(apiParams, (chunk) => {
+        if (chunk.content) {
+          setStreamingAnswer((prev) => prev + chunk.content);
+        }
+        if (chunk.sources) {
+          setStreamingSources(chunk.sources);
+        }
+      });
+    },
+  });
+
+  return {
+    ...mutation,
+    streamingAnswer,
+    streamingSources,
+    clearStreaming,
+  };
+}
@@ -0,0 +1,61 @@
+import { useState, useCallback, useRef } from "react";
+import { useAiGenerateStreamMutation } from "@/ee/ai/queries/ai-query.ts";
+import { AiGenerateDto } from "@/ee/ai/types/ai.types.ts";
+
+export function useAiStream() {
+  const [content, setContent] = useState("");
+  const [isStreaming, setIsStreaming] = useState(false);
+  const abortControllerRef = useRef<AbortController | null>(null);
+  const mutation = useAiGenerateStreamMutation();
+
+  const startStream = useCallback(
+    async (data: AiGenerateDto) => {
+      setContent("");
+      setIsStreaming(true);
+
+      try {
+        const controller = await mutation.mutateAsync({
+          ...data,
+          onChunk: (chunk) => {
+            setContent((prev) => prev + chunk.content);
+          },
+          onError: (error) => {
+            console.error("AI stream error:", error);
+            setIsStreaming(false);
+          },
+          onComplete: () => {
+            setIsStreaming(false);
+          },
+        });
+
+        abortControllerRef.current = controller;
+      } catch (error) {
+        console.error("Failed to start stream:", error);
+        setIsStreaming(false);
+      }
+    },
+    [mutation]
+  );
+
+  const stopStream = useCallback(() => {
+    if (abortControllerRef.current) {
+      abortControllerRef.current.abort();
+      abortControllerRef.current = null;
+      setIsStreaming(false);
+    }
+  }, []);
+
+  const resetContent = useCallback(() => {
+    setContent("");
+  }, []);
+
+  return {
+    content,
+    isStreaming,
+    startStream,
+    stopStream,
+    resetContent,
+    isLoading: mutation.isPending,
+    error: mutation.error,
+  };
+}