From e0af6d9fbf57a232e40703820fca78deb69f120a Mon Sep 17 00:00:00 2001 From: Logan Yang Date: Thu, 21 Nov 2024 17:36:52 -0800 Subject: [PATCH] Introduce Copilot Plus Alpha to testers (#835) - Implement Copilot Plus mode - Use broca endpoint and add license key setting - Implement context menu for Plus mode - Implement url mention - Implement cmd shift enter for vault search - Handle embedding model change at note reindex - Implement image support in chat - Fix vault qa and remove relevance threshold score - Implement brevilabs client, reranker when max score from orama is low - Implement chat history to plus mode - Pass only original user message for intent analysis and processing - Implement pdf in chat context - Implement web search - Implement youtube tool - Enable index on mode switch for copilot plus mode --- README.md | 2 +- package-lock.json | 33 ++ package.json | 3 + src/LLMProviders/brevilabsClient.ts | 149 ++++++ src/LLMProviders/chainManager.ts | 321 ++++------- src/LLMProviders/chainRunner.ts | 502 ++++++++++++++++++ src/LLMProviders/intentAnalyzer.ts | 161 ++++++ src/LLMProviders/memoryManager.ts | 30 +- src/LLMProviders/promptManager.ts | 33 ++ src/VectorStoreManager.ts | 233 ++++++-- src/aiParams.ts | 2 + src/components/AddContextNoteModal.tsx | 33 ++ src/components/AddImageModal.tsx | 29 + src/components/BaseNoteModal.tsx | 51 ++ src/components/Chat.tsx | 135 ++++- src/components/ChatComponents/ChatButtons.tsx | 19 +- .../ChatComponents/ChatContextMenu.tsx | 80 +++ .../ChatComponents/ChatControls.tsx | 196 ++++--- src/components/ChatComponents/ChatInput.tsx | 290 ++++++++-- .../ChatComponents/ChatMessages.tsx | 8 +- .../ChatComponents/ChatSingleMessage.tsx | 87 ++- src/components/CopilotView.tsx | 6 +- src/components/NoteTitleModal.tsx | 15 +- src/components/OramaSearchModal.tsx | 93 ++++ src/components/SourcesModal.tsx | 61 +++ src/constants.ts | 19 +- src/contextProcessor.ts | 131 +++++ src/customPromptProcessor.ts | 24 + src/encryptionService.ts | 5 +- src/langchainStream.ts | 3 +- src/main.ts | 68 ++- src/mentions/Mention.ts | 94 ++++ src/search/hybridRetriever.ts | 243 +++++++-- src/settings/SettingsPage.tsx | 1 + .../components/CopilotPlusSettings.tsx | 31 ++ src/settings/components/GeneralSettings.tsx | 1 + src/settings/components/SettingsMain.tsx | 2 + src/sharedState.ts | 4 + src/tools/FileParserManager.ts | 76 +++ src/tools/SearchTools.ts | 161 ++++++ src/tools/TimeTools.ts | 281 ++++++++++ src/tools/YoutubeTools.ts | 46 ++ src/tools/toolManager.ts | 10 + src/utils.ts | 18 + src/vectorDBManager.ts | 12 +- styles.css | 173 +++++- 46 files changed, 3505 insertions(+), 470 deletions(-) create mode 100644 src/LLMProviders/brevilabsClient.ts create mode 100644 src/LLMProviders/chainRunner.ts create mode 100644 src/LLMProviders/intentAnalyzer.ts create mode 100644 src/components/AddContextNoteModal.tsx create mode 100644 src/components/AddImageModal.tsx create mode 100644 src/components/BaseNoteModal.tsx create mode 100644 src/components/ChatComponents/ChatContextMenu.tsx create mode 100644 src/components/OramaSearchModal.tsx create mode 100644 src/components/SourcesModal.tsx create mode 100644 src/contextProcessor.ts create mode 100644 src/mentions/Mention.ts create mode 100644 src/settings/components/CopilotPlusSettings.tsx create mode 100644 src/tools/FileParserManager.ts create mode 100644 src/tools/SearchTools.ts create mode 100644 src/tools/TimeTools.ts create mode 100644 src/tools/YoutubeTools.ts create mode 100644 src/tools/toolManager.ts diff --git a/README.md b/README.md index 8a8b074c..8378746b 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,7 @@ If you enjoy Copilot for Obsidian, please consider [sponsoring this project](htt Buy Me A Coffee SPECIAL THANKS TO OUR TOP SPONSORS: -@pedramamini, @Arlorean, @dashinja, @azagore, @MTGMAD, @gpythomas, @emaynard, @scmarinelli, @borthwick, @adamhill, @gluecode, @rusi, @timgrote, @JiaruiYu-Consilium, @ddocta, @AMOz1, @agu3rra +@pedramamini, @Arlorean, @dashinja, @azagore, @MTGMAD, @gpythomas, @emaynard, @scmarinelli, @borthwick, @adamhill, @gluecode, @rusi, @timgrote, @JiaruiYu-Consilium, @ddocta, @AMOz1, @chchwy, @pborenstein, @GitTom, @kazukgw, @mjluser1, @joesfer, @rwaal [Changelog](https://github.com/logancyang/obsidian-copilot/releases) diff --git a/package-lock.json b/package-lock.json index 60a8c797..81f89ab1 100644 --- a/package-lock.json +++ b/package-lock.json @@ -22,6 +22,7 @@ "@radix-ui/react-tooltip": "^1.1.3", "@tabler/icons-react": "^2.14.0", "axios": "^1.3.4", + "chrono-node": "^2.7.7", "cohere-ai": "^7.13.0", "crypto-js": "^4.1.1", "esbuild-plugin-svg": "^0.1.0", @@ -30,6 +31,7 @@ "koa-proxies": "^0.12.3", "langchain": "^0.3.2", "lucide-react": "^0.454.0", + "luxon": "^3.5.0", "next-i18next": "^13.2.2", "prop-types": "^15.8.1", "react": "^18.2.0", @@ -47,6 +49,7 @@ "@types/jest": "^29.5.11", "@types/koa": "^2.13.7", "@types/koa__cors": "^4.0.0", + "@types/luxon": "^3.4.2", "@types/node": "^16.11.6", "@types/react": "^18.0.33", "@types/react-dom": "^18.0.11", @@ -6125,6 +6128,12 @@ "@types/koa": "*" } }, + "node_modules/@types/luxon": { + "version": "3.4.2", + "resolved": "https://registry.npmjs.org/@types/luxon/-/luxon-3.4.2.tgz", + "integrity": "sha512-TifLZlFudklWlMBfhubvgqTXRzLDI5pCbGa4P8a3wPyUQSW+1xQ5eDsreP9DWHX3tjq1ke96uYG/nwundroWcA==", + "dev": true + }, "node_modules/@types/mdast": { "version": "4.0.4", "resolved": "https://registry.npmjs.org/@types/mdast/-/mdast-4.0.4.tgz", @@ -7199,6 +7208,17 @@ "url": "https://github.com/sponsors/wooorm" } }, + "node_modules/chrono-node": { + "version": "2.7.7", + "resolved": "https://registry.npmjs.org/chrono-node/-/chrono-node-2.7.7.tgz", + "integrity": "sha512-p3S7gotuTPu5oqhRL2p1fLwQXGgdQaRTtWR3e8Di9P1Pa9mzkK5DWR5AWBieMUh2ZdOnPgrK+zCrbbtyuA+D/Q==", + "dependencies": { + "dayjs": "^1.10.0" + }, + "engines": { + "node": "^12.20.0 || ^14.13.1 || >=16.0.0" + } + }, "node_modules/ci-info": { "version": "3.8.0", "resolved": "https://registry.npmjs.org/ci-info/-/ci-info-3.8.0.tgz", @@ -7642,6 +7662,11 @@ "node": ">=12" } }, + "node_modules/dayjs": { + "version": "1.11.13", + "resolved": "https://registry.npmjs.org/dayjs/-/dayjs-1.11.13.tgz", + "integrity": "sha512-oaMBel6gjolK862uaPQOVTA7q3TZhuSvuMQAAglQDOWYO9A91IrAOUJEyKVlqJlHE0vq5p5UXxzdPfMH/x6xNg==" + }, "node_modules/debug": { "version": "4.3.6", "resolved": "https://registry.npmjs.org/debug/-/debug-4.3.6.tgz", @@ -12259,6 +12284,14 @@ "react": "^16.5.1 || ^17.0.0 || ^18.0.0 || ^19.0.0-rc" } }, + "node_modules/luxon": { + "version": "3.5.0", + "resolved": "https://registry.npmjs.org/luxon/-/luxon-3.5.0.tgz", + "integrity": "sha512-rh+Zjr6DNfUYR3bPwJEnuwDdqMbxZW7LOQfUN4B54+Cl+0o5zaU9RJ6bcidfDtC1cWCZXQ+nvX8bf6bAji37QQ==", + "engines": { + "node": ">=12" + } + }, "node_modules/lz-string": { "version": "1.5.0", "resolved": "https://registry.npmjs.org/lz-string/-/lz-string-1.5.0.tgz", diff --git a/package.json b/package.json index b496bc67..521b0e5f 100644 --- a/package.json +++ b/package.json @@ -37,6 +37,7 @@ "@types/jest": "^29.5.11", "@types/koa": "^2.13.7", "@types/koa__cors": "^4.0.0", + "@types/luxon": "^3.4.2", "@types/node": "^16.11.6", "@types/react": "^18.0.33", "@types/react-dom": "^18.0.11", @@ -73,6 +74,7 @@ "@radix-ui/react-tooltip": "^1.1.3", "@tabler/icons-react": "^2.14.0", "axios": "^1.3.4", + "chrono-node": "^2.7.7", "cohere-ai": "^7.13.0", "crypto-js": "^4.1.1", "esbuild-plugin-svg": "^0.1.0", @@ -81,6 +83,7 @@ "koa-proxies": "^0.12.3", "langchain": "^0.3.2", "lucide-react": "^0.454.0", + "luxon": "^3.5.0", "next-i18next": "^13.2.2", "prop-types": "^15.8.1", "react": "^18.2.0", diff --git a/src/LLMProviders/brevilabsClient.ts b/src/LLMProviders/brevilabsClient.ts new file mode 100644 index 00000000..b34e5838 --- /dev/null +++ b/src/LLMProviders/brevilabsClient.ts @@ -0,0 +1,149 @@ +import { BREVILABS_API_BASE_URL } from "@/constants"; +import { Notice } from "obsidian"; + +export interface BrocaResponse { + response: { + tool_calls: Array<{ + tool: string; + args: { + [key: string]: any; + }; + }>; + salience_terms: string[]; + }; + elapsed_time_ms: number; +} + +export interface RerankResponse { + response: { + object: string; + data: Array<{ + relevance_score: number; + index: number; + }>; + model: string; + usage: { + total_tokens: number; + }; + }; + elapsed_time_ms: number; +} + +export interface ToolCall { + tool: any; + args: any; +} + +export interface Url4llmResponse { + response: any; + elapsed_time_ms: number; +} + +export interface Pdf4llmResponse { + response: any; + elapsed_time_ms: number; +} + +export interface WebSearchResponse { + response: any; + elapsed_time_ms: number; +} + +export interface Youtube4llmResponse { + response: { + transcript: string; + }; + elapsed_time_ms: number; +} + +export class BrevilabsClient { + private static instance: BrevilabsClient; + private licenseKey: string; + private options: any; + + private constructor(licenseKey: string, options?: { debug?: boolean }) { + this.licenseKey = licenseKey; + this.options = options; + } + + static getInstance(licenseKey: string, options?: { debug?: boolean }): BrevilabsClient { + if (!BrevilabsClient.instance) { + BrevilabsClient.instance = new BrevilabsClient(licenseKey, options); + } + return BrevilabsClient.instance; + } + + private checkLicenseKey() { + if (!this.licenseKey) { + new Notice( + "Copilot Plus license key not found. Please enter your license key in the settings." + ); + throw new Error("License key not initialized"); + } + } + + private async makeRequest(endpoint: string, body: any, method = "POST"): Promise { + this.checkLicenseKey(); + + const url = new URL(`${BREVILABS_API_BASE_URL}${endpoint}`); + if (method === "GET") { + // Add query parameters for GET requests + Object.entries(body).forEach(([key, value]) => { + url.searchParams.append(key, value as string); + }); + } + + const response = await fetch(url.toString(), { + method, + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${this.licenseKey}`, + }, + ...(method === "POST" && { body: JSON.stringify(body) }), + }); + + const data = await response.json(); + if (this.options?.debug) { + console.log(`==== ${endpoint} request ====:`, data); + } + + return data; + } + + async broca(userMessage: string): Promise { + const brocaResponse = await this.makeRequest("/broca", { + message: userMessage, + }); + + return brocaResponse; + } + + async rerank(query: string, documents: string[]): Promise { + return this.makeRequest("/rerank", { + query, + documents, + model: "rerank-2", + }); + } + + async url4llm(url: string): Promise { + return this.makeRequest("/url4llm", { url }); + } + + async pdf4llm(binaryContent: ArrayBuffer): Promise { + // Convert ArrayBuffer to base64 string + const base64Content = Buffer.from(binaryContent).toString("base64"); + + return this.makeRequest("/pdf4llm", { + pdf: base64Content, + }); + } + + async webSearch(query: string): Promise { + return this.makeRequest("/websearch", { q: query }, "GET"); + } + + async youtube4llm(url: string): Promise { + return this.makeRequest("/youtube4llm", { url }); + } +} diff --git a/src/LLMProviders/chainManager.ts b/src/LLMProviders/chainManager.ts index 75b3dd41..cadf69be 100644 --- a/src/LLMProviders/chainManager.ts +++ b/src/LLMProviders/chainManager.ts @@ -1,17 +1,17 @@ import { CustomModel, LangChainParams, SetChainOptions } from "@/aiParams"; import ChainFactory, { ChainType, Document } from "@/chainFactory"; -import { ABORT_REASON, AI_SENDER, BUILTIN_CHAT_MODELS, USER_SENDER } from "@/constants"; +import { BUILTIN_CHAT_MODELS, USER_SENDER } from "@/constants"; import EncryptionService from "@/encryptionService"; +import { + ChainRunner, + CopilotPlusChainRunner, + LLMChainRunner, + VaultQAChainRunner, +} from "@/LLMProviders/chainRunner"; import { HybridRetriever } from "@/search/hybridRetriever"; import { CopilotSettings } from "@/settings/SettingsPage"; import { ChatMessage } from "@/sharedState"; -import { - extractChatHistory, - extractUniqueTitlesFromDocs, - formatDateTime, - isSupportedChain, -} from "@/utils"; -import VectorDBManager from "@/vectorDBManager"; +import { isSupportedChain } from "@/utils"; import VectorStoreManager from "@/VectorStoreManager"; import { ChatPromptTemplate, @@ -19,9 +19,8 @@ import { MessagesPlaceholder, } from "@langchain/core/prompts"; import { RunnableSequence } from "@langchain/core/runnables"; -import { BaseChatMemory } from "langchain/memory"; -import { MemoryVectorStore } from "langchain/vectorstores/memory"; import { App, Notice } from "obsidian"; +import { BrevilabsClient } from "./brevilabsClient"; import ChatModelManager from "./chatModelManager"; import EmbeddingsManager from "./embeddingManager"; import MemoryManager from "./memoryManager"; @@ -30,35 +29,34 @@ import PromptManager from "./promptManager"; export default class ChainManager { private static chain: RunnableSequence; private static retrievalChain: RunnableSequence; - private static retrievedDocuments: Document[] = []; - private app: App; private settings: CopilotSettings; - private vectorStore: MemoryVectorStore; - private vectorStoreManager: VectorStoreManager; - private promptManager: PromptManager; private encryptionService: EncryptionService; - private langChainParams: LangChainParams; + public app: App; + public vectorStoreManager: VectorStoreManager; public chatModelManager: ChatModelManager; public memoryManager: MemoryManager; public embeddingsManager: EmbeddingsManager; + public promptManager: PromptManager; + public brevilabsClient: BrevilabsClient; + public static retrievedDocuments: Document[] = []; constructor( app: App, getLangChainParams: () => LangChainParams, encryptionService: EncryptionService, settings: CopilotSettings, - // Ensure ChainManager always has the up-to-date dbVectorStores - vectorStoreManager: VectorStoreManager + vectorStoreManager: VectorStoreManager, + brevilabsClient: BrevilabsClient ) { // Instantiate singletons this.app = app; this.langChainParams = getLangChainParams(); this.settings = settings; this.vectorStoreManager = vectorStoreManager; - this.memoryManager = MemoryManager.getInstance(this.getLangChainParams()); + this.memoryManager = MemoryManager.getInstance(this.getLangChainParams(), settings.debug); this.encryptionService = encryptionService; this.chatModelManager = ChatModelManager.getInstance( () => this.getLangChainParams(), @@ -67,6 +65,7 @@ export default class ChainManager { ); this.embeddingsManager = this.vectorStoreManager.getEmbeddingsManager(); this.promptManager = PromptManager.getInstance(this.getLangChainParams()); + this.brevilabsClient = brevilabsClient; this.createChainWithNewModel(this.getLangChainParams().modelKey); } @@ -81,14 +80,37 @@ export default class ChainManager { this.langChainParams[key] = value; } - private setNoteFile(noteFile: any): void { - this.setLangChainParam("options", { ...this.getLangChainParams().options, noteFile }); + static getChain(): RunnableSequence { + return ChainManager.chain; + } + + static getRetrievalChain(): RunnableSequence { + return ChainManager.retrievalChain; } private validateChainType(chainType: ChainType): void { if (chainType === undefined || chainType === null) throw new Error("No chain type set"); } + private validateChatModel() { + 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); + throw new Error(errorMsg); + } + } + + private validateChainInitialization() { + if (!ChainManager.chain || !isSupportedChain(ChainManager.chain)) { + console.error( + "Chain is not initialized properly, re-initializing chain: ", + this.getLangChainParams().chainType + ); + this.setChain(this.getLangChainParams().chainType, this.getLangChainParams().options); + } + } + private findCustomModel(modelKey: string): CustomModel | undefined { const [name, provider] = modelKey.split("|"); return this.settings.activeModels.find( @@ -206,9 +228,11 @@ export default class ChainManager { this.app.vault, chatModel, embeddingsAPI, + this.brevilabsClient, { - minSimilarityScore: 0.3, + minSimilarityScore: 0.01, maxK: this.settings.maxSourceChunks, + salientTerms: [], }, options.debug ); @@ -232,14 +256,51 @@ export default class ChainManager { break; } + case ChainType.COPILOT_PLUS_CHAIN: { + // TODO: Create new copilotPlusChain with retriever + // For initial load of the plugin + if (options.forceNewCreation) { + ChainManager.chain = ChainFactory.createNewLLMChain({ + llm: chatModel, + memory: memory, + prompt: options.prompt || chatPrompt, + abortController: options.abortController, + }) as RunnableSequence; + } else { + // For navigating back to the plugin view + ChainManager.chain = ChainFactory.getLLMChainFromMap({ + llm: chatModel, + memory: memory, + prompt: options.prompt || chatPrompt, + abortController: options.abortController, + }) as RunnableSequence; + } + + this.setLangChainParam("chainType", ChainType.COPILOT_PLUS_CHAIN); + break; + } + default: this.validateChainType(chainType); break; } } + private getChainRunner(): ChainRunner { + switch (this.getLangChainParams().chainType) { + case ChainType.LLM_CHAIN: + return new LLMChainRunner(this); + case ChainType.VAULT_QA_CHAIN: + return new VaultQAChainRunner(this); + case ChainType.COPILOT_PLUS_CHAIN: + return new CopilotPlusChainRunner(this); + default: + throw new Error(`Unsupported chain type: ${this.getLangChainParams().chainType}`); + } + } + async runChain( - userMessage: string, + userMessage: ChatMessage, abortController: AbortController, updateCurrentAiMessage: (message: string) => void, addMessage: (message: ChatMessage) => void, @@ -251,211 +312,31 @@ export default class ChainManager { ) { 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."; - new Notice(errorMsg); - console.error(errorMsg); - return; - } - // Check if chain is initialized properly - if (!ChainManager.chain || !isSupportedChain(ChainManager.chain)) { - console.error( - "Chain is not initialized properly, re-initializing chain: ", - this.getLangChainParams().chainType - ); - this.setChain(this.getLangChainParams().chainType, this.getLangChainParams().options); - } + if (debug) console.log("==== Step 0: Initial user message ====\n", userMessage); - const { temperature, maxTokens, systemMessage, chatContextTurns, chainType } = - this.getLangChainParams(); + this.validateChatModel(); + this.validateChainInitialization(); - const memory = this.memoryManager.getMemory(); - const chatPrompt = this.promptManager.getChatPrompt(); - const systemPrompt = ignoreSystemMessage ? "" : systemMessage; - // Whether to ignore system prompt (for commands) + // Handle ignoreSystemMessage if (ignoreSystemMessage) { - const effectivePrompt = ignoreSystemMessage - ? ChatPromptTemplate.fromMessages([ - new MessagesPlaceholder("history"), - HumanMessagePromptTemplate.fromTemplate("{input}"), - ]) - : chatPrompt; - - this.setChain(chainType, { + const effectivePrompt = ChatPromptTemplate.fromMessages([ + new MessagesPlaceholder("history"), + HumanMessagePromptTemplate.fromTemplate("{input}"), + ]); + this.setChain(this.getLangChainParams().chainType, { ...this.getLangChainParams().options, prompt: effectivePrompt, }); - } else { - this.setChain(chainType, this.getLangChainParams().options); } - 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, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - } 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` - ); - } - - for await (const chunk of chatStream) { - if (abortController.signal.aborted) break; - fullAIResponse += chunk.content; - updateCurrentAiMessage(fullAIResponse); - } - break; - 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("chain RunnableSequence:", ChainManager.chain); - console.log("embedding model:", this.getLangChainParams().embeddingModelKey); - } - fullAIResponse = await this.runRetrievalChain( - userMessage, - memory, - updateCurrentAiMessage, - abortController, - { debug } - ); - break; - default: - console.error("Chain type not supported:", this.getLangChainParams().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."; - new Notice(modelNotFoundMsg); - console.error(modelNotFoundMsg); - } else { - new Notice(`LangChain error: ${errorCode}`); - console.error(errorData); - } - } finally { - // note: pls Check Abort Reason - if (fullAIResponse && abortController.signal.reason !== ABORT_REASON.NEW_CHAT) { - // This line is a must for memory to work with RunnableSequence! - await memory.saveContext({ input: userMessage }, { output: fullAIResponse }); - addMessage({ - message: fullAIResponse, - sender: AI_SENDER, - isVisible: true, - timestamp: formatDateTime(new Date()), - }); - } - updateCurrentAiMessage(""); - } - if (debug) { - console.log( - "Chat memory:", - (memory as any).chatHistory.messages.map((msg: any) => msg.content) - ); - } - return fullAIResponse; - } - - private async runRetrievalChain( - userMessage: string, - memory: BaseChatMemory, - updateCurrentAiMessage: (message: string) => void, - abortController: AbortController, - options: { - debug?: boolean; - } = {} - ): Promise { - const memoryVariables = await memory.loadMemoryVariables({}); - const chatHistory = extractChatHistory(memoryVariables); - const qaStream = await ChainManager.retrievalChain.stream({ - question: userMessage, - chat_history: chatHistory, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - } as any); - - let fullAIResponse = ""; - - for await (const chunk of qaStream) { - if (abortController.signal.aborted) break; - fullAIResponse += chunk.content; - updateCurrentAiMessage(fullAIResponse); - } - - if (options.debug) { - 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. - // That means multiple chunks can be from the same note. A more advanced logic is needed - // to show specific chunks in the future. E.g. collapsed note title when clicked, - // expand and reveal the chunk - if (this.getLangChainParams().chainType === ChainType.VAULT_QA_CHAIN) { - const docTitles = extractUniqueTitlesFromDocs(ChainManager.retrievedDocuments); - if (docTitles.length > 0) { - const markdownLinks = docTitles - .map( - (title) => - `- [${title}](obsidian://open?vault=${encodeURIComponent(this.app.vault.getName())}&file=${encodeURIComponent( - title - )})` - ) - .join("\n"); - fullAIResponse += "\n\n#### Sources:\n" + markdownLinks; - } - } - - return fullAIResponse; - } - - async indexFile(noteFile: any): Promise { - const embeddingsAPI = this.embeddingsManager.getEmbeddingsAPI(); - const db = this.vectorStoreManager.getDb(); - - if (!embeddingsAPI) { - const errorMsg = - "Failed to load file, embedding API is not set correctly, please check your settings."; - new Notice(errorMsg); - console.error(errorMsg); - return; - } - - if (!db) { - const errorMsg = "Vector database is not initialized or disabled on mobile."; - new Notice(errorMsg); - console.error(errorMsg); - return; - } - - return await VectorDBManager.indexFile(db, embeddingsAPI, noteFile); + const chainRunner = this.getChainRunner(); + return await chainRunner.run( + userMessage, + abortController, + updateCurrentAiMessage, + addMessage, + options + ); } async updateMemoryWithLoadedMessages(messages: ChatMessage[]) { diff --git a/src/LLMProviders/chainRunner.ts b/src/LLMProviders/chainRunner.ts new file mode 100644 index 00000000..ad3cb21f --- /dev/null +++ b/src/LLMProviders/chainRunner.ts @@ -0,0 +1,502 @@ +import { ABORT_REASON, AI_SENDER, LOADING_MESSAGES } from "@/constants"; +import { ChatMessage } from "@/sharedState"; +import { ToolManager } from "@/tools/toolManager"; +import { extractChatHistory, extractUniqueTitlesFromDocs, formatDateTime } from "@/utils"; +import { Notice } from "obsidian"; +import ChainManager from "./chainManager"; +import { COPILOT_TOOL_NAMES, IntentAnalyzer } from "./intentAnalyzer"; + +export interface ChainRunner { + run( + userMessage: ChatMessage, + abortController: AbortController, + updateCurrentAiMessage: (message: string) => void, + addMessage: (message: ChatMessage) => void, + options: { + debug?: boolean; + ignoreSystemMessage?: boolean; + updateLoading?: (loading: boolean) => void; + } + ): Promise; +} + +abstract class BaseChainRunner implements ChainRunner { + constructor(protected chainManager: ChainManager) {} + + abstract run( + userMessage: ChatMessage, + abortController: AbortController, + updateCurrentAiMessage: (message: string) => void, + addMessage: (message: ChatMessage) => void, + options: { + debug?: boolean; + ignoreSystemMessage?: boolean; + updateLoading?: (loading: boolean) => void; + } + ): Promise; + + protected async handleResponse( + fullAIResponse: string, + userMessage: ChatMessage, + abortController: AbortController, + addMessage: (message: ChatMessage) => void, + updateCurrentAiMessage: (message: string) => void, + debug: boolean, + sources?: { title: string; score: number }[] + ) { + if (fullAIResponse && abortController.signal.reason !== ABORT_REASON.NEW_CHAT) { + await this.chainManager.memoryManager + .getMemory() + .saveContext({ input: userMessage.message }, { output: fullAIResponse }); + + addMessage({ + message: fullAIResponse, + sender: AI_SENDER, + isVisible: true, + timestamp: formatDateTime(new Date()), + sources: sources, + }); + } + updateCurrentAiMessage(""); + if (debug) { + console.log( + "==== Chat Memory ====\n", + (this.chainManager.memoryManager.getMemory().chatHistory as any).messages.map( + (m: any) => m.content + ) + ); + console.log("==== Final AI Response ====\n", fullAIResponse); + } + return fullAIResponse; + } + + protected async handleError(error: any, debug: boolean) { + if (debug) console.error("Error during LLM invocation:", 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."; + new Notice(modelNotFoundMsg); + console.error(modelNotFoundMsg); + } else { + new Notice(`LangChain error: ${errorCode}`); + console.error(errorData); + } + } +} + +class LLMChainRunner extends BaseChainRunner { + async run( + userMessage: ChatMessage, + abortController: AbortController, + updateCurrentAiMessage: (message: string) => void, + addMessage: (message: ChatMessage) => void, + options: { + debug?: boolean; + ignoreSystemMessage?: boolean; + updateLoading?: (loading: boolean) => void; + } + ): Promise { + const { debug = false } = options; + let fullAIResponse = ""; + + try { + const chain = ChainManager.getChain(); + const chatStream = await chain.stream({ + input: userMessage.message, + } as any); + + for await (const chunk of chatStream) { + if (abortController.signal.aborted) break; + fullAIResponse += chunk.content; + updateCurrentAiMessage(fullAIResponse); + } + } catch (error) { + await this.handleError(error, debug); + } + + return this.handleResponse( + fullAIResponse, + userMessage, + abortController, + addMessage, + updateCurrentAiMessage, + debug + ); + } +} + +class VaultQAChainRunner extends BaseChainRunner { + async run( + userMessage: ChatMessage, + abortController: AbortController, + updateCurrentAiMessage: (message: string) => void, + addMessage: (message: ChatMessage) => void, + options: { + debug?: boolean; + ignoreSystemMessage?: boolean; + updateLoading?: (loading: boolean) => void; + } + ): Promise { + const { debug = false } = options; + let fullAIResponse = ""; + + try { + const memory = this.chainManager.memoryManager.getMemory(); + const memoryVariables = await memory.loadMemoryVariables({}); + const chatHistory = extractChatHistory(memoryVariables); + const qaStream = await ChainManager.getRetrievalChain().stream({ + question: userMessage.message, + chat_history: chatHistory, + } as any); + + for await (const chunk of qaStream) { + if (abortController.signal.aborted) break; + fullAIResponse += chunk.content; + updateCurrentAiMessage(fullAIResponse); + } + + fullAIResponse = this.addSourcestoResponse(fullAIResponse); + } catch (error) { + await this.handleError(error, debug); + } + + return this.handleResponse( + fullAIResponse, + userMessage, + abortController, + addMessage, + updateCurrentAiMessage, + debug + ); + } + + private addSourcestoResponse(response: string): string { + const docTitles = extractUniqueTitlesFromDocs(ChainManager.retrievedDocuments); + if (docTitles.length > 0) { + const markdownLinks = docTitles + .map( + (title) => + `- [${title}](obsidian://open?vault=${encodeURIComponent(this.chainManager.app.vault.getName())}&file=${encodeURIComponent( + title + )})` + ) + .join("\n"); + response += "\n\n#### Sources:\n" + markdownLinks; + } + return response; + } +} + +class CopilotPlusChainRunner extends BaseChainRunner { + private async streamMultimodalResponse( + textContent: string, + userMessage: ChatMessage, + abortController: AbortController, + updateCurrentAiMessage: (message: string) => void, + debug: boolean + ): Promise { + // Get chat history + const memory = this.chainManager.memoryManager.getMemory(); + const memoryVariables = await memory.loadMemoryVariables({}); + const chatHistory = extractChatHistory(memoryVariables); + + // Create messages array starting with system message + const messages: any[] = []; + + // Add system message if available + const systemMessage = this.chainManager.getLangChainParams().systemMessage; + if (systemMessage) { + messages.push({ + role: "system", + content: `${systemMessage}\nIMPORTANT: Maintain consistency with previous responses in the conversation. If you've provided information about a person or topic before, use that same information in follow-up questions.`, + }); + } + + // Add chat history with explicit instruction to maintain context + if (chatHistory.length > 0) { + messages.push({ + role: "system", + content: + "The following is the relevant conversation history. Use this context to maintain consistency in your responses:", + }); + } + + // Add chat history + for (const [human, ai] of chatHistory) { + messages.push({ role: "user", content: human }); + messages.push({ role: "assistant", content: ai }); + } + + // Create content array for current message + const content = [ + { + type: "text", + text: textContent, + }, + ]; + + // Add image content if present + if (userMessage.content && userMessage.content.length > 0) { + const imageContent = userMessage.content.filter( + (item) => item.type === "image_url" && item.image_url?.url + ); + content.push(...imageContent); + } + + // Add current user message + messages.push({ + role: "user", + content, + }); + + let fullAIResponse = ""; + const chatStream = await this.chainManager.chatModelManager.getChatModel().stream(messages); + + for await (const chunk of chatStream) { + if (abortController.signal.aborted) break; + fullAIResponse += chunk.content; + updateCurrentAiMessage(fullAIResponse); + } + + return fullAIResponse; + } + + async run( + userMessage: ChatMessage, + abortController: AbortController, + updateCurrentAiMessage: (message: string) => void, + addMessage: (message: ChatMessage) => void, + options: { + debug?: boolean; + ignoreSystemMessage?: boolean; + updateLoading?: (loading: boolean) => void; + updateLoadingMessage?: (message: string) => void; + } + ): Promise { + const { debug = false, updateLoadingMessage } = options; + let fullAIResponse = ""; + let sources: { title: string; score: number }[] = []; + + try { + if (debug) console.log("==== Step 1: Analyzing intent ===="); + // Use the original message for intent analysis + const messageForAnalysis = userMessage.originalMessage || userMessage.message; + const toolCalls = await IntentAnalyzer.analyzeIntent( + messageForAnalysis, + this.chainManager.vectorStoreManager, + this.chainManager.chatModelManager, + this.chainManager.brevilabsClient + ); + if (debug) + console.log( + "Tool identification result:", + toolCalls.map((call) => call.tool.name) + ); + + // Use the same removeAtCommands logic as IntentAnalyzer + const cleanedUserMessage = userMessage.message + .split(" ") + .filter((word) => !COPILOT_TOOL_NAMES.includes(word.toLowerCase())) + .join(" ") + .trim(); + + const toolOutputs = await this.executeToolCalls(toolCalls, debug, updateLoadingMessage); + const localSearchResult = toolOutputs.find( + (output) => output.tool === "localSearch" && output.output && output.output.length > 0 + ); + + if (localSearchResult) { + if (debug) console.log("==== Step 2: Processing local search results ===="); + const documents = JSON.parse(localSearchResult.output); + + // Format chat history from memory + const memory = this.chainManager.memoryManager.getMemory(); + const memoryVariables = await memory.loadMemoryVariables({}); + const chatHistory = extractChatHistory(memoryVariables); + + if (debug) console.log("==== Step 3: Condensing Question ===="); + const standaloneQuestion = await this.getStandaloneQuestion( + cleanedUserMessage, + chatHistory + ); + if (debug) console.log("Condensed standalone question: ", standaloneQuestion); + + if (debug) console.log("==== Step 4: Preparing context ===="); + const timeExpression = this.getTimeExpression(toolCalls); + const context = this.formatLocalSearchResult(documents, timeExpression); + + if (debug) console.log(context); + if (debug) console.log("==== Step 5: Invoking QA Chain ===="); + const qaPrompt = await this.chainManager.promptManager.getQAPrompt({ + question: standaloneQuestion, + context: context, + systemMessage: this.chainManager.getLangChainParams().systemMessage, + }); + + fullAIResponse = await this.streamMultimodalResponse( + qaPrompt, + userMessage, + abortController, + updateCurrentAiMessage, + debug + ); + + // Append sources to the response + sources = this.getSources(documents); + } else { + const enhancedUserMessage = this.prepareEnhancedUserMessage( + cleanedUserMessage, + toolOutputs + ); + // If no results, default to LLM Chain + if (debug) { + console.log("No local search results. Using standard LLM Chain."); + console.log("Enhanced user message:", enhancedUserMessage); + } + + fullAIResponse = await this.streamMultimodalResponse( + enhancedUserMessage, + userMessage, + abortController, + updateCurrentAiMessage, + debug + ); + } + } catch (error) { + // Reset loading message to default + updateLoadingMessage?.(LOADING_MESSAGES.DEFAULT); + await this.handleError(error, debug); + } + + return this.handleResponse( + fullAIResponse, + userMessage, + abortController, + addMessage, + updateCurrentAiMessage, + debug, + sources + ); + } + + private async getStandaloneQuestion( + question: string, + chatHistory: [string, string][] + ): Promise { + const condenseQuestionTemplate = `Given the following conversation and a follow up question, + summarize the conversation as context and keep the follow up question unchanged, in its original language. + If the follow up question is unrelated to its preceding messages, return this follow up question directly. + If it is related, then combine the summary and the follow up question to construct a standalone question. + Make sure to keep any [[]] wrapped note titles in the question unchanged. + If there's nothing in the chat history, just return the follow up question. + + Chat History: + {chat_history} + Follow Up Input: {question} + Standalone question:`; + + const formattedChatHistory = chatHistory + .map(([human, ai]) => `Human: ${human}\nAssistant: ${ai}`) + .join("\n"); + + const response = await this.chainManager.chatModelManager.getChatModel().invoke([ + { role: "system", content: condenseQuestionTemplate }, + { + role: "user", + content: condenseQuestionTemplate + .replace("{chat_history}", formattedChatHistory) + .replace("{question}", question), + }, + ]); + + return response.content as string; + } + + private getSources(documents: any): { title: string; score: number }[] { + if (!documents || !Array.isArray(documents)) { + console.warn("No valid documents provided to getSources"); + return []; + } + return this.sortUniqueDocsByScore(documents); + } + + private sortUniqueDocsByScore(documents: any[]): any[] { + const uniqueDocs = new Map(); + + // Iterate through all documents + for (const doc of documents) { + if (!doc.title || (!doc?.score && !doc?.rerank_score)) { + console.warn("Invalid document structure:", doc); + continue; + } + + const currentDoc = uniqueDocs.get(doc.title); + const isReranked = doc && "rerank_score" in doc; + const docScore = isReranked ? doc.rerank_score : doc.score; + + // If the title doesn't exist in the map, or if the new doc has a higher score, update the map + if (!currentDoc || docScore > (currentDoc.score ?? 0)) { + uniqueDocs.set(doc.title, { + title: doc.title, + score: docScore, + isReranked: isReranked, + }); + } + } + + // Convert the map values back to an array and sort by score in descending order + return Array.from(uniqueDocs.values()).sort((a, b) => (b.score ?? 0) - (a.score ?? 0)); + } + + private async executeToolCalls( + toolCalls: any[], + debug: boolean, + updateLoadingMessage?: (message: string) => void + ) { + const toolOutputs = []; + for (const toolCall of toolCalls) { + if (debug) { + console.log(`==== Step 2: Calling tool: ${toolCall.tool.name} ====`); + } + if (toolCall.tool.name === "localSearch") { + updateLoadingMessage?.(LOADING_MESSAGES.READING_FILES); + } else if (toolCall.tool.name === "webSearch") { + updateLoadingMessage?.(LOADING_MESSAGES.SEARCHING_WEB); + } + const output = await ToolManager.callTool(toolCall.tool, toolCall.args); + toolOutputs.push({ tool: toolCall.tool.name, output }); + } + return toolOutputs; + } + + private prepareEnhancedUserMessage(userMessage: string, toolOutputs: any[]) { + let context = ""; + if (toolOutputs.length > 0) { + context = + "\n# Additional context:\n\n" + + toolOutputs + .map((output) => `# ${output.tool}\n${JSON.stringify(output.output)}`) + .join("\n\n"); + } + return `User message: ${userMessage}\n${context}`; + } + + private getTimeExpression(toolCalls: any[]): string { + const timeRangeCall = toolCalls.find((call) => call.tool.name === "getTimeRangeMs"); + return timeRangeCall ? timeRangeCall.args.timeExpression : ""; + } + + private formatLocalSearchResult(documents: any[], timeExpression: string): string { + const formattedDocs = documents + .filter((doc) => doc.includeInContext) + .map((doc: any) => `Note in Vault: ${doc.content}`) + .join("\n\n"); + return timeExpression + ? `Local Search Result for ${timeExpression}:\n${formattedDocs}` + : `Local Search Result:\n${formattedDocs}`; + } +} + +export { CopilotPlusChainRunner, LLMChainRunner, VaultQAChainRunner }; diff --git a/src/LLMProviders/intentAnalyzer.ts b/src/LLMProviders/intentAnalyzer.ts new file mode 100644 index 00000000..4dbfd33c --- /dev/null +++ b/src/LLMProviders/intentAnalyzer.ts @@ -0,0 +1,161 @@ +import ChatModelManager from "@/LLMProviders/chatModelManager"; +import VectorStoreManager from "@/VectorStoreManager"; +import { indexTool, localSearchTool, webSearchTool } from "@/tools/SearchTools"; +import { + getCurrentTimeTool, + getTimeInfoByEpochTool, + getTimeRangeMsTool, + pomodoroTool, + TimeInfo, +} from "@/tools/TimeTools"; +import { simpleYoutubeTranscriptionTool } from "@/tools/YoutubeTools"; +import { ToolManager } from "@/tools/toolManager"; +import { extractYoutubeUrl } from "@/utils"; +import { BrevilabsClient } from "./brevilabsClient"; + +// TODO: Add @index with explicit pdf files in chat context menu +export const COPILOT_TOOL_NAMES = ["@vault", "@web", "@youtube", "@pomodoro"]; + +type ToolCall = { + tool: any; + args: any; +}; + +export class IntentAnalyzer { + private static tools = [ + getCurrentTimeTool, + getTimeInfoByEpochTool, + getTimeRangeMsTool, + localSearchTool, + indexTool, + pomodoroTool, + webSearchTool, + ]; + + static async analyzeIntent( + originalMessage: string, + vectorStoreManager: VectorStoreManager, + chatModelManager: ChatModelManager, + brevilabsClient: BrevilabsClient + ): Promise { + try { + // Only analyze the original message + const brocaResponse = await brevilabsClient.broca(originalMessage); + const brocaToolCalls = brocaResponse.response.tool_calls; + const salientTerms = brocaResponse.response.salience_terms; + + const processedToolCalls: ToolCall[] = []; + let timeRange: { startTime: TimeInfo; endTime: TimeInfo } | undefined; + + // Process tool calls from broca + for (const brocaToolCall of brocaToolCalls) { + const tool = this.tools.find((t) => t.name === brocaToolCall.tool); + if (tool) { + const args = brocaToolCall.args || {}; + + if (tool.name === "getTimeRangeMs") { + timeRange = await ToolManager.callTool(tool, args); + } + + if (tool.name === "indexVault") { + args.vectorStoreManager = vectorStoreManager; + } + + processedToolCalls.push({ tool, args }); + } + } + + // Process @ commands from original message only + await this.processAtCommands(originalMessage, processedToolCalls, { + timeRange, + salientTerms, + vectorStoreManager, + chatModelManager, + brevilabsClient, + }); + + return processedToolCalls; + } catch (error) { + console.error("Error in intent analysis:", error); + return []; + } + } + + private static async processAtCommands( + originalMessage: string, + processedToolCalls: ToolCall[], + context: { + timeRange?: { startTime: TimeInfo; endTime: TimeInfo }; + salientTerms: string[]; + vectorStoreManager: VectorStoreManager; + chatModelManager: ChatModelManager; + brevilabsClient: BrevilabsClient; + } + ): Promise { + const message = originalMessage.toLowerCase(); + const { timeRange, salientTerms, vectorStoreManager, chatModelManager, brevilabsClient } = + context; + + // Handle @vault command + if (message.includes("@vault") && (salientTerms.length > 0 || timeRange)) { + // Remove all @commands from the query + const cleanQuery = this.removeAtCommands(originalMessage); + + processedToolCalls.push({ + tool: localSearchTool, + args: { + timeRange: timeRange || undefined, + query: cleanQuery, + salientTerms, + vectorStoreManager, + chatModelManager, + brevilabsClient, + }, + }); + } + + // Handle @web command + if (message.includes("@web")) { + const cleanQuery = this.removeAtCommands(originalMessage); + processedToolCalls.push({ + tool: webSearchTool, + args: { + query: cleanQuery, + brevilabsClient, + }, + }); + } + + // Handle @pomodoro command + if (message.includes("@pomodoro")) { + const pomodoroMatch = originalMessage.match(/@pomodoro\s+(\S+)/i); + const interval = pomodoroMatch ? pomodoroMatch[1] : "25min"; + processedToolCalls.push({ + tool: pomodoroTool, + args: { interval }, + }); + } + + // Handle @youtube command + if (message.includes("@youtube")) { + const youtubeUrl = extractYoutubeUrl(originalMessage); + if (youtubeUrl) { + processedToolCalls.push({ + tool: simpleYoutubeTranscriptionTool, + args: { + url: youtubeUrl, + brevilabsClient, + }, + }); + } + } + } + + private static removeAtCommands(message: string): string { + return message + .split(" ") + .filter((word) => !COPILOT_TOOL_NAMES.includes(word.toLowerCase())) + .join(" ") + .trim(); + } +} diff --git a/src/LLMProviders/memoryManager.ts b/src/LLMProviders/memoryManager.ts index 0e410544..a8c3ee54 100644 --- a/src/LLMProviders/memoryManager.ts +++ b/src/LLMProviders/memoryManager.ts @@ -4,14 +4,19 @@ import { BaseChatMemory, BufferWindowMemory } from "langchain/memory"; export default class MemoryManager { private static instance: MemoryManager; private memory: BaseChatMemory; + private debug: boolean; - private constructor(private langChainParams: LangChainParams) { + private constructor( + private langChainParams: LangChainParams, + debug = false + ) { + this.debug = debug; this.initMemory(); } - static getInstance(langChainParams: LangChainParams): MemoryManager { + static getInstance(langChainParams: LangChainParams, debug = false): MemoryManager { if (!MemoryManager.instance) { - MemoryManager.instance = new MemoryManager(langChainParams); + MemoryManager.instance = new MemoryManager(langChainParams, debug); } return MemoryManager.instance; } @@ -23,14 +28,27 @@ export default class MemoryManager { inputKey: "input", returnMessages: true, }); + if (this.debug) + console.log("Memory initialized with context turns:", this.langChainParams.chatContextTurns); } getMemory(): BaseChatMemory { return this.memory; } - clearChatMemory(): void { - console.log("clearing chat memory"); - this.memory.clear(); + async clearChatMemory(): Promise { + if (this.debug) console.log("Clearing chat memory"); + await this.memory.clear(); + } + + async loadMemoryVariables(): Promise { + const variables = await this.memory.loadMemoryVariables({}); + if (this.debug) console.log("Loaded memory variables:", variables); + return variables; + } + + async saveContext(input: any, output: any): Promise { + if (this.debug) console.log("Saving to memory - Input:", input, "Output:", output); + await this.memory.saveContext(input, output); } } diff --git a/src/LLMProviders/promptManager.ts b/src/LLMProviders/promptManager.ts index c508d63f..2c90063d 100644 --- a/src/LLMProviders/promptManager.ts +++ b/src/LLMProviders/promptManager.ts @@ -9,9 +9,11 @@ import { export default class PromptManager { private static instance: PromptManager; private chatPrompt: ChatPromptTemplate; + private qaPrompt: ChatPromptTemplate; private constructor(private langChainParams: LangChainParams) { this.initChatPrompt(); + this.initQAPrompt(); } static getInstance(langChainParams: LangChainParams): PromptManager { @@ -32,6 +34,20 @@ export default class PromptManager { ]); } + private initQAPrompt(): void { + const qaTemplate = `{system_message} + +Answer the question with as detailed as possible based only on the following context: +{context} + +Question: {question} +`; + + this.qaPrompt = ChatPromptTemplate.fromMessages([ + SystemMessagePromptTemplate.fromTemplate(qaTemplate), + ]); + } + // Add this new method to escape curly braces private escapeTemplateString(str: string): string { return str.replace(/\{/g, "{{").replace(/\}/g, "}}"); @@ -40,4 +56,21 @@ export default class PromptManager { getChatPrompt(): ChatPromptTemplate { return this.chatPrompt; } + + async getQAPrompt({ + question, + context, + systemMessage, + }: { + question: string; + context: string; + systemMessage: string; + }): Promise { + const promptResult = await this.qaPrompt.format({ + question, + context, + system_message: systemMessage, + }); + return promptResult; + } } diff --git a/src/VectorStoreManager.ts b/src/VectorStoreManager.ts index d9acd039..47b788a5 100644 --- a/src/VectorStoreManager.ts +++ b/src/VectorStoreManager.ts @@ -1,4 +1,3 @@ -import { LangChainParams } from "@/aiParams"; import EncryptionService from "@/encryptionService"; import { CustomError } from "@/error"; import EmbeddingsManager from "@/LLMProviders/embeddingManager"; @@ -8,7 +7,9 @@ import VectorDBManager from "@/vectorDBManager"; import { Embeddings } from "@langchain/core/embeddings"; import { create, load, Orama, remove, removeMultiple, save, search } from "@orama/orama"; import { MD5 } from "crypto-js"; -import { App, Notice, Platform } from "obsidian"; +import { App, Notice, Platform, TAbstractFile, TFile, Vault } from "obsidian"; +import { LangChainParams } from "./aiParams"; +import { ChainType } from "./chainFactory"; import { VAULT_VECTOR_STORE_STRATEGY } from "./constants"; class VectorStoreManager { @@ -28,6 +29,10 @@ class VectorStoreManager { private totalFilesToIndex = 0; private initializationPromise: Promise; private isIndexLoaded = false; + private excludedFiles: Set = new Set(); + + private debounceDelay = 5000; // 5 seconds + private debounceTimer: number | null = null; constructor( app: App, @@ -69,6 +74,8 @@ class VectorStoreManager { getEmbeddingRequestsPerSecond: () => this.settings.embeddingRequestsPerSecond, debug: this.settings.debug, }); + + this.updateExcludedFiles(); } private async performPostInitializationTasks() { @@ -161,6 +168,12 @@ class VectorStoreManager { } const vectorLength = await this.getVectorLength(embeddingInstance); + if (!vectorLength || vectorLength === 0) { + throw new CustomError( + "Invalid vector length detected. Please check if your embedding model is working." + ); + } + const schema = this.createDynamicSchema(vectorLength); const db = await create({ @@ -180,10 +193,33 @@ class VectorStoreManager { return db; } + public getVault(): Vault { + return this.app.vault; + } + + public getSettings(): CopilotSettings { + return this.settings; + } + private async getVectorLength(embeddingInstance: Embeddings): Promise { - const sampleText = "Sample text for embedding"; - const sampleEmbedding = await embeddingInstance.embedQuery(sampleText); - return sampleEmbedding.length; + try { + const sampleText = "Sample text for embedding"; + const sampleEmbedding = await embeddingInstance.embedQuery(sampleText); + + if (!sampleEmbedding || sampleEmbedding.length === 0) { + throw new CustomError("Failed to get valid embedding vector length"); + } + + console.log( + `Detected vector length: ${sampleEmbedding.length} for model: ${EmbeddingsManager.getModelName(embeddingInstance)}` + ); + return sampleEmbedding.length; + } catch (error) { + console.error("Error getting vector length:", error); + throw new CustomError( + "Failed to determine embedding vector length. Please check your embedding model settings." + ); + } } private async ensureCorrectSchema(db: Orama, embeddingInstance: Embeddings): Promise { @@ -320,6 +356,49 @@ class VectorStoreManager { return allContent; } + private async checkAndHandleEmbeddingModelChange( + db: Orama, + embeddingInstance: Embeddings + ): Promise { + const singleDoc = await search(db, { + term: "", + limit: 1, + }); + + let prevEmbeddingModel: string | undefined; + + if (singleDoc.hits.length > 0) { + const oramaDocSample = singleDoc.hits[0]; + if ( + typeof oramaDocSample === "object" && + oramaDocSample !== null && + "document" in oramaDocSample + ) { + const document = oramaDocSample.document as { embeddingModel?: string }; + prevEmbeddingModel = document.embeddingModel; + } + } + + if (prevEmbeddingModel) { + const currEmbeddingModel = EmbeddingsManager.getModelName(embeddingInstance); + + if (!areEmbeddingModelsSame(prevEmbeddingModel, currEmbeddingModel)) { + // Model has changed, notify user and rebuild DB + new Notice("New embedding model detected. Rebuilding vector store from scratch."); + console.log("Detected change in embedding model. Rebuilding vector store from scratch."); + + // Create new DB with new model + this.oramaDb = await this.createNewDb(); + await this.saveDB(); + return true; + } + } else { + console.log("No previous embedding model found in the database."); + } + + return false; + } + public async indexVaultToVectorStore(overwrite?: boolean): Promise { // Add check at the start of the method if ((Platform.isMobile && this.settings.disableIndexOnMobile) || !this.oramaDb) { @@ -337,38 +416,13 @@ class VectorStoreManager { } await this.ensureCorrectSchema(this.oramaDb, embeddingInstance); - const singleDoc = await search(this.oramaDb, { - term: "", - limit: 1, - }); - - let prevEmbeddingModel: string | undefined; - - if (singleDoc.hits.length > 0) { - const oramaDocSample = singleDoc.hits[0]; - if ( - typeof oramaDocSample === "object" && - oramaDocSample !== null && - "document" in oramaDocSample - ) { - const document = oramaDocSample.document as { embeddingModel?: string }; - prevEmbeddingModel = document.embeddingModel; - } - } - - if (prevEmbeddingModel) { - const currEmbeddingModel = EmbeddingsManager.getModelName(embeddingInstance); - - if (!areEmbeddingModelsSame(prevEmbeddingModel, currEmbeddingModel)) { - // Model has changed, reinitialize DB - this.oramaDb = await this.createNewDb(); - overwrite = true; - new Notice("Detected change in embedding model. Rebuilding vector store from scratch."); - console.log("Detected change in embedding model. Rebuilding vector store from scratch."); - await this.saveDB(); - } - } else { - console.log("No previous embedding model found in the database."); + // Check for model change + const modelChanged = await this.checkAndHandleEmbeddingModelChange( + this.oramaDb, + embeddingInstance + ); + if (modelChanged) { + overwrite = true; } const latestMtime = await VectorDBManager.getLatestFileMtime(this.oramaDb); @@ -432,7 +486,7 @@ class VectorStoreManager { embeddingModel: EmbeddingsManager.getModelName(embeddingInstance), ctime: file.stat.ctime, mtime: file.stat.mtime, - tags: fileMetadatas[index]?.tags ?? [], // Assuming tags are in the metadata + tags: fileMetadatas[index]?.tags?.map((tag) => tag.tag) ?? [], extension: file.extension, metadata: fileMetadatas[index]?.frontmatter ?? {}, }; @@ -510,9 +564,14 @@ class VectorStoreManager { public async clearVectorStore(): Promise { try { - // Create a new, empty database instance + // Create a new, empty database instance with fresh schema this.oramaDb = await this.createNewDb(); + // Delete the existing database file + if (await this.app.vault.adapter.exists(this.dbPath)) { + await this.app.vault.adapter.remove(this.dbPath); + } + // Save the new, empty database await this.saveDB(); new Notice("Local vector store cleared successfully."); @@ -580,7 +639,6 @@ class VectorStoreManager { const searchResult = await search(this.oramaDb, { term: filePath, properties: ["path"], - tolerance: 1, }); if (searchResult.hits.length > 0) { await removeMultiple( @@ -588,6 +646,9 @@ class VectorStoreManager { searchResult.hits.map((hit) => hit.id), 500 ); + if (this.settings.debug) { + console.log(`Deleted document from local Copilot index: ${filePath}`); + } } } catch (err) { console.error("Error deleting document from local Copilotindex:", err); @@ -611,6 +672,98 @@ class VectorStoreManager { }); return result.hits[0]?.document; } + + public async initializeEventListeners() { + if (this.settings.debug) { + console.log("Copilot Plus: Initializing event listeners"); + } + this.app.vault.on("modify", this.handleFileModify); + this.app.vault.on("delete", this.handleFileDelete); + } + + private debouncedReindexFile = (file: TFile) => { + if (this.debounceTimer !== null) { + window.clearTimeout(this.debounceTimer); + } + this.debounceTimer = window.setTimeout(() => { + this.reindexFile(file); + this.debounceTimer = null; + }, this.debounceDelay); + }; + + private handleFileModify = async (file: TAbstractFile) => { + await this.updateExcludedFiles(); + const currentChainType = this.getLangChainParams().chainType; + if ( + file instanceof TFile && + file.extension === "md" && + !this.excludedFiles.has(file.path) && + currentChainType === ChainType.COPILOT_PLUS_CHAIN + ) { + if (this.settings.debug) { + console.log("Copilot Plus: Triggering reindex for file ", file.path); + } + await this.removeDocs(file.path); + this.debouncedReindexFile(file); + } + }; + + private handleFileDelete = async (file: TAbstractFile) => { + if (file instanceof TFile) { + await this.removeDocs(file.path); + } + }; + + private updateExcludedFiles = async () => { + this.excludedFiles = await this.getFilePathsForQA("exclusions"); + }; + + private async reindexFile(file: TFile) { + try { + const embeddingInstance = this.embeddingsManager.getEmbeddingsAPI(); + if (!embeddingInstance) { + throw new CustomError("Embedding instance not found."); + } + + const db = this.oramaDb; + if (!db) { + throw new CustomError("Copilot index is not loaded."); + } + + // Check for model change + const modelChanged = await this.checkAndHandleEmbeddingModelChange(db, embeddingInstance); + if (modelChanged) { + // Trigger full reindex and exit + await this.indexVaultToVectorStore(true); + return; + } + + // Proceed with single file reindex + const content = await this.app.vault.cachedRead(file); + const fileCache = this.app.metadataCache.getFileCache(file); + + const fileToSave = { + title: file.basename, + path: file.path, + content: content, + embeddingModel: EmbeddingsManager.getModelName(embeddingInstance), + ctime: file.stat.ctime, + mtime: file.stat.mtime, + tags: fileCache?.tags?.map((tag) => tag.tag) ?? [], + extension: file.extension, + metadata: fileCache?.frontmatter ?? {}, + }; + + await VectorDBManager.indexFile(db, embeddingInstance, fileToSave); + await this.saveDB(); + + if (this.settings.debug) { + console.log(`Reindexed file: ${file.path}`); + } + } catch (error) { + console.error(`Error reindexing file ${file.path}:`, error); + } + } } export default VectorStoreManager; diff --git a/src/aiParams.ts b/src/aiParams.ts index 051f449b..2c028a13 100644 --- a/src/aiParams.ts +++ b/src/aiParams.ts @@ -1,4 +1,5 @@ import { ChainType } from "@/chainFactory"; +import { BaseChatModel } from "@langchain/core/language_models/chat_models"; import { ChatPromptTemplate } from "@langchain/core/prompts"; export interface ModelConfig { @@ -54,6 +55,7 @@ export interface LangChainParams { export interface SetChainOptions { prompt?: ChatPromptTemplate; + chatModel?: BaseChatModel; noteFile?: any; forceNewCreation?: boolean; abortController?: AbortController; diff --git a/src/components/AddContextNoteModal.tsx b/src/components/AddContextNoteModal.tsx new file mode 100644 index 00000000..af615277 --- /dev/null +++ b/src/components/AddContextNoteModal.tsx @@ -0,0 +1,33 @@ +import { App, TFile } from "obsidian"; +import { BaseNoteModal } from "./BaseNoteModal"; + +interface AddContextNoteModalProps { + app: App; + onNoteSelect: (note: TFile) => void; + excludeNotes: string[]; +} + +export class AddContextNoteModal extends BaseNoteModal { + private onNoteSelect: (note: TFile) => void; + private excludeNotes: string[]; + + constructor({ app, onNoteSelect, excludeNotes }: AddContextNoteModalProps) { + super(app); + this.onNoteSelect = onNoteSelect; + this.excludeNotes = excludeNotes; + this.availableNotes = this.getOrderedNotes(excludeNotes); + } + + getItems(): TFile[] { + return this.availableNotes; + } + + getItemText(note: TFile): string { + const isActive = note.path === this.activeNote?.path; + return this.formatNoteTitle(note.basename, isActive, note.extension); + } + + onChooseItem(note: TFile, evt: MouseEvent | KeyboardEvent) { + this.onNoteSelect(note); + } +} diff --git a/src/components/AddImageModal.tsx b/src/components/AddImageModal.tsx new file mode 100644 index 00000000..4c3b6224 --- /dev/null +++ b/src/components/AddImageModal.tsx @@ -0,0 +1,29 @@ +import { App } from "obsidian"; + +export class AddImageModal { + private app: App; + private onImagesSelected: (files: File[]) => void; + + constructor(app: App, onImagesSelected: (files: File[]) => void) { + this.app = app; + this.onImagesSelected = onImagesSelected; + } + + open() { + const input = document.createElement("input"); + input.type = "file"; + input.accept = "image/*"; + input.multiple = true; + input.style.display = "none"; + + input.addEventListener("change", () => { + const files = Array.from(input.files || []); + this.onImagesSelected(files); + // Clean up + document.body.removeChild(input); + }); + + document.body.appendChild(input); + input.click(); + } +} diff --git a/src/components/BaseNoteModal.tsx b/src/components/BaseNoteModal.tsx new file mode 100644 index 00000000..dbb23830 --- /dev/null +++ b/src/components/BaseNoteModal.tsx @@ -0,0 +1,51 @@ +import { App, FuzzySuggestModal, TFile } from "obsidian"; + +export abstract class BaseNoteModal extends FuzzySuggestModal { + protected activeNote: TFile | null; + protected availableNotes: T[]; + + constructor(app: App) { + super(app); + this.activeNote = app.workspace.getActiveFile(); + } + + protected getOrderedNotes(excludeNotes: string[] = []): TFile[] { + // Get recently opened files first + const recentFiles = this.app.workspace + .getLastOpenFiles() + .map((filePath) => this.app.vault.getAbstractFileByPath(filePath)) + .filter( + (file): file is TFile => + file instanceof TFile && + (file.extension === "md" || file.extension === "pdf") && + !excludeNotes.includes(file.path) && + file.path !== this.activeNote?.path + ); + + // Get all other files that weren't recently opened + const allFiles = this.app.vault + .getFiles() + .filter((file) => file.extension === "md" || file.extension === "pdf"); + + const otherFiles = allFiles.filter( + (file) => + !recentFiles.some((recent) => recent.path === file.path) && + !excludeNotes.includes(file.path) && + file.path !== this.activeNote?.path + ); + + // Combine active note (if exists) with recent files and other files + return [...(this.activeNote ? [this.activeNote] : []), ...recentFiles, ...otherFiles]; + } + + protected formatNoteTitle(basename: string, isActive: boolean, extension?: string): string { + let title = basename; + if (isActive) { + title += " (current)"; + } + if (extension === "pdf") { + title += " (PDF)"; + } + return title; + } +} diff --git a/src/components/Chat.tsx b/src/components/Chat.tsx index cfb368ab..440f4e57 100644 --- a/src/components/Chat.tsx +++ b/src/components/Chat.tsx @@ -7,15 +7,19 @@ import { ABORT_REASON, AI_SENDER, EVENT_NAMES, + LOADING_MESSAGES, USER_SENDER, VAULT_VECTOR_STORE_STRATEGY, } from "@/constants"; import { AppContext } from "@/context"; +import { ContextProcessor } from "@/contextProcessor"; import { CustomPromptProcessor } from "@/customPromptProcessor"; import { getAIResponse } from "@/langchainStream"; import CopilotPlugin from "@/main"; +import { Mention } from "@/mentions/Mention"; import { CopilotSettings } from "@/settings/SettingsPage"; import SharedState, { ChatMessage, useSharedState } from "@/sharedState"; +import { FileParserManager } from "@/tools/FileParserManager"; import { createChangeToneSelectionPrompt, createTranslateSelectionPrompt, @@ -51,6 +55,7 @@ interface ChatProps { defaultSaveFolder: string; onSaveChat: (saveAsNote: () => Promise) => void; updateUserMessageHistory: (newMessage: string) => void; + fileParserManager: FileParserManager; plugin: CopilotPlugin; debug: boolean; } @@ -63,6 +68,7 @@ const Chat: React.FC = ({ defaultSaveFolder, onSaveChat, updateUserMessageHistory, + fileParserManager, plugin, debug, }) => { @@ -73,8 +79,16 @@ const Chat: React.FC = ({ const [inputMessage, setInputMessage] = useState(""); const [abortController, setAbortController] = useState(null); const [loading, setLoading] = useState(false); + const [loadingMessage, setLoadingMessage] = useState(LOADING_MESSAGES.DEFAULT); const [historyIndex, setHistoryIndex] = useState(-1); const [chatIsVisible, setChatIsVisible] = useState(false); + const [contextNotes, setContextNotes] = useState([]); + const [includeActiveNote, setIncludeActiveNote] = useState(false); + const [selectedImages, setSelectedImages] = useState([]); + + const mention = Mention.getInstance(plugin.settings.plusLicenseKey); + + const contextProcessor = ContextProcessor.getInstance(); useEffect(() => { const handleChatVisibility = (evt: CustomEvent<{ chatIsVisible: boolean }>) => { @@ -90,54 +104,118 @@ const Chat: React.FC = ({ const app = plugin.app || useContext(AppContext); - const handleSendMessage = async () => { - if (!inputMessage) return; + const processContextNotes = async ( + customPromptProcessor: CustomPromptProcessor, + fileParserManager: FileParserManager + ) => { + const activeNote = app.workspace.getActiveFile(); + return await contextProcessor.processContextNotes( + customPromptProcessor, + fileParserManager, + app.vault, + contextNotes, + includeActiveNote, + activeNote + ); + }; + const handleSendMessage = async (toolCalls?: string[]) => { + if (!inputMessage && selectedImages.length === 0) return; + + const timestamp = formatDateTime(new Date()); + + // Create message content array + const content: any[] = []; + + // Add text content if present + if (inputMessage) { + content.push({ + type: "text", + text: inputMessage, + }); + } + + // Add images if present + for (const image of selectedImages) { + const imageData = await image.arrayBuffer(); + const base64Image = Buffer.from(imageData).toString("base64"); + content.push({ + type: "image_url", + image_url: { + url: `data:${image.type};base64,${base64Image}`, + }, + }); + } + + const userMessage: ChatMessage = { + message: inputMessage || "Image message", + originalMessage: inputMessage, + sender: USER_SENDER, + isVisible: true, + timestamp: timestamp, + content: content, + }; + + // Clear input and images + setInputMessage(""); + setSelectedImages([]); + + // Add messages to chat history + addMessage(userMessage); + setLoading(true); + setLoadingMessage(LOADING_MESSAGES.DEFAULT); + + // First, process the original user message for custom prompts const customPromptProcessor = CustomPromptProcessor.getInstance(app.vault, settings); - const processedUserMessage = await customPromptProcessor.processCustomPrompt( - inputMessage, + let processedUserMessage = await customPromptProcessor.processCustomPrompt( + inputMessage || "", "", app.workspace.getActiveFile() as TFile | undefined ); - const timestamp = formatDateTime(new Date()); + // Extract Mentions (such as URLs) from original input message only + const urlContextAddition = await mention.processUrls(inputMessage || ""); - const userMessage: ChatMessage = { - message: inputMessage, - sender: USER_SENDER, - isVisible: true, - timestamp: timestamp, - }; + // Add context notes + const noteContextAddition = await processContextNotes(customPromptProcessor, fileParserManager); + + // Combine everything + processedUserMessage = processedUserMessage + urlContextAddition + noteContextAddition; + + let messageWithToolCalls = inputMessage; + // Add tool calls last + if (toolCalls) { + messageWithToolCalls += " " + toolCalls.join("\n"); + } const promptMessageHidden: ChatMessage = { message: processedUserMessage, + originalMessage: messageWithToolCalls, sender: USER_SENDER, isVisible: false, timestamp: timestamp, + content: content, }; - // Add user message to chat history - addMessage(userMessage); + // Add hidden user message to chat history addMessage(promptMessageHidden); - // Add to user message history - updateUserMessageHistory(inputMessage); - setHistoryIndex(-1); + // Add to user message history if there's text + if (inputMessage) { + updateUserMessageHistory(inputMessage); + setHistoryIndex(-1); + } - // Clear input - setInputMessage(""); - - // Display running dots to indicate loading - setLoading(true); await getAIResponse( promptMessageHidden, chainManager, addMessage, setCurrentAiMessage, setAbortController, - { debug } + { debug, updateLoadingMessage: setLoadingMessage } ); setLoading(false); + setLoadingMessage(LOADING_MESSAGES.DEFAULT); }; const navigateHistory = (direction: "up" | "down"): string => { @@ -301,7 +379,7 @@ ${chatContent}`; setLoading(true); try { const regeneratedResponse = await chainManager.runChain( - lastUserMessage.message, + lastUserMessage, new AbortController(), setCurrentAiMessage, addMessage, @@ -535,6 +613,7 @@ ${chatContent}`; currentAiMessage={currentAiMessage} indexVaultToVectorStore={settings.indexVaultToVectorStore as VAULT_VECTOR_STORE_STRATEGY} loading={loading} + loadingMessage={loadingMessage} app={app} onInsertAtCursor={handleInsertAtCursor} onRegenerate={handleRegenerate} @@ -571,10 +650,18 @@ ${chatContent}`; onSaveAsNote={() => handleSaveAsNote(true)} onRefreshVaultContext={refreshVaultContext} vault_qa_strategy={plugin.settings.indexVaultToVectorStore} - debug={debug} addMessage={addMessage} vault={app.vault} isIndexLoadedPromise={plugin.vectorStoreManager.getIsIndexLoaded()} + contextNotes={contextNotes} + setContextNotes={setContextNotes} + includeActiveNote={includeActiveNote} + setIncludeActiveNote={setIncludeActiveNote} + mention={mention} + selectedImages={selectedImages} + onAddImage={(files: File[]) => setSelectedImages((prev) => [...prev, ...files])} + setSelectedImages={setSelectedImages} + debug={debug} /> diff --git a/src/components/ChatComponents/ChatButtons.tsx b/src/components/ChatComponents/ChatButtons.tsx index 68be0596..47b14a37 100644 --- a/src/components/ChatComponents/ChatButtons.tsx +++ b/src/components/ChatComponents/ChatButtons.tsx @@ -1,6 +1,14 @@ import { USER_SENDER } from "@/constants"; import { ChatMessage } from "@/sharedState"; -import { Check, Copy, MessageSquarePlus, PenSquare, RotateCw, Trash2 } from "lucide-react"; +import { + Check, + Copy, + LibraryBig, + MessageSquarePlus, + PenSquare, + RotateCw, + Trash2, +} from "lucide-react"; import React from "react"; interface ChatButtonsProps { @@ -11,6 +19,8 @@ interface ChatButtonsProps { onRegenerate?: () => void; onEdit?: () => void; onDelete: () => void; + onShowSources?: () => void; + hasSources: boolean; } export const ChatButtons: React.FC = ({ @@ -21,6 +31,8 @@ export const ChatButtons: React.FC = ({ onRegenerate, onEdit, onDelete, + onShowSources, + hasSources, }) => { return (
@@ -38,6 +50,11 @@ export const ChatButtons: React.FC = ({ ) : ( <> + {hasSources && ( + + )} +
+ ); + + return ( +
+ }> + Add Note to Context + +
+ {activeNote && renderNote(activeNote, true)} + {uniqueNotes.map((note) => renderNote(note))} + {uniqueUrls.map((url) => ( +
+ + {(() => { + const hostname = new URL(url).hostname; + const cleanUrl = url.replace(/\/+$/, ""); + const path = cleanUrl.slice(cleanUrl.indexOf(hostname) + hostname.length); + if (path.length <= 1) return hostname; + return `${hostname}/...${cleanUrl.slice(-2)}`; + })()} + + url + +
+ ))} +
+
+ ); +}; diff --git a/src/components/ChatComponents/ChatControls.tsx b/src/components/ChatComponents/ChatControls.tsx index b5e2f590..ced40ab8 100644 --- a/src/components/ChatComponents/ChatControls.tsx +++ b/src/components/ChatComponents/ChatControls.tsx @@ -1,17 +1,20 @@ import { SetChainOptions } from "@/aiParams"; -import { CopilotPlusModal } from "@/components/CopilotPlusModal"; import { VAULT_VECTOR_STORE_STRATEGY } from "@/constants"; import { CustomError } from "@/error"; import { CopilotSettings } from "@/settings/SettingsPage"; -import { Notice } from "obsidian"; +import { App, Notice } from "obsidian"; import React, { useEffect, useState } from "react"; import { ChainType } from "@/chainFactory"; +import { AddContextNoteModal } from "@/components/AddContextNoteModal"; import { TooltipActionButton } from "@/components/ChatComponents/TooltipActionButton"; import { stringToChainType } from "@/utils"; import * as DropdownMenu from "@radix-ui/react-dropdown-menu"; import { ChevronDown, Download, Puzzle, RefreshCw } from "lucide-react"; +import { TFile } from "obsidian"; +import { ChatContextMenu } from "./ChatContextMenu"; + interface ChatControlsProps { currentChain: ChainType; setCurrentChain: (chain: ChainType, options?: SetChainOptions) => void; @@ -21,6 +24,13 @@ interface ChatControlsProps { settings: CopilotSettings; vault_qa_strategy: string; isIndexLoadedPromise: Promise; + app: App; + contextNotes: TFile[]; + setContextNotes: React.Dispatch>; + includeActiveNote: boolean; + setIncludeActiveNote: React.Dispatch>; + contextUrls: string[]; + onRemoveUrl: (url: string) => void; debug?: boolean; } @@ -32,11 +42,19 @@ const ChatControls: React.FC = ({ onRefreshVaultContext, settings, vault_qa_strategy, - debug, isIndexLoadedPromise, + app, + contextNotes, + setContextNotes, + includeActiveNote, + setIncludeActiveNote, + contextUrls, + onRemoveUrl, + debug, }) => { const [selectedChain, setSelectedChain] = useState(currentChain); const [isIndexLoaded, setIsIndexLoaded] = useState(false); + const activeNote = app.workspace.getActiveFile(); useEffect(() => { isIndexLoadedPromise.then((loaded) => { @@ -48,13 +66,14 @@ const ChatControls: React.FC = ({ const newChain = stringToChainType(value); setSelectedChain(newChain); - if (newChain === ChainType.COPILOT_PLUS_CHAIN) { - new CopilotPlusModal(app).open(); - // Reset the selected chain to the previous value - setSelectedChain(currentChain); - } else { - setCurrentChain(newChain, { debug }); - } + // TODO: Update Copilot Plus Modal to check the license key when ready to ship + // if (newChain === ChainType.COPILOT_PLUS_CHAIN) { + // new CopilotPlusModal(app).open(); + // // Reset the selected chain to the previous value + // setSelectedChain(currentChain); + // } else { + // setCurrentChain(newChain, { debug }); + // } }; useEffect(() => { @@ -64,10 +83,12 @@ const ChatControls: React.FC = ({ return; } - if (selectedChain === ChainType.VAULT_QA_CHAIN) { - if (vault_qa_strategy === VAULT_VECTOR_STORE_STRATEGY.ON_MODE_SWITCH) { - await onRefreshVaultContext(); - } + if ( + (selectedChain === ChainType.VAULT_QA_CHAIN || + selectedChain === ChainType.COPILOT_PLUS_CHAIN) && + vault_qa_strategy === VAULT_VECTOR_STORE_STRATEGY.ON_MODE_SWITCH + ) { + await onRefreshVaultContext(); } try { @@ -104,59 +125,102 @@ const ChatControls: React.FC = ({ // new SimilarNotesModal(app, similarChunks).open(); // }; - return ( -
- { - onNewChat(false); - }} - Icon={} - > -
New Chat
- {!settings.autosaveChat &&
(Unsaved history will be lost)
} - - }> - Save as Note - - {selectedChain === "vault_qa" && ( - } - > - Refresh Index for Vault - - )} -
- - - {currentChain === "llm_chain" && "chat"} - {currentChain === "vault_qa" && "vault QA (basic)"} - {currentChain === "copilot_plus" && "copilot plus (alpha)"} - - + const handleAddContext = () => { + const excludeNotes = [ + ...contextNotes.map((note) => note.path), + ...(includeActiveNote && activeNote ? [activeNote.path] : []), + ].filter(Boolean) as string[]; - - - handleChainChange({ value: "llm_chain" })}> - chat - - handleChainChange({ value: "vault_qa" })} - disabled={!isIndexLoaded} - className={!isIndexLoaded ? "disabled-menu-item" : ""} - > - vault QA (basic) {!isIndexLoaded && "(index not loaded)"} - - handleChainChange({ value: "copilot_plus" })} - disabled={!isIndexLoaded} - className={!isIndexLoaded ? "disabled-menu-item" : ""} - > - copilot plus (alpha) {!isIndexLoaded && "(index not loaded)"} - - - - + new AddContextNoteModal({ + app, + onNoteSelect: (note) => { + if (activeNote && note.path === activeNote.path) { + setIncludeActiveNote(true); + // Remove the note from contextNotes if it exists there + setContextNotes((prev) => prev.filter((n) => n.path !== note.path)); + } else { + setContextNotes((prev) => [...prev, note]); + } + }, + excludeNotes, + }).open(); + }; + + const handleRemoveContext = (path: string) => { + if (activeNote && path === activeNote.path) { + setIncludeActiveNote(false); + } else { + setContextNotes((prev) => prev.filter((note) => note.path !== path)); + } + }; + + return ( +
+
+ {currentChain === ChainType.COPILOT_PLUS_CHAIN && ( + + )} +
+ { + onNewChat(false); + }} + Icon={} + > +
New Chat
+ {!settings.autosaveChat &&
(Unsaved history will be lost)
} + + }> + Save as Note + + {selectedChain === "vault_qa" && ( + } + > + Refresh Index for Vault + + )} +
+ + + {currentChain === "llm_chain" && "chat"} + {currentChain === "vault_qa" && "vault QA (basic)"} + {currentChain === "copilot_plus" && "copilot plus (alpha)"} + + + + + + handleChainChange({ value: "llm_chain" })}> + chat + + handleChainChange({ value: "vault_qa" })} + disabled={!isIndexLoaded} + className={!isIndexLoaded ? "disabled-menu-item" : ""} + > + vault QA (basic) {!isIndexLoaded && "(index not loaded)"} + + handleChainChange({ value: "copilot_plus" })} + disabled={!isIndexLoaded} + className={!isIndexLoaded ? "disabled-menu-item" : ""} + > + copilot plus (alpha) {!isIndexLoaded && "(index not loaded)"} + + + + +
+
); diff --git a/src/components/ChatComponents/ChatInput.tsx b/src/components/ChatComponents/ChatInput.tsx index 17ee634f..7df2c410 100644 --- a/src/components/ChatComponents/ChatInput.tsx +++ b/src/components/ChatComponents/ChatInput.tsx @@ -2,19 +2,25 @@ import { CustomModel, SetChainOptions } from "@/aiParams"; import { ChainType } from "@/chainFactory"; import { ListPromptModal } from "@/components/ListPromptModal"; import { NoteTitleModal } from "@/components/NoteTitleModal"; +import { ContextProcessor } from "@/contextProcessor"; import { CustomPromptProcessor } from "@/customPromptProcessor"; +import { COPILOT_TOOL_NAMES } from "@/LLMProviders/intentAnalyzer"; +import { Mention } from "@/mentions/Mention"; import { CopilotSettings } from "@/settings/SettingsPage"; import { ChatMessage } from "@/sharedState"; +import { extractNoteTitles } from "@/utils"; import * as DropdownMenu from "@radix-ui/react-dropdown-menu"; -import { ChevronUp, Command, CornerDownLeft, StopCircle } from "lucide-react"; +import { ArrowBigUp, ChevronUp, Command, CornerDownLeft, Image, StopCircle } from "lucide-react"; import { App, Platform, TFile, Vault } from "obsidian"; import React, { useEffect, useRef, useState } from "react"; +import { AddImageModal } from "../AddImageModal"; import ChatControls from "./ChatControls"; +import { TooltipActionButton } from "./TooltipActionButton"; interface ChatInputProps { inputMessage: string; setInputMessage: (message: string) => void; - handleSendMessage: () => void; + handleSendMessage: (toolCalls?: string[]) => void; isGenerating: boolean; chatIsVisible: boolean; onStopGenerating: () => void; @@ -32,6 +38,14 @@ interface ChatInputProps { vault: Vault; vault_qa_strategy: string; isIndexLoadedPromise: Promise; + contextNotes: TFile[]; + setContextNotes: React.Dispatch>; + includeActiveNote: boolean; + setIncludeActiveNote: (include: boolean) => void; + mention: Mention; + selectedImages: File[]; + onAddImage: (files: File[]) => void; + setSelectedImages: React.Dispatch>; debug?: boolean; } @@ -58,15 +72,70 @@ const ChatInput: React.FC = ({ vault, vault_qa_strategy, isIndexLoadedPromise, + contextNotes, + setContextNotes, + includeActiveNote, + setIncludeActiveNote, + mention, + selectedImages, + onAddImage, + setSelectedImages, debug, }) => { const [shouldFocus, setShouldFocus] = useState(false); const [historyIndex, setHistoryIndex] = useState(-1); const [tempInput, setTempInput] = useState(""); const [isModelDropdownOpen, setIsModelDropdownOpen] = useState(false); + const [contextUrls, setContextUrls] = useState([]); const textAreaRef = useRef(null); const containerRef = useRef(null); + const debounce = any>( + fn: T, + delay: number + ): ((...args: Parameters) => void) => { + let timeoutId: NodeJS.Timeout; + return (...args: Parameters) => { + clearTimeout(timeoutId); + timeoutId = setTimeout(() => fn(...args), delay); + }; + }; + + // Debounce the context update to prevent excessive re-renders + const debouncedUpdateContext = debounce( + async ( + inputValue: string, + setContextNotes: React.Dispatch>, + currentContextNotes: TFile[], + app: App + ) => { + const noteTitles = extractNoteTitles(inputValue); + + const notesToAdd = await Promise.all( + noteTitles.map(async (title) => { + const files = app.vault.getMarkdownFiles(); + const file = files.find((file) => file.basename === title); + if (file) { + return Object.assign(file, { wasAddedViaReference: true }) as TFile & { + wasAddedViaReference: boolean; + }; + } + return undefined; + }) + ); + + const validNotes = notesToAdd.filter( + (note): note is TFile & { wasAddedViaReference: boolean } => + note !== undefined && !currentContextNotes.some((existing) => existing.path === note.path) + ); + + if (validNotes.length > 0) { + setContextNotes((prev) => [...prev, ...validNotes]); + } + }, + 50 + ); + const handleInputChange = async (event: React.ChangeEvent) => { const inputValue = event.target.value; const cursorPos = event.target.selectionStart; @@ -74,10 +143,26 @@ const ChatInput: React.FC = ({ setInputMessage(inputValue); adjustTextareaHeight(); + // Extract URLs and update mentions + const urls = mention.extractAllUrls(inputValue); + + // Update URLs in context, ensuring uniqueness + const newUrls = urls.filter((url) => !contextUrls.includes(url)); + if (newUrls.length > 0) { + // Use Set to ensure uniqueness + setContextUrls((prev) => Array.from(new Set([...prev, ...newUrls]))); + } + + // Update context with debouncing + debouncedUpdateContext(inputValue, setContextNotes, contextNotes, app); + + // Handle other input triggers if (cursorPos >= 2 && inputValue.slice(cursorPos - 2, cursorPos) === "[[") { showNoteTitleModal(cursorPos); } else if (inputValue === "/") { showCustomPromptModal(); + } else if (inputValue.slice(-1) === "@" && currentChain === ChainType.COPILOT_PLUS_CHAIN) { + showCopilotPlusOptionsModal(); } }; @@ -95,11 +180,31 @@ const ChatInput: React.FC = ({ const showNoteTitleModal = (cursorPos: number) => { const fetchNoteTitles = async () => { const noteTitles = app.vault.getMarkdownFiles().map((file: TFile) => file.basename); + const contextProcessor = ContextProcessor.getInstance(); - new NoteTitleModal(app, noteTitles, (noteTitle: string) => { + new NoteTitleModal(app, noteTitles, async (noteTitle: string) => { const before = inputMessage.slice(0, cursorPos - 2); const after = inputMessage.slice(cursorPos - 1); - setInputMessage(`${before}[[${noteTitle}]]${after}`); + const newInputMessage = `${before}[[${noteTitle}]]${after}`; + setInputMessage(newInputMessage); + + // Manually invoke debouncedUpdateContext + debouncedUpdateContext(newInputMessage, setContextNotes, contextNotes, app); + + const activeNote = app.workspace.getActiveFile(); + const noteFile = app.vault.getMarkdownFiles().find((file) => file.basename === noteTitle); + + if (noteFile) { + await contextProcessor.addNoteToContext( + noteFile, + vault, + contextNotes, + activeNote, + setContextNotes, + setIncludeActiveNote + ); + } + // Add a delay to ensure the cursor is set after inputMessage is updated setTimeout(() => { if (textAreaRef.current) { @@ -109,7 +214,6 @@ const ChatInput: React.FC = ({ }, 0); }).open(); }; - fetchNoteTitles(); }; @@ -127,6 +231,13 @@ const ChatInput: React.FC = ({ }).open(); }; + const showCopilotPlusOptionsModal = () => { + const options = COPILOT_TOOL_NAMES; + new ListPromptModal(app, options, (selectedOption: string) => { + setInputMessage(inputMessage + selectedOption + " "); + }).open(); + }; + useEffect(() => { setShouldFocus(chatIsVisible); }, [chatIsVisible]); @@ -147,6 +258,21 @@ const ChatInput: React.FC = ({ const lines = value.split("\n"); const currentLineIndex = value.substring(0, selectionStart).split("\n").length - 1; + // Check for Cmd+Shift+Enter (Mac) or Ctrl+Shift+Enter (Windows) + if (e.key === "Enter" && e.shiftKey && (Platform.isMacOS ? e.metaKey : e.ctrlKey)) { + e.preventDefault(); + e.stopPropagation(); + + if (currentChain === ChainType.COPILOT_PLUS_CHAIN) { + handleSendMessage(["@vault"]); + } else { + handleSendMessage(); + } + setHistoryIndex(-1); + setTempInput(""); + return; + } + if (e.key === "Enter" && !e.shiftKey) { e.preventDefault(); handleSendMessage(); @@ -196,6 +322,53 @@ const ChatInput: React.FC = ({ } }; + useEffect(() => { + // Get all note titles that are referenced using [[note]] syntax in the input + const currentTitles = new Set(extractNoteTitles(inputMessage)); + // Get all URLs mentioned in the input + const currentUrls = mention.extractAllUrls(inputMessage); + // Get the currently open note in the editor + const activeNote = app.workspace.getActiveFile(); + + setContextNotes((prev) => + prev.filter((note) => { + // Check if this note was added by typing [[note]] in the input + // as opposed to being added via the "Add Note to Context" button + const wasAddedViaReference = (note as any).wasAddedViaReference === true; + + // Special handling for the active note (currently open in editor) + if (note.path === activeNote?.path) { + if (wasAddedViaReference) { + // Case 1: Active note was added by typing [[note]] + // Keep it only if its title is still in the input + // This ensures it's removed when you delete the [[note]] reference + return currentTitles.has(note.basename); + } else { + // Case 2: Active note was NOT added by [[note]], but by the includeActiveNote toggle + // Keep it only if includeActiveNote is true + // This handles the "Include active note" toggle in the UI + return includeActiveNote; + } + } else { + // Handling for all other notes (not the active note) + if (wasAddedViaReference) { + // Case 3: Other note was added by typing [[note]] + // Keep it only if its title is still in the input + // This ensures it's removed when you delete the [[note]] reference + return currentTitles.has(note.basename); + } else { + // Case 4: Other note was added via "Add Note to Context" button + // Always keep these notes as they were manually added + return true; + } + } + }) + ); + + // Remove any URLs that are no longer present in the input + setContextUrls((prev) => prev.filter((url) => currentUrls.includes(url))); + }, [inputMessage, includeActiveNote]); + return (
= ({ settings={settings} vault_qa_strategy={vault_qa_strategy} isIndexLoadedPromise={isIndexLoadedPromise} + app={app} + contextNotes={contextNotes} + setContextNotes={setContextNotes} + includeActiveNote={includeActiveNote} + setIncludeActiveNote={setIncludeActiveNote} + contextUrls={contextUrls} + onRemoveUrl={(url: string) => setContextUrls((prev) => prev.filter((u) => u !== url))} debug={debug} /> + {selectedImages.length > 0 && ( +
+ {selectedImages.map((file, index) => ( +
+ {file.name} + +
+ ))} +
+ )} +