feat(kb): Use document as attachment

package.json

@@ -3,7 +3,11 @@
   "version": "0.0.1",
   "description": "A UI for decentralized AI",
   "productName": "Libertai UI",
-  "author": "[email protected]",
+  "author": "LibertAI Team",
+  "contributors": [
+    "David Amelekh <[email protected]>",
+    "Reza Rahemtola <[email protected]>"
+  ],
   "private": true,
   "license": "MIT",
   "engines": {

src/components/MessageInput.vue

@@ -63,7 +63,7 @@
 </template>
 
 <script lang="ts" setup>
-import { ref } from 'vue';
+import { PropType, ref, watch } from 'vue';
 import { MessageAttachment, SendMessageParams } from 'src/types/chats';
 import { processAttachment } from 'src/utils/knowledge/attachments';
 import { useQuasar } from 'quasar';
@@ -79,8 +79,27 @@ const props = defineProps({
     type: String,
     default: '',
   },
+  additionalAttachment: {
+    type: Object as PropType<File>,
+    default: undefined,
+  },
 });
 
+watch(
+  () => props.additionalAttachment,
+  async (additionalAttachmentFile: File | undefined) => {
+    if (!additionalAttachmentFile) {
+      return;
+    }
+    const attachmentData = await processAttachmentFile(additionalAttachmentFile);
+    if (!attachmentData) {
+      return;
+    }
+    attachments.value = attachments.value.concat([attachmentData]);
+  },
+  { immediate: true },
+);
+
 const $q = useQuasar();
 
 const emit = defineEmits<{ sendMessage: [value: SendMessageParams] }>();
@@ -95,20 +114,27 @@ const processMessageAttachments = async (event: any) => {
 
   await Promise.all(
     Array.from(target.files as FileList).map(async (file) => {
-      try {
-        const fileData = await processAttachment(file);
-        attachmentsData.push(fileData);
-      } catch (error) {
-        $q.notify({
-          message: (error as Error)?.message ?? 'File processing failed, please try again',
-          color: 'negative',
-        });
+      const attachment = await processAttachmentFile(file);
+      if (!attachment) {
+        return;
       }
+      attachmentsData.push(attachment);
     }),
   );
   attachments.value = attachments.value.concat(attachmentsData);
 };
 
+const processAttachmentFile = async (file: File): Promise<MessageAttachment | undefined> => {
+  try {
+    return await processAttachment(file);
+  } catch (error) {
+    $q.notify({
+      message: (error as Error)?.message ?? 'File processing failed, please try again',
+      color: 'negative',
+    });
+  }
+};
+
 const removeAttachment = (attachmentId: string) => {
   attachments.value = attachments.value.filter((a) => a.id !== attachmentId);
 };

src/pages/Chat.vue

@@ -220,7 +220,7 @@ async function generatePersonaMessage() {
 
     // Finding related knowledge document chunks
     if (knowledgeBaseIds.length > 0) {
-      const documents = knowledgeStore.getDocumentsFrom(knowledgeBaseIds);
+      const documents = knowledgeStore.getDocumentsFromBases(knowledgeBaseIds);
       const lastUserMessage = messages.findLast((message) => message.author === 'user')!;
       knowledgeSearchResults = await searchDocuments(lastUserMessage.content, documents);
     }

src/pages/KnowledgeBase.vue

@@ -63,8 +63,16 @@
         <p class="tw-font-bold tw-text-base">{{ document.name }}</p>
         <div class="tw-ml-auto tw-flex tw-items-center tw-gap-4">
           <p class="max-sm:tw-hidden">{{ filesize(document.size, { round: 0 }) }}</p>
-          <q-btn class="tw-w-10 tw-h-10" disable unelevated>
+          <q-btn
+            :disable="document.size > MAX_ATTACHMENT_SIZE"
+            class="tw-w-10 tw-h-10"
+            unelevated
+            @click="chatWithDocument(document)"
+          >
             <ltai-icon name="svguse:icons.svg#chat" />
+            <q-tooltip v-if="document.size > MAX_ATTACHMENT_SIZE">
+              Document is too big to be used as chat attachment
+            </q-tooltip>
           </q-btn>
 
           <q-btn class="tw-w-10 tw-h-10" unelevated @click="downloadDocument(document)">
@@ -163,6 +171,7 @@ import { processDocument } from 'src/utils/knowledge/document';
 import { decryptFile, encryptFile } from 'src/utils/encryption';
 import KnowledgeBaseRenameDialog from 'components/dialog/KnowledgeBaseRenameDialog.vue';
 import { supportedInputFiles } from 'src/utils/knowledge/parsing';
+import { MAX_ATTACHMENT_SIZE } from 'src/utils/knowledge/attachments';
 
 const $q = useQuasar();
 const route = useRoute();
@@ -364,4 +373,17 @@ const deleteKnowledgeBase = async () => {
 
   await router.push({ path: '/knowledge-base' });
 };
+
+const chatWithDocument = async (document: KnowledgeDocument) => {
+  if (knowledgeBaseIdentifierRef.value === undefined) {
+    return;
+  }
+
+  await router.push({
+    path: '/new',
+    query: {
+      knowledgeDocumentAttachment: `${knowledgeBaseIdentifierRef.value.id},${document.id}`,
+    },
+  });
+};
 </script>

src/pages/NewChat.vue

@@ -43,6 +43,7 @@
       </div>
       <div class="fixed-bottom absolute q-mb-xl tw-pb-1">
         <message-input
+          :additional-attachment="knowledgeAttachmentFile"
           hint="Disclaimer: This chat bot uses personas for entertainment and informational purposes only. The
               chat bot's responses are not a reflection of any real person or organization's views or opinions, and should not
               be used as a substitute for professional advice. The accuracy and reliability of the chat bot's responses cannot
@@ -56,13 +57,13 @@
   </q-page>
 </template>
 <script lang="ts" setup>
-import { defaultChatTopic } from 'src/utils/chat'; // Import State
+import { defaultChatTopic } from 'src/utils/chat';
 import { useModelsStore } from 'stores/models';
 import { useChatsStore } from 'stores/chats';
 import { usePersonasStore } from 'stores/personas';
 import { useSettingsStore } from 'stores/settings';
 import { ref, watch } from 'vue';
-import { useRoute, useRouter } from 'vue-router'; // Import components
+import { useRoute, useRouter } from 'vue-router';
 import MessageInput from 'src/components/MessageInput.vue';
 import PersonaDropdown from 'components/select/PersonaSelector.vue';
 import { getPersonaAvatarUrl } from 'src/utils/personas';
@@ -71,14 +72,16 @@ import { UIModel } from 'src/utils/models';
 import { UIPersona } from 'src/types/personas';
 import { SendMessageParams } from 'src/types/chats';
 import KnowledgeBasesSelector from 'components/select/KnowledgeBasesSelector.vue';
+import { useKnowledgeStore } from 'stores/knowledge';
 
 const router = useRouter();
 
-// Stored State
+// Store state
 const modelsStore = useModelsStore();
 const chatsStore = useChatsStore();
 const personasStore = usePersonasStore();
 const settingsStore = useSettingsStore();
+const knowledgeStore = useKnowledgeStore();
 
 const route = useRoute();
 
@@ -88,21 +91,36 @@ const selectedPersona = ref<UIPersona>(personasStore.personas[0]);
 const username = ref(settingsStore.username);
 const selectedKnowledgeBases = ref<string[]>([]);
 
+const knowledgeAttachmentFile = ref<File | undefined>(undefined);
+
 watch(
   () => route.query.persona as string | undefined,
   (personaId: string | undefined) => {
-    if (personaId) {
-      const persona = personasStore.personas.find((p) => p.id === personaId);
-      if (persona) {
-        selectedPersona.value = persona;
-      }
+    if (!personaId) {
+      return;
+    }
+    const persona = personasStore.personas.find((p) => p.id === personaId);
+    if (persona) {
+      selectedPersona.value = persona;
     }
   },
   {
     immediate: true,
   },
 );
 
+watch(
+  () => route.query.knowledgeDocumentAttachment as string | undefined,
+  async (knowledgeDocumentAttachmentInfo: string | undefined) => {
+    if (!knowledgeDocumentAttachmentInfo) {
+      return;
+    }
+    const [kbIdentifierId, documentId] = knowledgeDocumentAttachmentInfo.split(',');
+    knowledgeAttachmentFile.value = await knowledgeStore.getDocumentFile(kbIdentifierId, documentId);
+  },
+  { immediate: true },
+);
+
 async function sendMessage({ content, attachments }: SendMessageParams) {
   // Extract the values out of our relevant refs
   const title = defaultChatTopic;

src/stores/knowledge.ts

@@ -3,6 +3,7 @@ import { v4 as uuidv4 } from 'uuid';
 
 import { KnowledgeBase, KnowledgeBaseIdentifier, KnowledgeDocument } from 'src/types/knowledge';
 import { useAccountStore } from 'stores/account';
+import { decryptFile } from 'src/utils/encryption';
 
 type KnowledgeStoreState = {
   knowledgeBases: KnowledgeBase[];
@@ -17,14 +18,41 @@ export const useKnowledgeStore = defineStore('knowledge', {
     isLoaded: false,
   }),
   getters: {
-    getDocumentsFrom: (state) => {
+    getDocumentsFromBases: (state) => {
       return (ids: string[]): KnowledgeDocument[] => {
         return state.knowledgeBases
           .filter((kb) => ids.includes(kb.id))
           .map((kb) => kb.documents)
           .flat();
       };
     },
+    getDocumentFile: (state) => {
+      return async (kbIdentifierId: string, documentId: string): Promise<File | undefined> => {
+        const { alephStorage } = useAccountStore();
+        if (alephStorage === null) {
+          return;
+        }
+
+        const document = state.knowledgeBases
+          .map((kb) => kb.documents)
+          .flat()
+          .find((document) => document.id === documentId);
+        if (!document) {
+          return;
+        }
+        const kbIdentifier = state.knowledgeBaseIdentifiers.find((kbi) => kbi.id === kbIdentifierId);
+        if (!kbIdentifier) {
+          return;
+        }
+
+        const encryptionKey = Buffer.from(kbIdentifier.encryption.key);
+        const encryptionIv = Buffer.from(kbIdentifier.encryption.iv);
+
+        const downloadedFile = await alephStorage.downloadFile(document.store.ipfs_hash);
+        const decryptedContent = decryptFile(downloadedFile, encryptionKey, encryptionIv);
+        return new File([decryptedContent], document.name, { type: document.type });
+      };
+    },
   },
   actions: {
     async load() {

src/utils/knowledge/attachments.ts

@@ -2,11 +2,13 @@ import { v4 as uuidv4 } from 'uuid';
 import { MessageAttachment } from 'src/types/chats';
 import { extractFileContent } from 'src/utils/knowledge/parsing';
 
+export const MAX_ATTACHMENT_SIZE = 4 * 1024; // 4 KiB
+
 export const processAttachment = async (file: File): Promise<MessageAttachment> => {
   const title = file.name;
   const fileInfo = await extractFileContent(file);
 
-  if (fileInfo.content.length > 4 * 1024) {
+  if (fileInfo.content.length > MAX_ATTACHMENT_SIZE) {
     // File is too big to be inlined, rejecting it.
     // Later we'll use a knowledge db to fix this.
     throw new Error('File is too big, please use a file of 4 KB of content or less.');

src/utils/knowledge/parsing.ts

@@ -2,7 +2,7 @@ import mime from 'mime';
 import * as pdfjs from 'pdfjs-dist';
 import { TextItem } from 'pdfjs-dist/types/src/display/api';
 
-export const supportedInputFiles = ['.txt', '.md', '.pdf'].join(',');
+export const supportedInputFiles = ['.txt', '.md', '.pdf', '.py'].join(',');
 
 const extractTextFromPdfFile = async (file: File): Promise<string> => {
   const pdfUrl = URL.createObjectURL(file);