ENG-1734 Add Roam semantic search to Advanced Node Search

apps/roam/src/components/AdvancedNodeSearchDialog/useAdvancedNodeSearchResults.ts

@@ -1,5 +1,10 @@
-import { useMemo } from "react";
+import { useEffect, useState } from "react";
 import MiniSearch from "minisearch";
+import getDiscourseNodes from "~/utils/getDiscourseNodes";
+import {
+  searchDiscourseNodesWithSemanticFallback,
+  type DiscourseSearchHit,
+} from "~/utils/discourseNodeSemanticSearch";
 import {
   searchIndexedNodes,
   sortSearchResults,
@@ -23,6 +28,57 @@ type UseAdvancedNodeSearchResultsArgs = {
   dockedResults?: SearchResult[];
 };
 
+const toKeywordHits = (
+  scoredHits: ReturnType<typeof searchIndexedNodes>,
+): DiscourseSearchHit[] =>
+  scoredHits.map((hit) => ({
+    uid: hit.result.uid,
+    text: hit.result.title,
+    type: hit.result.type,
+    nodeTypeLabel: hit.result.nodeTypeLabel,
+    score: hit.score,
+    source: "keyword" as const,
+  }));
+
+const hitsToScoredSearchHits = ({
+  hits,
+  resultsByUid,
+}: {
+  hits: DiscourseSearchHit[];
+  resultsByUid: Map<string, SearchResult>;
+}) =>
+  hits
+    .map((hit) => {
+      const indexedResult = resultsByUid.get(hit.uid);
+      if (indexedResult) {
+        return { result: indexedResult, score: hit.score, source: hit.source };
+      }
+
+      return {
+        result: {
+          uid: hit.uid,
+          title: hit.text,
+          type: hit.type || "",
+          nodeTypeLabel: hit.nodeTypeLabel || "",
+          excerpt: "",
+          createdAt: "",
+          lastModified: "",
+          authorName: "Unknown",
+        },
+        score: hit.score,
+        source: hit.source,
+      };
+    })
+    .filter(
+      (
+        hit,
+      ): hit is {
+        result: SearchResult;
+        score: number;
+        source: DiscourseSearchHit["source"];
+      } => !!hit,
+    );
+
 export const useAdvancedNodeSearchResults = ({
   debouncedSearchTerm,
   selectedNodeTypeIds,
@@ -32,30 +88,80 @@ export const useAdvancedNodeSearchResults = ({
   searchIndex,
   dockedQuery,
   dockedResults,
-}: UseAdvancedNodeSearchResultsArgs): SearchResult[] =>
-  useMemo(() => {
-    if (!debouncedSearchTerm) return [];
+}: UseAdvancedNodeSearchResultsArgs): SearchResult[] => {
+  const [results, setResults] = useState<SearchResult[]>([]);
+
+  useEffect(() => {
+    if (!debouncedSearchTerm) {
+      setResults([]);
+      return;
+    }
 
     const isDockedQuery =
       dockedQuery !== undefined &&
       debouncedSearchTerm.trim() === dockedQuery.trim();
 
     if (isDockedQuery && dockedResults) {
-      return dockedResults;
+      setResults(dockedResults);
+      return;
     }
 
     if (isIndexLoading || indexError || !searchIndex) {
-      return [];
+      setResults([]);
+      return;
     }
 
-    const scoredHits = searchIndexedNodes({
-      miniSearch: searchIndex.miniSearch,
-      allResults: searchIndex.allResults,
-      searchTerm: debouncedSearchTerm,
-      typeFilter: selectedNodeTypeIds.length ? selectedNodeTypeIds : undefined,
-    });
+    let cancelled = false;
+    const typeFilter = selectedNodeTypeIds.length
+      ? selectedNodeTypeIds
+      : undefined;
+    const discourseNodes = getDiscourseNodes().filter(
+      (node) =>
+        node.backedBy === "user" &&
+        (!typeFilter || typeFilter.includes(node.type)),
+    );
+    const resultsByUid = new Map(
+      searchIndex.allResults.map((result) => [result.uid, result]),
+    );
 
-    return sortSearchResults({ hits: scoredHits, sort });
+    const runKeywordSearch = (): DiscourseSearchHit[] =>
+      toKeywordHits(
+        searchIndexedNodes({
+          miniSearch: searchIndex.miniSearch,
+          allResults: searchIndex.allResults,
+          searchTerm: debouncedSearchTerm,
+          typeFilter,
+        }),
+      );
+
+    void searchDiscourseNodesWithSemanticFallback({
+      nodeTypes: discourseNodes,
+      query: debouncedSearchTerm,
+      runKeywordSearch,
+    })
+      .then((hits) => {
+        if (cancelled) return;
+
+        const scoredHits = hitsToScoredSearchHits({ hits, resultsByUid });
+        setResults(sortSearchResults({ hits: scoredHits, sort }));
+      })
+      .catch((error) => {
+        console.error("Advanced node search failed:", error);
+        if (cancelled) return;
+        setResults(
+          sortSearchResults({
+            hits: hitsToScoredSearchHits({
+              hits: runKeywordSearch(),
+              resultsByUid,
+            }),
+            sort,
+          }),
+        );
+      });
+
+    return () => {
+      cancelled = true;
+    };
   }, [
     debouncedSearchTerm,
     dockedQuery,
@@ -66,3 +172,6 @@ export const useAdvancedNodeSearchResults = ({
     selectedNodeTypeIds,
     sort,
   ]);
+
+  return results;
+};

apps/roam/src/components/AdvancedNodeSearchDialog/utils.ts

@@ -59,6 +59,7 @@ export type DockedSearchState = {
 export type ScoredSearchHit = {
   result: SearchResult;
   score: number;
+  source?: "semantic" | "keyword";
 };
 
 type MiniSearchDocument = SearchResult & {
@@ -229,6 +230,17 @@ export const sortSearchResults = ({
 
     switch (sort.field) {
       case "relevance":
+        if (aHit.source && bHit.source && aHit.source !== bHit.source) {
+          comparison =
+            sort.direction === "desc"
+              ? aHit.source === "semantic"
+                ? -1
+                : 1
+              : aHit.source === "semantic"
+                ? 1
+                : -1;
+          break;
+        }
         comparison = compareNumbers(aHit.score, bHit.score, sort.direction);
         break;
       case "alphabetical":

apps/roam/src/utils/discourseNodeSearchProviders.ts

@@ -249,7 +249,7 @@ export const searchSemanticNodeTitles = async ({
   );
 };
 
-const runRoamSemanticSearch = async ({
+export const runRoamSemanticSearch = async ({
   nodeTypes,
   query,
 }: {

apps/roam/src/utils/discourseNodeSemanticSearch.ts

@@ -0,0 +1,88 @@
+import type { DiscourseNode } from "~/utils/getDiscourseNodes";
+import {
+  runRoamSemanticSearch,
+  type AdminSearchResultItem,
+} from "~/utils/discourseNodeSearchProviders";
+
+export const SEMANTIC_SEARCH_MIN_DISCOURSE_RESULTS = 5;
+
+export type DiscourseSearchHitSource = "semantic" | "keyword";
+
+export type DiscourseSearchHit = {
+  uid: string;
+  text: string;
+  type?: string;
+  nodeTypeLabel?: string;
+  score: number;
+  source: DiscourseSearchHitSource;
+};
+
+export const shouldUseRoamSemanticSearch = (): boolean =>
+  window.roamAlphaAPI.data.semanticSearchEnabled();
+
+export const combineDiscourseSearchResults = ({
+  semantic,
+  keyword,
+}: {
+  semantic: DiscourseSearchHit[];
+  keyword: DiscourseSearchHit[];
+}): DiscourseSearchHit[] => {
+  const seenUids = new Set(semantic.map((hit) => hit.uid));
+  const combined = [...semantic];
+
+  keyword.forEach((hit) => {
+    if (seenUids.has(hit.uid)) return;
+    seenUids.add(hit.uid);
+    combined.push(hit);
+  });
+
+  return combined;
+};
+
+const toSemanticHit = (item: AdminSearchResultItem): DiscourseSearchHit => ({
+  uid: item.uid,
+  text: item.text,
+  type: item.type,
+  nodeTypeLabel: item.nodeTypeLabel,
+  score: item.score ?? 0,
+  source: "semantic",
+});
+
+export const searchDiscourseNodesWithSemanticFallback = async ({
+  nodeTypes,
+  query,
+  runKeywordSearch,
+}: {
+  nodeTypes: DiscourseNode[];
+  query: string;
+  runKeywordSearch: () => DiscourseSearchHit[];
+}): Promise<DiscourseSearchHit[]> => {
+  const trimmedQuery = query.trim();
+  if (!trimmedQuery) return [];
+
+  if (!shouldUseRoamSemanticSearch()) {
+    return runKeywordSearch();
+  }
+
+  try {
+    const providerResult = await runRoamSemanticSearch({
+      nodeTypes,
+      query: trimmedQuery,
+    });
+    const semanticHits = providerResult.filteredResults.map(toSemanticHit);
+
+    if (
+      providerResult.filteredResultCount >=
+      SEMANTIC_SEARCH_MIN_DISCOURSE_RESULTS
+    ) {
+      return semanticHits;
+    }
+
+    return combineDiscourseSearchResults({
+      semantic: semanticHits,
+      keyword: runKeywordSearch(),
+    });
+  } catch {
+    return runKeywordSearch();
+  }
+};