feature: index sort keys, metadata, and tag_set. Setup tag_set filtering

for the search component
devflowinc · Dec 19, 2024 · 2fff219 · 2fff219
1 parent 031fb4d
commit 2fff219
Show file tree

Hide file tree

Showing 7 changed files with 180 additions and 67 deletions.
diff --git a/clients/search-component/example/.env.dist b/clients/search-component/example/.env.dist
@@ -9,3 +9,5 @@ VITE_SHOW_FLOATING_BTN="true"
 VITE_FLOATING_BTN_POSITION="bottom-right"
 VITE_USE_PAGEFIND="false"
 VITE_PAGEFIND_URL="https://pagefind-testing-index.trieve.ai/pagfind-index-west/pagefind"
+VITE_USE_GROUP_SEARCH="false"
+VITE_DEFAULT_TAGS='[{"tag": "pink", "label": "Pink"}]'
diff --git a/clients/search-component/example/index.html b/clients/search-component/example/index.html
@@ -8,13 +8,13 @@
     <link rel="icon" type="image/png" sizes="32x32" href="https://cdn.trieve.ai/favicon-32x32.png" />
     <link rel="icon" type="image/png" sizes="16x16" href="https://cdn.trieve.ai/favicon-16x16.png" />
     <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-    <meta property="og:title" content="Trieve Drop-In RAG and Search Component">
+    <meta property="og:title" content="Search Component By Trieve">
     <meta property="og:site_name" content="advanced relevance semantic search and RAG API">
     <meta property="og:url" content="https://docsearch.trieve.ai">
     <meta property="og:description" content="Build better, faster, and more relevant search and RAG with our open source API. Date recency biasing, re-ranker models, dense vector search, sub-sentence highlighting, and more all on one endpoint that you can host yourself." >
     <meta property="og:type" content="">
     <meta property="og:image" content="https://cdn.trieve.ai/trieve-og.png">
-    <title>Trieve Drop-In RAG and Search Component</title>
+    <title>Search Component By Trieve</title>
   </head>
   <body>
     <div id="root"></div>

diff --git a/clients/search-component/example/src/routes/ecommerce.tsx b/clients/search-component/example/src/routes/ecommerce.tsx
@@ -19,6 +19,8 @@ export default function ECommerce() {
   const useGroupSearch = import.meta.env.VITE_USE_GROUP_SEARCH == "true";
   const showFloatingButton = import.meta.env.VITE_SHOW_FLOATING_BTN == "true";
   const floatingButtonPosition = import.meta.env.VITE_FLOATING_BTN_POSITION;
+  const usePagefind = import.meta.env.VITE_USE_PAGEFIND == "true";
+  const pagefindUrl = import.meta.env.VITE_PAGEFIND_URL;
 
   const defaultSearchQueries: string[] = (
     import.meta.env.VITE_DEFAULT_SEARCH_QUERIES ?? ""
@@ -33,9 +35,8 @@ export default function ECommerce() {
   return (
     <>
       <div
-        className={`p-12 flex flex-col items-center justify-center w-screen h-screen relative ${
-          theme === "dark" ? "bg-zinc-900 text-zinc-50" : ""
-        }`}
+        className={`p-12 flex flex-col items-center justify-center w-screen h-screen relative ${theme === "dark" ? "bg-zinc-900 text-zinc-50" : ""
+          }`}
       >
         <div className="absolute top-6 right-6">
           <ul>
@@ -81,6 +82,16 @@ export default function ECommerce() {
                 use_autocomplete: false,
                 search_type: "fulltext",
               }}
+              buttonTriggers={[
+                {
+                  selector: ".random-trigger-location",
+                  mode: "chat",
+                },
+              ]}
+              pagefindOptions={{
+                usePagefind: usePagefind,
+                cdnBaseUrl: pagefindUrl,
+              }}
               defaultSearchQueries={defaultSearchQueries}
               tags={defaultTags}
               floatingButtonPosition={floatingButtonPosition}

diff --git a/clients/search-component/src/utils/hooks/modal-context.tsx b/clients/search-component/src/utils/hooks/modal-context.tsx
@@ -14,6 +14,7 @@ import {
 } from "trieve-ts-sdk";
 import {
   countChunks,
+  countChunksWithPagefind,
   groupSearchWithTrieve,
   searchWithPagefind,
   searchWithTrieve,
@@ -110,7 +111,7 @@ const defaultProps = {
   analytics: true,
   chat: true,
   suggestedQueries: true,
-  trieve: (() => {}) as unknown as TrieveSDK,
+  trieve: (() => { }) as unknown as TrieveSDK,
   openKeyCombination: [{ ctrl: true }, { key: "k", label: "K" }],
   type: "docs" as ModalTypes,
   useGroupSearch: false,
@@ -244,7 +245,12 @@ const ModalProvider = ({
         setResults(Array.from(groupMap.values()));
         setRequestID(results.requestID);
       } else if (props.pagefindOptions?.usePagefind) {
-        const results = await searchWithPagefind(query, props.pagefindOptions, props.datasetId);
+        const results = await searchWithPagefind(
+          query,
+          props.pagefindOptions,
+          props.datasetId,
+          currentTag !== "all" ? currentTag : undefined
+        );
         setResults(results);
       } else {
         const results = await searchWithTrieve({
@@ -276,24 +282,35 @@ const ModalProvider = ({
       return;
     }
     if (props.tags?.length) {
-      try {
-        const numberOfRecords = await Promise.all(
-          [ALL_TAG, ...props.tags].map((tag) =>
-            countChunks({
-              query: query,
-              trieve: trieve,
-              abortController,
-              ...(tag.tag !== "all" && { tag: tag.tag }),
-            })
-          )
+      if (props.pagefindOptions?.usePagefind) {
+        let filterCounts = await countChunksWithPagefind(
+          query,
+          props.pagefindOptions,
+          props.datasetId,
+          props.tags
         );
-        setTagCounts(numberOfRecords);
-      } catch (e) {
-        if (
-          e != "AbortError" &&
-          e != "AbortError: signal is aborted without reason"
-        ) {
-          console.error(e);
+        console.log("fildas", filterCounts, props.tags);
+        setTagCounts(filterCounts);
+      } else {
+        try {
+          const numberOfRecords = await Promise.all(
+            [ALL_TAG, ...props.tags].map((tag) =>
+              countChunks({
+                query: query,
+                trieve: trieve,
+                abortController,
+                ...(tag.tag !== "all" && { tag: tag.tag }),
+              })
+            )
+          );
+          setTagCounts(numberOfRecords);
+        } catch (e) {
+          if (
+            e != "AbortError" &&
+            e != "AbortError: signal is aborted without reason"
+          ) {
+            console.error(e);
+          }
         }
       }
     }

diff --git a/clients/search-component/src/utils/trieve.ts b/clients/search-component/src/utils/trieve.ts
@@ -1,6 +1,7 @@
 import {
   ChunkMetadata,
   ChunkMetadataStringTagSet,
+  CountChunkQueryResponseBody,
   SearchResponseBody,
   TrieveSDK,
 } from "trieve-ts-sdk";
@@ -36,7 +37,7 @@ export const searchWithTrieve = async ({
   const scoreThreshold =
     searchOptions.score_threshold ??
     ((searchOptions.search_type ?? "fulltext") === "fulltext" ||
-    searchOptions.search_type == "bm25"
+      searchOptions.search_type == "bm25"
       ? 2
       : 0.3);
 
@@ -129,7 +130,7 @@ export const groupSearchWithTrieve = async ({
   const scoreThreshold =
     searchOptions.score_threshold ??
     ((searchOptions.search_type ?? "fulltext") === "fulltext" ||
-    searchOptions.search_type == "bm25"
+      searchOptions.search_type == "bm25"
       ? 2
       : 0.3);
 
@@ -191,7 +192,7 @@ export const countChunks = async ({
   const scoreThreshold =
     searchOptions?.score_threshold ??
     ((searchOptions?.search_type ?? "fulltext") === "fulltext" ||
-    searchOptions?.search_type == "bm25"
+      searchOptions?.search_type == "bm25"
       ? 2
       : 0.3);
 
@@ -310,11 +311,19 @@ export const searchWithPagefind = async (
   query: string,
   pagefindOptions: PagefindOptions,
   datasetId: string,
+  tag?: string,
 ) => {
   const pagefind_base_url = `${pagefindOptions.cdnBaseUrl}/${datasetId}`;
 
   const pagefind = await import(`${pagefind_base_url}/pagefind.js`);
-  const response = await pagefind.search(query);
+  console.log("searchWithPagefind B");
+  const response = await pagefind.search(query,
+    tag && {
+      filters: {
+        tag_set: tag
+      }
+    }
+  );
 
   const results = await Promise.all(response.results.map(async (result: any) => {
     return await result.data();
@@ -330,10 +339,10 @@ export const searchWithPagefind = async (
         created_at: "",
         dataset_id: datasetId,
         id: i.toString(),
-        image_urls: null,
+        image_urls: result.meta.image_urls.split(", "),
         location: null,
         num_value: null,
-        tag_set: null,
+        tag_set: result.meta.tag_set,
         time_stamp: null,
         tracking_id: null,
         updated_at: "",
@@ -346,3 +355,42 @@ export const searchWithPagefind = async (
 
   return pagefindResultsMappedToTrieve;
 };
+
+export const countChunksWithPagefind = async (
+  query: string,
+  pagefindOptions: PagefindOptions,
+  datasetId: string,
+  tags: {
+    tag: string;
+    label?: string;
+    selected?: boolean;
+    iconClassName?: string;
+    icon?: () => JSX.Element;
+  }[]
+): Promise<CountChunkQueryResponseBody[]> => {
+  const pagefind_base_url = `${pagefindOptions.cdnBaseUrl}/${datasetId}`;
+
+  const pagefind = await import(`${pagefind_base_url}/pagefind.js`);
+  console.log("filters load");
+  console.log("filters");
+
+  await pagefind.filters();
+  const response = await pagefind.search(query);
+
+  const tag_set = response.filters.tag_set;
+  console.log("tag_set", response);
+
+  const counts: CountChunkQueryResponseBody[] = tags.map((tag) => {
+    if (tag.tag in tag_set) {
+      return {
+        count: tag_set[tag.tag] as number,
+      };
+    }
+  })
+
+  counts.unshift({
+    count: response.unfilteredResultCount as number,
+  });
+
+  return counts;
+};
diff --git a/server/src/operators/clickhouse_operator.rs b/server/src/operators/clickhouse_operator.rs
@@ -200,7 +200,9 @@ impl EventQueue {
 
     pub async fn send(&self, event: ClickHouseEvent) {
         if let Some(sender) = &self.sender {
-            sender.send(event).await.unwrap();
+            let _ = sender.send(event).await.map_err(|e| {
+                log::error!("Error sending event to clickhouse: {:?}", e);
+            });
         }
     }
 }