feat(kb): Chunks generation and basic search results

package.json

@@ -38,6 +38,7 @@
     "marked": "^13.0.3",
     "marked-highlight": "^2.1.3",
     "mime": "^4.0.4",
+    "ml-distance": "^4.0.1",
     "pdfjs-dist": "^4.5.136",
     "pinia": "^2.2.1",
     "pinia-plugin-persistedstate": "^3.2.1",

src/pages/Chat.vue

@@ -119,6 +119,8 @@ import { useSettingsStore } from 'stores/settings';
 import { Chat, SendMessageParams, UIMessage } from 'src/types/chats';
 import dayjs from 'dayjs';
 import LtaiIcon from 'components/libertai/LtaiIcon.vue';
+import { searchDocuments } from 'src/utils/knowledge/embedding';
+import { useKnowledgeStore } from 'stores/knowledge';
 
 const $q = useQuasar();
 const route = useRoute();
@@ -128,6 +130,7 @@ const router = useRouter();
 const chatsStore = useChatsStore();
 const modelsStore = useModelsStore();
 const settingsStore = useSettingsStore();
+const knowledgeStore = useKnowledgeStore();
 
 // Local page state
 const isLoadingRef = ref(false);
@@ -185,7 +188,8 @@ async function generatePersonaMessage() {
 
   const chatId = chatRef.value.id;
   const username = chatRef.value.username;
-  const messages = JSON.parse(JSON.stringify(chatRef.value.messages));
+  const messages: UIMessage[] = JSON.parse(JSON.stringify(chatRef.value.messages));
+  const knowledgeBaseIds = chatRef.value.knowledgeBases;
   const persona = chatRef.value.persona;
 
   const modelId = chatRef.value.modelId;
@@ -211,61 +215,42 @@ async function generatePersonaMessage() {
     // Set loading state
     isLoadingRef.value = true;
 
-    // NOTE: assuming last message is guaranteed to be non-empty and the user's last message
-    // Get the last message from the user
-    // const lastMessage = messages[messages.length - 1];
-    // const searchResultMessages: Message[] = [];
-    // const searchResults = await knowledgeStore.searchDocuments(lastMessage.content, chatTags);
-    // searchResults.forEach((result) => {
-    //   searchResultMessages.push({
-    //     role: 'search-result',
-    //     content: result.content,
-    //   });
-    // });
+    let searchResultMessages: Message[] = [];
+
+    // Finding related knowledge document chunks
+    if (knowledgeBaseIds.length > 0) {
+      const documents = knowledgeStore.getDocumentsFrom(knowledgeBaseIds);
+      const lastUserMessage = messages.findLast((message) => message.author === 'user')!;
+      const searchResults = await searchDocuments(lastUserMessage.content, documents);
+      console.log(searchResults);
+      searchResultMessages = searchResults.map(
+        (result): Message => ({
+          role: 'search-result',
+          content: result.content,
+        }),
+      );
+    }
 
     // Expand all the messages to inline any compatible attachments
     const expandedMessages = messages
       .map((message: UIMessage): Message[] => {
         const ret = [];
         // Push any attachments as messages ahead of the message itself
         message.attachments?.forEach((attachment) => {
-          if (attachment.content) {
-            ret.push({
-              role: 'attachment',
-              content: `[${attachment.title}](${attachment.content})`,
-            });
-          }
-          // else if (attachment.documentId) {
-          //   ret.push({
-          //     role: 'attachment',
-          //     content: `[${attachment.title}](document-id-${attachment.documentId})`,
-          //   });
-          // }
+          ret.push({
+            role: 'attachment',
+            content: `[${attachment.title}](${attachment.content})`,
+          });
         });
 
-        // Push what search results we found based on the message
-        // TODO: this should probably be a more generic tool-call or llm-chain-link
-        // TODO: this should probably link back to the document id
-        // TODO: I should probably write these below messages in the log
-        //  Really these search results should get attached to the message that
-        //   lead to them being queried
-        // if (message.searchResults) {
-        //   message.searchResults.forEach((result: Message) => {
-        //     ret.push({
-        //       role: 'search-result',
-        //       content: result.content,
-        //     });
-        //   });
-        // }
-
         // Push the message itself
         ret.push(message);
         return ret;
       })
       .flat();
 
     // Append the search results to the messages
-    const allMessages: Message[] = [...expandedMessages /*...searchResultMessages */];
+    const allMessages: Message[] = [...expandedMessages, ...searchResultMessages];
 
     // Generate a stream of responses from the AI
     for await (const output of inferenceEngine.generateAnswer(allMessages, model, persona, username, false)) {

src/pages/KnowledgeBase.vue

@@ -43,15 +43,33 @@
 
           <q-btn-dropdown class="tw-p-1" dropdown-icon="more_horiz" unelevated>
             <q-list>
-              <q-item v-close-popup clickable @click="showRenameDocument = true">
+              <q-item
+                v-close-popup
+                clickable
+                @click="
+                  () => {
+                    selectedDocument = document;
+                    showRenameDocument = true;
+                  }
+                "
+              >
                 <q-item-section avatar>
                   <ltai-icon class="tw-mx-auto" name="svguse:icons.svg#pencil" />
                 </q-item-section>
                 <q-item-section>
                   <q-item-label>Rename</q-item-label>
                 </q-item-section>
               </q-item>
-              <q-item v-close-popup clickable @click="showDeleteDocumentConfirmation = true">
+              <q-item
+                v-close-popup
+                clickable
+                @click="
+                  () => {
+                    selectedDocument = document;
+                    showDeleteDocumentConfirmation = true;
+                  }
+                "
+              >
                 <q-item-section avatar>
                   <ltai-icon class="tw-mx-auto" name="svguse:icons.svg#delete" />
                 </q-item-section>
@@ -62,19 +80,22 @@
             </q-list>
           </q-btn-dropdown>
         </div>
-
-        <!-- Dialogs-->
-        <knowledge-base-rename-document-dialog
-          v-model="showRenameDocument"
-          :name="document.name"
-          @save="(newName: string) => renameDocument(document, newName)"
-        />
-        <ltai-dialog v-model="showDeleteDocumentConfirmation" title="Delete document" @save="deleteDocument(document)">
-          <q-card-section class="row">
-            <span>Are you sure you want to delete the the document {{ document.name }}?</span>
-          </q-card-section>
-        </ltai-dialog>
       </div>
+      <!-- Dialogs-->
+      <knowledge-base-rename-document-dialog
+        v-model="showRenameDocument"
+        :name="selectedDocument?.name ?? ''"
+        @save="(newName: string) => renameDocument(selectedDocument!, newName)"
+      />
+      <ltai-dialog
+        v-model="showDeleteDocumentConfirmation"
+        title="Delete document"
+        @save="deleteDocument(selectedDocument!)"
+      >
+        <q-card-section class="row">
+          <span>Are you sure you want to delete the the document {{ selectedDocument!.name }}?</span>
+        </q-card-section>
+      </ltai-dialog>
     </div>
   </section>
 </template>
@@ -103,6 +124,7 @@ const knowledgeStore = useKnowledgeStore();
 
 const knowledgeBaseRef = ref<KnowledgeBase | undefined>(undefined);
 const knowledgeBaseIdentifierRef = ref<KnowledgeBaseIdentifier | undefined>(undefined);
+const selectedDocument = ref<KnowledgeDocument | undefined>(undefined);
 const showRenameDocument = ref(false);
 const showDeleteDocumentConfirmation = ref(false);
 

src/pages/KnowledgeBasesList.vue

@@ -22,7 +22,7 @@
           <p class="tw-font-bold tw-text-base">{{ knowledgeBase.name }}</p>
           <div class="tw-ml-auto tw-flex tw-gap-4">
             <p>{{ knowledgeBase.documents.length }} File{{ knowledgeBase.documents.length !== 1 ? 's' : '' }}</p>
-            <p>Last updated: {{ dayjs(knowledgeBase.lastUpdatedAt).format('LL') }}</p>
+            <p class="max-sm:tw-hidden">Last updated: {{ dayjs(knowledgeBase.lastUpdatedAt).format('LL') }}</p>
             <ltai-icon class="tw-w-5 tw-h-5" name="svguse:icons.svg#chevron-right" />
           </div>
         </div>

src/stores/chats.ts

@@ -8,17 +8,6 @@ import localforage from 'localforage';
 const CHATS_STORE_NAME = 'chats-store';
 const CHATS_STORE_PINIA_KEY = 'chats-store-pinia-key';
 
-// TODO: Search results are not yet implemented
-/**
- * Representation of a search result:
- * interface SearchResult {
- *  // embedding document id
- *  documentId: string;
- *  // embedding content
- *  content: string;
- * }
- */
-
 type ChatsStoreState = {
   version: number;
   chats: Chat[];

src/stores/knowledge.ts

@@ -1,7 +1,7 @@
 import { defineStore } from 'pinia';
 import { v4 as uuidv4 } from 'uuid';
 
-import { KnowledgeBase, KnowledgeBaseIdentifier } from 'src/types/knowledge';
+import { KnowledgeBase, KnowledgeBaseIdentifier, KnowledgeDocument } from 'src/types/knowledge';
 import { useAccountStore } from 'stores/account';
 
 type KnowledgeStoreState = {
@@ -16,6 +16,16 @@ export const useKnowledgeStore = defineStore('knowledge', {
     knowledgeBaseIdentifiers: [],
     isLoaded: false,
   }),
+  getters: {
+    getDocumentsFrom: (state) => {
+      return (ids: string[]): KnowledgeDocument[] => {
+        return state.knowledgeBases
+          .filter((kb) => ids.includes(kb.id))
+          .map((kb) => kb.documents)
+          .flat();
+      };
+    },
+  },
   actions: {
     async load() {
       const { alephStorage } = useAccountStore();

src/types/knowledge.ts

@@ -1,5 +1,10 @@
 import { z } from 'zod';
 
+export type KnowledgeSearchResult = {
+  content: string;
+  distance: number;
+};
+
 const knowledgeDocumentChunk = z.object({
   content: z.string(),
   vector: z.array(z.number()),

src/utils/knowledge/document.ts

@@ -7,7 +7,7 @@ import { generateChunks } from 'src/utils/knowledge/embedding';
 export const processDocument = async (file: File): Promise<Omit<KnowledgeDocument, 'store'>> => {
   const fileInfo = await extractFileContent(file);
 
-  const chunks = await generateChunks(file.name, fileInfo.content);
+  const chunks = await generateChunks(fileInfo.content);
 
   return { ...fileInfo, id: uuidv4(), name: file.name, size: file.size, chunks };
 };