Chat improv (#401)

* Allow full files to be passed to context. * feat: tools and chat history/metadata clarify. * fix: loading state enum * error handling * fix: generation cleanup * rename: update->savechat * v0.5: tool calling * refactor: separate message components. * feat: tool call rendering * feat: tool call execution. * fix:saving chat. * fix:tool calls. * fix:tool calls. * rm vector field * c * results rendering * custom tool result render. * Add shadcn. * Button. * Markdown comp. * a * Respond to tool result.
reorproject · Sep 15, 2024 · 3e15f0b · 3e15f0b
1 parent cc972cc
commit 3e15f0b
Show file tree

Hide file tree

Showing 45 changed files with 1,431 additions and 894 deletions.
diff --git a/components.json b/components.json
@@ -0,0 +1,20 @@
+{
+  "$schema": "https://ui.shadcn.com/schema.json",
+  "style": "new-york",
+  "rsc": false,
+  "tsx": true,
+  "tailwind": {
+    "config": "tailwind.config.js",
+    "css": "src/styles/global.css",
+    "baseColor": "slate",
+    "cssVariables": true,
+    "prefix": ""
+  },
+  "aliases": {
+    "components": "@/components",
+    "utils": "@/lib/utils",
+    "ui": "@/components/ui",
+    "lib": "@/lib",
+    "hooks": "@/hooks"
+  }
+}
diff --git a/electron/main/electron-store/ipcHandlers.ts b/electron/main/electron-store/ipcHandlers.ts
@@ -3,7 +3,6 @@ import path from 'path'
 import { ipcMain } from 'electron'
 import Store from 'electron-store'
 import {
-  Tab,
   EmbeddingModelConfig,
   EmbeddingModelWithLocalPath,
   EmbeddingModelWithRepo,
@@ -14,7 +13,7 @@ import {
 import WindowsManager from '../common/windowManager'
 
 import { initializeAndMaybeMigrateStore } from './storeSchemaMigrator'
-import { Chat } from '@/components/Chat/types'
+import { Chat, ChatMetadata } from '@/components/Chat/types'
 
 export const registerStoreHandlers = (store: Store<StoreSchema>, windowsManager: WindowsManager) => {
   initializeAndMaybeMigrateStore(store)
@@ -131,7 +130,7 @@ export const registerStoreHandlers = (store: Store<StoreSchema>, windowsManager:
     store.set(StoreKeys.hasUserOpenedAppBefore, true)
   })
 
-  ipcMain.handle('get-all-chats', (event) => {
+  ipcMain.handle('get-all-chats-metadata', (event) => {
     const vaultDir = windowsManager.getVaultDirectoryForWinContents(event.sender)
 
     if (!vaultDir) {
@@ -140,33 +139,34 @@ export const registerStoreHandlers = (store: Store<StoreSchema>, windowsManager:
 
     const allHistories = store.get(StoreKeys.ChatHistories)
     const chatHistoriesCorrespondingToVault = allHistories?.[vaultDir] ?? []
-    return chatHistoriesCorrespondingToVault
+    return chatHistoriesCorrespondingToVault.map(({ messages, ...rest }) => rest) as ChatMetadata[]
   })
 
-  ipcMain.handle('update-chat', (event, newChat: Chat) => {
+  ipcMain.handle('save-chat', (event, newChat: Chat) => {
     const vaultDir = windowsManager.getVaultDirectoryForWinContents(event.sender)
-    const allChatHistories = store.get(StoreKeys.ChatHistories)
     if (!vaultDir) {
       return
     }
+
+    const allChatHistories = store.get(StoreKeys.ChatHistories)
     const chatHistoriesCorrespondingToVault = allChatHistories?.[vaultDir] ?? []
-    // check if chat history already exists. if it does, update it. if it doesn't append it
+
     const existingChatIndex = chatHistoriesCorrespondingToVault.findIndex((chat) => chat.id === newChat.id)
     if (existingChatIndex !== -1) {
       chatHistoriesCorrespondingToVault[existingChatIndex] = newChat
     } else {
       chatHistoriesCorrespondingToVault.push(newChat)
     }
-    // store.set(StoreKeys.ChatHistories, allChatHistories);
     store.set(StoreKeys.ChatHistories, {
       ...allChatHistories,
       [vaultDir]: chatHistoriesCorrespondingToVault,
     })
-
-    event.sender.send('update-chat-histories', chatHistoriesCorrespondingToVault)
   })
 
-  ipcMain.handle('get-chat', (event, chatId: string) => {
+  ipcMain.handle('get-chat', (event, chatId: string | undefined) => {
+    if (!chatId) {
+      return undefined
+    }
     const vaultDir = windowsManager.getVaultDirectoryForWinContents(event.sender)
     if (!vaultDir) {
       return null
@@ -176,7 +176,7 @@ export const registerStoreHandlers = (store: Store<StoreSchema>, windowsManager:
     return vaultChatHistories.find((chat) => chat.id === chatId)
   })
 
-  ipcMain.handle('delete-chat-at-id', (event, chatID: string | undefined) => {
+  ipcMain.handle('delete-chat', (event, chatID: string | undefined) => {
     if (!chatID) return
     const vaultDir = windowsManager.getVaultDirectoryForWinContents(event.sender)
 
@@ -188,56 +188,9 @@ export const registerStoreHandlers = (store: Store<StoreSchema>, windowsManager:
     const allChatHistories = chatHistoriesMap[vaultDir] || []
     const filteredChatHistories = allChatHistories.filter((item) => item.id !== chatID)
 
-    chatHistoriesMap[vaultDir] = filteredChatHistories.reverse()
+    chatHistoriesMap[vaultDir] = filteredChatHistories
     store.set(StoreKeys.ChatHistories, chatHistoriesMap)
   })
-
-  ipcMain.handle('get-current-open-files', () => store.get(StoreKeys.OpenTabs) || [])
-
-  ipcMain.handle('add-current-open-files', (event, tab: Tab) => {
-    if (tab === null) return
-    const openTabs: Tab[] = store.get(StoreKeys.OpenTabs) || []
-    const existingTab = openTabs.findIndex((item) => item.path === tab.path)
-
-    /* If tab is already open, do not do anything */
-    if (existingTab !== -1) return
-    openTabs.push(tab)
-    store.set(StoreKeys.OpenTabs, openTabs)
-  })
-
-  ipcMain.handle('remove-current-open-files', (event, tabId: string, idx: number, newIndex: number) => {
-    // Ensure indices are within range
-    const openTabs: Tab[] = store.get(StoreKeys.OpenTabs) || []
-    if (idx < 0 || idx >= openTabs.length || newIndex < 0 || newIndex >= openTabs.length) return
-    openTabs[idx].lastAccessed = false
-    openTabs[newIndex].lastAccessed = true
-    const updatedTabs = openTabs.filter((tab) => tab.id !== tabId)
-    store.set(StoreKeys.OpenTabs, updatedTabs)
-  })
-
-  ipcMain.handle('clear-current-open-files', () => {
-    store.set(StoreKeys.OpenTabs, [])
-  })
-
-  ipcMain.handle('update-current-open-files', (event, draggedIndex: number, targetIndex: number) => {
-    const openTabs: Tab[] = store.get(StoreKeys.OpenTabs) || []
-    if (draggedIndex < 0 || draggedIndex >= openTabs.length || targetIndex < 0 || targetIndex >= openTabs.length) return
-    ;[openTabs[draggedIndex], openTabs[targetIndex]] = [openTabs[targetIndex], openTabs[draggedIndex]]
-    store.set(StoreKeys.OpenTabs, openTabs)
-  })
-
-  ipcMain.handle('set-current-open-files', (event, tabs: Tab[]) => {
-    if (tabs) store.set(StoreKeys.OpenTabs, tabs)
-  })
-
-  ipcMain.handle('remove-current-open-files-by-path', (event, filePath: string) => {
-    if (!filePath) return
-    const openTabs: Tab[] = store.get(StoreKeys.OpenTabs) || []
-    // Filter out selected tab
-    const updatedTabs = openTabs.filter((tab) => tab.path !== filePath)
-    store.set(StoreKeys.OpenTabs, updatedTabs)
-    event.sender.send('remove-tab-after-deletion', updatedTabs)
-  })
 }
 
 export function getDefaultEmbeddingModelConfig(store: Store<StoreSchema>): EmbeddingModelConfig {

diff --git a/electron/main/electron-store/storeSchemaMigrator.ts b/electron/main/electron-store/storeSchemaMigrator.ts
@@ -1,5 +1,6 @@
 import Store from 'electron-store'
 
+import getDisplayableChatName from '../../../shared/utils'
 import { StoreKeys, StoreSchema } from './storeConfig'
 import { defaultEmbeddingModelRepos } from '../vector-database/embeddings'
 import { defaultOllamaAPI } from '../llm/models/ollama'
@@ -50,7 +51,7 @@ export function setupDefaultStoreValues(store: Store<StoreSchema>) {
   }
 
   if (!store.get(StoreKeys.ChunkSize)) {
-    store.set(StoreKeys.ChunkSize, 500)
+    store.set(StoreKeys.ChunkSize, 1000)
   }
 
   setupDefaultAnalyticsValue(store)
@@ -84,6 +85,29 @@ function ensureChatHistoryIsCorrectProperty(store: Store<StoreSchema>) {
   store.set(StoreKeys.ChatHistories, chatHistories)
 }
 
+function ensureChatHistoryHasDisplayNameAndTimestamp(store: Store<StoreSchema>) {
+  const chatHistories = store.get(StoreKeys.ChatHistories)
+  if (!chatHistories) {
+    return
+  }
+
+  Object.keys(chatHistories).forEach((vaultDir) => {
+    const chats = chatHistories[vaultDir]
+    chats.map((chat) => {
+      const outputChat = chat
+      if (!outputChat.displayName || outputChat.displayName === chat.id) {
+        outputChat.displayName = getDisplayableChatName(chat.messages)
+      }
+      if (!outputChat.timeOfLastMessage) {
+        outputChat.timeOfLastMessage = Date.now()
+      }
+      return outputChat
+    })
+  })
+
+  store.set(StoreKeys.ChatHistories, chatHistories)
+}
+
 export const initializeAndMaybeMigrateStore = (store: Store<StoreSchema>) => {
   const storeSchemaVersion = store.get(StoreKeys.SchemaVersion)
   if (storeSchemaVersion !== currentSchemaVersion) {
@@ -93,6 +117,6 @@ export const initializeAndMaybeMigrateStore = (store: Store<StoreSchema>) => {
   }
 
   ensureChatHistoryIsCorrectProperty(store)
-
+  ensureChatHistoryHasDisplayNameAndTimestamp(store)
   setupDefaultStoreValues(store)
 }
diff --git a/electron/main/filesystem/filesystem.ts b/electron/main/filesystem/filesystem.ts
@@ -6,6 +6,7 @@ import chokidar from 'chokidar'
 import { BrowserWindow } from 'electron'
 
 import { FileInfo, FileInfoTree, isFileNodeDirectory } from './types'
+import addExtensionToFilenameIfNoExtensionPresent from '../path/path'
 
 export const markdownExtensions = ['.md', '.markdown', '.mdown', '.mkdn', '.mkd']
 
@@ -97,11 +98,9 @@ export function GetFilesInfoList(directory: string): FileInfo[] {
 }
 
 export function GetFilesInfoListForListOfPaths(paths: string[]): FileInfo[] {
-  // so perhaps for this function, all we maybe need to do is remove
   const fileInfoTree = paths.map((_path) => GetFilesInfoTree(_path)).flat()
   const fileInfoList = flattenFileInfoTree(fileInfoTree)
 
-  // remove duplicates:
   const uniquePaths = new Set()
   const fileInfoListWithoutDuplicates = fileInfoList.filter((fileInfo) => {
     if (uniquePaths.has(fileInfo.path)) {
@@ -123,8 +122,9 @@ export function createFileRecursive(filePath: string, content: string, charset?:
   if (fs.existsSync(filePath)) {
     return
   }
+  const filePathWithExtension = addExtensionToFilenameIfNoExtensionPresent(filePath, markdownExtensions, '.md')
 
-  fs.writeFileSync(filePath, content, charset)
+  fs.writeFileSync(filePathWithExtension, content, charset)
 }
 
 export function updateFileListForRenderer(win: BrowserWindow, directory: string): void {

diff --git a/electron/main/filesystem/ipcHandlers.ts b/electron/main/filesystem/ipcHandlers.ts
@@ -6,12 +6,7 @@ import Store from 'electron-store'
 
 import WindowsManager from '../common/windowManager'
 import { StoreSchema } from '../electron-store/storeConfig'
-import { DBEntry } from '../vector-database/schema'
-import {
-  convertFileInfoListToDBItems,
-  orchestrateEntryMove,
-  updateFileInTable,
-} from '../vector-database/tableHelperFunctions'
+import { orchestrateEntryMove, updateFileInTable } from '../vector-database/tableHelperFunctions'
 
 import {
   GetFilesInfoTree,
@@ -21,7 +16,7 @@ import {
   startWatchingDirectory,
   updateFileListForRenderer,
 } from './filesystem'
-import { FileInfoTree, WriteFileProps, RenameFileProps } from './types'
+import { FileInfoTree, WriteFileProps, RenameFileProps, FileInfoWithContent } from './types'
 
 const registerFileHandlers = (store: Store<StoreSchema>, _windowsManager: WindowsManager) => {
   const windowsManager = _windowsManager
@@ -164,12 +159,10 @@ const registerFileHandlers = (store: Store<StoreSchema>, _windowsManager: Window
     orchestrateEntryMove(windowInfo.dbTableClient, sourcePath, destinationPath)
   })
 
-  ipcMain.handle('get-filesystem-paths-as-db-items', async (_event, filePaths: string[]): Promise<DBEntry[]> => {
+  ipcMain.handle('get-files', async (_event, filePaths: string[]): Promise<FileInfoWithContent[]> => {
     const fileItems = GetFilesInfoListForListOfPaths(filePaths)
-
-    const dbItems = await convertFileInfoListToDBItems(fileItems)
-
-    return dbItems.flat()
+    const fileContents = fileItems.map((fileItem) => fs.readFileSync(fileItem.path, 'utf-8'))
+    return fileItems.map((fileItem, index) => ({ ...fileItem, content: fileContents[index] }))
   })
 
   ipcMain.handle('get-files-in-directory', (event, dirName: string) => {

diff --git a/electron/main/filesystem/types.ts b/electron/main/filesystem/types.ts
@@ -6,6 +6,10 @@ export type FileInfo = {
   dateCreated: Date
 }
 
+export type FileInfoWithContent = FileInfo & {
+  content: string
+}
+
 export type FileInfoNode = FileInfo & {
   children?: FileInfoNode[]
 }

diff --git a/electron/main/vector-database/lanceTableWrapper.ts b/electron/main/vector-database/lanceTableWrapper.ts
@@ -11,7 +11,8 @@ export function unsanitizePathForFileSystem(dbPath: string): string {
 }
 
 export function convertRecordToDBType<T extends DBEntry | DBQueryResult>(record: Record<string, unknown>): T | null {
-  const recordWithType = record as T
+  const { vector, ...recordWithoutVector } = record
+  const recordWithType = recordWithoutVector as unknown as T
   recordWithType.notepath = unsanitizePathForFileSystem(recordWithType.notepath)
   return recordWithType
 }
@@ -100,12 +101,7 @@ class LanceDBTableWrapper {
     }
   }
 
-  async search(
-    query: string,
-    //   metricType: string,
-    limit: number,
-    filter?: string,
-  ): Promise<DBQueryResult[]> {
+  async search(query: string, limit: number, filter?: string): Promise<DBQueryResult[]> {
     const lanceQuery = await this.lanceTable.search(query).metricType(MetricType.Cosine).limit(limit)
 
     if (filter) {

diff --git a/electron/main/vector-database/schema.ts b/electron/main/vector-database/schema.ts
@@ -2,7 +2,6 @@ import { Schema, Field, Utf8, FixedSizeList, Float32, Float64, DateUnit, Date_ a
 
 export interface DBEntry {
   notepath: string
-  vector?: Float32Array
   content: string
   subnoteindex: number
   timeadded: Date

diff --git a/electron/preload/index.ts b/electron/preload/index.ts
@@ -6,12 +6,11 @@ import {
   LLMConfig,
   LLMAPIConfig,
   LLMGenerationParameters,
-  Tab,
 } from 'electron/main/electron-store/storeConfig'
-import { FileInfoTree, RenameFileProps, WriteFileProps } from 'electron/main/filesystem/types'
-import { DBEntry, DBQueryResult } from 'electron/main/vector-database/schema'
+import { FileInfoTree, FileInfoWithContent, RenameFileProps, WriteFileProps } from 'electron/main/filesystem/types'
+import { DBQueryResult } from 'electron/main/vector-database/schema'
 
-import { Chat } from '@/components/Chat/types'
+import { Chat, ChatMetadata } from '@/components/Chat/types'
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 type IPCHandler<T extends (...args: any[]) => any> = (...args: Parameters<T>) => Promise<ReturnType<T>>
@@ -23,10 +22,6 @@ function createIPCHandler<T extends (...args: any[]) => any>(channel: string): I
 
 const database = {
   search: createIPCHandler<(query: string, limit: number, filter?: string) => Promise<DBQueryResult[]>>('search'),
-  searchWithReranking:
-    createIPCHandler<(query: string, limit: number, filter?: string) => Promise<DBQueryResult[]>>(
-      'search-with-reranking',
-    ),
   deleteLanceDBEntriesByFilePath: createIPCHandler<(filePath: string) => Promise<void>>(
     'delete-lance-db-entries-by-filepath',
   ),
@@ -68,24 +63,14 @@ const electronStore = {
   setSpellCheckMode: createIPCHandler<(isSpellCheck: boolean) => Promise<void>>('set-spellcheck-mode'),
   getHasUserOpenedAppBefore: createIPCHandler<() => Promise<boolean>>('has-user-opened-app-before'),
   setHasUserOpenedAppBefore: createIPCHandler<() => Promise<void>>('set-user-has-opened-app-before'),
-  getAllChats: createIPCHandler<() => Promise<Chat[]>>('get-all-chats'),
-  updateChat: createIPCHandler<(chat: Chat) => Promise<void>>('update-chat'),
-  deleteChatAtID: createIPCHandler<(chatID: string) => Promise<void>>('delete-chat-at-id'),
-  getChat: createIPCHandler<(chatID: string) => Promise<Chat>>('get-chat'),
+  getAllChatsMetadata: createIPCHandler<() => Promise<ChatMetadata[]>>('get-all-chats-metadata'),
+  saveChat: createIPCHandler<(chat: Chat) => Promise<void>>('save-chat'),
+  deleteChat: createIPCHandler<(chatID: string) => Promise<void>>('delete-chat'),
+  getChat: createIPCHandler<(chatID: string | undefined) => Promise<Chat | undefined>>('get-chat'),
   getSBCompact: createIPCHandler<() => Promise<boolean>>('get-sb-compact'),
   setSBCompact: createIPCHandler<(isSBCompact: boolean) => Promise<void>>('set-sb-compact'),
   getEditorFlexCenter: createIPCHandler<() => Promise<boolean>>('get-editor-flex-center'),
   setEditorFlexCenter: createIPCHandler<(editorFlexCenter: boolean) => Promise<void>>('set-editor-flex-center'),
-  getCurrentOpenTabs: createIPCHandler<() => Promise<Tab[]>>('get-current-open-files'),
-  setCurrentOpenTabs: createIPCHandler<(action: string, args: any) => Promise<void>>('set-current-open-files'),
-  addOpenTabs: createIPCHandler<(tab: Tab) => Promise<void>>('add-current-open-files'),
-  removeOpenTabs:
-    createIPCHandler<(tabId: string, idx: number, newIndex: number) => Promise<void>>('remove-current-open-files'),
-  clearOpenTabs: createIPCHandler<() => Promise<void>>('clear-current-open-files'),
-  updateOpenTabs:
-    createIPCHandler<(draggedIndex: number, targetIndex: number) => Promise<void>>('update-current-open-files'),
-  selectOpenTabs: createIPCHandler<(tabs: Tab[]) => Promise<void>>('set-current-open-files'),
-  removeOpenTabsByPath: createIPCHandler<(path: string) => Promise<void>>('remove-current-open-files-by-path'),
 }
 
 const fileSystem = {
@@ -102,10 +87,8 @@ const fileSystem = {
   checkFileExists: createIPCHandler<(filePath: string) => Promise<boolean>>('check-file-exists'),
   deleteFile: createIPCHandler<(filePath: string) => Promise<void>>('delete-file'),
   moveFileOrDir: createIPCHandler<(sourcePath: string, destinationPath: string) => Promise<void>>('move-file-or-dir'),
-  getFilesystemPathsAsDBItems: createIPCHandler<(paths: string[]) => Promise<DBEntry[]>>(
-    'get-filesystem-paths-as-db-items',
-  ),
   getAllFilenamesInDirectory: createIPCHandler<(dirName: string) => Promise<string[]>>('get-files-in-directory'),
+  getFiles: createIPCHandler<(filePaths: string[]) => Promise<FileInfoWithContent[]>>('get-files'),
 }
 
 const path = {