diff --git a/src/LLMProviders/embeddingManager.ts b/src/LLMProviders/embeddingManager.ts index b98143d8..d4448a5a 100644 --- a/src/LLMProviders/embeddingManager.ts +++ b/src/LLMProviders/embeddingManager.ts @@ -1,6 +1,6 @@ /* eslint-disable @typescript-eslint/no-explicit-any */ import { CustomModel } from "@/aiParams"; -import { BREVILABS_API_BASE_URL, EMBEDDING_BATCH_SIZE, EmbeddingModelProviders } from "@/constants"; +import { BREVILABS_API_BASE_URL, EmbeddingModelProviders } from "@/constants"; import { getDecryptedKey } from "@/encryptionService"; import { CustomError } from "@/error"; import { getModelKeyFromModel, getSettings, subscribeToSettingsChange } from "@/settings/model"; @@ -189,7 +189,7 @@ export default class EmbeddingManager { modelName, apiKey: await getDecryptedKey(settings.plusLicenseKey), timeout: 10000, - batchSize: EMBEDDING_BATCH_SIZE, + batchSize: getSettings().embeddingBatchSize, configuration: { baseURL: BREVILABS_API_BASE_URL, fetch: customModel.enableCors ? safeFetch : undefined, @@ -199,7 +199,7 @@ export default class EmbeddingManager { model: modelName, apiKey: await getDecryptedKey(settings.plusLicenseKey), timeout: 10000, - batchSize: 128, + batchSize: getSettings().embeddingBatchSize, dimensions: customModel.dimensions, baseUrl: BREVILABS_API_BASE_URL + "/embeddings", configuration: { @@ -210,6 +210,7 @@ export default class EmbeddingManager { modelName, apiKey: await getDecryptedKey(customModel.apiKey || settings.openAIApiKey), timeout: 10000, + batchSize: getSettings().embeddingBatchSize, configuration: { baseURL: customModel.baseUrl, fetch: customModel.enableCors ? safeFetch : undefined, @@ -252,6 +253,7 @@ export default class EmbeddingManager { [EmbeddingModelProviders.OPENAI_FORMAT]: { modelName, openAIApiKey: await getDecryptedKey(customModel.apiKey || ""), + batchSize: getSettings().embeddingBatchSize, configuration: { baseURL: customModel.baseUrl, fetch: customModel.enableCors ? safeFetch : undefined, diff --git a/src/constants.ts b/src/constants.ts index 3515005d..2037b3dd 100644 --- a/src/constants.ts +++ b/src/constants.ts @@ -20,8 +20,7 @@ export const DEFAULT_SYSTEM_PROMPT = `You are Obsidian Copilot, a helpful assist 11. Do NOT mention the additional context provided such as getCurrentTime and getTimeRangeMs if it's irrelevant to the user message.`; export const EMPTY_INDEX_ERROR_MESSAGE = "Copilot index does not exist. Please index your vault first!\n\n1. Set a working embedding model in QA settings. If it's not a local model, don't forget to set the API key. \n\n2. Click 'Refresh Index for Vault' and wait for indexing to complete. If you encounter the rate limiting error, please turn your request per second down in QA setting."; -export const CHUNK_SIZE = 4000; -export const EMBEDDING_BATCH_SIZE = 16; +export const CHUNK_SIZE = 6000; export const CONTEXT_SCORE_THRESHOLD = 0.4; export const TEXT_WEIGHT = 0.4; export const PLUS_MODE_DEFAULT_SOURCE_CHUNKS = 15; @@ -538,7 +537,8 @@ export const DEFAULT_SETTINGS: CopilotSettings = { groqApiKey: "", activeModels: BUILTIN_CHAT_MODELS, activeEmbeddingModels: BUILTIN_EMBEDDING_MODELS, - embeddingRequestsPerSecond: 10, + embeddingRequestsPerMin: 90, + embeddingBatchSize: 16, disableIndexOnMobile: true, showSuggestedPrompts: true, showRelevantNotes: true, diff --git a/src/rateLimiter.ts b/src/rateLimiter.ts index 58cc46fd..81a6ce6a 100644 --- a/src/rateLimiter.ts +++ b/src/rateLimiter.ts @@ -1,51 +1,28 @@ export class RateLimiter { - private queue: (() => void)[] = []; private lastRequestTime = 0; - private processing = false; - private requestsPerSecond: number; + private requestsPerMin: number; - constructor(requestsPerSecond: number) { - this.requestsPerSecond = requestsPerSecond; + constructor(requestsPerMin: number) { + this.requestsPerMin = requestsPerMin; } - setRequestsPerSecond(requestsPerSecond: number) { - this.requestsPerSecond = requestsPerSecond; + setRequestsPerMin(requestsPerMin: number) { + this.requestsPerMin = requestsPerMin; } - getRequestsPerSecond(): number { - return this.requestsPerSecond; + getRequestsPerMin(): number { + return this.requestsPerMin; } async wait(): Promise { - return new Promise((resolve) => { - this.queue.push(resolve); - this.process(); - }); - } + const now = Date.now(); + const timeSinceLastRequest = now - this.lastRequestTime; + const timeToWait = Math.max(0, 60000 / this.requestsPerMin - timeSinceLastRequest); - private async process(): Promise { - // Use an atomic compare-and-swap operation - if (this.processing) return; - this.processing = true; - - try { - while (this.queue.length > 0) { - const now = Date.now(); - // Calculate the time to wait until the next request can be made - const timeToWait = Math.max(0, this.lastRequestTime + 1000 / this.requestsPerSecond - now); - - if (timeToWait > 0) { - await new Promise((resolve) => setTimeout(resolve, timeToWait)); - } - - const resolve = this.queue.shift(); - if (resolve) { - this.lastRequestTime = Date.now(); - resolve(); - } - } - } finally { - this.processing = false; + if (timeToWait > 0) { + await new Promise((resolve) => setTimeout(resolve, timeToWait)); } + + this.lastRequestTime = Date.now(); } } diff --git a/src/search/dbOperations.ts b/src/search/dbOperations.ts index 7249ca3b..09db626a 100644 --- a/src/search/dbOperations.ts +++ b/src/search/dbOperations.ts @@ -1,5 +1,6 @@ import EmbeddingsManager from "@/LLMProviders/embeddingManager"; import { CustomError } from "@/error"; +import { logError, logInfo } from "@/logger"; import { getSettings, subscribeToSettingsChange } from "@/settings/model"; import { areEmbeddingModelsSame } from "@/utils"; import { Embeddings } from "@langchain/core/embeddings"; @@ -52,11 +53,11 @@ export class DBOperations { const newPath = await this.getDbPath(); if (this.dbPath && newPath !== this.dbPath) { - console.log("Path change detected, reinitializing database..."); + logInfo("Path change detected, reinitializing database..."); this.dbPath = newPath; await this.initializeChunkedStorage(); await this.initializeDB(await EmbeddingsManager.getInstance().getEmbeddingsAPI()); - console.log("Database reinitialized with new path:", newPath); + logInfo("Database reinitialized with new path:", newPath); } }); } @@ -90,12 +91,12 @@ export class DBOperations { try { if (await this.chunkedStorage.exists()) { this.oramaDb = await this.chunkedStorage.loadDatabase(); - console.log("Loaded existing chunked Orama database from disk."); + logInfo("Loaded existing chunked Orama database from disk."); return this.oramaDb; } } catch (error) { // If loading fails, we'll create a new database - console.log("Failed to load existing database, creating new one:", error); + logError("Failed to load existing database, creating new one:", error); } // Create new database if none exists or loading failed @@ -103,7 +104,7 @@ export class DBOperations { this.oramaDb = newDb; return newDb; } catch (error) { - console.error(`Error initializing Orama database:`, error); + logError(`Error initializing Orama database:`, error); new Notice("Failed to initialize Copilot database. Some features may be limited."); return undefined; } @@ -124,7 +125,7 @@ export class DBOperations { throw new CustomError("Orama database not found."); } } catch (error) { - console.error("Failed to initialize database during save:", error); + logError("Failed to initialize database during save:", error); throw new CustomError("Failed to initialize and save database."); } } @@ -134,10 +135,10 @@ export class DBOperations { this.hasUnsavedChanges = false; if (getSettings().debug) { - console.log("Orama database saved successfully at:", this.dbPath); + logInfo("Orama database saved successfully at:", this.dbPath); } } catch (error) { - console.error(`Error saving Orama database:`, error); + logError(`Error saving Orama database:`, error); throw error; } } @@ -162,9 +163,9 @@ export class DBOperations { await this.saveDB(); new Notice("Local Copilot index cleared successfully."); - console.log("Local Copilot index cleared successfully, new instance created."); + logInfo("Local Copilot index cleared successfully, new instance created."); } catch (err) { - console.error("Error clearing the local Copilot index:", err); + logError("Error clearing the local Copilot index:", err); new Notice("An error occurred while clearing the local Copilot index."); throw err; } @@ -186,12 +187,12 @@ export class DBOperations { 500 ); if (getSettings().debug) { - console.log(`Deleted document from local Copilot index: ${filePath}`); + logInfo(`Deleted document from local Copilot index: ${filePath}`); } } this.markUnsavedChanges(); } catch (err) { - console.error("Error deleting document from local Copilotindex:", err); + logError("Error deleting document from local Copilotindex:", err); } } @@ -237,7 +238,7 @@ export class DBOperations { // Ensure the directory exists if (!(await this.app.vault.adapter.exists(baseDir))) { await this.app.vault.adapter.mkdir(baseDir); - console.log("Created directory:", baseDir); + logInfo("Created directory:", baseDir); } } @@ -276,7 +277,7 @@ export class DBOperations { }, }, }); - console.log( + logInfo( `Created new Orama database for ${this.dbPath}. ` + `Embedding model: ${EmbeddingsManager.getModelName(embeddingInstance)} with vector length ${vectorLength}.` ); @@ -334,7 +335,7 @@ export class DBOperations { return 0; // Return 0 if no documents found } catch (err) { - console.error("Error getting latest file mtime from VectorDB:", err); + logError("Error getting latest file mtime from VectorDB:", err); return 0; } } @@ -380,14 +381,14 @@ export class DBOperations { try { await insert(this.oramaDb, docToSave); if (getSettings().debug) { - console.log( + logInfo( `${existingDoc.hits.length > 0 ? "Updated" : "Inserted"} document ${docToSave.id} in partition ${partition}` ); } this.markUnsavedChanges(); return docToSave; } catch (insertErr) { - console.error( + logError( `Failed to ${existingDoc.hits.length > 0 ? "update" : "insert"} document ${docToSave.id}:`, insertErr ); @@ -397,13 +398,13 @@ export class DBOperations { try { await insert(this.oramaDb, existingDoc.hits[0].document); } catch (restoreErr) { - console.error("Failed to restore previous document version:", restoreErr); + logError("Failed to restore previous document version:", restoreErr); } } return undefined; } } catch (err) { - console.error(`Error upserting document ${docToSave.id}:`, err); + logError(`Error upserting document ${docToSave.id}:`, err); return undefined; } } @@ -428,18 +429,27 @@ export class DBOperations { return 0; // Return 0 if no documents found } catch (err) { - console.error("Error getting latest file mtime from VectorDB:", err); + logError("Error getting latest file mtime from VectorDB:", err); return 0; } } async checkAndHandleEmbeddingModelChange(embeddingInstance: Embeddings): Promise { if (!this.oramaDb) { - console.error( - "Orama database not found. Please make sure you have a working embedding model." + logInfo( + "Embedding model change detected. Orama database not found. Initializing new database..." ); - return false; + try { + await this.initializeDB(embeddingInstance); + return true; + } catch (error) { + logError("Failed to initialize database:", error); + throw new CustomError( + "Failed to initialize Orama database. Please check your embedding model settings." + ); + } } + const singleDoc = await search(this.oramaDb, { term: "", limit: 1, @@ -465,7 +475,7 @@ export class DBOperations { if (!areEmbeddingModelsSame(prevEmbeddingModel, currEmbeddingModel)) { // Model has changed, notify user and rebuild DB new Notice("New embedding model detected. Rebuilding Copilot index from scratch."); - console.log("Detected change in embedding model. Rebuilding Copilot index from scratch."); + logInfo("Detected change in embedding model. Rebuilding Copilot index from scratch."); // Create new DB with new model this.oramaDb = await this.createNewDb(embeddingInstance); @@ -473,7 +483,7 @@ export class DBOperations { return true; } } else { - console.log("No previous embedding model found in the database."); + logInfo("No previous embedding model found in the database."); } return false; @@ -490,8 +500,24 @@ export class DBOperations { public async garbageCollect(): Promise { if (!this.oramaDb) { - throw new CustomError("Orama database not found."); + logInfo("Orama database not found during garbage collection. Attempting to initialize..."); + try { + const embeddingInstance = await EmbeddingsManager.getInstance().getEmbeddingsAPI(); + if (!embeddingInstance) { + throw new CustomError("No embedding model available."); + } + await this.initializeDB(embeddingInstance); + if (!this.oramaDb) { + throw new CustomError("Failed to initialize database after attempt."); + } + } catch (error) { + logError("Failed to initialize database during garbage collection:", error); + throw new CustomError( + "Failed to initialize database. Please check your embedding model settings." + ); + } } + try { const files = this.app.vault.getMarkdownFiles(); const filePaths = new Set(files.map((file) => file.path)); @@ -505,7 +531,7 @@ export class DBOperations { return 0; } - console.log( + logInfo( "Copilot index: Docs to remove during garbage collection:", Array.from(new Set(docsToRemove.map((doc) => doc.path))).join(", ") ); @@ -523,7 +549,7 @@ export class DBOperations { await this.saveDB(); return docsToRemove.length; } catch (err) { - console.error("Error garbage collecting the Copilot index:", err); + logError("Error garbage collecting the Copilot index:", err); throw new CustomError("Failed to garbage collect the Copilot index."); } } @@ -546,7 +572,7 @@ export class DBOperations { // Convert Set to sorted array return Array.from(uniquePaths).sort(); } catch (err) { - console.error("Error getting indexed files:", err); + logError("Error getting indexed files:", err); throw new CustomError("Failed to retrieve indexed files."); } } @@ -563,7 +589,7 @@ export class DBOperations { }); return result.hits.length === 0; } catch (err) { - console.error("Error checking if database is empty:", err); + logError("Error checking if database is empty:", err); throw new CustomError("Failed to check if database is empty."); } } diff --git a/src/search/indexOperations.ts b/src/search/indexOperations.ts index 4e22e719..f4739b96 100644 --- a/src/search/indexOperations.ts +++ b/src/search/indexOperations.ts @@ -1,4 +1,4 @@ -import { CHUNK_SIZE, EMBEDDING_BATCH_SIZE } from "@/constants"; +import { CHUNK_SIZE } from "@/constants"; import EmbeddingsManager from "@/LLMProviders/embeddingManager"; import { RateLimiter } from "@/rateLimiter"; import { getSettings, subscribeToSettingsChange } from "@/settings/model"; @@ -9,8 +9,6 @@ import { App, Notice, TFile } from "obsidian"; import { DBOperations } from "./dbOperations"; import { extractAppIgnoreSettings, getMatchingPatterns, shouldIndexFile } from "./searchUtils"; -const CHECKPOINT_INTERVAL = 8 * EMBEDDING_BATCH_SIZE; - export interface IndexingState { isIndexingPaused: boolean; isIndexingCancelled: boolean; @@ -23,6 +21,8 @@ export interface IndexingState { export class IndexOperations { private rateLimiter: RateLimiter; + private checkpointInterval: number; + private embeddingBatchSize: number; private state: IndexingState = { isIndexingPaused: false, isIndexingCancelled: false, @@ -38,12 +38,17 @@ export class IndexOperations { private dbOps: DBOperations, private embeddingsManager: EmbeddingsManager ) { - this.rateLimiter = new RateLimiter(getSettings().embeddingRequestsPerSecond); + const settings = getSettings(); + this.rateLimiter = new RateLimiter(settings.embeddingRequestsPerMin); + this.embeddingBatchSize = settings.embeddingBatchSize; + this.checkpointInterval = 8 * this.embeddingBatchSize; // Subscribe to settings changes subscribeToSettingsChange(async () => { const settings = getSettings(); - this.rateLimiter = new RateLimiter(settings.embeddingRequestsPerSecond); + this.rateLimiter = new RateLimiter(settings.embeddingRequestsPerMin); + this.embeddingBatchSize = settings.embeddingBatchSize; + this.checkpointInterval = 8 * this.embeddingBatchSize; }); } @@ -90,11 +95,11 @@ export class IndexOperations { // Process chunks in batches const errors: string[] = []; - for (let i = 0; i < allChunks.length; i += EMBEDDING_BATCH_SIZE) { + for (let i = 0; i < allChunks.length; i += this.embeddingBatchSize) { if (this.state.isIndexingCancelled) break; await this.handlePause(); - const batch = allChunks.slice(i, i + EMBEDDING_BATCH_SIZE); + const batch = allChunks.slice(i, i + this.embeddingBatchSize); try { await this.rateLimiter.wait(); const embeddings = await embeddingInstance.embedDocuments( @@ -122,9 +127,9 @@ export class IndexOperations { // Calculate if we've crossed a checkpoint threshold const previousCheckpoint = Math.floor( - (this.state.indexedCount - batch.length) / CHECKPOINT_INTERVAL + (this.state.indexedCount - batch.length) / this.checkpointInterval ); - const currentCheckpoint = Math.floor(this.state.indexedCount / CHECKPOINT_INTERVAL); + const currentCheckpoint = Math.floor(this.state.indexedCount / this.checkpointInterval); if (currentCheckpoint > previousCheckpoint) { await this.dbOps.saveDB(); diff --git a/src/settings/model.ts b/src/settings/model.ts index e687be0a..0160b365 100644 --- a/src/settings/model.ts +++ b/src/settings/model.ts @@ -54,7 +54,8 @@ export interface CopilotSettings { activeModels: Array; activeEmbeddingModels: Array; promptUsageTimestamps: Record; - embeddingRequestsPerSecond: number; + embeddingRequestsPerMin: number; + embeddingBatchSize: number; defaultOpenArea: DEFAULT_OPEN_AREA; disableIndexOnMobile: boolean; showSuggestedPrompts: boolean; @@ -158,6 +159,16 @@ export function sanitizeSettings(settings: CopilotSettings): CopilotSettings { ? DEFAULT_SETTINGS.contextTurns : contextTurns; + const embeddingRequestsPerMin = Number(settingsToSanitize.embeddingRequestsPerMin); + sanitizedSettings.embeddingRequestsPerMin = isNaN(embeddingRequestsPerMin) + ? DEFAULT_SETTINGS.embeddingRequestsPerMin + : embeddingRequestsPerMin; + + const embeddingBatchSize = Number(settingsToSanitize.embeddingBatchSize); + sanitizedSettings.embeddingBatchSize = isNaN(embeddingBatchSize) + ? DEFAULT_SETTINGS.embeddingBatchSize + : embeddingBatchSize; + return sanitizedSettings; } diff --git a/src/settings/v2/components/QASettings.tsx b/src/settings/v2/components/QASettings.tsx index 64d343ac..6a383321 100644 --- a/src/settings/v2/components/QASettings.tsx +++ b/src/settings/v2/components/QASettings.tsx @@ -130,16 +130,28 @@ const QASettings: React.FC = ({ indexVaultToVectorStore }) => { onChange={(value) => updateSetting("maxSourceChunks", value)} /> - {/* Requests per second */} + {/* Requests per Minute */} updateSetting("embeddingRequestsPerMin", value)} + /> + + {/* Embedding batch size */} + updateSetting("embeddingRequestsPerSecond", value)} + value={settings.embeddingBatchSize} + onChange={(value) => updateSetting("embeddingBatchSize", value)} /> {/* Number of Partitions */}