From b2f28af93fcbc700069122d6c7d3d85601e53865 Mon Sep 17 00:00:00 2001
From: skeptrune <nick.k@trieve.ai>
Date: Mon, 20 May 2024 17:01:10 -0700
Subject: [PATCH] bugfix: checkbox for chat feature in search UI for group
 search

---
 search/src/components/ChatPopup.tsx       | 18 ++++++++++++++----
 search/src/components/GroupPage.tsx       |  2 +-
 search/src/components/ResultsPage.tsx     | 11 +++++++----
 search/src/components/SingleChunkPage.tsx |  2 +-
 4 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/search/src/components/ChatPopup.tsx b/search/src/components/ChatPopup.tsx
index 9e2369babd..165cec6e3d 100644
--- a/search/src/components/ChatPopup.tsx
+++ b/search/src/components/ChatPopup.tsx
@@ -13,6 +13,7 @@ import {
 } from "solid-js";
 import { FiSend, FiStopCircle } from "solid-icons/fi";
 import {
+  GroupScoreChunkDTO,
   type Message,
   messageRoleFromIndex,
   ScoreChunkDTO,
@@ -23,10 +24,11 @@ import { DatasetAndUserContext } from "./Contexts/DatasetAndUserContext";
 export interface LayoutProps {
   selectedIds: Accessor<string[]>;
   chunks: Accessor<ScoreChunkDTO[]>;
+  groupChunks?: Accessor<GroupScoreChunkDTO[]>;
   setOpenChat: Setter<boolean>;
 }
 
-const ChatPopup = (props: LayoutProps) => {
+export const ChatPopup = (props: LayoutProps) => {
   const api_host = import.meta.env.VITE_API_HOST as unknown as string;
   const datasetAndUserContext = useContext(DatasetAndUserContext);
   const $dataset = datasetAndUserContext.currentDataset;
@@ -166,7 +168,17 @@ const ChatPopup = (props: LayoutProps) => {
 
   const messageChunks = createMemo(() => {
     const selectedIds = props.selectedIds();
-    const chunks = props.chunks();
+    let chunks = props.chunks();
+    if (!chunks.length) {
+      chunks =
+        props
+          .groupChunks?.()
+          ?.flatMap((group) =>
+            group.metadata.flatMap((metadata) => metadata),
+          ) ?? [];
+
+      console.log("groups", chunks);
+    }
 
     return chunks.filter((chunk) => selectedIds.includes(chunk.metadata[0].id));
   });
@@ -273,5 +285,3 @@ const ChatPopup = (props: LayoutProps) => {
     </div>
   );
 };
-
-export default ChatPopup;
diff --git a/search/src/components/GroupPage.tsx b/search/src/components/GroupPage.tsx
index bd15ca7487..17dc7d842c 100644
--- a/search/src/components/GroupPage.tsx
+++ b/search/src/components/GroupPage.tsx
@@ -28,7 +28,7 @@ import { ScoreChunkArray } from "./ScoreChunkArray";
 import SearchForm from "./SearchForm";
 import ChunkMetadataDisplay from "./ChunkMetadataDisplay";
 import { Portal } from "solid-js/web";
-import ChatPopup from "./ChatPopup";
+import { ChatPopup } from "./ChatPopup";
 import { AiOutlineRobot } from "solid-icons/ai";
 import { IoDocumentOutline, IoDocumentsOutline } from "solid-icons/io";
 import { useLocation, useNavigate } from "@solidjs/router";
diff --git a/search/src/components/ResultsPage.tsx b/search/src/components/ResultsPage.tsx
index f0b2df42c7..f82582acc1 100644
--- a/search/src/components/ResultsPage.tsx
+++ b/search/src/components/ResultsPage.tsx
@@ -25,7 +25,7 @@ import { ConfirmModal } from "./Atoms/ConfirmModal";
 import { ScoreChunkArray } from "./ScoreChunkArray";
 import { Portal } from "solid-js/web";
 import { AiOutlineRobot } from "solid-icons/ai";
-import ChatPopup from "./ChatPopup";
+import { ChatPopup } from "./ChatPopup";
 import { IoDocumentOutline, IoDocumentsOutline } from "solid-icons/io";
 import { DatasetAndUserContext } from "./Contexts/DatasetAndUserContext";
 import { FaSolidChevronDown, FaSolidChevronUp } from "solid-icons/fa";
@@ -283,6 +283,7 @@ const ResultsPage = (props: ResultsPageProps) => {
             <div class="max-h-[75vh] min-h-[75vh] min-w-[75vw] max-w-[75vw] overflow-y-auto rounded-md scrollbar-thin">
               <ChatPopup
                 chunks={resultChunks}
+                groupChunks={groupResultChunks}
                 selectedIds={selectedIds}
                 setOpenChat={setOpenChat}
               />
@@ -501,11 +502,13 @@ const ResultsPage = (props: ResultsPageProps) => {
                     class="relative h-[52px] w-[52px] items-center justify-center rounded-lg border border-gray-200 bg-white text-gray-500 shadow-sm hover:bg-gray-50 hover:text-gray-900 focus:outline-none focus:ring-4 focus:ring-gray-300 dark:border-gray-600 dark:bg-gray-700 dark:text-gray-400 dark:hover:bg-gray-600 dark:hover:text-white dark:focus:ring-gray-400"
                     onClick={() => {
                       setSelectedIds(
-                        resultChunks()
+                        groupResultChunks()
+                          .flatMap((g) => {
+                            return g.metadata;
+                          })
                           .flatMap((c) => {
                             return c.metadata.map((m) => m.id);
-                          })
-                          .slice(0, 10),
+                          }),
                       );
                       setOpenChat(true);
                     }}
diff --git a/search/src/components/SingleChunkPage.tsx b/search/src/components/SingleChunkPage.tsx
index f5f4279987..cd1e57b477 100644
--- a/search/src/components/SingleChunkPage.tsx
+++ b/search/src/components/SingleChunkPage.tsx
@@ -20,7 +20,7 @@ import { FullScreenModal } from "./Atoms/FullScreenModal";
 import { ConfirmModal } from "./Atoms/ConfirmModal";
 import ChunkMetadataDisplay from "./ChunkMetadataDisplay";
 import { Portal } from "solid-js/web";
-import ChatPopup from "./ChatPopup";
+import { ChatPopup } from "./ChatPopup";
 import { AiOutlineRobot } from "solid-icons/ai";
 import { IoDocumentOutline } from "solid-icons/io";
 import { DatasetAndUserContext } from "./Contexts/DatasetAndUserContext";