obsidian-plugin/logancyang_obsidian-copilot
1
0
Fork 0
mirror of https://github.com/logancyang/obsidian-copilot.git synced 2026-07-22 07:50:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Create vector store per vault (#348)

Browse source
This commit is contained in:
Logan Yang 2024-03-08 21:50:45 -08:00 committed by GitHub
parent 08eb5c8fcb
commit e30c91fd57
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 351 additions and 275 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.editorconfig
View file

@ -8,3 +8,4 @@ insert_final_newline = true
indent_style = space
indent_size = 2
tab_width = 2
trim_trailing_whitespace = true

335
src/LLMProviders/chainManager.ts
View file

@ -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<VectorStoreDocument>;
/**
* 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<VectorStoreDocument>,
) {
// 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<void> {
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<string> {
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<VectorStoreDocument | undefined> {
async indexFile(
noteFile: NoteFile,
): Promise<VectorStoreDocument | undefined> {
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,
);
}
}

156
src/main.ts
View file

@ -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<void> | null = null;
chatIsVisible = false;
dbPrompts: PouchDB.Database;
dbVectorStores: PouchDB.Database;
dbVectorStores: PouchDB.Database<VectorStoreDocument>;
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<VectorStoreDocument>(
`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<VectorStoreDocument>(
"copilot_vector_stores"
);
this.embeddingsManager = EmbeddingsManager.getInstance(
langChainParams,
this.encryptionService
this.encryptionService,
);
this.dbPrompts = new PouchDB<CustomPrompt>("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<VectorStoreDocument>(
"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<void> {
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<number> {
@ -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<VectorStoreDocument>(
"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<number> {
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;
}
}

134
src/vectorDBManager.ts
View file

@ -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<MemoryVectorStore> {
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<void> {
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<VectorStoreDocument | undefined> {
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<NoteFile[]> {
if (!this.db) throw new Error("DB not initialized");
public static async getNoteFiles(db: PouchDB.Database): Promise<NoteFile[]> {
if (!db) throw new Error("DB not initialized");
try {
const allDocsResponse = await this.db.allDocs<VectorStoreDocument>({
const allDocsResponse = await db.allDocs<VectorStoreDocument>({
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<void> {
if (!this.db) throw new Error("DB not initialized");
public static async removeMemoryVectors(
db: PouchDB.Database,
docHash: string,
): Promise<void> {
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<number> {
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<number> {
if (!db) throw new Error("DB not initialized");
try {
const allDocsResponse = await this.db.allDocs<VectorStoreDocument>({
const allDocsResponse = await db.allDocs<VectorStoreDocument>({
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<string | undefined> {
if (!this.db) throw new Error("DB not initialized");
public static async checkEmbeddingModel(
db: PouchDB.Database,
): Promise<string | undefined> {
if (!db) throw new Error("DB not initialized");
try {
// Fetch all documents
const allDocsResponse = await this.db.allDocs<VectorStoreDocument>({
const allDocsResponse = await db.allDocs<VectorStoreDocument>({
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<MemoryVector[] | undefined> {
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);
}