ENG-1734 Add Roam semantic search to Advanced Node Search

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

@@ -1,8 +1,11 @@
-import { useMemo } from "react";
+import { useEffect, useMemo, useState } from "react";
 import MiniSearch from "minisearch";
+import getDiscourseNodes from "~/utils/getDiscourseNodes";
+import { searchDiscourseNodes } from "~/utils/searchDiscourseNodes";
 import {
-  searchIndexedNodes,
+  searchDiscourseNodesWithMiniSearch,
   sortSearchResults,
+  type ScoredSearchResult,
   type SearchResult,
   type SortConfig,
 } from "./utils";
@@ -32,37 +35,91 @@ export const useAdvancedNodeSearchResults = ({
   searchIndex,
   dockedQuery,
   dockedResults,
-}: UseAdvancedNodeSearchResultsArgs): SearchResult[] =>
-  useMemo(() => {
-    if (!debouncedSearchTerm) return [];
-
-    const isDockedQuery =
+}: UseAdvancedNodeSearchResultsArgs): SearchResult[] => {
+  const isDockedQuery = useMemo(
+    () =>
       dockedQuery !== undefined &&
-      debouncedSearchTerm.trim() === dockedQuery.trim();
+      debouncedSearchTerm.trim() === dockedQuery.trim(),
+    [debouncedSearchTerm, dockedQuery],
+  );
+
+  const [unsortedScoredResults, setUnsortedScoredResults] = useState<
+    ScoredSearchResult[]
+  >([]);
 
-    if (isDockedQuery && dockedResults) {
-      return dockedResults;
+  useEffect(() => {
+    if (isDockedQuery) {
+      setUnsortedScoredResults([]);
+      return;
+    }
+
+    if (!debouncedSearchTerm) {
+      setUnsortedScoredResults([]);
+      return;
     }
 
     if (isIndexLoading || indexError || !searchIndex) {
-      return [];
+      setUnsortedScoredResults([]);
+      return;
     }
 
-    const scoredHits = searchIndexedNodes({
-      miniSearch: searchIndex.miniSearch,
-      allResults: searchIndex.allResults,
-      searchTerm: debouncedSearchTerm,
-      typeFilter: selectedNodeTypeIds.length ? selectedNodeTypeIds : undefined,
-    });
+    setUnsortedScoredResults([]);
+    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]),
+    );
+
+    const runMiniSearch = (): ScoredSearchResult[] =>
+      searchDiscourseNodesWithMiniSearch({
+        miniSearch: searchIndex.miniSearch,
+        allResults: searchIndex.allResults,
+        searchTerm: debouncedSearchTerm,
+        typeFilter,
+      });
 
-    return sortSearchResults({ hits: scoredHits, sort });
+    void searchDiscourseNodes({
+      nodeTypes: discourseNodes,
+      query: debouncedSearchTerm,
+      resultsByUid,
+      runMiniSearch,
+    })
+      .then((results) => {
+        if (cancelled) return;
+        setUnsortedScoredResults(results);
+      })
+      .catch(() => {
+        if (cancelled) return;
+        setUnsortedScoredResults(runMiniSearch());
+      });
+
+    return () => {
+      cancelled = true;
+    };
   }, [
     debouncedSearchTerm,
-    dockedQuery,
-    dockedResults,
     indexError,
+    isDockedQuery,
     isIndexLoading,
     searchIndex,
     selectedNodeTypeIds,
-    sort,
   ]);
+
+  const results = useMemo(
+    () => sortSearchResults({ scoredResults: unsortedScoredResults, sort }),
+    [unsortedScoredResults, sort],
+  );
+
+  if (isDockedQuery && dockedResults) {
+    return dockedResults;
+  }
+
+  return results;
+};

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

@@ -9,6 +9,15 @@ import {
   getPulledDiscourseNodeUid,
   queryDiscourseNodesByFormat,
 } from "~/utils/discourseNodeSearch";
+import type {
+  ScoredSearchResult,
+  SearchResult,
+} from "~/utils/discourseNodeSearchTypes";
+
+export type {
+  ScoredSearchResult,
+  SearchResult,
+} from "~/utils/discourseNodeSearchTypes";
 
 export const DEBOUNCE_MS = 250;
 export const MAX_RESULTS = 50;
@@ -36,17 +45,6 @@ export const SORT_FIELD_LABELS: Record<SortField, string> = {
   author: "Author",
 };
 
-export type SearchResult = {
-  uid: string;
-  title: string;
-  type: string;
-  nodeTypeLabel: string;
-  excerpt: string;
-  createdAt: string;
-  lastModified: string;
-  authorName: string;
-};
-
 export type DockedSearchState = {
   query: string;
   results: SearchResult[];
@@ -56,11 +54,6 @@ export type DockedSearchState = {
   dgSearchId?: string;
 };
 
-export type ScoredSearchHit = {
-  result: SearchResult;
-  score: number;
-};
-
 type MiniSearchDocument = SearchResult & {
   id: string;
 };
@@ -214,22 +207,26 @@ export const isNonDefaultSort = (sort: SortConfig): boolean =>
   sort.direction !== DEFAULT_SORT_CONFIG.direction;
 
 export const sortSearchResults = ({
-  hits,
+  scoredResults,
   sort,
 }: {
-  hits: ScoredSearchHit[];
+  scoredResults: ScoredSearchResult[];
   sort: SortConfig;
 }): SearchResult[] => {
-  const sorted = [...hits];
+  const sorted = [...scoredResults];
 
-  sorted.sort((aHit, bHit) => {
-    const a = aHit.result;
-    const b = bHit.result;
+  sorted.sort((aEntry, bEntry) => {
+    const a = aEntry.result;
+    const b = bEntry.result;
     let comparison = 0;
 
     switch (sort.field) {
       case "relevance":
-        comparison = compareNumbers(aHit.score, bHit.score, sort.direction);
+        if (aEntry.source !== bEntry.source) {
+          comparison = aEntry.source === "semantic" ? -1 : 1;
+          break;
+        }
+        comparison = compareNumbers(aEntry.score, bEntry.score, sort.direction);
         break;
       case "alphabetical":
         comparison = compareStrings(
@@ -259,10 +256,10 @@ export const sortSearchResults = ({
     return comparison || a.uid.localeCompare(b.uid);
   });
 
-  return sorted.map((hit) => hit.result);
+  return sorted.map((entry) => entry.result);
 };
 
-export const searchIndexedNodes = ({
+export const searchDiscourseNodesWithMiniSearch = ({
   miniSearch,
   allResults,
   searchTerm,
@@ -272,7 +269,7 @@ export const searchIndexedNodes = ({
   allResults: SearchResult[];
   searchTerm: string;
   typeFilter?: string[];
-}): ScoredSearchHit[] => {
+}): ScoredSearchResult[] => {
   const resultsByUid = new Map(
     allResults.map((result) => [result.uid, result]),
   );
@@ -291,7 +288,11 @@ export const searchIndexedNodes = ({
     .map((result) => {
       const searchResult = resultsByUid.get(String(result.id));
       if (!searchResult) return null;
-      return { result: searchResult, score: result.score };
+      return {
+        result: searchResult,
+        score: result.score,
+        source: "miniSearch",
+      };
     })
-    .filter((hit): hit is ScoredSearchHit => !!hit);
+    .filter((entry): entry is ScoredSearchResult => !!entry);
 };

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/discourseNodeSearchTypes.ts

@@ -0,0 +1,73 @@
+export type DiscourseNodeSearchSource = "semantic" | "miniSearch";
+
+export type SearchResult = {
+  uid: string;
+  title: string;
+  type: string;
+  nodeTypeLabel: string;
+  excerpt: string;
+  createdAt: string;
+  lastModified: string;
+  authorName: string;
+};
+
+export type ScoredSearchResult = {
+  result: SearchResult;
+  score: number;
+  source: DiscourseNodeSearchSource;
+};
+
+export const combineSemanticAndMiniSearchResults = ({
+  semantic,
+  miniSearch,
+}: {
+  semantic: ScoredSearchResult[];
+  miniSearch: ScoredSearchResult[];
+}): ScoredSearchResult[] => {
+  const seenUids = new Set(semantic.map((entry) => entry.result.uid));
+  const combined = [...semantic];
+
+  miniSearch.forEach((entry) => {
+    if (seenUids.has(entry.result.uid)) return;
+    seenUids.add(entry.result.uid);
+    combined.push(entry);
+  });
+
+  return combined;
+};
+
+export const toScoredSearchResultFromSemantic = ({
+  uid,
+  title,
+  type,
+  nodeTypeLabel,
+  score,
+  resultsByUid,
+}: {
+  uid: string;
+  title: string;
+  type?: string;
+  nodeTypeLabel?: string;
+  score: number;
+  resultsByUid: Map<string, SearchResult>;
+}): ScoredSearchResult => {
+  const indexedResult = resultsByUid.get(uid);
+  if (indexedResult) {
+    return { result: indexedResult, score, source: "semantic" };
+  }
+
+  return {
+    result: {
+      uid,
+      title,
+      type: type || "",
+      nodeTypeLabel: nodeTypeLabel || "",
+      excerpt: "",
+      createdAt: "",
+      lastModified: "",
+      authorName: "Unknown",
+    },
+    score,
+    source: "semantic",
+  };
+};

apps/roam/src/utils/searchDiscourseNodes.ts

@@ -0,0 +1,69 @@
+import type { DiscourseNode } from "~/utils/getDiscourseNodes";
+import {
+  combineSemanticAndMiniSearchResults,
+  toScoredSearchResultFromSemantic,
+  type ScoredSearchResult,
+  type SearchResult,
+} from "~/utils/discourseNodeSearchTypes";
+import { runRoamSemanticSearch } from "~/utils/discourseNodeSearchProviders";
+
+export const SEMANTIC_SEARCH_MIN_DISCOURSE_RESULTS = 5;
+
+export type {
+  DiscourseNodeSearchSource,
+  ScoredSearchResult,
+  SearchResult,
+} from "~/utils/discourseNodeSearchTypes";
+
+export const isRoamSemanticSearchEnabled = (): boolean =>
+  window.roamAlphaAPI.data.semanticSearchEnabled();
+
+export const searchDiscourseNodes = async ({
+  nodeTypes,
+  query,
+  resultsByUid,
+  runMiniSearch,
+}: {
+  nodeTypes: DiscourseNode[];
+  query: string;
+  resultsByUid: Map<string, SearchResult>;
+  runMiniSearch: () => ScoredSearchResult[];
+}): Promise<ScoredSearchResult[]> => {
+  const trimmedQuery = query.trim();
+  if (!trimmedQuery) return [];
+
+  if (!isRoamSemanticSearchEnabled()) {
+    return runMiniSearch();
+  }
+
+  try {
+    const providerResult = await runRoamSemanticSearch({
+      nodeTypes,
+      query: trimmedQuery,
+    });
+    const semanticResults = providerResult.filteredResults.map((item) =>
+      toScoredSearchResultFromSemantic({
+        uid: item.uid,
+        title: item.text,
+        type: item.type,
+        nodeTypeLabel: item.nodeTypeLabel,
+        score: item.score ?? 0,
+        resultsByUid,
+      }),
+    );
+
+    if (
+      providerResult.filteredResultCount >=
+      SEMANTIC_SEARCH_MIN_DISCOURSE_RESULTS
+    ) {
+      return semanticResults;
+    }
+
+    return combineSemanticAndMiniSearchResults({
+      semantic: semanticResults,
+      miniSearch: runMiniSearch(),
+    });
+  } catch {
+    return runMiniSearch();
+  }
+};