From e30c91fd57727d5b1cb65ff32c4ca2d1b09eb396 Mon Sep 17 00:00:00 2001 From: Logan Yang Date: Fri, 8 Mar 2024 21:50:45 -0800 Subject: [PATCH] Create vector store per vault (#348) --- .editorconfig | 1 + src/LLMProviders/chainManager.ts | 335 ++++++++++++++++++------------- src/main.ts | 156 +++++++------- src/vectorDBManager.ts | 134 +++++++------ 4 files changed, 351 insertions(+), 275 deletions(-) diff --git a/.editorconfig b/.editorconfig index 72201e7b..8d8cada3 100644 --- a/.editorconfig +++ b/.editorconfig @@ -8,3 +8,4 @@ insert_final_newline = true indent_style = space indent_size = 2 tab_width = 2 +trim_trailing_whitespace = true diff --git a/src/LLMProviders/chainManager.ts b/src/LLMProviders/chainManager.ts index ee03c59f..7afa8b9c 100644 --- a/src/LLMProviders/chainManager.ts +++ b/src/LLMProviders/chainManager.ts @@ -1,39 +1,37 @@ -import { LangChainParams, SetChainOptions } from '@/aiParams'; -import ChainFactory, { - ChainType, - Document, -} from '@/chainFactory'; -import { - AI_SENDER, - ChatModelDisplayNames -} from '@/constants'; -import EncryptionService from '@/encryptionService'; -import { ProxyChatOpenAI } from '@/langchainWrappers'; +import { LangChainParams, SetChainOptions } from "@/aiParams"; +import ChainFactory, { ChainType, Document } from "@/chainFactory"; +import { AI_SENDER, ChatModelDisplayNames } from "@/constants"; +import EncryptionService from "@/encryptionService"; +import { ProxyChatOpenAI } from "@/langchainWrappers"; import { CopilotSettings } from "@/settings/SettingsPage"; -import { ChatMessage } from '@/sharedState'; +import { ChatMessage } from "@/sharedState"; import { extractChatHistory, extractUniqueTitlesFromDocs, getModelName, - isSupportedChain -} from '@/utils'; -import VectorDBManager, { MemoryVector, NoteFile, VectorStoreDocument } from '@/vectorDBManager'; + isSupportedChain, +} from "@/utils"; +import VectorDBManager, { + MemoryVector, + NoteFile, + VectorStoreDocument, +} from "@/vectorDBManager"; import { ChatOllama } from "@langchain/community/chat_models/ollama"; import { RunnableSequence } from "@langchain/core/runnables"; import { BaseChatMemory } from "langchain/memory"; import { ChatPromptTemplate, HumanMessagePromptTemplate, - MessagesPlaceholder + MessagesPlaceholder, } from "langchain/prompts"; import { MultiQueryRetriever } from "langchain/retrievers/multi_query"; import { ScoreThresholdRetriever } from "langchain/retrievers/score_threshold"; import { MemoryVectorStore } from "langchain/vectorstores/memory"; -import { App, Notice } from 'obsidian'; -import ChatModelManager from './chatModelManager'; -import EmbeddingsManager from './embeddingManager'; -import MemoryManager from './memoryManager'; -import PromptManager from './promptManager'; +import { App, Notice } from "obsidian"; +import ChatModelManager from "./chatModelManager"; +import EmbeddingsManager from "./embeddingManager"; +import MemoryManager from "./memoryManager"; +import PromptManager from "./promptManager"; export default class ChainManager { private static chain: RunnableSequence; @@ -52,6 +50,7 @@ export default class ChainManager { public chatModelManager: ChatModelManager; public langChainParams: LangChainParams; public memoryManager: MemoryManager; + private getDbVectorStores: () => PouchDB.Database; /** * Constructor for initializing langChainParams and instantiating singletons. @@ -64,6 +63,8 @@ export default class ChainManager { langChainParams: LangChainParams, encryptionService: EncryptionService, settings: CopilotSettings, + // Ensure ChainManager always has the up-to-date dbVectorStores + getDbVectorStores: () => PouchDB.Database, ) { // Instantiate singletons this.app = app; @@ -71,8 +72,12 @@ export default class ChainManager { this.settings = settings; this.memoryManager = MemoryManager.getInstance(this.langChainParams); this.encryptionService = encryptionService; - this.chatModelManager = ChatModelManager.getInstance(this.langChainParams, encryptionService); + this.chatModelManager = ChatModelManager.getInstance( + this.langChainParams, + encryptionService, + ); this.promptManager = PromptManager.getInstance(this.langChainParams); + this.getDbVectorStores = getDbVectorStores; this.createChainWithNewModel(this.langChainParams.modelDisplayName); } @@ -81,7 +86,8 @@ export default class ChainManager { } private validateChainType(chainType: ChainType): void { - if (chainType === undefined || chainType === null) throw new Error('No chain type set'); + if (chainType === undefined || chainType === null) + throw new Error("No chain type set"); } static storeRetrieverDocuments(documents: Document[]) { @@ -96,8 +102,10 @@ export default class ChainManager { * @return {void} */ createChainWithNewModel(newModelDisplayName: string): void { - ChainManager.isOllamaModelActive = newModelDisplayName === ChatModelDisplayNames.OLLAMA; - ChainManager.isOpenRouterModelActive = newModelDisplayName === ChatModelDisplayNames.OPENROUTERAI; + ChainManager.isOllamaModelActive = + newModelDisplayName === ChatModelDisplayNames.OLLAMA; + ChainManager.isOpenRouterModelActive = + newModelDisplayName === ChatModelDisplayNames.OPENROUTERAI; // model and model display name must be update at the same time! let newModel = getModelName(newModelDisplayName); @@ -106,7 +114,7 @@ export default class ChainManager { newModel = this.langChainParams.ollamaModel; break; case ChatModelDisplayNames.LM_STUDIO: - newModel = 'check_model_in_lm_studio_ui'; + newModel = "check_model_in_lm_studio_ui"; break; case ChatModelDisplayNames.OPENROUTERAI: newModel = this.langChainParams.openRouterModel; @@ -120,10 +128,10 @@ export default class ChainManager { // Must update the chatModel for chain because ChainFactory always // retrieves the old chain without the chatModel change if it exists! // Create a new chain with the new chatModel - this.createChain( - this.langChainParams.chainType, - { ...this.langChainParams.options, forceNewCreation: true }, - ) + this.createChain(this.langChainParams.chainType, { + ...this.langChainParams.options, + forceNewCreation: true, + }); console.log(`Setting model to ${newModelDisplayName}: ${newModel}`); } catch (error) { console.error("createChainWithNewModel failed: ", error); @@ -132,16 +140,13 @@ export default class ChainManager { } /* Create a new chain, or update chain with new model */ - createChain( - chainType: ChainType, - options?: SetChainOptions, - ): void { + createChain(chainType: ChainType, options?: SetChainOptions): void { this.validateChainType(chainType); try { this.setChain(chainType, options); } catch (error) { - new Notice('Error creating chain:', error); - console.error('Error creating chain:', error); + new Notice("Error creating chain:", error); + console.error("Error creating chain:", error); } } @@ -149,15 +154,25 @@ export default class ChainManager { chainType: ChainType, options: SetChainOptions = {}, ): Promise { - if (!this.chatModelManager.validateChatModel(this.chatModelManager.getChatModel())) { + if ( + !this.chatModelManager.validateChatModel( + this.chatModelManager.getChatModel(), + ) + ) { // No need to throw error and trigger multiple Notices to user - console.error('setChain failed: No chat model set.'); + console.error("setChain failed: No chat model set."); return; } this.validateChainType(chainType); // MUST set embeddingsManager when switching to QA mode - if (chainType === ChainType.LONG_NOTE_QA_CHAIN || chainType === ChainType.VAULT_QA_CHAIN) { - this.embeddingsManager = EmbeddingsManager.getInstance(this.langChainParams, this.encryptionService); + if ( + chainType === ChainType.LONG_NOTE_QA_CHAIN || + chainType === ChainType.VAULT_QA_CHAIN + ) { + this.embeddingsManager = EmbeddingsManager.getInstance( + this.langChainParams, + this.encryptionService, + ); } // Get chatModel, memory, prompt, and embeddingAPI from respective managers @@ -173,7 +188,8 @@ export default class ChainManager { if (ChainManager.isOllamaModelActive) { (chatModel as ChatOllama).model = this.langChainParams.ollamaModel; } else if (ChainManager.isOpenRouterModelActive) { - (chatModel as ProxyChatOpenAI).modelName = this.langChainParams.openRouterModel; + (chatModel as ProxyChatOpenAI).modelName = + this.langChainParams.openRouterModel; } ChainManager.chain = ChainFactory.createNewLLMChain({ @@ -197,16 +213,22 @@ export default class ChainManager { } case ChainType.LONG_NOTE_QA_CHAIN: { if (!options.noteFile) { - new Notice('No note content provided'); - throw new Error('No note content provided'); + new Notice("No note content provided"); + throw new Error("No note content provided"); } this.setNoteFile(options.noteFile); const docHash = VectorDBManager.getDocumentHash(options.noteFile.path); - const parsedMemoryVectors: MemoryVector[] | undefined = await VectorDBManager.getMemoryVectors(docHash); + const parsedMemoryVectors: MemoryVector[] | undefined = + await VectorDBManager.getMemoryVectors( + this.getDbVectorStores(), + docHash, + ); const embeddingsAPI = this.embeddingsManager.getEmbeddingsAPI(); if (!embeddingsAPI) { - console.error('Error getting embeddings API. Please check your settings.'); + console.error( + "Error getting embeddings API. Please check your settings.", + ); return; } if (parsedMemoryVectors) { @@ -217,67 +239,70 @@ export default class ChainManager { ); // Create new conversational retrieval chain - ChainManager.retrievalChain = ChainFactory.createConversationalRetrievalChain( - { - llm: chatModel, - retriever: vectorStore.asRetriever( - undefined, - (doc) => { + ChainManager.retrievalChain = + ChainFactory.createConversationalRetrievalChain( + { + llm: chatModel, + retriever: vectorStore.asRetriever(undefined, (doc) => { return doc.metadata.path === options.noteFile?.path; - } - ), - }, - ChainManager.storeRetrieverDocuments.bind(ChainManager), - ) - console.log('Existing vector store for document hash: ', docHash); + }), + }, + ChainManager.storeRetrieverDocuments.bind(ChainManager), + ); + console.log("Existing vector store for document hash: ", docHash); } else { // Index doesn't exist const vectorStoreDoc = await this.indexFile(options.noteFile); this.vectorStore = await VectorDBManager.getMemoryVectorStore( + this.getDbVectorStores(), embeddingsAPI, vectorStoreDoc?._id, ); if (!this.vectorStore) { - console.error('Error creating vector store.'); + console.error("Error creating vector store."); return; } const retriever = MultiQueryRetriever.fromLLM({ llm: chatModel, - retriever: this.vectorStore.asRetriever( - undefined, - (doc) => { - return doc.metadata.path === options.noteFile?.path; - } - ), + retriever: this.vectorStore.asRetriever(undefined, (doc) => { + return doc.metadata.path === options.noteFile?.path; + }), verbose: false, }); - ChainManager.retrievalChain = ChainFactory.createConversationalRetrievalChain( - { - llm: chatModel, - retriever: retriever, - }, - ChainManager.storeRetrieverDocuments.bind(ChainManager) - ) + ChainManager.retrievalChain = + ChainFactory.createConversationalRetrievalChain( + { + llm: chatModel, + retriever: retriever, + }, + ChainManager.storeRetrieverDocuments.bind(ChainManager), + ); console.log( - 'New conversational retrieval QA chain with multi-query retriever created for ' - + 'document hash: ', docHash + "New conversational retrieval QA chain with multi-query retriever created for " + + "document hash: ", + docHash, ); } this.langChainParams.chainType = ChainType.LONG_NOTE_QA_CHAIN; - console.log('Set chain:', ChainType.LONG_NOTE_QA_CHAIN); + console.log("Set chain:", ChainType.LONG_NOTE_QA_CHAIN); break; } case ChainType.VAULT_QA_CHAIN: { const embeddingsAPI = this.embeddingsManager.getEmbeddingsAPI(); if (!embeddingsAPI) { - console.error('Error getting embeddings API. Please check your settings.'); + console.error( + "Error getting embeddings API. Please check your settings.", + ); return; } - const vectorStore = await VectorDBManager.getMemoryVectorStore(embeddingsAPI) + const vectorStore = await VectorDBManager.getMemoryVectorStore( + this.getDbVectorStores(), + embeddingsAPI, + ); const retriever = ScoreThresholdRetriever.fromVectorStore(vectorStore, { minSimilarityScore: 0.3, // TODO: Make this a user setting maxK: this.settings.maxSourceChunks, // The maximum number of docs (chunks) to retrieve @@ -285,19 +310,20 @@ export default class ChainManager { }); // Create new conversational retrieval chain - ChainManager.retrievalChain = ChainFactory.createConversationalRetrievalChain( - { - llm: chatModel, - retriever: retriever, - }, - ChainManager.storeRetrieverDocuments.bind(ChainManager) - ) + ChainManager.retrievalChain = + ChainFactory.createConversationalRetrievalChain( + { + llm: chatModel, + retriever: retriever, + }, + ChainManager.storeRetrieverDocuments.bind(ChainManager), + ); console.log( - 'New conversational retrieval QA chain with multi-query retriever created for entire vault' + "New conversational retrieval QA chain with multi-query retriever created for entire vault", ); this.langChainParams.chainType = ChainType.VAULT_QA_CHAIN; - console.log('Set chain:', ChainType.VAULT_QA_CHAIN); + console.log("Set chain:", ChainType.VAULT_QA_CHAIN); break; } @@ -313,19 +339,21 @@ export default class ChainManager { updateCurrentAiMessage: (message: string) => void, addMessage: (message: ChatMessage) => void, options: { - debug?: boolean, - ignoreSystemMessage?: boolean, - updateLoading?: (loading: boolean) => void + debug?: boolean; + ignoreSystemMessage?: boolean; + updateLoading?: (loading: boolean) => void; } = {}, ) { - const { - debug = false, - ignoreSystemMessage = false, - } = options; + const { debug = false, ignoreSystemMessage = false } = options; // Check if chat model is initialized - if (!this.chatModelManager.validateChatModel(this.chatModelManager.getChatModel())) { - const errorMsg = 'Chat model is not initialized properly, check your API key in Copilot setting and make sure you have API access.'; + if ( + !this.chatModelManager.validateChatModel( + this.chatModelManager.getChatModel(), + ) + ) { + const errorMsg = + "Chat model is not initialized properly, check your API key in Copilot setting and make sure you have API access."; new Notice(errorMsg); console.error(errorMsg); return; @@ -333,10 +361,13 @@ export default class ChainManager { // Check if chain is initialized properly if (!ChainManager.chain || !isSupportedChain(ChainManager.chain)) { console.error( - 'Chain is not initialized properly, re-initializing chain: ', - this.langChainParams.chainType + "Chain is not initialized properly, re-initializing chain: ", + this.langChainParams.chainType, + ); + this.setChain( + this.langChainParams.chainType, + this.langChainParams.options, ); - this.setChain(this.langChainParams.chainType, this.langChainParams.options); } const { @@ -349,14 +380,14 @@ export default class ChainManager { const memory = this.memoryManager.getMemory(); const chatPrompt = this.promptManager.getChatPrompt(); - const systemPrompt = ignoreSystemMessage ? '' : systemMessage; + const systemPrompt = ignoreSystemMessage ? "" : systemMessage; // Whether to ignore system prompt (for commands) if (ignoreSystemMessage) { const effectivePrompt = ignoreSystemMessage ? ChatPromptTemplate.fromMessages([ - new MessagesPlaceholder("history"), - HumanMessagePromptTemplate.fromTemplate("{input}"), - ]) + new MessagesPlaceholder("history"), + HumanMessagePromptTemplate.fromTemplate("{input}"), + ]) : chatPrompt; this.setChain(chainType, { @@ -367,28 +398,31 @@ export default class ChainManager { this.setChain(chainType, this.langChainParams.options); } - let fullAIResponse = ''; + let fullAIResponse = ""; // eslint-disable-next-line @typescript-eslint/no-explicit-any const chatModel = (ChainManager.chain as any).last.bound; // eslint-disable-next-line @typescript-eslint/no-explicit-any - const chatStream = await ChainManager.chain.stream({ input: userMessage } as any); + const chatStream = await ChainManager.chain.stream({ + input: userMessage, + } as any); try { switch (chainType) { case ChainType.LLM_CHAIN: if (debug) { - console.log(`*** DEBUG INFO ***\n` - + `user message: ${userMessage}\n` - // ChatOpenAI has modelName, some other ChatModels like ChatOllama have model - + `model: ${chatModel.modelName || chatModel.model}\n` - + `chain type: ${chainType}\n` - + `temperature: ${temperature}\n` - + `maxTokens: ${maxTokens}\n` - + `system prompt: ${systemPrompt}\n` - + `chat context turns: ${chatContextTurns}\n`, + console.log( + `*** DEBUG INFO ***\n` + + `user message: ${userMessage}\n` + + // ChatOpenAI has modelName, some other ChatModels like ChatOllama have model + `model: ${chatModel.modelName || chatModel.model}\n` + + `chain type: ${chainType}\n` + + `temperature: ${temperature}\n` + + `maxTokens: ${maxTokens}\n` + + `system prompt: ${systemPrompt}\n` + + `chat context turns: ${chatContextTurns}\n`, ); - console.log('chain RunnableSequence:', ChainManager.chain); - console.log('Chat memory:', memory); + console.log("chain RunnableSequence:", ChainManager.chain); + console.log("Chat memory:", memory); } for await (const chunk of chatStream) { @@ -400,30 +434,42 @@ export default class ChainManager { case ChainType.LONG_NOTE_QA_CHAIN: case ChainType.VAULT_QA_CHAIN: if (debug) { - console.log(`*** DEBUG INFO ***\n` - + `user message: ${userMessage}\n` - + `model: ${chatModel.modelName || chatModel.model}\n` - + `chain type: ${chainType}\n` - + `temperature: ${temperature}\n` - + `maxTokens: ${maxTokens}\n` - + `system prompt: ${systemPrompt}\n` - + `chat context turns: ${chatContextTurns}\n`, + console.log( + `*** DEBUG INFO ***\n` + + `user message: ${userMessage}\n` + + `model: ${chatModel.modelName || chatModel.model}\n` + + `chain type: ${chainType}\n` + + `temperature: ${temperature}\n` + + `maxTokens: ${maxTokens}\n` + + `system prompt: ${systemPrompt}\n` + + `chat context turns: ${chatContextTurns}\n`, + ); + console.log("chain RunnableSequence:", ChainManager.chain); + console.log( + "embedding model:", + this.langChainParams.embeddingModel, ); - console.log('chain RunnableSequence:', ChainManager.chain); - console.log('embedding model:', this.langChainParams.embeddingModel); } fullAIResponse = await this.runRetrievalChain( - userMessage, memory, updateCurrentAiMessage, abortController, { debug }, + userMessage, + memory, + updateCurrentAiMessage, + abortController, + { debug }, ); break; default: - console.error('Chain type not supported:', this.langChainParams.chainType); + console.error( + "Chain type not supported:", + this.langChainParams.chainType, + ); } } catch (error) { const errorData = error?.response?.data?.error || error; const errorCode = errorData?.code || error; - if (errorCode === 'model_not_found') { - const modelNotFoundMsg = "You do not have access to this model or the model does not exist, please check with your API provider."; + if (errorCode === "model_not_found") { + const modelNotFoundMsg = + "You do not have access to this model or the model does not exist, please check with your API provider."; new Notice(modelNotFoundMsg); console.error(modelNotFoundMsg); } else { @@ -435,7 +481,7 @@ export default class ChainManager { // This line is a must for memory to work with RunnableSequence! await memory.saveContext( { input: userMessage }, - { output: fullAIResponse } + { output: fullAIResponse }, ); addMessage({ message: fullAIResponse, @@ -443,7 +489,7 @@ export default class ChainManager { isVisible: true, }); } - updateCurrentAiMessage(''); + updateCurrentAiMessage(""); } return fullAIResponse; } @@ -454,7 +500,7 @@ export default class ChainManager { updateCurrentAiMessage: (message: string) => void, abortController: AbortController, options: { - debug?: boolean, + debug?: boolean; } = {}, ): Promise { const memoryVariables = await memory.loadMemoryVariables({}); @@ -465,7 +511,7 @@ export default class ChainManager { // eslint-disable-next-line @typescript-eslint/no-explicit-any } as any); - let fullAIResponse = ''; + let fullAIResponse = ""; for await (const chunk of qaStream) { if (abortController.signal.aborted) break; @@ -474,8 +520,8 @@ export default class ChainManager { } if (options.debug) { - console.log('Max source chunks:', this.settings.maxSourceChunks); - console.log('Retrieved chunks:', ChainManager.retrievedDocuments); + console.log("Max source chunks:", this.settings.maxSourceChunks); + console.log("Retrieved chunks:", ChainManager.retrievedDocuments); } // TODO: This only returns unique note titles, but actual retrieved docs are chunks. @@ -483,13 +529,15 @@ export default class ChainManager { // to show specific chunks in the future. E.g. collapsed note title when clicked, // expand and reveal the chunk if (this.langChainParams.chainType === ChainType.VAULT_QA_CHAIN) { - const docTitles = extractUniqueTitlesFromDocs(ChainManager.retrievedDocuments); + const docTitles = extractUniqueTitlesFromDocs( + ChainManager.retrievedDocuments, + ); const markdownLinks = docTitles .map( (title) => `[${title}](obsidian://open?vault=${this.app.vault.getName()}&file=${encodeURIComponent( - title - )})` + title, + )})`, ) .join("\n"); fullAIResponse += "\n\n**Sources**:\n" + markdownLinks; @@ -498,14 +546,21 @@ export default class ChainManager { return fullAIResponse; } - async indexFile(noteFile: NoteFile): Promise { + async indexFile( + noteFile: NoteFile, + ): Promise { const embeddingsAPI = this.embeddingsManager.getEmbeddingsAPI(); if (!embeddingsAPI) { - const errorMsg = 'Failed to load file, embedding API is not set correctly, please check your settings.'; + const errorMsg = + "Failed to load file, embedding API is not set correctly, please check your settings."; new Notice(errorMsg); console.error(errorMsg); return; } - return await VectorDBManager.indexFile(noteFile, embeddingsAPI); + return await VectorDBManager.indexFile( + this.getDbVectorStores(), + embeddingsAPI, + noteFile, + ); } } diff --git a/src/main.ts b/src/main.ts index c8b831a8..ea735b07 100644 --- a/src/main.ts +++ b/src/main.ts @@ -12,14 +12,19 @@ import { CHAT_VIEWTYPE, DEFAULT_SETTINGS, DEFAULT_SYSTEM_PROMPT, - VAULT_VECTOR_STORE_STRATEGY + VAULT_VECTOR_STORE_STRATEGY, } from "@/constants"; import { CustomPrompt } from "@/customPromptProcessor"; import EncryptionService from "@/encryptionService"; import { CopilotSettingTab, CopilotSettings } from "@/settings/SettingsPage"; import SharedState from "@/sharedState"; -import { areEmbeddingModelsSame, getAllNotesContent, sanitizeSettings } from "@/utils"; +import { + areEmbeddingModelsSame, + getAllNotesContent, + sanitizeSettings, +} from "@/utils"; import VectorDBManager, { VectorStoreDocument } from "@/vectorDBManager"; +import { MD5 } from "crypto-js"; import { Server } from "http"; import { Editor, Notice, Plugin, TFile, WorkspaceLeaf } from "obsidian"; import PouchDB from "pouchdb"; @@ -33,7 +38,7 @@ export default class CopilotPlugin extends Plugin { activateViewPromise: Promise | null = null; chatIsVisible = false; dbPrompts: PouchDB.Database; - dbVectorStores: PouchDB.Database; + dbVectorStores: PouchDB.Database; embeddingsManager: EmbeddingsManager; encryptionService: EncryptionService; server: Server | null = null; @@ -47,32 +52,30 @@ export default class CopilotPlugin extends Plugin { this.sharedState = new SharedState(); const langChainParams = this.getChainManagerParams(); this.encryptionService = new EncryptionService(this.settings); + this.dbVectorStores = new PouchDB( + `copilot_vector_stores_${this.getVaultIdentifier()}`, + ); this.chainManager = new ChainManager( this.app, langChainParams, this.encryptionService, this.settings, + () => this.dbVectorStores, ); if (this.settings.enableEncryption) { await this.saveSettings(); } - this.dbVectorStores = new PouchDB( - "copilot_vector_stores" - ); this.embeddingsManager = EmbeddingsManager.getInstance( langChainParams, - this.encryptionService + this.encryptionService, ); this.dbPrompts = new PouchDB("copilot_custom_prompts"); - VectorDBManager.initializeDB(this.dbVectorStores); - VectorDBManager.setEmbeddingModel(this.settings.embeddingModel); - this.registerView( CHAT_VIEWTYPE, - (leaf: WorkspaceLeaf) => new CopilotView(leaf, this) + (leaf: WorkspaceLeaf) => new CopilotView(leaf, this), ); this.addCommand({ @@ -108,7 +111,7 @@ export default class CopilotPlugin extends Plugin { new Notice("Custom prompt saved successfully."); } catch (e) { new Notice( - "Error saving custom prompt. Please check if the title already exists." + "Error saving custom prompt. Please check if the title already exists.", ); console.error(e); } @@ -131,30 +134,30 @@ export default class CopilotPlugin extends Plugin { } try { const doc = (await this.dbPrompts.get( - promptTitle + promptTitle, )) as CustomPrompt; if (!doc.prompt) { new Notice( - `No prompt found with the title "${promptTitle}".` + `No prompt found with the title "${promptTitle}".`, ); return; } this.processCustomPrompt( editor, "applyCustomPrompt", - doc.prompt + doc.prompt, ); } catch (err) { if (err.name === "not_found") { new Notice( - `No prompt found with the title "${promptTitle}".` + `No prompt found with the title "${promptTitle}".`, ); } else { console.error(err); new Notice("An error occurred."); } } - } + }, ).open(); }); }, @@ -173,7 +176,7 @@ export default class CopilotPlugin extends Plugin { console.error(err); new Notice("An error occurred."); } - } + }, ); modal.open(); @@ -202,25 +205,25 @@ export default class CopilotPlugin extends Plugin { const doc = await this.dbPrompts.get(promptTitle); if (doc._rev) { await this.dbPrompts.remove( - doc as PouchDB.Core.RemoveDocument + doc as PouchDB.Core.RemoveDocument, ); new Notice(`Prompt "${promptTitle}" has been deleted.`); } else { new Notice( - `Failed to delete prompt "${promptTitle}": No revision found.` + `Failed to delete prompt "${promptTitle}": No revision found.`, ); } } catch (err) { if (err.name === "not_found") { new Notice( - `No prompt found with the title "${promptTitle}".` + `No prompt found with the title "${promptTitle}".`, ); } else { console.error(err); new Notice("An error occurred while deleting the prompt."); } } - } + }, ).open(); }); @@ -248,7 +251,7 @@ export default class CopilotPlugin extends Plugin { try { const doc = (await this.dbPrompts.get( - promptTitle + promptTitle, )) as CustomPrompt; if (doc.prompt) { new AddPromptModal( @@ -262,24 +265,24 @@ export default class CopilotPlugin extends Plugin { }, doc._id, doc.prompt, - true + true, ).open(); } else { new Notice( - `No prompt found with the title "${promptTitle}".` + `No prompt found with the title "${promptTitle}".`, ); } } catch (err) { if (err.name === "not_found") { new Notice( - `No prompt found with the title "${promptTitle}".` + `No prompt found with the title "${promptTitle}".`, ); } else { console.error(err); new Notice("An error occurred."); } } - } + }, ).open(); }); @@ -296,18 +299,16 @@ export default class CopilotPlugin extends Plugin { await this.dbVectorStores.destroy(); // Reinitialize the database this.dbVectorStores = new PouchDB( - "copilot_vector_stores" + `copilot_vector_stores_${this.getVaultIdentifier()}`, ); - // Make sure to update the instance with VectorDBManager - VectorDBManager.updateDBInstance(this.dbVectorStores); new Notice("Local vector store cleared successfully."); console.log( - "Local vector store cleared successfully, new instance created." + "Local vector store cleared successfully, new instance created.", ); } catch (err) { console.error("Error clearing the local vector store:", err); new Notice( - "An error occurred while clearing the local vector store." + "An error occurred while clearing the local vector store.", ); } }, @@ -320,15 +321,18 @@ export default class CopilotPlugin extends Plugin { try { const files = this.app.vault.getMarkdownFiles(); const filePaths = files.map((file) => file.path); - const indexedFiles = await VectorDBManager.getNoteFiles(); + const indexedFiles = await VectorDBManager.getNoteFiles( + this.dbVectorStores, + ); const indexedFilePaths = indexedFiles.map((file) => file.path); const filesToDelete = indexedFilePaths.filter( - (filePath) => !filePaths.includes(filePath) + (filePath) => !filePaths.includes(filePath), ); const deletePromises = filesToDelete.map(async (filePath) => { VectorDBManager.removeMemoryVectors( - VectorDBManager.getDocumentHash(filePath) + this.dbVectorStores, + VectorDBManager.getDocumentHash(filePath), ); }); @@ -336,12 +340,12 @@ export default class CopilotPlugin extends Plugin { new Notice("Local vector store garbage collected successfully."); console.log( - "Local vector store garbage collected successfully, new instance created." + "Local vector store garbage collected successfully, new instance created.", ); } catch (err) { console.error("Error clearing the local vector store:", err); new Notice( - "An error occurred while clearing the local vector store." + "An error occurred while clearing the local vector store.", ); } }, @@ -355,10 +359,10 @@ export default class CopilotPlugin extends Plugin { const indexedFileCount = await this.indexVaultToVectorStore(); new Notice( - `${indexedFileCount} vault files indexed to vector store.` + `${indexedFileCount} vault files indexed to vector store.`, ); console.log( - `${indexedFileCount} vault files indexed to vector store.` + `${indexedFileCount} vault files indexed to vector store.`, ); } catch (err) { console.error("Error indexing vault to vector store:", err); @@ -375,15 +379,15 @@ export default class CopilotPlugin extends Plugin { const indexedFileCount = await this.indexVaultToVectorStore(true); new Notice( - `${indexedFileCount} vault files indexed to vector store.` + `${indexedFileCount} vault files indexed to vector store.`, ); console.log( - `${indexedFileCount} vault files indexed to vector store.` + `${indexedFileCount} vault files indexed to vector store.`, ); } catch (err) { console.error("Error re-indexing vault to vector store:", err); new Notice( - "An error occurred while re-indexing vault to vector store." + "An error occurred while re-indexing vault to vector store.", ); } }, @@ -401,7 +405,7 @@ export default class CopilotPlugin extends Plugin { this.settings.chatNoteContextPath = path; this.settings.chatNoteContextTags = tags; await this.saveSettings(); - } + }, ).open(); }, }); @@ -409,15 +413,15 @@ export default class CopilotPlugin extends Plugin { this.registerEvent( this.app.vault.on("delete", (file) => { const docHash = VectorDBManager.getDocumentHash(file.path); - VectorDBManager.removeMemoryVectors(docHash); - }) + VectorDBManager.removeMemoryVectors(this.dbVectorStores, docHash); + }), ); // Index vault to vector store on startup and after loading all commands // This can take a while, so we don't want to block the startup process if ( this.settings.indexVaultToVectorStore === - VAULT_VECTOR_STORE_STRATEGY.ON_STARTUP + VAULT_VECTOR_STORE_STRATEGY.ON_STARTUP ) { try { await this.indexVaultToVectorStore(); @@ -428,6 +432,11 @@ export default class CopilotPlugin extends Plugin { } } + private getVaultIdentifier(): string { + const vaultName = this.app.vault.getName(); + return MD5(vaultName).toString(); + } + async saveFileToVectorStore(file: TFile): Promise { const embeddingInstance = this.embeddingsManager.getEmbeddingsAPI(); if (!embeddingInstance) { @@ -443,7 +452,7 @@ export default class CopilotPlugin extends Plugin { content: fileContent, metadata: fileMetadata?.frontmatter ?? {}, }; - VectorDBManager.indexFile(noteFile, embeddingInstance); + VectorDBManager.indexFile(this.dbVectorStores, embeddingInstance, noteFile); } async indexVaultToVectorStore(overwrite?: boolean): Promise { @@ -453,12 +462,17 @@ export default class CopilotPlugin extends Plugin { } // Check if embedding model has changed - const prevEmbeddingModel = await VectorDBManager.checkEmbeddingModel(); + const prevEmbeddingModel = await VectorDBManager.checkEmbeddingModel( + this.dbVectorStores, + ); // TODO: Remove this when Ollama model is dynamically set - const currEmbeddingModel = EmbeddingsManager.getModelName(embeddingInstance); + const currEmbeddingModel = + EmbeddingsManager.getModelName(embeddingInstance); console.log( - 'Prev vs Current embedding models:', prevEmbeddingModel, currEmbeddingModel + "Prev vs Current embedding models:", + prevEmbeddingModel, + currEmbeddingModel, ); if (!areEmbeddingModelsSame(prevEmbeddingModel, currEmbeddingModel)) { @@ -470,33 +484,33 @@ export default class CopilotPlugin extends Plugin { await this.dbVectorStores.destroy(); // Reinitialize the database this.dbVectorStores = new PouchDB( - "copilot_vector_stores" + `copilot_vector_stores_${this.getVaultIdentifier()}`, + ); + new Notice( + "Detected change in embedding model. Rebuild vector store from scratch.", ); - // Make sure to update the instance with VectorDBManager - VectorDBManager.updateDBInstance(this.dbVectorStores); - new Notice("Detected change in embedding model. Rebuild vector store from scratch."); console.log( - "Detected change in embedding model. Rebuild vector store from scratch." + "Detected change in embedding model. Rebuild vector store from scratch.", ); } catch (err) { console.error("Error clearing vector store for reindexing:", err); - new Notice( - "Error clearing vector store for reindexing." - ); + new Notice("Error clearing vector store for reindexing."); } } - const latestMtime = await VectorDBManager.getLatestFileMtime(); + const latestMtime = await VectorDBManager.getLatestFileMtime( + this.dbVectorStores, + ); const files = this.app.vault.getMarkdownFiles().filter((file) => { if (!latestMtime || overwrite) return true; return file.stat.mtime > latestMtime; }); const fileContents: string[] = await Promise.all( - files.map((file) => this.app.vault.cachedRead(file)) + files.map((file) => this.app.vault.cachedRead(file)), ); const fileMetadatas = files.map((file) => - this.app.metadataCache.getFileCache(file) + this.app.metadataCache.getFileCache(file), ); const totalFiles = files.length; @@ -507,8 +521,8 @@ export default class CopilotPlugin extends Plugin { let indexedCount = 0; const indexNotice = new Notice( - `Copilot is indexing your vault... 0/${totalFiles} files processed.`, - 0 + `Copilot is indexing your vault...\n0/${totalFiles} files processed.`, + 0, ); const loadPromises = files.map(async (file, index) => { @@ -520,12 +534,13 @@ export default class CopilotPlugin extends Plugin { metadata: fileMetadatas[index]?.frontmatter ?? {}, }; const result = await VectorDBManager.indexFile( + this.dbVectorStores, + embeddingInstance, noteFile, - embeddingInstance ); indexedCount++; indexNotice.setMessage( - `Copilot is indexing your vault... ${indexedCount}/${totalFiles} files processed.` + `Copilot is indexing your vault...\n${indexedCount}/${totalFiles} files processed.`, ); return result; }); @@ -541,7 +556,7 @@ export default class CopilotPlugin extends Plugin { editor: Editor, eventType: string, eventSubtype?: string, - checkSelectedText = true + checkSelectedText = true, ) { const selectedText = editor.getSelection(); @@ -586,7 +601,7 @@ export default class CopilotPlugin extends Plugin { }); await this.activateViewPromise; this.app.workspace.revealLeaf( - this.app.workspace.getLeavesOfType(CHAT_VIEWTYPE)[0] + this.app.workspace.getLeavesOfType(CHAT_VIEWTYPE)[0], ); this.chatIsVisible = true; } @@ -609,7 +624,7 @@ export default class CopilotPlugin extends Plugin { }); await this.activateViewPromise; this.app.workspace.revealLeaf( - this.app.workspace.getLeavesOfType(CHAT_VIEWTYPE)[0] + this.app.workspace.getLeavesOfType(CHAT_VIEWTYPE)[0], ); this.chatIsVisible = true; } @@ -636,10 +651,11 @@ export default class CopilotPlugin extends Plugin { async countTotalTokens(): Promise { try { const allContent = await getAllNotesContent(this.app.vault); - const totalTokens = await this.chainManager.chatModelManager.countTokens(allContent); + const totalTokens = + await this.chainManager.chatModelManager.countTokens(allContent); return totalTokens; } catch (error) { - console.error('Error counting tokens: ', error); + console.error("Error counting tokens: ", error); return 0; } } diff --git a/src/vectorDBManager.ts b/src/vectorDBManager.ts index b0b63880..c5c0f686 100644 --- a/src/vectorDBManager.ts +++ b/src/vectorDBManager.ts @@ -1,8 +1,9 @@ -import { MD5 } from 'crypto-js'; +import { MD5 } from "crypto-js"; import { Document } from "langchain/document"; -import { Embeddings } from 'langchain/embeddings/base'; -import { RecursiveCharacterTextSplitter } from 'langchain/text_splitter'; +import { Embeddings } from "langchain/embeddings/base"; +import { RecursiveCharacterTextSplitter } from "langchain/text_splitter"; import { MemoryVectorStore } from "langchain/vectorstores/memory"; +import EmbeddingManager from "./LLMProviders/embeddingManager"; // TODOs // 1. Use embeddingModel rather than embeddingProvider @@ -34,42 +35,27 @@ export interface NoteFile { } class VectorDBManager { - public static db: PouchDB.Database | null = null; - public static embeddingModel: string | null = null; - - public static initializeDB(db: PouchDB.Database): void { - this.db = db; - } - - public static updateDBInstance(newDb: PouchDB.Database): void { - this.db = newDb; - } - - public static setEmbeddingModel(embeddingModel: string): void { - this.embeddingModel = embeddingModel; - } - public static getDocumentHash(sourceDocument: string): string { return MD5(sourceDocument).toString(); } public static async rebuildMemoryVectorStore( memoryVectors: MemoryVector[], - embeddingsAPI: Embeddings + embeddingsAPI: Embeddings, ) { if (!Array.isArray(memoryVectors)) { throw new TypeError("Expected memoryVectors to be an array"); } // Extract the embeddings and documents from the deserialized memoryVectors const embeddingsArray: number[][] = memoryVectors.map( - (memoryVector) => memoryVector.embedding + (memoryVector) => memoryVector.embedding, ); const documentsArray = memoryVectors.map( (memoryVector) => new Document({ pageContent: memoryVector.content, metadata: memoryVector.metadata, - }) + }), ); // Create a new MemoryVectorStore instance @@ -79,37 +65,40 @@ class VectorDBManager { } public static async getMemoryVectorStore( + db: PouchDB.Database, embeddingsAPI: Embeddings, - docHash?: string + docHash?: string, ): Promise { - if (!this.db) throw new Error("DB not initialized"); - if (!this.embeddingModel) throw new Error("Embedding model not set"); + if (!db) throw new Error("DB not initialized"); let allDocsResponse; if (docHash) { // Fetch a single document by its _id - const doc = await this.db.get(docHash); + const doc = await db.get(docHash); allDocsResponse = { rows: [{ doc }] }; } else { // Fetch all documents - allDocsResponse = await this.db.allDocs({ include_docs: true }); + allDocsResponse = await db.allDocs({ include_docs: true }); } + const embeddingModel = EmbeddingManager.getModelName(embeddingsAPI); + if (!embeddingModel) + console.error("EmbeddingManager could not determine model name!"); const allDocs = allDocsResponse.rows .map((row) => row.doc as VectorStoreDocument) - .filter((doc) => doc.embeddingModel === this.embeddingModel); + .filter((doc) => doc.embeddingModel === embeddingModel); const memoryVectors = allDocs .map((doc) => JSON.parse(doc.memory_vectors) as MemoryVector[]) .flat(); const embeddingsArray: number[][] = memoryVectors.map( - (memoryVector) => memoryVector.embedding + (memoryVector) => memoryVector.embedding, ); const documentsArray = memoryVectors.map( (memoryVector) => new Document({ pageContent: memoryVector.content, metadata: memoryVector.metadata, - }) + }), ); // Create a new MemoryVectorStore instance const memoryVectorStore = new MemoryVectorStore(embeddingsAPI); @@ -118,17 +107,18 @@ class VectorDBManager { } public static async setMemoryVectors( + db: PouchDB.Database, memoryVectors: MemoryVector[], - docHash: string + docHash: string, ): Promise { - if (!this.db) throw new Error("DB not initialized"); + if (!db) throw new Error("DB not initialized"); if (!Array.isArray(memoryVectors)) { throw new TypeError("Expected memoryVectors to be an array"); } const serializedMemoryVectors = JSON.stringify(memoryVectors); try { // Attempt to fetch the existing document, if it exists. - const existingDoc = await this.db.get(docHash).catch((err) => null); + const existingDoc = await db.get(docHash).catch((err) => null); // Prepare the document to be saved. const docToSave = { @@ -139,22 +129,29 @@ class VectorDBManager { }; // Save the document. - await this.db.put(docToSave); + await db.put(docToSave); } catch (err) { console.error("Error storing vectors in VectorDB:", err); } } - public static async indexFile(noteFile: NoteFile, embeddingsAPI: Embeddings) { - if (!this.db) throw new Error("DB not initialized"); - if (!this.embeddingModel) throw new Error("Embedding model not set"); + public static async indexFile( + db: PouchDB.Database, + embeddingsAPI: Embeddings, + noteFile: NoteFile, + ): Promise { + if (!db) throw new Error("DB not initialized"); + + const embeddingModel = EmbeddingManager.getModelName(embeddingsAPI); + if (!embeddingModel) + console.error("EmbeddingManager could not determine model name!"); // Markdown splitter: https://js.langchain.com/docs/modules/data_connection/document_transformers/code_splitter#markdown const textSplitter = RecursiveCharacterTextSplitter.fromLanguage( "markdown", { chunkSize: 5000, - } + }, ); // Add note title as contextual chunk headers // https://js.langchain.com/docs/modules/data_connection/document_transformers/contextual_chunk_headers @@ -164,10 +161,10 @@ class VectorDBManager { { chunkHeader: "[[" + noteFile.basename + "]]" + "\n\n---\n\n", appendChunkOverlapHeader: true, - } + }, ); const docVectors = await embeddingsAPI.embedDocuments( - splitDocument.map((doc) => doc.pageContent) + splitDocument.map((doc) => doc.pageContent), ); const memoryVectors = docVectors.map((docVector, i) => ({ content: splitDocument[i].pageContent, @@ -183,34 +180,34 @@ class VectorDBManager { const serializedMemoryVectors = JSON.stringify(memoryVectors); try { // Attempt to fetch the existing document, if it exists. - const existingDoc = await this.db.get(docHash).catch((err) => null); + const existingDoc = await db.get(docHash).catch((err) => null); // Prepare the document to be saved. const docToSave: VectorStoreDocument = { _id: docHash, memory_vectors: serializedMemoryVectors, file_mtime: noteFile.mtime, - embeddingModel: this.embeddingModel, + embeddingModel: embeddingModel, created_at: Date.now(), _rev: existingDoc?._rev, // Add the current revision if the document exists. }; // Save the document. - await this.db.put(docToSave); + await db.put(docToSave); return docToSave; } catch (err) { console.error("Error storing vectors in VectorDB:", err); } } - public static async getNoteFiles(): Promise { - if (!this.db) throw new Error("DB not initialized"); + public static async getNoteFiles(db: PouchDB.Database): Promise { + if (!db) throw new Error("DB not initialized"); try { - const allDocsResponse = await this.db.allDocs({ + const allDocsResponse = await db.allDocs({ include_docs: true, }); const allDocs = allDocsResponse.rows.map( - (row) => row.doc as VectorStoreDocument + (row) => row.doc as VectorStoreDocument, ); const memoryVectors = allDocs .map((doc) => JSON.parse(doc.memory_vectors) as MemoryVector[]) @@ -236,29 +233,33 @@ class VectorDBManager { } } - public static async removeMemoryVectors(docHash: string): Promise { - if (!this.db) throw new Error("DB not initialized"); + public static async removeMemoryVectors( + db: PouchDB.Database, + docHash: string, + ): Promise { + if (!db) throw new Error("DB not initialized"); try { - const doc = await this.db.get(docHash); + const doc = await db.get(docHash); if (doc) { - await this.db.remove(doc); + await db.remove(doc); } } catch (err) { console.error("Error removing file from VectorDB:", err); } } - public static async getLatestFileMtime(): Promise { - if (!this.db) throw new Error("DB not initialized"); - if (!this.embeddingModel) throw new Error("Embedding provider not set"); + public static async getLatestFileMtime( + db: PouchDB.Database, + ): Promise { + if (!db) throw new Error("DB not initialized"); try { - const allDocsResponse = await this.db.allDocs({ + const allDocsResponse = await db.allDocs({ include_docs: true, }); - const allDocs = allDocsResponse.rows - .map((row) => row.doc as VectorStoreDocument) - .filter((doc) => doc.embeddingModel === this.embeddingModel); + const allDocs = allDocsResponse.rows.map( + (row) => row.doc as VectorStoreDocument, + ); const newestFileMtime = allDocs .map((doc) => doc.file_mtime) .sort((a, b) => b - a)[0]; @@ -269,16 +270,18 @@ class VectorDBManager { } } - public static async checkEmbeddingModel(): Promise { - if (!this.db) throw new Error("DB not initialized"); + public static async checkEmbeddingModel( + db: PouchDB.Database, + ): Promise { + if (!db) throw new Error("DB not initialized"); try { // Fetch all documents - const allDocsResponse = await this.db.allDocs({ + const allDocsResponse = await db.allDocs({ include_docs: true, }); const allDocs = allDocsResponse.rows.map( - (row) => row.doc as VectorStoreDocument + (row) => row.doc as VectorStoreDocument, ); // Check if there are any documents @@ -291,7 +294,7 @@ class VectorDBManager { // Check if all documents have the same embeddingModel const allSame = embeddingModels.every( - (model) => model === embeddingModels[0] + (model) => model === embeddingModels[0], ); return allSame ? embeddingModels[0] : undefined; } catch (err) { @@ -301,11 +304,12 @@ class VectorDBManager { } public static async getMemoryVectors( - docHash: string + db: PouchDB.Database, + docHash: string, ): Promise { - if (!this.db) throw new Error("DB not initialized"); + if (!db) throw new Error("DB not initialized"); try { - const doc: VectorStoreDocument = await this.db.get(docHash); + const doc: VectorStoreDocument = await db.get(docHash); if (doc && doc.memory_vectors) { return JSON.parse(doc.memory_vectors); }