Implement rag prototpye

package.json

@@ -53,6 +53,7 @@
     "@tanstack/react-query": "^5.56.2",
     "diff": "^7.0.0",
     "drizzle-orm": "^0.35.2",
+    "exponential-backoff": "^3.1.1",
     "fuzzysort": "^3.1.0",
     "groq-sdk": "^0.7.0",
     "js-tiktoken": "^1.0.15",

src/ChatView.tsx

@@ -13,12 +13,14 @@ import { SettingsProvider } from './contexts/settings-context'
 import SmartCopilotPlugin from './main'
 import { MentionableBlockData } from './types/mentionable'
 import { SmartCopilotSettings } from './types/settings'
+import { RAGEngine } from './utils/ragEngine'
 
 export class ChatView extends ItemView {
   private root: Root | null = null
   private settings: SmartCopilotSettings
   private initialChatProps?: ChatProps
   private chatRef: React.RefObject<ChatRef> = React.createRef()
+  private onRAGEngineChange: (ragEngine: RAGEngine) => void
 
   constructor(
     leaf: WorkspaceLeaf,
@@ -27,6 +29,9 @@ export class ChatView extends ItemView {
     super(leaf)
     this.settings = plugin.settings
     this.initialChatProps = plugin.initialChatProps
+    this.onRAGEngineChange = (ragEngine) => {
+      plugin.ragEngine = ragEngine
+    }
   }
 
   getViewType() {
@@ -70,11 +75,7 @@ export class ChatView extends ItemView {
         >
           <DarkModeProvider>
             <LLMProvider>
-              <RAGProvider
-                setPluginRAGEngine={(ragEngine) =>
-                  (this.plugin.ragEngine = ragEngine)
-                }
-              >
+              <RAGProvider onRAGEngineChange={this.onRAGEngineChange}>
                 <QueryClientProvider client={queryClient}>
                   <React.StrictMode>
                     <Chat ref={this.chatRef} {...this.initialChatProps} />

src/components/chat-view/Chat.tsx

@@ -45,7 +45,7 @@ const getNewInputMessage = (app: App): ChatUserMessage => {
   return {
     role: 'user',
     content: null,
-    parsedContent: null,
+    promptContent: null,
     id: uuidv4(),
     mentionables: [
       {
@@ -199,7 +199,7 @@ const Chat = forwardRef<ChatRef, ChatProps>((props, ref) => {
           await promptGenerator.generateRequestMessages({
             messages: newChatHistory,
             useVaultSearch,
-            setQueryProgress,
+            onQueryProgressChange: setQueryProgress,
           })
         setQueryProgress({
           type: 'idle',

src/components/chat-view/chat-input/plugins/mention/MentionPlugin.tsx

@@ -15,7 +15,7 @@ import { createPortal } from 'react-dom'
 import { serializeMentionable } from 'src/utils/mentionable'
 
 import { Mentionable } from '../../../../../types/mentionable'
-import { SearchResultItem } from '../../../../../utils/fuzzy-search'
+import { SearchableMentionable } from '../../../../../utils/fuzzy-search'
 import { getMentionableName } from '../../../../../utils/mentionable'
 import { getMentionableIcon } from '../../utils/get-metionable-icon'
 import { MenuOption, MenuTextMatch } from '../shared/LexicalMenu'
@@ -105,7 +105,7 @@ class MentionTypeaheadOption extends MenuOption {
   mentionable: Mentionable
   icon: React.ReactNode
 
-  constructor(result: SearchResultItem) {
+  constructor(result: SearchableMentionable) {
     switch (result.type) {
       case 'file':
         super(result.file.path)
@@ -167,7 +167,7 @@ export default function NewMentionsPlugin({
   searchResultByQuery,
   onAddMention,
 }: {
-  searchResultByQuery: (query: string) => SearchResultItem[]
+  searchResultByQuery: (query: string) => SearchableMentionable[]
   onAddMention: (mentionable: Mentionable) => void
 }): JSX.Element | null {
   const [editor] = useLexicalComposerContext()

src/contexts/rag-context.tsx

@@ -3,10 +3,9 @@ import {
   createContext,
   useContext,
   useEffect,
-  useMemo,
+  useState,
 } from 'react'
 
-import { getEmbeddingModel } from '../utils/embedding'
 import { RAGEngine } from '../utils/ragEngine'
 
 import { useApp } from './app-context'
@@ -19,40 +18,28 @@ export type RAGContextType = {
 const RAGContext = createContext<RAGContextType | null>(null)
 
 export function RAGProvider({
-  setPluginRAGEngine,
   children,
-}: PropsWithChildren<{ setPluginRAGEngine: (ragEngine: RAGEngine) => void }>) {
+  onRAGEngineChange,
+}: PropsWithChildren<{ onRAGEngineChange: (ragEngine: RAGEngine) => void }>) {
   const app = useApp()
   const { settings } = useSettings()
+  const [ragEngine, setRagEngine] = useState<RAGEngine | null>(null)
 
-  const ragEngine = useMemo(() => {
-    return new RAGEngine(app, settings)
+  useEffect(() => {
+    RAGEngine.create(app, settings).then(setRagEngine)
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [app])
 
-  const embeddingModel = useMemo(() => {
-    return getEmbeddingModel(settings.embeddingModel, {
-      openAIApiKey: settings.openAIApiKey,
-    })
-  }, [settings.embeddingModel, settings.openAIApiKey])
-
   useEffect(() => {
-    void ragEngine.initialize(embeddingModel)
-    setPluginRAGEngine(ragEngine)
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [ragEngine])
+    if (ragEngine) {
+      onRAGEngineChange(ragEngine)
+    }
+  }, [ragEngine, onRAGEngineChange])
 
   useEffect(() => {
     ragEngine?.setSettings(settings)
   }, [ragEngine, settings])
 
-  useEffect(() => {
-    if (!ragEngine || !embeddingModel) {
-      return
-    }
-    ragEngine.setEmbeddingModel(embeddingModel)
-  }, [ragEngine, embeddingModel])
-
   return (
     <RAGContext.Provider value={{ ragEngine }}>{children}</RAGContext.Provider>
   )

src/hooks/useChatHistory.ts

@@ -23,7 +23,7 @@ const serializeChatMessage = (message: ChatMessage): SerializedChatMessage => {
       return {
         role: 'user',
         content: message.content,
-        parsedContent: message.parsedContent,
+        promptContent: message.promptContent,
         id: message.id,
         mentionables: message.mentionables.map(serializeMentionable),
       }
@@ -45,7 +45,7 @@ const deserializeChatMessage = (
       return {
         role: 'user',
         content: message.content,
-        parsedContent: message.parsedContent,
+        promptContent: message.promptContent,
         id: message.id,
         mentionables: message.mentionables
           .map((m) => deserializeMentionable(m, app))

src/main.ts

@@ -51,8 +51,9 @@ export default class SmartCopilotPlugin extends Plugin {
       callback: async () => {
         const notice = new Notice('Re-indexing vault...', 0)
         try {
-          await this.ragEngine?.updateVaultIndex(
-            { overwrite: true },
+          const ragEngine = await this.getRAGEngine()
+          await ragEngine.updateVaultIndex(
+            { reindexAll: true },
             (queryProgress) => {
               if (queryProgress.type === 'indexing') {
                 const { completedChunks, totalChunks } =
@@ -81,8 +82,9 @@ export default class SmartCopilotPlugin extends Plugin {
       callback: async () => {
         const notice = new Notice('Updating vault index...', 0)
         try {
-          await this.ragEngine?.updateVaultIndex(
-            { overwrite: false },
+          const ragEngine = await this.getRAGEngine()
+          await ragEngine.updateVaultIndex(
+            { reindexAll: false },
             (queryProgress) => {
               if (queryProgress.type === 'indexing') {
                 const { completedChunks, totalChunks } =
@@ -181,4 +183,11 @@ export default class SmartCopilotPlugin extends Plugin {
     chatView.addSelectionToChat(data)
     chatView.focusMessage()
   }
+
+  async getRAGEngine(): Promise<RAGEngine> {
+    if (!this.ragEngine) {
+      this.ragEngine = await RAGEngine.create(this.app, this.settings)
+    }
+    return this.ragEngine
+  }
 }

src/types/chat.ts

@@ -5,7 +5,7 @@ import { Mentionable, SerializedMentionable } from './mentionable'
 export type ChatUserMessage = {
   role: 'user'
   content: SerializedEditorState | null
-  parsedContent: string | null
+  promptContent: string | null
   id: string
   mentionables: Mentionable[]
 }
@@ -19,7 +19,7 @@ export type ChatMessage = ChatUserMessage | ChatAssistantMessage
 export type SerializedChatUserMessage = {
   role: 'user'
   content: SerializedEditorState | null
-  parsedContent: string | null
+  promptContent: string | null
   id: string
   mentionables: SerializedMentionable[]
 }

src/utils/fuzzy-search.ts

@@ -9,13 +9,11 @@ import {
 
 import { calculateFileDistance, getOpenFiles } from './obsidian'
 
-export type SearchResultItem =
+export type SearchableMentionable =
   | MentionableFile
   | MentionableFolder
   | MentionableVault
 
-type SearchItem = FolderWithMetadata | FileWithMetadata | VaultSearchItem
-
 type VaultSearchItem = {
   type: 'vault'
   path: string
@@ -36,6 +34,8 @@ type FolderWithMetadata = {
   folder: TFolder
 }
 
+type SearchItem = FolderWithMetadata | FileWithMetadata | VaultSearchItem
+
 function scoreFnWithBoost(score: number, searchItem: SearchItem) {
   switch (searchItem.type) {
     case 'file': {
@@ -76,7 +76,7 @@ function scoreFnWithBoost(score: number, searchItem: SearchItem) {
 function getEmptyQueryResult(
   searchItems: SearchItem[],
   limit: number,
-): SearchResultItem[] {
+): SearchableMentionable[] {
   // Sort files based on a custom scoring function
   const sortedFiles = searchItems.sort((a, b) => {
     const scoreA = scoreFnWithBoost(0.5, a) // Use 0.5 as a base score
@@ -85,27 +85,12 @@ function getEmptyQueryResult(
   })
 
   // Return only the top 'limit' files
-  return sortedFiles.slice(0, limit).map((item) => {
-    switch (item.type) {
-      case 'file':
-        return {
-          type: 'file',
-          file: item.file,
-        }
-      case 'folder':
-        return {
-          type: 'folder',
-          folder: item.folder,
-        }
-      case 'vault':
-        return {
-          type: 'vault',
-        }
-    }
-  })
+  return sortedFiles
+    .slice(0, limit)
+    .map((item) => searchItemToMentionable(item))
 }
 
-export function fuzzySearch(app: App, query: string): SearchResultItem[] {
+export function fuzzySearch(app: App, query: string): SearchableMentionable[] {
   const currentFile = app.workspace.getActiveFile()
   const openFiles = getOpenFiles(app)
 
@@ -158,22 +143,24 @@ export function fuzzySearch(app: App, query: string): SearchResultItem[] {
     scoreFn: (result) => scoreFnWithBoost(result.score, result.obj),
   })
 
-  return results.map((result) => {
-    switch (result.obj.type) {
-      case 'file':
-        return {
-          type: 'file',
-          file: result.obj.file,
-        }
-      case 'folder':
-        return {
-          type: 'folder',
-          folder: result.obj.folder,
-        }
-      case 'vault':
-        return {
-          type: 'vault',
-        }
-    }
-  })
+  return results.map((result) => searchItemToMentionable(result.obj))
+}
+
+function searchItemToMentionable(item: SearchItem): SearchableMentionable {
+  switch (item.type) {
+    case 'file':
+      return {
+        type: 'file',
+        file: item.file,
+      }
+    case 'folder':
+      return {
+        type: 'folder',
+        folder: item.folder,
+      }
+    case 'vault':
+      return {
+        type: 'vault',
+      }
+  }
 }