diff --git a/src/renderer/components/Experiment/Generate/GenerateModal.tsx b/src/renderer/components/Experiment/Generate/GenerateModal.tsx
index c30c45bb4bf76ccd1ac0a6fb2973d51d38d39e38..fd19420ada37c7e5010d139c8218ade4c18533f4 100644
--- a/src/renderer/components/Experiment/Generate/GenerateModal.tsx
+++ b/src/renderer/components/Experiment/Generate/GenerateModal.tsx
@@ -22,6 +22,7 @@ import {
} from '@mui/joy';
import DynamicPluginForm from '../DynamicPluginForm';
import TrainingModalDataTab from '../Train/TraningModalDataTab';
+import PickADocumentMenu from '../Rag/PickADocumentMenu';
import { generateFriendlyName } from 'renderer/lib/utils';
import exp from 'node:constants';
@@ -70,7 +71,8 @@ export default function GenerateModal({
const [hasDatasetKey, setHasDatasetKey] = useState(false);
const [hasDocumentsKey, setHasDocumentsKey] = useState(false);
const [hasContextKey, setHasContextKey] = useState(false);
- const [selectedFiles, setSelectedFiles] = useState<File[]>([]);
+ const [selectedFiles, setSelectedFiles] = useState<string[]>([]);
+ const [selectedFileNames, setSelectedFileNames] = useState<string[]>([]);
const [nameInput, setNameInput] = useState('');
const [currentTab, setCurrentTab] = useState(0);
const [contextInput, setContextInput] = useState('');
@@ -110,8 +112,11 @@ export default function GenerateModal({
useEffect(() => {
if (open) {
+ setSelectedFiles([]);
+ setSelectedFileNames([]);
if (!currentEvalName || currentEvalName === '') {
setNameInput(generateFriendlyName());
+
}
}
}, [open]);
@@ -148,8 +153,10 @@ export default function GenerateModal({
if (docsKeyExists && evalConfig.script_parameters.docs.length > 0) {
// const docstemp = evalConfig.script_parameters.docs.split(',').map((path) => ({ path }));
const docPaths = evalConfig.script_parameters.docs.split(',');
- const docFiles = docPaths.map((path) => new File([], path));
- setSelectedFiles(docFiles);
+ const docNames = evalConfig.script_parameters.doc_names.split(',');
+ // const docFiles = docPaths.map((path) => new File([], path));
+ setSelectedFiles(docPaths);
+ setSelectedFileNames(docNames);
delete evalConfig.script_parameters.docs;
setConfig(evalConfig.script_parameters)
@@ -292,41 +299,72 @@ export default function GenerateModal({
);
}
- function DocsTab() {
- const handleFileChange = (event: React.ChangeEvent<HTMLInputElement>) => {
- if (event.target.files) {
- setSelectedFiles(Array.from(event.target.files));
- }
- };
+ // function DocsTab() {
+ // const handleFileChange = (event: React.ChangeEvent<HTMLInputElement>) => {
+ // if (event.target.files) {
+ // setSelectedFiles(Array.from(event.target.files));
+ // }
+ // };
+
+ // return (
+ // <Stack spacing={2}>
+ // <FormControl>
+ // <FormLabel>Upload Documents</FormLabel>
+ // <Input
+ // type="file"
+ // multiple={true}
+ // onChange={handleFileChange}
+ // name="docs"
+ // />
+ // <FormHelperText>
+ // Select multiple documents to upload
+ // </FormHelperText>
+ // </FormControl>
+ // {selectedFiles.length > 0 && (
+ // <Stack spacing={1} mt={2}>
+ // <FormLabel>Selected Documents:</FormLabel>
+ // {selectedFiles.map((file, index) => (
+ // <Sheet key={index} variant="outlined" p={1}>
+ // {file.name}
+ // </Sheet>
+ // ))}
+ // </Stack>
+ // )}
+ // </Stack>
+ // );
+ // }
+
+ function DocsTab({ experimentInfo }) {
return (
<Stack spacing={2}>
<FormControl>
- <FormLabel>Upload Documents</FormLabel>
- <Input
- type="file"
- multiple={true}
- onChange={handleFileChange}
- name="docs"
+ <FormLabel>Pick Documents</FormLabel>
+ <PickADocumentMenu
+ experimentInfo={experimentInfo}
+ showFoldersOnly={false}
+ setSelectedFiles={setSelectedFiles}
+ setSelectedFileNames={setSelectedFileNames}
/>
- <FormHelperText>
- Select multiple documents to upload
- </FormHelperText>
+ <FormHelperText>Select documents to upload</FormHelperText>
</FormControl>
- {selectedFiles.length > 0 && (
- <Stack spacing={1} mt={2}>
- <FormLabel>Selected Documents:</FormLabel>
- {selectedFiles.map((file, index) => (
- <Sheet key={index} variant="outlined" p={1}>
- {file.name}
- </Sheet>
- ))}
- </Stack>
- )}
+ {selectedFileNames.length > 0 && (
+ <Stack spacing={1} mt={2}>
+ <FormLabel>Selected Documents:</FormLabel>
+ {selectedFileNames.map((file, index) => (
+ <Sheet key={index} variant="outlined" p={1}>
+ {file}
+ </Sheet>
+ ))}
+ </Stack>
+ )}
</Stack>
);
}
+
+
+
function ContextTab({contextInput, setContextInput}) {
return (
@@ -361,8 +399,13 @@ export default function GenerateModal({
formJson.run_name = formJson.template_name;
}
// Add the selected file paths to the formJson as comma separated string
+ // if (hasDocumentsKey && selectedFiles.length > 0) {
+ // formJson.docs = selectedFiles.map((file) => file.path).join(',');
+ // formJson.generation_type = 'docs';
+ // }
if (hasDocumentsKey && selectedFiles.length > 0) {
- formJson.docs = selectedFiles.map((file) => file.path).join(',');
+ formJson.docs = selectedFiles.join(',');;
+ formJson.doc_names = selectedFileNames.join(',');
formJson.generation_type = 'docs';
}
// Add context to the formJson
@@ -386,6 +429,7 @@ export default function GenerateModal({
setNameInput(generateFriendlyName());
setContextInput('');
setSelectedFiles([]);
+ setSelectedFileNames([]);
} else {
const template_name = formJson.template_name;
delete formJson.template_name;
@@ -471,8 +515,14 @@ export default function GenerateModal({
</TabPanel>
{hasDocumentsKey && (
<TabPanel value={3} sx={{ p: 2, overflow: 'auto' }} keepMounted>
- <DocsTab />
+ <DocsTab
+ experimentInfo={experimentInfo}
+ />
+ {/* <PickADocumentMenu
+ experimentInfo={experimentInfo}
+ /> */}
</TabPanel>
+ // <DocsTab />
)}
{hasContextKey && (
<TabPanel value={3} sx={{ p: 2, overflow: 'auto' }} keepMounted>
diff --git a/src/renderer/components/Experiment/Rag/PickADocumentMenu.tsx b/src/renderer/components/Experiment/Rag/PickADocumentMenu.tsx
index 47b10356d74132c96d4f8546a89a542d7ab47e4b..bd1919ac01480594865f6e361900db4e80cd8fbb 100644
--- a/src/renderer/components/Experiment/Rag/PickADocumentMenu.tsx
+++ b/src/renderer/components/Experiment/Rag/PickADocumentMenu.tsx
@@ -1,59 +1,173 @@
-import {
- Dropdown,
- List,
- ListItem,
- Menu,
- MenuButton,
- MenuItem,
- MenuList,
- Typography,
-} from '@mui/joy';
-import React, { useState, useEffect } from 'react';
+import React, { useState } from 'react';
+import { Box, Typography } from '@mui/joy';
import * as chatAPI from '../../../lib/transformerlab-api-sdk';
import useSWR from 'swr';
-import { FolderIcon } from 'lucide-react';
-const fetcher = (url) => fetch(url).then((res) => res.json());
+import { FolderIcon, Check } from 'lucide-react';
+
+const fetcher = (url: string) => fetch(url).then((res) => res.json());
+
+interface FolderChildrenProps {
+ experimentId: string;
+ folderId: string;
+ toggleSelect: (id: string, filePath: string, fileName: string) => void;
+ selectedIds: Set<string>;
+}
+
+function FolderChildren({ experimentId, folderId, toggleSelect, selectedIds }: FolderChildrenProps) {
+ const { data, isLoading } = useSWR(
+ chatAPI.Endpoints.Documents.List(experimentId, folderId),
+ fetcher
+ );
+
+ if (isLoading) {
+ return <Typography sx={{ ml: 4 }}>Loading...</Typography>;
+ }
+ if (!data || data.length === 0) {
+ return <Typography sx={{ ml: 4 }}>No documents</Typography>;
+ }
+ return (
+ <Box sx={{ ml: 4 }}>
+ {data.map((child: any, idx: number) => {
+ const childId = child.id ? child.id.toString() : `child-index-${idx}`;
+ const isFolder = child?.type === 'folder';
+ return (
+ <Box
+ key={childId}
+ onClick={!isFolder ? () => toggleSelect(childId, child?.path || '', child?.name || 'Unknown') : undefined}
+ sx={{
+ display: 'flex',
+ alignItems: 'center',
+ p: 1,
+ my: 0.5,
+ borderRadius: 'sm',
+ cursor: !isFolder ? 'pointer' : 'default',
+ backgroundColor: !isFolder && selectedIds.has(childId)
+ ? 'primary.softHoverBg'
+ : 'transparent',
+ '&:hover': !isFolder ? { backgroundColor: 'primary.softHoverBg' } : undefined,
+ }}
+ >
+ {isFolder && <FolderIcon size="14px" />}
+ <Typography ml={isFolder ? 1 : 0}>
+ {child?.name || 'Unnamed'}
+ </Typography>
+ {!isFolder && selectedIds.has(childId) && (
+ <Check size="16px" style={{ marginLeft: 'auto', color: 'green' }} />
+ )}
+ </Box>
+ );
+ })}
+ </Box>
+ );
+}
+
+interface PickADocumentMenuProps {
+ experimentInfo: any;
+ showFoldersOnly?: boolean;
+ setSelectedFiles: React.Dispatch<React.SetStateAction<string[]>>;
+ setSelectedFileNames: React.Dispatch<React.SetStateAction<string[]>>;
+}
export default function PickADocumentMenu({
experimentInfo,
showFoldersOnly = false,
-}) {
- const {
- data: rows,
- isLoading,
- mutate,
- } = useSWR(chatAPI.Endpoints.Documents.List(experimentInfo?.id, ''), fetcher);
+ setSelectedFiles,
+ setSelectedFileNames,
+}: PickADocumentMenuProps) {
+ const { data: rows, isLoading } = useSWR(
+ chatAPI.Endpoints.Documents.List(experimentInfo?.id, ''),
+ fetcher
+ );
+
+ // State for expanded folders (by unique id)
+ const [expandedFolders, setExpandedFolders] = useState<Set<string>>(new Set());
+
+ // State for selected non-folder items (by unique id)
+ const [selectedIds, setSelectedIds] = useState<Set<string>>(new Set());
+
+ const toggleExpand = (id: string) => {
+ const newExpanded = new Set(expandedFolders);
+ if (newExpanded.has(id)) {
+ newExpanded.delete(id);
+ } else {
+ newExpanded.add(id);
+ }
+ setExpandedFolders(newExpanded);
+ };
+
+ const toggleSelect = (id: string, filePath: string, fileName: string) => {
+ const isSelected = selectedIds.has(id);
+ const newSelected = new Set(selectedIds);
+ if (isSelected) {
+ newSelected.delete(id);
+ setSelectedFiles((prevFiles) => prevFiles.filter((path) => path !== filePath));
+ setSelectedFileNames((prevNames) => prevNames.filter((name) => name !== fileName));
+ } else {
+ newSelected.add(id);
+ setSelectedFiles((prevFiles) => [...prevFiles, filePath]);
+ setSelectedFileNames((prevNames) => [...prevNames, fileName]);
+ }
+ setSelectedIds(newSelected);
+ };
return (
- <Dropdown>
- <MenuButton>Pick {showFoldersOnly ? 'Folder' : 'File'}</MenuButton>
- <Menu>
- {isLoading ? (
- <MenuItem>Loading...</MenuItem>
- ) : (
- rows?.map((row) =>
- showFoldersOnly ? (
- row?.type == 'folder' && (
- <MenuItem key={row.id} onClick={() => console.log(row)}>
- <Typography sx={{ display: 'flex', alignItems: 'center' }}>
- {row?.type == 'folder' ? <FolderIcon size="14px" /> : null}
-
- {row.name}
- </Typography>
- </MenuItem>
- )
- ) : (
- <MenuItem key={row.id} onClick={() => console.log(row)}>
- <Typography sx={{ display: 'flex', alignItems: 'center' }}>
- {row?.type == 'folder' ? <FolderIcon size="14px" /> : null}
-
- {row.name}
+ <Box sx={{ border: '1px solid', borderColor: 'neutral.outlinedBorder', borderRadius: 'sm', p: 1 }}>
+ <Typography level="h6" mb={1}>
+ Pick {showFoldersOnly ? 'Folder' : 'File'}
+ </Typography>
+ {isLoading ? (
+ <Typography>Loading...</Typography>
+ ) : (
+ rows?.map((row: any, index: number) => {
+ // Use row.id if available; fallback to index.
+ const uniqueId = row.name ? row.name.toString() : `index-${index}`;
+ const isFolder = row?.type === 'folder';
+ return (
+ <Box key={uniqueId}>
+ <Box
+ onClick={() => {
+ if (isFolder) {
+ toggleExpand(uniqueId);
+ } else {
+ // We use row.path for file identifier; adjust if necessary.
+ toggleSelect(uniqueId, row?.path || '', row?.name || '');
+ }
+ }}
+ sx={{
+ display: 'flex',
+ alignItems: 'center',
+ p: 1,
+ my: 0.5,
+ borderRadius: 'sm',
+ cursor: 'pointer',
+ backgroundColor: isFolder
+ ? expandedFolders.has(uniqueId)
+ ? 'primary.softHoverBg'
+ : 'transparent'
+ : selectedIds.has(uniqueId)
+ ? 'primary.softHoverBg'
+ : 'transparent',
+ '&:hover': { backgroundColor: 'primary.softHoverBg' },
+ }}
+ >
+ {isFolder && <FolderIcon size="14px" />}
+ <Typography ml={isFolder ? 1 : 0}>
+ {row?.name || 'Unnamed'}
</Typography>
- </MenuItem>
- )
- )
- )}
- </Menu>
- </Dropdown>
+ {!isFolder && selectedIds.has(uniqueId) && (
+ <Check size="16px" style={{ marginLeft: 'auto', color: 'green' }} />
+ )}
+ </Box>
+ {isFolder && expandedFolders.has(uniqueId) && (
+ <FolderChildren experimentId={experimentInfo?.id} folderId={uniqueId} toggleSelect={toggleSelect} selectedIds={selectedIds} />
+ )}
+ </Box>
+ );
+ })
+ )}
+ {!isLoading && (!rows || rows.length === 0) && (
+ <Typography>No documents found</Typography>
+ )}
+ </Box>
);
}