From 4907d87d3aa9a4057de2f97413a74bb29c78e211 Mon Sep 17 00:00:00 2001 From: Logan Yang Date: Mon, 4 Nov 2024 17:28:29 -0800 Subject: [PATCH] Implement QA Inclusions filter (#777) --- src/VectorStoreManager.ts | 80 +++++++++--------- src/constants.ts | 1 + src/main.ts | 2 +- src/promptUsageStrategy.ts | 2 +- src/settings/SettingsPage.tsx | 1 + src/settings/components/QASettings.tsx | 9 ++- src/utils.ts | 32 ++++++++ tests/customPromptProcessor.test.ts | 12 ++- tests/utils.test.ts | 108 +++++++++++++++++++++++++ 9 files changed, 203 insertions(+), 44 deletions(-) diff --git a/src/VectorStoreManager.ts b/src/VectorStoreManager.ts index 25ea350b..c7ad4571 100644 --- a/src/VectorStoreManager.ts +++ b/src/VectorStoreManager.ts @@ -3,7 +3,7 @@ import EncryptionService from "@/encryptionService"; import { CustomError } from "@/error"; import EmbeddingsManager from "@/LLMProviders/embeddingManager"; import { CopilotSettings } from "@/settings/SettingsPage"; -import { areEmbeddingModelsSame, getNotesFromTags, isPathInList } from "@/utils"; +import { areEmbeddingModelsSame, getFilePathsFromPatterns } from "@/utils"; import VectorDBManager from "@/vectorDBManager"; import { Embeddings } from "@langchain/core/embeddings"; import { create, load, Orama, remove, removeMultiple, save, search } from "@orama/orama"; @@ -214,9 +214,14 @@ class VectorStoreManager { private updateIndexingNoticeMessage() { if (this.indexNoticeMessage) { const status = this.isIndexingPaused ? " (Paused)" : ""; - this.indexNoticeMessage.textContent = `Copilot is indexing your vault...\n${this.indexedCount}/${this.totalFilesToIndex} files processed.${status}\nExclusions: ${ - this.settings.qaExclusions ? this.settings.qaExclusions : "None" - }`; + const filterType = this.settings.qaInclusions + ? `Inclusions: ${this.settings.qaInclusions}` + : `Exclusions: ${this.settings.qaExclusions || "None"}`; + + this.indexNoticeMessage.textContent = + `Copilot is indexing your vault...\n` + + `${this.indexedCount}/${this.totalFilesToIndex} files processed.${status}\n` + + filterType; } } @@ -247,49 +252,36 @@ class VectorStoreManager { return new Notice(frag, 0); } - private async getExcludedFiles(): Promise> { - const excludedFiles = new Set(); + private async getFilePathsForQA(filterType: "exclusions" | "inclusions"): Promise> { + const targetFiles = new Set(); - if (this.settings.qaExclusions) { + if (filterType === "exclusions" && this.settings.qaExclusions) { const exclusions = this.settings.qaExclusions.split(",").map((item) => item.trim()); - - for (const exclusion of exclusions) { - if (exclusion.startsWith("#")) { - // Tag-based exclusion - const tagName = exclusion.slice(1); - const taggedFiles = await getNotesFromTags(this.app.vault, [tagName]); - taggedFiles.forEach((file) => excludedFiles.add(file.path)); - } else if (exclusion.startsWith("*")) { - // File-extension-based exclusion - const extensionName = exclusion.slice(1); - this.app.vault.getFiles().forEach((file) => { - if (file.name.toLowerCase().contains(extensionName.toLowerCase())) { - excludedFiles.add(file.path); - } - }); - } else { - // Path-based exclusion - this.app.vault.getFiles().forEach((file) => { - if (isPathInList(file.path, exclusion)) { - excludedFiles.add(file.path); - } - }); - } - } + const excludedFilePaths = await getFilePathsFromPatterns(exclusions, this.app.vault); + excludedFilePaths.forEach((filePath) => targetFiles.add(filePath)); + } else if (filterType === "inclusions" && this.settings.qaInclusions) { + const inclusions = this.settings.qaInclusions.split(",").map((item) => item.trim()); + const includedFilePaths = await getFilePathsFromPatterns(inclusions, this.app.vault); + includedFilePaths.forEach((filePath) => targetFiles.add(filePath)); } - return excludedFiles; + return targetFiles; } - public async getAllNonExcludedNotesContent(): Promise { + public async getAllQAMarkdownContent(): Promise { let allContent = ""; - const excludedFiles = await this.getExcludedFiles(); - const nonExcludedFiles = this.app.vault - .getMarkdownFiles() - .filter((file) => !excludedFiles.has(file.path)); + const includedFiles = await this.getFilePathsForQA("inclusions"); + const excludedFiles = await this.getFilePathsForQA("exclusions"); - await Promise.all(nonExcludedFiles.map((file) => this.app.vault.cachedRead(file))).then( + const filteredFiles = this.app.vault.getMarkdownFiles().filter((file) => { + if (includedFiles.size > 0) { + return includedFiles.has(file.path); + } + return !excludedFiles.has(file.path); + }); + + await Promise.all(filteredFiles.map((file) => this.app.vault.cachedRead(file))).then( (contents) => contents.map((c) => (allContent += c + " ")) ); @@ -345,7 +337,8 @@ class VectorStoreManager { this.isIndexingPaused = false; this.isIndexingCancelled = false; - const excludedFiles = await this.getExcludedFiles(); + const includedFiles = await this.getFilePathsForQA("inclusions"); + const excludedFiles = await this.getFilePathsForQA("exclusions"); const files = this.app.vault .getMarkdownFiles() @@ -353,7 +346,14 @@ class VectorStoreManager { if (!latestMtime || overwrite) return true; return file.stat.mtime > latestMtime; }) - .filter((file) => !excludedFiles.has(file.path)); + .filter((file) => { + if (includedFiles.size > 0) { + // If inclusions are specified, only include files in the inclusion set + return includedFiles.has(file.path); + } + // Otherwise, use exclusion filter + return !excludedFiles.has(file.path); + }); const fileContents: string[] = await Promise.all( files.map((file) => this.app.vault.cachedRead(file)) diff --git a/src/constants.ts b/src/constants.ts index 46be4165..f7a55860 100644 --- a/src/constants.ts +++ b/src/constants.ts @@ -227,6 +227,7 @@ export const DEFAULT_SETTINGS: CopilotSettings = { customPromptsFolder: "copilot-custom-prompts", indexVaultToVectorStore: VAULT_VECTOR_STORE_STRATEGY.ON_MODE_SWITCH, qaExclusions: "", + qaInclusions: "", chatNoteContextPath: "", chatNoteContextTags: [], debug: false, diff --git a/src/main.ts b/src/main.ts index 05c250b6..2513c309 100644 --- a/src/main.ts +++ b/src/main.ts @@ -549,7 +549,7 @@ export default class CopilotPlugin extends Plugin { async countTotalTokens(): Promise { try { - const allContent = await this.vectorStoreManager.getAllNonExcludedNotesContent(); + const allContent = await this.vectorStoreManager.getAllQAMarkdownContent(); const totalTokens = await this.chainManager.chatModelManager.countTokens(allContent); return totalTokens; } catch (error) { diff --git a/src/promptUsageStrategy.ts b/src/promptUsageStrategy.ts index eb73926c..c4f1fee5 100644 --- a/src/promptUsageStrategy.ts +++ b/src/promptUsageStrategy.ts @@ -35,7 +35,7 @@ export class TimestampUsageStrategy implements PromptUsageStrategy { removeUnusedPrompts(existingPromptTitles: Array): PromptUsageStrategy { for (const key in this.usageData) { - if (!existingPromptTitles.contains(key)) { + if (!existingPromptTitles.includes(key)) { delete this.usageData[key]; } } diff --git a/src/settings/SettingsPage.tsx b/src/settings/SettingsPage.tsx index fc7cb3c9..4dd6c319 100644 --- a/src/settings/SettingsPage.tsx +++ b/src/settings/SettingsPage.tsx @@ -43,6 +43,7 @@ export interface CopilotSettings { enableEncryption: boolean; maxSourceChunks: number; qaExclusions: string; + qaInclusions: string; groqApiKey: string; enabledCommands: Record; activeModels: Array; diff --git a/src/settings/components/QASettings.tsx b/src/settings/components/QASettings.tsx index fa1c10e4..93ed34e2 100644 --- a/src/settings/components/QASettings.tsx +++ b/src/settings/components/QASettings.tsx @@ -131,12 +131,19 @@ const QASettings: React.FC = ({ onChange={(value) => updateSettings({ embeddingRequestsPerSecond: value })} /> updateSettings({ qaExclusions: value })} /> + updateSettings({ qaInclusions: value })} + /> ); }; diff --git a/src/utils.ts b/src/utils.ts index b3a5c898..17385c30 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -499,3 +499,35 @@ export function extractUniqueTitlesFromDocs(docs: Document[]): string[] { return Array.from(titlesSet); } + +export async function getFilePathsFromPatterns( + patterns: string[], + vault: Vault +): Promise { + const filePaths = new Set(); + + for (const pattern of patterns) { + if (pattern.startsWith("#")) { + // Tag-based pattern + const taggedFiles = await getNotesFromTags(vault, [pattern]); + taggedFiles.forEach((file) => filePaths.add(file.path)); + } else if (pattern.startsWith("*")) { + // File-extension-based pattern + const extensionName = pattern.slice(1); + vault.getFiles().forEach((file) => { + if (file.name.toLowerCase().endsWith(extensionName.toLowerCase())) { + filePaths.add(file.path); + } + }); + } else { + // Path-based pattern + vault.getFiles().forEach((file) => { + if (isPathInList(file.path, pattern)) { + filePaths.add(file.path); + } + }); + } + } + + return Array.from(filePaths); +} diff --git a/tests/customPromptProcessor.test.ts b/tests/customPromptProcessor.test.ts index dea02903..43fa7e97 100644 --- a/tests/customPromptProcessor.test.ts +++ b/tests/customPromptProcessor.test.ts @@ -21,6 +21,16 @@ jest.mock("@/utils", () => ({ processVariableNameForNotePath: jest.fn(), })); +const mockUsageStrategy = { + incrementPromptUsage: jest.fn(), + getPromptUsage: jest.fn(), + recordUsage: jest.fn(), + updateUsage: jest.fn(), + removeUnusedPrompts: jest.fn(), + compare: jest.fn(), + save: jest.fn(), +}; + describe("CustomPromptProcessor", () => { let processor: CustomPromptProcessor; let mockVault: Vault; @@ -40,7 +50,7 @@ describe("CustomPromptProcessor", () => { } as TFile; // Create an instance of CustomPromptProcessor with mocked dependencies - processor = CustomPromptProcessor.getInstance(mockVault, mockSettings); + processor = CustomPromptProcessor.getInstance(mockVault, mockSettings, mockUsageStrategy); }); it("should add 1 context and selectedText", async () => { diff --git a/tests/utils.test.ts b/tests/utils.test.ts index 572fbc12..77814135 100644 --- a/tests/utils.test.ts +++ b/tests/utils.test.ts @@ -2,6 +2,7 @@ import * as Obsidian from "obsidian"; import { TFile } from "obsidian"; import { extractNoteTitles, + getFilePathsFromPatterns, getNotesFromPath, getNotesFromTags, isFolderMatch, @@ -312,3 +313,110 @@ describe("extractNoteTitles", () => { expect(result).toEqual(expected); }); }); + +describe("getFilePathsFromPatterns", () => { + it("should return files matching tag patterns", async () => { + const mockVault = new Obsidian.Vault(); + const patterns = ["#tag1"]; + + const result = await getFilePathsFromPatterns(patterns, mockVault); + expect(result).toEqual(["test/test2/note1.md", "note4.md"]); + }); + + it("should return files matching file extension patterns", async () => { + const mockVault = new Obsidian.Vault(); + mockVault.getFiles = jest.fn().mockReturnValue([ + { path: "test/image1.jpg", name: "image1.jpg" }, + { path: "test/doc.md", name: "doc.md" }, + { path: "image2.jpg", name: "image2.jpg" }, + { path: "test/image3.JPG", name: "image3.JPG" }, + ]); + + const patterns = ["*.jpg"]; + + const result = await getFilePathsFromPatterns(patterns, mockVault); + expect(result).toEqual(["test/image1.jpg", "image2.jpg", "test/image3.JPG"]); + }); + + it("should return files matching path patterns", async () => { + const mockVault = new Obsidian.Vault(); + mockVault.getFiles = jest.fn().mockReturnValue([ + { path: "test/test2/note1.md", name: "note1.md" }, + { path: "test/note2.md", name: "note2.md" }, + { path: "test2/note3.md", name: "note3.md" }, + ]); + + const patterns = ["test/test2"]; + + const result = await getFilePathsFromPatterns(patterns, mockVault); + expect(result).toEqual(["test/test2/note1.md"]); + }); + + it("should handle multiple patterns of different types", async () => { + const mockVault = new Obsidian.Vault(); + mockVault.getFiles = jest.fn().mockReturnValue([ + { path: "test/test2/note1.md", name: "note1.md" }, + { path: "test/image1.jpg", name: "image1.jpg" }, + { path: "note4.md", name: "note4.md" }, + { path: "image2.jpg", name: "image2.jpg" }, + ]); + + const patterns = ["#tag1", "*.jpg", "test/test2"]; + + const result = await getFilePathsFromPatterns(patterns, mockVault); + expect(result).toEqual([ + "test/test2/note1.md", // From tag1 and test/test2 + "note4.md", // From tag1 + "test/image1.jpg", // From *.jpg + "image2.jpg", // From *.jpg + ]); + }); + + it("should handle note title patterns with [[]]", async () => { + const mockVault = new Obsidian.Vault(); + mockVault.getFiles = jest.fn().mockReturnValue([ + { path: "test/test2/note1.md", name: "note1.md" }, + { path: "test/note2.md", name: "note2.md" }, + ]); + + const patterns = ["[[note1]]"]; + + const result = await getFilePathsFromPatterns(patterns, mockVault); + expect(result).toEqual(["test/test2/note1.md"]); + }); + + it("should return empty array for non-matching patterns", async () => { + const mockVault = new Obsidian.Vault(); + mockVault.getFiles = jest.fn().mockReturnValue([ + { path: "test/test2/note1.md", name: "note1.md" }, + { path: "test/note2.md", name: "note2.md" }, + ]); + + const patterns = ["#nonexistenttag", "*.xyz", "nonexistentfolder"]; + + const result = await getFilePathsFromPatterns(patterns, mockVault); + expect(result).toEqual([]); + }); + + it("should handle empty pattern array", async () => { + const mockVault = new Obsidian.Vault(); + const patterns: string[] = []; + + const result = await getFilePathsFromPatterns(patterns, mockVault); + expect(result).toEqual([]); + }); + + it("should be case insensitive for file extensions", async () => { + const mockVault = new Obsidian.Vault(); + mockVault.getFiles = jest.fn().mockReturnValue([ + { path: "test/image1.JPG", name: "image1.JPG" }, + { path: "test/image2.jpg", name: "image2.jpg" }, + { path: "image3.Jpg", name: "image3.Jpg" }, + ]); + + const patterns = ["*.jpg"]; + + const result = await getFilePathsFromPatterns(patterns, mockVault); + expect(result).toEqual(["test/image1.JPG", "test/image2.jpg", "image3.Jpg"]); + }); +});