Merge pull request #453 from l3vels/fix/datasource

Fix/datasource
l3vels · Mar 21, 2024 · 27456b9 · 27456b9
2 parents 2eeb0a2 + 4e71efd
commit 27456b9
Show file tree

Hide file tree

Showing 11 changed files with 76 additions and 35 deletions.
diff --git a/apps/server/controllers/configuration.py b/apps/server/controllers/configuration.py
@@ -56,14 +56,19 @@ def index_documents(value: str, datasource_id: UUID, account: AccountOutput):
             similarity_top_k,
         )
         retriever.index_documents(file_urls)
-
         datasource.status = DatasourceStatus.READY.value
         datasource.error = None
+
     except Exception as err:
-        print(err)
-        sentry_sdk.capture_exception(err)
         datasource.status = DatasourceStatus.FAILED.value
-        datasource.error = str(err)
+        try:
+            error_body = json.loads(err.body)
+            error_message = error_body.get("error", {}).get("message", "")
+            print("HERE: ", error_message)
+            datasource.error = str(error_message)
+        except (AttributeError, json.JSONDecodeError):
+            sentry_sdk.capture_exception(err)
+            datasource.error = str(err)
 
     session.add(datasource)
     session.commit()

diff --git a/apps/server/datasources/file/file_retriever.py b/apps/server/datasources/file/file_retriever.py
@@ -5,6 +5,7 @@
 from uuid import UUID, uuid4
 
 import s3fs
+from fastapi import HTTPException
 from llama_index.core import (ServiceContext, Settings, SimpleDirectoryReader,
                               StorageContext, SummaryIndex, TreeIndex,
                               VectorStoreIndex, load_index_from_storage)
@@ -15,6 +16,7 @@
 from llama_index.vector_stores.pinecone import PineconeVectorStore
 from llama_index.vector_stores.weaviate import WeaviateVectorStore
 from llama_index.vector_stores.zep import ZepVectorStore
+from pinecone import Pinecone, ServerlessSpec
 
 from config import Config
 from services.aws_s3 import AWSS3Service
@@ -93,20 +95,28 @@ def get_vector_store(self, is_retriever: bool = False):
                 embedding_dimensions=1536,
             )
         elif self.vector_store == VectorStoreProvider.PINECONE.value:
-            import pinecone
-
             # Pinecone only supports alphanumeric characters. Max length 40
             index_name = UUID(self.datasource_id).hex
 
-            pinecone.init(
+            # pinecone.init(
+            #     api_key=self.settings.pinecone_api_key,
+            #     environment=self.settings.pinecone_environment,
+            # )
+
+            pc = Pinecone(
                 api_key=self.settings.pinecone_api_key,
                 environment=self.settings.pinecone_environment,
             )
 
             if not is_retriever:
-                pinecone.create_index(index_name, dimension=1536, metric="cosine")
-
-            pinecone_index = pinecone.Index(index_name)
+                pc.create_index(
+                    index_name,
+                    dimension=1536,
+                    metric="cosine",
+                    spec=ServerlessSpec(cloud="aws", region="us-west-2"),
+                )
+
+            pinecone_index = pc.Index(index_name)
 
             vector_store = PineconeVectorStore(
                 pinecone_index=pinecone_index,

diff --git a/apps/ui/src/components/UploadedFile/UploadedFile.tsx b/apps/ui/src/components/UploadedFile/UploadedFile.tsx
@@ -7,15 +7,25 @@ import IconButton from 'share-ui/components/IconButton/IconButton'
 import Close from 'share-ui/components/Icon/Icons/components/Close'
 
 type UploadedFileProps = {
-  onClick: (id?: string) => void
-  id?: string
+  onClick?: () => void
   name: string
-  hasDeleteIcon?: boolean
+  onDeleteClick?: () => void
 }
 
-const UploadedFile = ({ id, name, onClick, hasDeleteIcon = false }: UploadedFileProps) => {
+const UploadedFile = ({ name, onClick, onDeleteClick }: UploadedFileProps) => {
+  const handleDelete = (event: any) => {
+    event.stopPropagation()
+    if (onDeleteClick) {
+      onDeleteClick()
+    }
+  }
+
   return (
-    <StyledUploadedFile onClick={() => onClick(id)}>
+    <StyledUploadedFile
+      onClick={() => {
+        if (onClick) onClick()
+      }}
+    >
       <StyledIconWrapper>
         <StyledDocIcon />
       </StyledIconWrapper>
@@ -27,8 +37,9 @@ const UploadedFile = ({ id, name, onClick, hasDeleteIcon = false }: UploadedFile
           customColor={'#000'}
         />
       </StyledTextWrapper>
-      {hasDeleteIcon && (
+      {handleDelete && (
         <IconButton
+          onClick={handleDelete}
           size={IconButton.sizes?.XS}
           icon={() => <StyledCloseIcon size='22' />}
           kind={IconButton.kinds?.TERTIARY}

diff --git a/apps/ui/src/modals/AIChatModal/components/ChatV2.tsx b/apps/ui/src/modals/AIChatModal/components/ChatV2.tsx
@@ -430,7 +430,7 @@ const ChatV2 = ({ chatSessionId }: { chatSessionId?: string }) => {
             {uploadedFileObject && (
               <StyledFileWrapper>
                 <UploadedFile
-                  onClick={() => setUploadedFileObject(null)}
+                  onDeleteClick={() => setUploadedFileObject(null)}
                   name={uploadedFileObject.fileName}
                 />
               </StyledFileWrapper>

diff --git a/apps/ui/src/modals/AIChatModal/fileTypes.ts b/apps/ui/src/modals/AIChatModal/fileTypes.ts
@@ -12,13 +12,13 @@ export const FILE_TYPES = [
   'application/epub+zip',
   'application/mbox',
   'application/x-ipynb+json',
-  'audio/mpeg',
-  'video/mp4',
-  'text/plain',
+  // 'audio/mpeg',
+  // 'video/mp4',
+  // 'text/plain',
 ]
 
 export const SUPPORTED_FILE_EXTENSIONS = [
-  'txt',
+  // 'txt',
   'pdf',
   'csv',
   'json',
@@ -31,6 +31,6 @@ export const SUPPORTED_FILE_EXTENSIONS = [
   'epub',
   'mbox',
   'ipynb',
-  'mp3',
-  'mp4',
+  // 'mp3',
+  // 'mp4',
 ]
diff --git a/apps/ui/src/pages/Datasource/Datasource.tsx b/apps/ui/src/pages/Datasource/Datasource.tsx
@@ -48,10 +48,10 @@ const Datasource = () => {
 
   const navigate = useNavigate()
 
-  useEffect(() => {
-    if (datasources?.length > 0) navigate(`/datasources/${datasources?.[0].id}/edit-datasource`)
-    else if (groups?.length > 0) navigate(`/datasources/${groups?.[0].id}/edit-group`)
-  }, [datasources])
+  // useEffect(() => {
+  //   if (datasources?.length > 0) navigate(`/datasources/${datasources?.[0].id}/edit-datasource`)
+  //   else if (groups?.length > 0) navigate(`/datasources/${groups?.[0].id}/edit-group`)
+  // }, [datasources])
 
   return (
     <>

diff --git a/apps/ui/src/pages/Datasource/DatasourceForm/DatasourceForm.tsx b/apps/ui/src/pages/Datasource/DatasourceForm/DatasourceForm.tsx
@@ -112,10 +112,10 @@ const DatasourceForm = ({ formik, isLoading, isEdit = false, type }: DatasourceF
                       files.map((file: any) => (
                         <UploadedFile
                           key={file.url}
-                          id={file.url}
-                          hasDeleteIcon
-                          onClick={id => {
-                            const filteredFiles = files.filter((file: any) => file.url !== id)
+                          onDeleteClick={() => {
+                            const filteredFiles = files.filter(
+                              (files: any) => files.url !== file.url,
+                            )
                             setFieldValue('files', filteredFiles)
                           }}
                           name={file.name}

diff --git a/apps/ui/src/pages/Datasource/DatasourceForm/useDatasourceForm.ts b/apps/ui/src/pages/Datasource/DatasourceForm/useDatasourceForm.ts
@@ -1,8 +1,11 @@
+import { ToastContext } from 'contexts'
 import useUploadFile from 'hooks/useUploadFile'
-import { ChangeEvent, useState } from 'react'
+import { ChangeEvent, useContext, useState } from 'react'
 import { useDataLoadersService } from 'services/datasource/useDataLoadersService'
 
 export const useDatasourceForm = (formik: any) => {
+  const { setToast } = useContext(ToastContext)
+
   const { data: dataLoaders } = useDataLoadersService()
 
   const { setFieldValue, values } = formik
@@ -26,6 +29,16 @@ export const useDatasourceForm = (formik: any) => {
     const promises = []
 
     for (const file of files) {
+      if (file.type.startsWith('video/')) {
+        setToast({
+          message: 'Cant upload video!',
+          type: 'warning',
+          open: true,
+        })
+
+        continue
+      }
+
       promises.push(
         uploadFile(
           {

diff --git a/apps/ui/src/pages/Datasource/useEditDatasource.ts b/apps/ui/src/pages/Datasource/useEditDatasource.ts
@@ -125,8 +125,6 @@ export const useEditDatasource = ({ incomingDatasourceId }: { incomingDatasource
     await Promise.all(promises)
     await Promise.all([refetchDatasources(), refetchConfigs()])
 
-    if (!incomingDatasourceId) navigate('/datasources')
-
     setToast({
       message: 'Datasource was updated!',
       type: 'positive',

diff --git a/apps/ui/src/pages/TeamOfAgents/TeamOfAgentsForm/TeamOfAgentsForm.tsx b/apps/ui/src/pages/TeamOfAgents/TeamOfAgentsForm/TeamOfAgentsForm.tsx
@@ -126,7 +126,7 @@ const TeamOfAgentsForm = ({ formik, isLoading }: TeamOfAgentsFormProps) => {
 
                   {config_value && (
                     <UploadedFile
-                      onClick={() => setFieldValue('config_value', null)}
+                      onDeleteClick={() => setFieldValue('config_value', null)}
                       name={'file'}
                     />
                   )}

diff --git a/apps/ui/src/share-ui/components/IconButton/IconButton.tsx b/apps/ui/src/share-ui/components/IconButton/IconButton.tsx
@@ -217,4 +217,8 @@ export default IconButton
 
 const StyledIconButtonWrapper = styled.div`
   max-height: 36px;
+
+  display: flex;
+  justify-content: center;
+  align-items: center;
 `