501 lines
25 KiB
JavaScript
501 lines
25 KiB
JavaScript
import { useEffect, useState, useRef } from 'react'
|
|
import { useDispatch, useSelector } from 'react-redux'
|
|
import { useNavigate, useParams } from 'react-router-dom'
|
|
import ReactJson from 'flowise-react-json-view'
|
|
import { cloneDeep } from 'lodash'
|
|
import { v4 as uuidv4 } from 'uuid'
|
|
|
|
// material-ui
|
|
import { Box, Card, Grid, Stack, Typography, OutlinedInput, IconButton, Button } from '@mui/material'
|
|
import Embeddings from '@mui/icons-material/DynamicFeed'
|
|
import { useTheme, styled } from '@mui/material/styles'
|
|
import CardContent from '@mui/material/CardContent'
|
|
import chunks_emptySVG from '@/assets/images/chunks_empty.svg'
|
|
import { IconSearch, IconFileStack, IconDeviceFloppy, IconX } from '@tabler/icons-react'
|
|
|
|
// project imports
|
|
import MainCard from '@/ui-component/cards/MainCard'
|
|
import { BackdropLoader } from '@/ui-component/loading/BackdropLoader'
|
|
import ConfirmDialog from '@/ui-component/dialog/ConfirmDialog'
|
|
import ExpandedChunkDialog from './ExpandedChunkDialog'
|
|
import ViewHeader from '@/layout/MainLayout/ViewHeader'
|
|
import DocStoreInputHandler from '@/views/docstore/DocStoreInputHandler'
|
|
import { PermissionButton } from '@/ui-component/button/RBACButtons'
|
|
|
|
// API
|
|
import documentsApi from '@/api/documentstore'
|
|
import nodesApi from '@/api/nodes'
|
|
|
|
// Hooks
|
|
import useApi from '@/hooks/useApi'
|
|
import { useAuth } from '@/hooks/useAuth'
|
|
import useNotifier from '@/utils/useNotifier'
|
|
import { baseURL } from '@/store/constant'
|
|
import { initNode, showHideInputParams } from '@/utils/genericHelper'
|
|
import { closeSnackbar as closeSnackbarAction, enqueueSnackbar as enqueueSnackbarAction } from '@/store/actions'
|
|
|
|
const CardWrapper = styled(MainCard)(({ theme }) => ({
|
|
background: theme.palette.card.main,
|
|
color: theme.darkTextPrimary,
|
|
overflow: 'auto',
|
|
position: 'relative',
|
|
boxShadow: '0 2px 14px 0 rgb(32 40 45 / 8%)',
|
|
cursor: 'pointer',
|
|
'&:hover': {
|
|
background: theme.palette.card.hover,
|
|
boxShadow: '0 2px 14px 0 rgb(32 40 45 / 20%)'
|
|
},
|
|
maxHeight: '250px',
|
|
minHeight: '250px',
|
|
maxWidth: '100%',
|
|
overflowWrap: 'break-word',
|
|
whiteSpace: 'pre-line',
|
|
padding: 1
|
|
}))
|
|
|
|
const VectorStoreQuery = () => {
|
|
const customization = useSelector((state) => state.customization)
|
|
const navigate = useNavigate()
|
|
const theme = useTheme()
|
|
const dispatch = useDispatch()
|
|
const inputRef = useRef(null)
|
|
const { hasAssignedWorkspace } = useAuth()
|
|
|
|
useNotifier()
|
|
|
|
const enqueueSnackbar = (...args) => dispatch(enqueueSnackbarAction(...args))
|
|
const closeSnackbar = (...args) => dispatch(closeSnackbarAction(...args))
|
|
|
|
const { storeId } = useParams()
|
|
|
|
const [documentChunks, setDocumentChunks] = useState([])
|
|
const [loading, setLoading] = useState(false)
|
|
const [showExpandedChunkDialog, setShowExpandedChunkDialog] = useState(false)
|
|
const [expandedChunkDialogProps, setExpandedChunkDialogProps] = useState({})
|
|
const [documentStore, setDocumentStore] = useState({})
|
|
const [query, setQuery] = useState('')
|
|
|
|
const [timeTaken, setTimeTaken] = useState(-1)
|
|
const [retrievalError, setRetrievalError] = useState(undefined)
|
|
|
|
const getSpecificDocumentStoreApi = useApi(documentsApi.getSpecificDocumentStore)
|
|
const queryVectorStoreApi = useApi(documentsApi.queryVectorStore)
|
|
|
|
const getVectorStoreNodeDetailsApi = useApi(nodesApi.getSpecificNode)
|
|
const [selectedVectorStoreProvider, setSelectedVectorStoreProvider] = useState({})
|
|
|
|
const handleVectorStoreProviderDataChange = ({ inputParam, newValue }) => {
|
|
setSelectedVectorStoreProvider((prevData) => {
|
|
const updatedData = { ...prevData }
|
|
updatedData.inputs[inputParam.name] = newValue
|
|
updatedData.inputParams = showHideInputParams(updatedData)
|
|
return updatedData
|
|
})
|
|
}
|
|
|
|
const chunkSelected = (chunkId, selectedChunkNumber) => {
|
|
const selectedChunk = documentChunks.find((chunk) => chunk.id === chunkId)
|
|
const dialogProps = {
|
|
data: {
|
|
selectedChunk,
|
|
selectedChunkNumber
|
|
}
|
|
}
|
|
setExpandedChunkDialogProps(dialogProps)
|
|
setShowExpandedChunkDialog(true)
|
|
}
|
|
|
|
const handleEnter = (e) => {
|
|
// Check if IME composition is in progress
|
|
const isIMEComposition = e.isComposing || e.keyCode === 229
|
|
if (e.key === 'Enter' && query && !isIMEComposition) {
|
|
if (!e.shiftKey && query) {
|
|
if (inputRef.current) {
|
|
inputRef.current.blur()
|
|
}
|
|
doQuery()
|
|
}
|
|
} else if (e.key === 'Enter') {
|
|
e.preventDefault()
|
|
}
|
|
}
|
|
|
|
const doQuery = () => {
|
|
setLoading(true)
|
|
const data = {
|
|
query: query,
|
|
storeId: storeId,
|
|
inputs: selectedVectorStoreProvider.inputs
|
|
}
|
|
queryVectorStoreApi.request(data)
|
|
}
|
|
|
|
const saveConfig = async () => {
|
|
setLoading(true)
|
|
const data = {
|
|
storeId: storeId
|
|
}
|
|
|
|
if (selectedVectorStoreProvider.inputs) {
|
|
data.vectorStoreConfig = {}
|
|
data.vectorStoreName = selectedVectorStoreProvider.name
|
|
Object.keys(selectedVectorStoreProvider.inputs).map((key) => {
|
|
if (key === 'FLOWISE_CREDENTIAL_ID') {
|
|
data.vectorStoreConfig['credential'] = selectedVectorStoreProvider.inputs[key]
|
|
} else {
|
|
data.vectorStoreConfig[key] = selectedVectorStoreProvider.inputs[key]
|
|
}
|
|
})
|
|
}
|
|
|
|
try {
|
|
const updateResp = await documentsApi.updateVectorStoreConfig(data)
|
|
setLoading(false)
|
|
if (updateResp.data) {
|
|
enqueueSnackbar({
|
|
message: 'Vector Store Config Successfully Updated',
|
|
options: {
|
|
key: new Date().getTime() + Math.random(),
|
|
variant: 'success',
|
|
action: (key) => (
|
|
<Button style={{ color: 'white' }} onClick={() => closeSnackbar(key)}>
|
|
<IconX />
|
|
</Button>
|
|
)
|
|
}
|
|
})
|
|
}
|
|
} catch (error) {
|
|
setLoading(false)
|
|
const errorData = error.response?.data || `${error.response?.status}: ${error.response?.statusText}`
|
|
enqueueSnackbar({
|
|
message: `Failed to update vector store config: ${errorData}`,
|
|
options: {
|
|
key: new Date().getTime() + Math.random(),
|
|
variant: 'error',
|
|
persist: true,
|
|
action: (key) => (
|
|
<Button style={{ color: 'white' }} onClick={() => closeSnackbar(key)}>
|
|
<IconX />
|
|
</Button>
|
|
)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
useEffect(() => {
|
|
if (queryVectorStoreApi.data) {
|
|
setDocumentChunks(queryVectorStoreApi.data.docs)
|
|
setTimeTaken(queryVectorStoreApi.data.timeTaken)
|
|
setRetrievalError(undefined)
|
|
setLoading(false)
|
|
if (inputRef.current) {
|
|
inputRef.current.focus()
|
|
}
|
|
}
|
|
|
|
// eslint-disable-next-line react-hooks/exhaustive-deps
|
|
}, [queryVectorStoreApi.data])
|
|
|
|
useEffect(() => {
|
|
if (queryVectorStoreApi.error) {
|
|
if (queryVectorStoreApi.error.response?.data?.message) {
|
|
const message = queryVectorStoreApi.error.response.data.message
|
|
// remove the text 'documentStoreServices.queryVectorStore - ' from the error message to make it readable
|
|
setRetrievalError(message.replace('documentStoreServices.queryVectorStore - ', ''))
|
|
setDocumentChunks([])
|
|
setTimeTaken(-1)
|
|
}
|
|
setLoading(false)
|
|
}
|
|
|
|
// eslint-disable-next-line react-hooks/exhaustive-deps
|
|
}, [queryVectorStoreApi.error])
|
|
|
|
useEffect(() => {
|
|
if (getVectorStoreNodeDetailsApi.data) {
|
|
const node = getVectorStoreNodeDetailsApi.data
|
|
fetchVectorStoreDetails(node)
|
|
}
|
|
|
|
// eslint-disable-next-line react-hooks/exhaustive-deps
|
|
}, [getVectorStoreNodeDetailsApi.data])
|
|
|
|
const fetchVectorStoreDetails = (component) => {
|
|
const nodeData = cloneDeep(initNode(component, uuidv4()))
|
|
if (documentStore.vectorStoreConfig) {
|
|
nodeData.inputs = documentStore.vectorStoreConfig.config
|
|
nodeData.credential = documentStore.vectorStoreConfig.config.credential
|
|
}
|
|
setSelectedVectorStoreProvider(nodeData)
|
|
}
|
|
|
|
useEffect(() => {
|
|
getSpecificDocumentStoreApi.request(storeId)
|
|
|
|
// eslint-disable-next-line react-hooks/exhaustive-deps
|
|
}, [])
|
|
|
|
useEffect(() => {
|
|
if (getSpecificDocumentStoreApi.data) {
|
|
if (!hasAssignedWorkspace(getSpecificDocumentStoreApi.data.workspaceId)) {
|
|
navigate('/unauthorized')
|
|
return
|
|
}
|
|
setDocumentStore(getSpecificDocumentStoreApi.data)
|
|
const vectorStoreConfig = getSpecificDocumentStoreApi.data.vectorStoreConfig
|
|
if (vectorStoreConfig) {
|
|
getVectorStoreNodeDetailsApi.request(vectorStoreConfig.name)
|
|
}
|
|
setLoading(false)
|
|
}
|
|
|
|
// eslint-disable-next-line react-hooks/exhaustive-deps
|
|
}, [getSpecificDocumentStoreApi.data])
|
|
|
|
return (
|
|
<>
|
|
<MainCard style={{ position: 'relative' }}>
|
|
<Stack flexDirection='column' sx={{ gap: 1 }}>
|
|
<ViewHeader
|
|
isBackButton={true}
|
|
search={false}
|
|
title={documentStore?.name || 'Document Store'}
|
|
description='Retrieval Playground - Test your vector store retrieval settings'
|
|
onBack={() => navigate(-1)}
|
|
>
|
|
<PermissionButton
|
|
permissionId={'documentStores:upsert-config'}
|
|
variant='outlined'
|
|
color='secondary'
|
|
sx={{ borderRadius: 2, height: '100%' }}
|
|
startIcon={<IconDeviceFloppy />}
|
|
onClick={saveConfig}
|
|
>
|
|
Save Config
|
|
</PermissionButton>
|
|
</ViewHeader>
|
|
<div style={{ width: '100%' }}></div>
|
|
<div>
|
|
<Grid container spacing={2}>
|
|
<Grid sx={{ ml: 1, mr: 1 }} item xs={12} sm={12} md={12} lg={12}>
|
|
<Box>
|
|
<div style={{ display: 'flex', flexDirection: 'row' }}>
|
|
<Typography variant='overline'>
|
|
Enter your Query<span style={{ color: 'red' }}> *</span>
|
|
</Typography>
|
|
|
|
<div style={{ flexGrow: 1 }}></div>
|
|
</div>
|
|
<OutlinedInput
|
|
size='small'
|
|
multiline={true}
|
|
rows={4}
|
|
sx={{ mt: 1 }}
|
|
type='string'
|
|
fullWidth
|
|
inputRef={inputRef}
|
|
onChange={(e) => setQuery(e.target.value)}
|
|
onKeyDown={handleEnter}
|
|
value={query ?? ''}
|
|
endAdornment={
|
|
<IconButton variant='contained' onClick={doQuery}>
|
|
<IconSearch />
|
|
</IconButton>
|
|
}
|
|
/>
|
|
</Box>
|
|
</Grid>
|
|
<Grid sx={{ ml: 1, mr: 1, mt: 1 }} container spacing={1}>
|
|
<Grid item xs={12} sm={4} md={4}>
|
|
<Box>
|
|
<Grid container spacing='2'>
|
|
<Grid item xs={12} md={12} lg={12} sm={12}>
|
|
<div
|
|
style={{
|
|
display: 'flex',
|
|
flexDirection: 'column',
|
|
paddingRight: 15,
|
|
paddingTop: 5
|
|
}}
|
|
>
|
|
<Box
|
|
sx={{
|
|
display: 'flex',
|
|
alignItems: 'center',
|
|
flexDirection: 'row',
|
|
p: 1
|
|
}}
|
|
>
|
|
<div
|
|
style={{
|
|
width: 40,
|
|
height: 40,
|
|
borderRadius: '50%',
|
|
backgroundColor: 'white',
|
|
display: 'flex',
|
|
alignItems: 'center',
|
|
justifyContent: 'center',
|
|
boxShadow: '0 2px 14px 0 rgb(32 40 45 / 25%)'
|
|
}}
|
|
>
|
|
{selectedVectorStoreProvider.label ? (
|
|
<img
|
|
style={{
|
|
width: '100%',
|
|
height: '100%',
|
|
padding: 7,
|
|
borderRadius: '50%',
|
|
objectFit: 'contain'
|
|
}}
|
|
alt={selectedVectorStoreProvider.label ?? 'embeddings'}
|
|
src={`${baseURL}/api/v1/node-icon/${selectedVectorStoreProvider?.name}`}
|
|
/>
|
|
) : (
|
|
<Embeddings color='black' />
|
|
)}
|
|
</div>
|
|
<Typography sx={{ ml: 2 }} variant='h3'>
|
|
{selectedVectorStoreProvider.label}
|
|
</Typography>
|
|
<div style={{ flex: 1 }}></div>
|
|
</Box>
|
|
{selectedVectorStoreProvider &&
|
|
Object.keys(selectedVectorStoreProvider).length > 0 &&
|
|
showHideInputParams(selectedVectorStoreProvider)
|
|
.filter((inputParam) => !inputParam.hidden && inputParam.display !== false)
|
|
.map((inputParam, index) => (
|
|
<DocStoreInputHandler
|
|
key={index}
|
|
data={selectedVectorStoreProvider}
|
|
inputParam={inputParam}
|
|
isAdditionalParams={inputParam.additionalParams}
|
|
onNodeDataChange={handleVectorStoreProviderDataChange}
|
|
/>
|
|
))}
|
|
</div>
|
|
</Grid>
|
|
</Grid>
|
|
</Box>
|
|
</Grid>
|
|
<Grid item xs={12} sm={8} md={8}>
|
|
<Box
|
|
sx={{
|
|
display: 'flex',
|
|
alignItems: 'center',
|
|
flexDirection: 'row',
|
|
p: 1,
|
|
paddingTop: 2,
|
|
marginBottom: 4
|
|
}}
|
|
>
|
|
<div
|
|
style={{
|
|
width: 40,
|
|
height: 40,
|
|
borderRadius: '50%',
|
|
backgroundColor: 'white',
|
|
display: 'flex',
|
|
alignItems: 'center',
|
|
justifyContent: 'center',
|
|
boxShadow: '0 2px 14px 0 rgb(32 40 45 / 25%)'
|
|
}}
|
|
>
|
|
<IconFileStack
|
|
style={{
|
|
width: '100%',
|
|
height: '100%',
|
|
padding: 7,
|
|
borderRadius: '50%',
|
|
objectFit: 'contain'
|
|
}}
|
|
/>
|
|
</div>
|
|
<Typography sx={{ ml: 2 }} variant='h3'>
|
|
Retrieved Documents
|
|
{timeTaken > -1 && (
|
|
<Typography variant='body2' sx={{ color: 'gray' }}>
|
|
Count: {documentChunks.length}. Time taken: {timeTaken} millis.
|
|
</Typography>
|
|
)}
|
|
{retrievalError && (
|
|
<Typography variant='body2' sx={{ color: 'gray' }}>
|
|
{retrievalError}
|
|
</Typography>
|
|
)}
|
|
</Typography>
|
|
<div style={{ flex: 1 }}></div>
|
|
</Box>
|
|
{!documentChunks.length && (
|
|
<div
|
|
style={{
|
|
display: 'flex',
|
|
flexDirection: 'column',
|
|
alignItems: 'center',
|
|
width: '100%'
|
|
}}
|
|
>
|
|
<Box sx={{ mt: 5, p: 2, height: 'auto' }}>
|
|
<img
|
|
style={{ objectFit: 'cover', height: '16vh', width: 'auto' }}
|
|
src={chunks_emptySVG}
|
|
alt='chunks_emptySVG'
|
|
/>
|
|
</Box>
|
|
<div>No Documents Retrieved</div>
|
|
</div>
|
|
)}
|
|
<Grid container spacing={2}>
|
|
{documentChunks?.length > 0 &&
|
|
documentChunks.map((row, index) => (
|
|
<Grid item lg={6} md={6} sm={6} xs={6} key={index}>
|
|
<CardWrapper
|
|
content={false}
|
|
onClick={() => chunkSelected(row.id, row.chunkNo)}
|
|
sx={{ border: 1, borderColor: theme.palette.grey[900] + 25, borderRadius: 2 }}
|
|
>
|
|
<Card>
|
|
<CardContent sx={{ p: 2 }}>
|
|
<Typography sx={{ wordWrap: 'break-word', mb: 1 }} variant='h5'>
|
|
{`#${row.chunkNo}. Characters: ${row.pageContent.length}`}
|
|
</Typography>
|
|
<Typography sx={{ wordWrap: 'break-word' }} variant='body2'>
|
|
{row.pageContent}
|
|
</Typography>
|
|
<ReactJson
|
|
theme={customization.isDarkMode ? 'ocean' : 'rjv-default'}
|
|
style={{ paddingTop: 10 }}
|
|
src={row.metadata || {}}
|
|
name={null}
|
|
quotesOnKeys={false}
|
|
enableClipboard={false}
|
|
displayDataTypes={false}
|
|
collapsed={true}
|
|
/>
|
|
</CardContent>
|
|
</Card>
|
|
</CardWrapper>
|
|
</Grid>
|
|
))}
|
|
</Grid>
|
|
</Grid>
|
|
</Grid>
|
|
</Grid>
|
|
</div>
|
|
</Stack>
|
|
</MainCard>
|
|
<ConfirmDialog />
|
|
<ExpandedChunkDialog
|
|
show={showExpandedChunkDialog}
|
|
dialogProps={expandedChunkDialogProps}
|
|
onCancel={() => setShowExpandedChunkDialog(false)}
|
|
isReadOnly={true}
|
|
></ExpandedChunkDialog>
|
|
{loading && <BackdropLoader open={loading} />}
|
|
</>
|
|
)
|
|
}
|
|
|
|
export default VectorStoreQuery
|