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

Add configurable batch size, update rps (#1240)

Browse source 
* Add configurable embedding batch size setting
* Enhance error handling in database initialization and garbage collection
* Refactor rate limiting to use requests per minute
This commit is contained in:
Logan Yang 2025-02-11 22:20:30 -08:00 committed by GitHub
parent 7845904e6c
commit 2ca4a4a955
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 122 additions and 89 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
src/LLMProviders/embeddingManager.ts
View file

@ -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,

6
src/constants.ts
View file

@ -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,

51
src/rateLimiter.ts
View file

@ -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<void> {
return new Promise<void>((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<void> {
// 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();
}
}

86
src/search/dbOperations.ts
View file

@ -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<boolean> {
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<number> {
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.");
}
}

23
src/search/indexOperations.ts
View file

@ -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();

13
src/settings/model.ts
View file

@ -54,7 +54,8 @@ export interface CopilotSettings {
activeModels: Array<CustomModel>;
activeEmbeddingModels: Array<CustomModel>;
promptUsageTimestamps: Record<string, number>;
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;
}

24
src/settings/v2/components/QASettings.tsx
View file

@ -130,16 +130,28 @@ const QASettings: React.FC<QASettingsProps> = ({ indexVaultToVectorStore }) => {
onChange={(value) => updateSetting("maxSourceChunks", value)}
/>
{/* Requests per second */}
{/* Requests per Minute */}
<SettingItem
type="slider"
title="Requests per second"
description="Default is 10. Decrease if you are rate limited by your embedding provider."
title="Requests per Minute"
description="Default is 90. Decrease if you are rate limited by your embedding provider."
min={10}
max={300}
step={10}
value={settings.embeddingRequestsPerMin}
onChange={(value) => updateSetting("embeddingRequestsPerMin", value)}
/>
{/* Embedding batch size */}
<SettingItem
type="slider"
title="Embedding Batch Size"
description="Default is 16. Increase if you are rate limited by your embedding provider."
min={1}
max={30}
max={128}
step={1}
value={settings.embeddingRequestsPerSecond}
onChange={(value) => updateSetting("embeddingRequestsPerSecond", value)}
value={settings.embeddingBatchSize}
onChange={(value) => updateSetting("embeddingBatchSize", value)}
/>
{/* Number of Partitions */}