Initial attempt at janhq#1995

Signed-off-by: Richard Tweed <RichardoC@users.noreply.github.com>
RichardoC · Apr 25, 2024 · 3c4dd22 · 3c4dd22
1 parent 63a2f22
commit 3c4dd22
Show file tree

Hide file tree

Showing 6 changed files with 133 additions and 26 deletions.
diff --git a/core/src/types/message/messageEntity.ts b/core/src/types/message/messageEntity.ts
@@ -97,6 +97,7 @@ export enum ContentType {
   Text = 'text',
   Image = 'image',
   Pdf = 'pdf',
+  PlainText = 'plainText',
 }
 
 /**

diff --git a/extensions/assistant-extension/src/node/retrieval.ts b/extensions/assistant-extension/src/node/retrieval.ts
@@ -1,6 +1,7 @@
 import { RecursiveCharacterTextSplitter } from 'langchain/text_splitter'
 import { formatDocumentsAsString } from 'langchain/util/document'
 import { PDFLoader } from 'langchain/document_loaders/fs/pdf'
+import { TextLoader } from 'langchain/document_loaders/fs/text'
 
 import { HNSWLib } from 'langchain/vectorstores/hnswlib'
 
@@ -50,9 +51,14 @@ export class Retrieval {
     filePath: string,
     memoryPath: string
   ): Promise<any> => {
-    const loader = new PDFLoader(filePath, {
+    var loader
+    if (filePath.endsWith(".pdf") ) {
+     loader = new PDFLoader(filePath, {
       splitPages: true,
     })
+    } else {
+      loader = new TextLoader(filePath)
+    }
     if (!this.embeddingModel) return Promise.reject()
     const doc = await loader.load()
     const docs = await this.textSplitter!.splitDocuments(doc)

diff --git a/extensions/conversational-extension/src/index.ts b/extensions/conversational-extension/src/index.ts
@@ -113,30 +113,49 @@ export default class JSONConversationalExtension extends ConversationalExtension
       ])
       if (!(await fs.existsSync(threadDirPath))) await fs.mkdir(threadDirPath)
 
-      if (message.content[0]?.type === 'image') {
-        const filesPath = await joinPath([threadDirPath, 'files'])
-        if (!(await fs.existsSync(filesPath))) await fs.mkdir(filesPath)
-
-        const imagePath = await joinPath([filesPath, `${message.id}.png`])
-        const base64 = message.content[0].text.annotations[0]
-        await this.storeImage(base64, imagePath)
-        if ((await fs.existsSync(imagePath)) && message.content?.length) {
-          // Use file path instead of blob
-          message.content[0].text.annotations[0] = `threads/${message.thread_id}/files/${message.id}.png`
+      switch (message.content[0]?.type){
+        case 'image': {
+          const filesPath = await joinPath([threadDirPath, 'files'])
+          if (!(await fs.existsSync(filesPath))) await fs.mkdir(filesPath)
+
+          const imagePath = await joinPath([filesPath, `${message.id}.png`])
+          const base64 = message.content[0].text.annotations[0]
+          await this.storeImage(base64, imagePath)
+          if ((await fs.existsSync(imagePath)) && message.content?.length) {
+            // Use file path instead of blob
+            message.content[0].text.annotations[0] = `threads/${message.thread_id}/files/${message.id}.png`
+          }
+          break
+        }
+
+        case 'pdf': {
+          const filesPath = await joinPath([threadDirPath, 'files'])
+          if (!(await fs.existsSync(filesPath))) await fs.mkdir(filesPath)
+
+          const filePath = await joinPath([filesPath, `${message.id}.pdf`])
+          const blob = message.content[0].text.annotations[0]
+          await this.storePdf(blob, filePath)
+
+          if ((await fs.existsSync(filePath)) && message.content?.length) {
+            // Use file path instead of blob
+            message.content[0].text.annotations[0] = `threads/${message.thread_id}/files/${message.id}.pdf`
+          }
+          break
         }
-      }
 
-      if (message.content[0]?.type === 'pdf') {
-        const filesPath = await joinPath([threadDirPath, 'files'])
-        if (!(await fs.existsSync(filesPath))) await fs.mkdir(filesPath)
+        case 'plainText': {
+          const filesPath = await joinPath([threadDirPath, 'files'])
+          if (!(await fs.existsSync(filesPath))) await fs.mkdir(filesPath)
 
-        const filePath = await joinPath([filesPath, `${message.id}.pdf`])
-        const blob = message.content[0].text.annotations[0]
-        await this.storeFile(blob, filePath)
+          const filePath = await joinPath([filesPath, `${message.id}.txt`])
+          const blob = message.content[0].text.annotations[0]
+          await this.storePlainFile(blob, filePath)
 
-        if ((await fs.existsSync(filePath)) && message.content?.length) {
-          // Use file path instead of blob
-          message.content[0].text.annotations[0] = `threads/${message.thread_id}/files/${message.id}.pdf`
+          if ((await fs.existsSync(filePath)) && message.content?.length) {
+            // Use file path instead of blob
+            message.content[0].text.annotations[0] = `threads/${message.thread_id}/files/${message.id}.txt`
+          }
+          break
         }
       }
       await fs.appendFileSync(threadMessagePath, JSON.stringify(message) + '\n')
@@ -156,7 +175,7 @@ export default class JSONConversationalExtension extends ConversationalExtension
     }
   }
 
-  async storeFile(base64: string, filePath: string): Promise<void> {
+  async storePdf(base64: string, filePath: string): Promise<void> {
     const base64Data = base64.replace(/^data:application\/pdf;base64,/, '')
     try {
       await fs.writeBlob(filePath, base64Data)
@@ -165,6 +184,16 @@ export default class JSONConversationalExtension extends ConversationalExtension
     }
   }
 
+  // Complete guess
+  async storePlainFile(base64: string, filePath: string): Promise<void> {
+    const base64Data = base64.replace(/^data:text\/plain;base64,/, '')
+    try {
+      await fs.writeBlob(filePath, base64Data)
+    } catch (err) {
+      console.error(err)
+    }
+  }
+
   async writeMessages(
     threadId: string,
     messages: ThreadMessage[]

diff --git a/web/containers/Providers/Jotai.tsx b/web/containers/Providers/Jotai.tsx
@@ -21,7 +21,7 @@ export default function JotaiWrapper({ children }: Props) {
   return <Provider>{children}</Provider>
 }
 
-export type FileType = 'image' | 'pdf'
+export type FileType = 'image' | 'pdf' | 'plain/text'
 
 export type FileInfo = {
   file: File

diff --git a/web/screens/Chat/ChatInput/index.tsx b/web/screens/Chat/ChatInput/index.tsx
@@ -116,11 +116,16 @@ const ChatInput: React.FC = () => {
    * Its to be used to display the extension file name of the selected file.
    * @param event - The change event object.
    */
-  const handleFileChange = (event: React.ChangeEvent<HTMLInputElement>) => {
+  const handlePDFFileChange = (event: React.ChangeEvent<HTMLInputElement>) => {
     const file = event.target.files?.[0]
     if (!file) return
     setFileUpload([{ file: file, type: 'pdf' }])
   }
+  const handleFileChange = (event: React.ChangeEvent<HTMLInputElement>) => {
+    const file = event.target.files?.[0]
+    if (!file) return
+    setFileUpload([{ file: file, type: 'plain/text' }])
+  }
 
   const handleImageChange = (event: React.ChangeEvent<HTMLInputElement>) => {
     const file = event.target.files?.[0]
@@ -262,7 +267,53 @@ const ChatInput: React.FC = () => {
                     }}
                   >
                     <FileTextIcon size={16} />
-                    <span className="font-medium">Document</span>
+                    <span className="font-medium">PDF Document</span>
+                  </li>
+                </TooltipTrigger>
+                <TooltipPortal>
+                  {(!activeThread?.assistants[0].tools ||
+                    !activeThread?.assistants[0].tools[0]?.enabled ||
+                    activeThread?.assistants[0].model.settings.text_model ===
+                      false) && (
+                    <TooltipContent side="top" className="max-w-[154px] px-3">
+                      {activeThread?.assistants[0].model.settings.text_model ===
+                      false ? (
+                        <span>
+                          This model does not support text-based retrieval.
+                        </span>
+                      ) : (
+                        <span>
+                          Turn on Retrieval in Assistant Settings to use this
+                          feature.
+                        </span>
+                      )}
+                      <TooltipArrow />
+                    </TooltipContent>
+                  )}
+                </TooltipPortal>
+              </Tooltip>
+              <Tooltip>
+                <TooltipTrigger asChild>
+                  <li
+                    className={twMerge(
+                      'flex w-full cursor-pointer items-center space-x-2 px-4 py-2 text-muted-foreground hover:bg-secondary',
+                      activeThread?.assistants[0].model.settings.text_model ===
+                        false
+                        ? 'cursor-not-allowed opacity-50'
+                        : 'cursor-pointer'
+                    )}
+                    onClick={() => {
+                      if (
+                        activeThread?.assistants[0].model.settings
+                          .text_model !== false
+                      ) {
+                        fileInputRef.current?.click()
+                        setShowAttacmentMenus(false)
+                      }
+                    }}
+                  >
+                    <FileTextIcon size={16} />
+                    <span className="font-medium">Plain Text Document</span>
                   </li>
                 </TooltipTrigger>
                 <TooltipPortal>
@@ -305,9 +356,17 @@ const ChatInput: React.FC = () => {
         className="hidden"
         ref={fileInputRef}
         value=""
-        onChange={handleFileChange}
+        onChange={handlePDFFileChange}
         accept="application/pdf"
       />
+      <input
+        type="file"
+        className="hidden"
+        ref={fileInputRef}
+        value=""
+        onChange={handleFileChange}
+        accept="plain/text"
+      />
 
       {messages[messages.length - 1]?.status !== MessageStatus.Pending &&
       !isGeneratingResponse &&

diff --git a/web/utils/threadMessageBuilder.ts b/web/utils/threadMessageBuilder.ts
@@ -60,6 +60,18 @@ export class ThreadMessageBuilder {
       })
     }
 
+    if (base64 && fileUpload[0]?.type === 'plain/text') {
+      this.content.push({
+        type: ContentType.PlainText,
+        text: {
+          value: prompt,
+          annotations: [base64],
+          name: fileUpload[0].file.name,
+          size: fileUpload[0].file.size,
+        },
+      })
+    }
+
     if (prompt && !base64) {
       this.content.push({
         type: ContentType.Text,