This is page 5 of 8. Use http://codebase.md/hangwin/mcp-chrome?page={x} to view the full context.
# Directory Structure
```
├── .gitattributes
├── .github
│ └── workflows
│ └── build-release.yml
├── .gitignore
├── .husky
│ ├── commit-msg
│ └── pre-commit
├── .prettierignore
├── .prettierrc.json
├── .vscode
│ └── extensions.json
├── app
│ ├── chrome-extension
│ │ ├── _locales
│ │ │ ├── de
│ │ │ │ └── messages.json
│ │ │ ├── en
│ │ │ │ └── messages.json
│ │ │ ├── ja
│ │ │ │ └── messages.json
│ │ │ ├── ko
│ │ │ │ └── messages.json
│ │ │ ├── zh_CN
│ │ │ │ └── messages.json
│ │ │ └── zh_TW
│ │ │ └── messages.json
│ │ ├── .env.example
│ │ ├── assets
│ │ │ └── vue.svg
│ │ ├── common
│ │ │ ├── constants.ts
│ │ │ ├── message-types.ts
│ │ │ └── tool-handler.ts
│ │ ├── entrypoints
│ │ │ ├── background
│ │ │ │ ├── index.ts
│ │ │ │ ├── native-host.ts
│ │ │ │ ├── semantic-similarity.ts
│ │ │ │ ├── storage-manager.ts
│ │ │ │ └── tools
│ │ │ │ ├── base-browser.ts
│ │ │ │ ├── browser
│ │ │ │ │ ├── bookmark.ts
│ │ │ │ │ ├── common.ts
│ │ │ │ │ ├── console.ts
│ │ │ │ │ ├── file-upload.ts
│ │ │ │ │ ├── history.ts
│ │ │ │ │ ├── index.ts
│ │ │ │ │ ├── inject-script.ts
│ │ │ │ │ ├── interaction.ts
│ │ │ │ │ ├── keyboard.ts
│ │ │ │ │ ├── network-capture-debugger.ts
│ │ │ │ │ ├── network-capture-web-request.ts
│ │ │ │ │ ├── network-request.ts
│ │ │ │ │ ├── screenshot.ts
│ │ │ │ │ ├── vector-search.ts
│ │ │ │ │ ├── web-fetcher.ts
│ │ │ │ │ └── window.ts
│ │ │ │ └── index.ts
│ │ │ ├── content.ts
│ │ │ ├── offscreen
│ │ │ │ ├── index.html
│ │ │ │ └── main.ts
│ │ │ └── popup
│ │ │ ├── App.vue
│ │ │ ├── components
│ │ │ │ ├── ConfirmDialog.vue
│ │ │ │ ├── icons
│ │ │ │ │ ├── BoltIcon.vue
│ │ │ │ │ ├── CheckIcon.vue
│ │ │ │ │ ├── DatabaseIcon.vue
│ │ │ │ │ ├── DocumentIcon.vue
│ │ │ │ │ ├── index.ts
│ │ │ │ │ ├── TabIcon.vue
│ │ │ │ │ ├── TrashIcon.vue
│ │ │ │ │ └── VectorIcon.vue
│ │ │ │ ├── ModelCacheManagement.vue
│ │ │ │ └── ProgressIndicator.vue
│ │ │ ├── index.html
│ │ │ ├── main.ts
│ │ │ └── style.css
│ │ ├── eslint.config.js
│ │ ├── inject-scripts
│ │ │ ├── click-helper.js
│ │ │ ├── fill-helper.js
│ │ │ ├── inject-bridge.js
│ │ │ ├── interactive-elements-helper.js
│ │ │ ├── keyboard-helper.js
│ │ │ ├── network-helper.js
│ │ │ ├── screenshot-helper.js
│ │ │ └── web-fetcher-helper.js
│ │ ├── LICENSE
│ │ ├── package.json
│ │ ├── public
│ │ │ ├── icon
│ │ │ │ ├── 128.png
│ │ │ │ ├── 16.png
│ │ │ │ ├── 32.png
│ │ │ │ ├── 48.png
│ │ │ │ └── 96.png
│ │ │ ├── libs
│ │ │ │ └── ort.min.js
│ │ │ └── wxt.svg
│ │ ├── README.md
│ │ ├── tsconfig.json
│ │ ├── utils
│ │ │ ├── content-indexer.ts
│ │ │ ├── i18n.ts
│ │ │ ├── image-utils.ts
│ │ │ ├── lru-cache.ts
│ │ │ ├── model-cache-manager.ts
│ │ │ ├── offscreen-manager.ts
│ │ │ ├── semantic-similarity-engine.ts
│ │ │ ├── simd-math-engine.ts
│ │ │ ├── text-chunker.ts
│ │ │ └── vector-database.ts
│ │ ├── workers
│ │ │ ├── ort-wasm-simd-threaded.jsep.mjs
│ │ │ ├── ort-wasm-simd-threaded.jsep.wasm
│ │ │ ├── ort-wasm-simd-threaded.mjs
│ │ │ ├── ort-wasm-simd-threaded.wasm
│ │ │ ├── simd_math_bg.wasm
│ │ │ ├── simd_math.js
│ │ │ └── similarity.worker.js
│ │ └── wxt.config.ts
│ └── native-server
│ ├── debug.sh
│ ├── install.md
│ ├── jest.config.js
│ ├── package.json
│ ├── README.md
│ ├── src
│ │ ├── cli.ts
│ │ ├── constant
│ │ │ └── index.ts
│ │ ├── file-handler.ts
│ │ ├── index.ts
│ │ ├── mcp
│ │ │ ├── mcp-server-stdio.ts
│ │ │ ├── mcp-server.ts
│ │ │ ├── register-tools.ts
│ │ │ └── stdio-config.json
│ │ ├── native-messaging-host.ts
│ │ ├── scripts
│ │ │ ├── browser-config.ts
│ │ │ ├── build.ts
│ │ │ ├── constant.ts
│ │ │ ├── postinstall.ts
│ │ │ ├── register-dev.ts
│ │ │ ├── register.ts
│ │ │ ├── run_host.bat
│ │ │ ├── run_host.sh
│ │ │ └── utils.ts
│ │ ├── server
│ │ │ ├── index.ts
│ │ │ └── server.test.ts
│ │ └── util
│ │ └── logger.ts
│ └── tsconfig.json
├── commitlint.config.cjs
├── docs
│ ├── ARCHITECTURE_zh.md
│ ├── ARCHITECTURE.md
│ ├── CHANGELOG.md
│ ├── CONTRIBUTING_zh.md
│ ├── CONTRIBUTING.md
│ ├── TOOLS_zh.md
│ ├── TOOLS.md
│ ├── TROUBLESHOOTING_zh.md
│ ├── TROUBLESHOOTING.md
│ └── WINDOWS_INSTALL_zh.md
├── eslint.config.js
├── LICENSE
├── package.json
├── packages
│ ├── shared
│ │ ├── package.json
│ │ ├── src
│ │ │ ├── constants.ts
│ │ │ ├── index.ts
│ │ │ ├── tools.ts
│ │ │ └── types.ts
│ │ └── tsconfig.json
│ └── wasm-simd
│ ├── .gitignore
│ ├── BUILD.md
│ ├── Cargo.toml
│ ├── package.json
│ ├── README.md
│ └── src
│ └── lib.rs
├── pnpm-lock.yaml
├── pnpm-workspace.yaml
├── prompt
│ ├── content-analize.md
│ ├── excalidraw-prompt.md
│ └── modify-web.md
├── README_zh.md
├── README.md
├── releases
│ ├── chrome-extension
│ │ └── latest
│ │ └── chrome-mcp-server-lastest.zip
│ └── README.md
└── test-inject-script.js
```
# Files
--------------------------------------------------------------------------------
/app/chrome-extension/utils/vector-database.ts:
--------------------------------------------------------------------------------
```typescript
/**
* Vector database manager
* Uses hnswlib-wasm for high-performance vector similarity search
* Implements singleton pattern to avoid duplicate WASM module initialization
*/
import { loadHnswlib } from 'hnswlib-wasm-static';
import type { TextChunk } from './text-chunker';
export interface VectorDocument {
id: string;
tabId: number;
url: string;
title: string;
chunk: TextChunk;
embedding: Float32Array;
timestamp: number;
}
export interface SearchResult {
document: VectorDocument;
similarity: number;
distance: number;
}
export interface VectorDatabaseConfig {
dimension: number;
maxElements: number;
efConstruction: number;
M: number;
efSearch: number;
indexFileName: string;
enableAutoCleanup?: boolean;
maxRetentionDays?: number;
}
let globalHnswlib: any = null;
let globalHnswlibInitPromise: Promise<any> | null = null;
let globalHnswlibInitialized = false;
let syncInProgress = false;
let pendingSyncPromise: Promise<void> | null = null;
const DB_NAME = 'VectorDatabaseStorage';
const DB_VERSION = 1;
const STORE_NAME = 'documentMappings';
/**
* IndexedDB helper functions
*/
class IndexedDBHelper {
private static dbPromise: Promise<IDBDatabase> | null = null;
static async getDB(): Promise<IDBDatabase> {
if (!this.dbPromise) {
this.dbPromise = new Promise((resolve, reject) => {
const request = indexedDB.open(DB_NAME, DB_VERSION);
request.onerror = () => reject(request.error);
request.onsuccess = () => resolve(request.result);
request.onupgradeneeded = (event) => {
const db = (event.target as IDBOpenDBRequest).result;
if (!db.objectStoreNames.contains(STORE_NAME)) {
const store = db.createObjectStore(STORE_NAME, { keyPath: 'id' });
store.createIndex('indexFileName', 'indexFileName', { unique: false });
}
};
});
}
return this.dbPromise;
}
static async saveData(indexFileName: string, data: any): Promise<void> {
const db = await this.getDB();
const transaction = db.transaction([STORE_NAME], 'readwrite');
const store = transaction.objectStore(STORE_NAME);
await new Promise<void>((resolve, reject) => {
const request = store.put({
id: indexFileName,
indexFileName,
data,
timestamp: Date.now(),
});
request.onsuccess = () => resolve();
request.onerror = () => reject(request.error);
});
}
static async loadData(indexFileName: string): Promise<any | null> {
const db = await this.getDB();
const transaction = db.transaction([STORE_NAME], 'readonly');
const store = transaction.objectStore(STORE_NAME);
return new Promise<any | null>((resolve, reject) => {
const request = store.get(indexFileName);
request.onsuccess = () => {
const result = request.result;
resolve(result ? result.data : null);
};
request.onerror = () => reject(request.error);
});
}
static async deleteData(indexFileName: string): Promise<void> {
const db = await this.getDB();
const transaction = db.transaction([STORE_NAME], 'readwrite');
const store = transaction.objectStore(STORE_NAME);
await new Promise<void>((resolve, reject) => {
const request = store.delete(indexFileName);
request.onsuccess = () => resolve();
request.onerror = () => reject(request.error);
});
}
/**
* Clear all IndexedDB data (for complete cleanup during model switching)
*/
static async clearAllData(): Promise<void> {
try {
const db = await this.getDB();
const transaction = db.transaction([STORE_NAME], 'readwrite');
const store = transaction.objectStore(STORE_NAME);
await new Promise<void>((resolve, reject) => {
const request = store.clear();
request.onsuccess = () => {
console.log('IndexedDBHelper: All data cleared from IndexedDB');
resolve();
};
request.onerror = () => reject(request.error);
});
} catch (error) {
console.error('IndexedDBHelper: Failed to clear all data:', error);
throw error;
}
}
/**
* Get all stored keys (for debugging)
*/
static async getAllKeys(): Promise<string[]> {
try {
const db = await this.getDB();
const transaction = db.transaction([STORE_NAME], 'readonly');
const store = transaction.objectStore(STORE_NAME);
return new Promise<string[]>((resolve, reject) => {
const request = store.getAllKeys();
request.onsuccess = () => resolve(request.result as string[]);
request.onerror = () => reject(request.error);
});
} catch (error) {
console.error('IndexedDBHelper: Failed to get all keys:', error);
return [];
}
}
}
/**
* Global hnswlib-wasm initialization function
* Ensures initialization only once across the entire application
*/
async function initializeGlobalHnswlib(): Promise<any> {
if (globalHnswlibInitialized && globalHnswlib) {
return globalHnswlib;
}
if (globalHnswlibInitPromise) {
return globalHnswlibInitPromise;
}
globalHnswlibInitPromise = (async () => {
try {
console.log('VectorDatabase: Initializing global hnswlib-wasm instance...');
globalHnswlib = await loadHnswlib();
globalHnswlibInitialized = true;
console.log('VectorDatabase: Global hnswlib-wasm instance initialized successfully');
return globalHnswlib;
} catch (error) {
console.error('VectorDatabase: Failed to initialize global hnswlib-wasm:', error);
globalHnswlibInitPromise = null;
throw error;
}
})();
return globalHnswlibInitPromise;
}
export class VectorDatabase {
private index: any = null;
private isInitialized = false;
private isInitializing = false;
private initPromise: Promise<void> | null = null;
private documents = new Map<number, VectorDocument>();
private tabDocuments = new Map<number, Set<number>>();
private nextLabel = 0;
private readonly config: VectorDatabaseConfig;
constructor(config?: Partial<VectorDatabaseConfig>) {
this.config = {
dimension: 384,
maxElements: 100000,
efConstruction: 200,
M: 48,
efSearch: 50,
indexFileName: 'tab_content_index.dat',
enableAutoCleanup: true,
maxRetentionDays: 30,
...config,
};
console.log('VectorDatabase: Initialized with config:', {
dimension: this.config.dimension,
efSearch: this.config.efSearch,
M: this.config.M,
efConstruction: this.config.efConstruction,
enableAutoCleanup: this.config.enableAutoCleanup,
maxRetentionDays: this.config.maxRetentionDays,
});
}
/**
* Initialize vector database
*/
public async initialize(): Promise<void> {
if (this.isInitialized) return;
if (this.isInitializing && this.initPromise) return this.initPromise;
this.isInitializing = true;
this.initPromise = this._doInitialize().finally(() => {
this.isInitializing = false;
});
return this.initPromise;
}
private async _doInitialize(): Promise<void> {
try {
console.log('VectorDatabase: Initializing...');
const hnswlib = await initializeGlobalHnswlib();
hnswlib.EmscriptenFileSystemManager.setDebugLogs(true);
this.index = new hnswlib.HierarchicalNSW(
'cosine',
this.config.dimension,
this.config.indexFileName,
);
await this.syncFileSystem('read');
const indexExists = hnswlib.EmscriptenFileSystemManager.checkFileExists(
this.config.indexFileName,
);
if (indexExists) {
console.log('VectorDatabase: Loading existing index...');
try {
await this.index.readIndex(this.config.indexFileName, this.config.maxElements);
this.index.setEfSearch(this.config.efSearch);
await this.loadDocumentMappings();
if (this.documents.size > 0) {
const maxLabel = Math.max(...Array.from(this.documents.keys()));
this.nextLabel = maxLabel + 1;
console.log(
`VectorDatabase: Loaded existing index with ${this.documents.size} documents, next label: ${this.nextLabel}`,
);
} else {
const indexCount = this.index.getCurrentCount();
if (indexCount > 0) {
console.warn(
`VectorDatabase: Index has ${indexCount} vectors but no document mappings found. This may cause label mismatch.`,
);
this.nextLabel = indexCount;
} else {
this.nextLabel = 0;
}
console.log(
`VectorDatabase: No document mappings found, starting with next label: ${this.nextLabel}`,
);
}
} catch (loadError) {
console.warn(
'VectorDatabase: Failed to load existing index, creating new one:',
loadError,
);
this.index.initIndex(
this.config.maxElements,
this.config.M,
this.config.efConstruction,
200,
);
this.index.setEfSearch(this.config.efSearch);
this.nextLabel = 0;
}
} else {
console.log('VectorDatabase: Creating new index...');
this.index.initIndex(
this.config.maxElements,
this.config.M,
this.config.efConstruction,
200,
);
this.index.setEfSearch(this.config.efSearch);
this.nextLabel = 0;
}
this.isInitialized = true;
console.log('VectorDatabase: Initialization completed successfully');
} catch (error) {
console.error('VectorDatabase: Initialization failed:', error);
this.isInitialized = false;
throw error;
}
}
/**
* Add document to vector database
*/
public async addDocument(
tabId: number,
url: string,
title: string,
chunk: TextChunk,
embedding: Float32Array,
): Promise<number> {
if (!this.isInitialized) {
await this.initialize();
}
const documentId = this.generateDocumentId(tabId, chunk.index);
const document: VectorDocument = {
id: documentId,
tabId,
url,
title,
chunk,
embedding,
timestamp: Date.now(),
};
try {
// Validate vector data
if (!embedding || embedding.length !== this.config.dimension) {
const errorMsg = `Invalid embedding dimension: expected ${this.config.dimension}, got ${embedding?.length || 0}`;
console.error('VectorDatabase: Dimension mismatch detected!', {
expectedDimension: this.config.dimension,
actualDimension: embedding?.length || 0,
documentId,
tabId,
url,
title: title.substring(0, 50) + '...',
});
// This might be caused by model switching, suggest reinitialization
console.warn(
'VectorDatabase: This might be caused by model switching. Consider reinitializing the vector database with the correct dimension.',
);
throw new Error(errorMsg);
}
// Check if vector data contains invalid values
for (let i = 0; i < embedding.length; i++) {
if (!isFinite(embedding[i])) {
throw new Error(`Invalid embedding value at index ${i}: ${embedding[i]}`);
}
}
// Ensure we have a clean Float32Array
let cleanEmbedding: Float32Array;
if (embedding instanceof Float32Array) {
cleanEmbedding = embedding;
} else {
cleanEmbedding = new Float32Array(embedding);
}
// Use current nextLabel as label
const label = this.nextLabel++;
console.log(
`VectorDatabase: Adding document with label ${label}, embedding dimension: ${embedding.length}`,
);
// Add vector to index
// According to hnswlib-wasm-static emscripten binding requirements, need to create VectorFloat type
console.log(`VectorDatabase: 🔧 DEBUGGING - About to call addPoint with:`, {
embeddingType: typeof cleanEmbedding,
isFloat32Array: cleanEmbedding instanceof Float32Array,
length: cleanEmbedding.length,
firstFewValues: Array.from(cleanEmbedding.slice(0, 3)),
label: label,
replaceDeleted: false,
});
// Method 1: Try using VectorFloat constructor (if available)
let vectorToAdd;
try {
// Check if VectorFloat constructor exists
if (globalHnswlib && globalHnswlib.VectorFloat) {
console.log('VectorDatabase: Using VectorFloat constructor');
vectorToAdd = new globalHnswlib.VectorFloat();
// Add elements to VectorFloat one by one
for (let i = 0; i < cleanEmbedding.length; i++) {
vectorToAdd.push_back(cleanEmbedding[i]);
}
} else {
// Method 2: Use plain JS array (fallback)
console.log('VectorDatabase: Using plain JS array as fallback');
vectorToAdd = Array.from(cleanEmbedding);
}
// Call addPoint with constructed vector
this.index.addPoint(vectorToAdd, label, false);
// Clean up VectorFloat object (if manually created)
if (vectorToAdd && typeof vectorToAdd.delete === 'function') {
vectorToAdd.delete();
}
} catch (vectorError) {
console.error(
'VectorDatabase: VectorFloat approach failed, trying alternatives:',
vectorError,
);
// Method 3: Try passing Float32Array directly
try {
console.log('VectorDatabase: Trying Float32Array directly');
this.index.addPoint(cleanEmbedding, label, false);
} catch (float32Error) {
console.error('VectorDatabase: Float32Array approach failed:', float32Error);
// Method 4: Last resort - use spread operator
console.log('VectorDatabase: Trying spread operator as last resort');
this.index.addPoint([...cleanEmbedding], label, false);
}
}
console.log(`VectorDatabase: ✅ Successfully added document with label ${label}`);
// Store document mapping
this.documents.set(label, document);
// Update tab document mapping
if (!this.tabDocuments.has(tabId)) {
this.tabDocuments.set(tabId, new Set());
}
this.tabDocuments.get(tabId)!.add(label);
// Save index and mappings
await this.saveIndex();
await this.saveDocumentMappings();
// Check if auto cleanup is needed
if (this.config.enableAutoCleanup) {
await this.checkAndPerformAutoCleanup();
}
console.log(`VectorDatabase: Successfully added document ${documentId} with label ${label}`);
return label;
} catch (error) {
console.error('VectorDatabase: Failed to add document:', error);
console.error('VectorDatabase: Embedding info:', {
type: typeof embedding,
constructor: embedding?.constructor?.name,
length: embedding?.length,
isFloat32Array: embedding instanceof Float32Array,
firstFewValues: embedding ? Array.from(embedding.slice(0, 5)) : null,
});
throw error;
}
}
/**
* Search similar documents
*/
public async search(queryEmbedding: Float32Array, topK: number = 10): Promise<SearchResult[]> {
if (!this.isInitialized) {
await this.initialize();
}
try {
// Validate query vector
if (!queryEmbedding || queryEmbedding.length !== this.config.dimension) {
throw new Error(
`Invalid query embedding dimension: expected ${this.config.dimension}, got ${queryEmbedding?.length || 0}`,
);
}
// Check if query vector contains invalid values
for (let i = 0; i < queryEmbedding.length; i++) {
if (!isFinite(queryEmbedding[i])) {
throw new Error(`Invalid query embedding value at index ${i}: ${queryEmbedding[i]}`);
}
}
console.log(
`VectorDatabase: Searching with query embedding dimension: ${queryEmbedding.length}, topK: ${topK}`,
);
// Check if index is empty
const currentCount = this.index.getCurrentCount();
if (currentCount === 0) {
console.log('VectorDatabase: Index is empty, returning no results');
return [];
}
console.log(`VectorDatabase: Index contains ${currentCount} vectors`);
// Check if document mapping and index are synchronized
const mappingCount = this.documents.size;
if (mappingCount === 0 && currentCount > 0) {
console.warn(
`VectorDatabase: Index has ${currentCount} vectors but document mapping is empty. Attempting to reload mappings...`,
);
await this.loadDocumentMappings();
if (this.documents.size === 0) {
console.error(
'VectorDatabase: Failed to load document mappings. Index and mappings are out of sync.',
);
return [];
}
console.log(
`VectorDatabase: Successfully reloaded ${this.documents.size} document mappings`,
);
}
// Process query vector according to hnswlib-wasm-static emscripten binding requirements
let queryVector;
let searchResult;
try {
// Method 1: Try using VectorFloat constructor (if available)
if (globalHnswlib && globalHnswlib.VectorFloat) {
console.log('VectorDatabase: Using VectorFloat for search query');
queryVector = new globalHnswlib.VectorFloat();
// Add elements to VectorFloat one by one
for (let i = 0; i < queryEmbedding.length; i++) {
queryVector.push_back(queryEmbedding[i]);
}
searchResult = this.index.searchKnn(queryVector, topK, undefined);
// Clean up VectorFloat object
if (queryVector && typeof queryVector.delete === 'function') {
queryVector.delete();
}
} else {
// Method 2: Use plain JS array (fallback)
console.log('VectorDatabase: Using plain JS array for search query');
const queryArray = Array.from(queryEmbedding);
searchResult = this.index.searchKnn(queryArray, topK, undefined);
}
} catch (vectorError) {
console.error(
'VectorDatabase: VectorFloat search failed, trying alternatives:',
vectorError,
);
// Method 3: Try passing Float32Array directly
try {
console.log('VectorDatabase: Trying Float32Array directly for search');
searchResult = this.index.searchKnn(queryEmbedding, topK, undefined);
} catch (float32Error) {
console.error('VectorDatabase: Float32Array search failed:', float32Error);
// Method 4: Last resort - use spread operator
console.log('VectorDatabase: Trying spread operator for search as last resort');
searchResult = this.index.searchKnn([...queryEmbedding], topK, undefined);
}
}
const results: SearchResult[] = [];
console.log(`VectorDatabase: Processing ${searchResult.neighbors.length} search neighbors`);
console.log(`VectorDatabase: Available documents in mapping: ${this.documents.size}`);
console.log(`VectorDatabase: Index current count: ${this.index.getCurrentCount()}`);
for (let i = 0; i < searchResult.neighbors.length; i++) {
const label = searchResult.neighbors[i];
const distance = searchResult.distances[i];
const similarity = 1 - distance; // Convert cosine distance to similarity
console.log(
`VectorDatabase: Processing neighbor ${i}: label=${label}, distance=${distance}, similarity=${similarity}`,
);
// Find corresponding document by label
const document = this.findDocumentByLabel(label);
if (document) {
console.log(`VectorDatabase: Found document for label ${label}: ${document.id}`);
results.push({
document,
similarity,
distance,
});
} else {
console.warn(`VectorDatabase: No document found for label ${label}`);
// Detailed debug information
if (i < 5) {
// Only show detailed info for first 5 neighbors to avoid log spam
console.warn(
`VectorDatabase: Available labels (first 20): ${Array.from(this.documents.keys()).slice(0, 20).join(', ')}`,
);
console.warn(`VectorDatabase: Total available labels: ${this.documents.size}`);
console.warn(
`VectorDatabase: Label type: ${typeof label}, Available label types: ${Array.from(
this.documents.keys(),
)
.slice(0, 3)
.map((k) => typeof k)
.join(', ')}`,
);
}
}
}
console.log(
`VectorDatabase: Found ${results.length} search results out of ${searchResult.neighbors.length} neighbors`,
);
// If no results found but index has data, indicates label mismatch
if (results.length === 0 && searchResult.neighbors.length > 0) {
console.error(
'VectorDatabase: Label mismatch detected! Index has vectors but no matching documents found.',
);
console.error(
'VectorDatabase: This usually indicates the index and document mappings are out of sync.',
);
console.error('VectorDatabase: Consider rebuilding the index to fix this issue.');
// Provide some diagnostic information
const sampleLabels = searchResult.neighbors.slice(0, 5);
const availableLabels = Array.from(this.documents.keys()).slice(0, 5);
console.error('VectorDatabase: Sample search labels:', sampleLabels);
console.error('VectorDatabase: Sample available labels:', availableLabels);
}
return results.sort((a, b) => b.similarity - a.similarity);
} catch (error) {
console.error('VectorDatabase: Search failed:', error);
console.error('VectorDatabase: Query embedding info:', {
type: typeof queryEmbedding,
constructor: queryEmbedding?.constructor?.name,
length: queryEmbedding?.length,
isFloat32Array: queryEmbedding instanceof Float32Array,
firstFewValues: queryEmbedding ? Array.from(queryEmbedding.slice(0, 5)) : null,
});
throw error;
}
}
/**
* Remove all documents for a tab
*/
public async removeTabDocuments(tabId: number): Promise<void> {
if (!this.isInitialized) {
await this.initialize();
}
const documentLabels = this.tabDocuments.get(tabId);
if (!documentLabels) {
return;
}
try {
// Remove documents from mapping (hnswlib-wasm doesn't support direct deletion, only mark as deleted)
for (const label of documentLabels) {
this.documents.delete(label);
}
// Clean up tab mapping
this.tabDocuments.delete(tabId);
// Save changes
await this.saveDocumentMappings();
console.log(`VectorDatabase: Removed ${documentLabels.size} documents for tab ${tabId}`);
} catch (error) {
console.error('VectorDatabase: Failed to remove tab documents:', error);
throw error;
}
}
/**
* Get database statistics
*/
public getStats(): {
totalDocuments: number;
totalTabs: number;
indexSize: number;
isInitialized: boolean;
} {
return {
totalDocuments: this.documents.size,
totalTabs: this.tabDocuments.size,
indexSize: this.calculateStorageSize(),
isInitialized: this.isInitialized,
};
}
/**
* Calculate actual storage size (bytes)
*/
private calculateStorageSize(): number {
let totalSize = 0;
try {
// 1. 计算文档映射的大小
const documentsSize = this.calculateDocumentMappingsSize();
totalSize += documentsSize;
// 2. 计算向量数据的大小
const vectorsSize = this.calculateVectorsSize();
totalSize += vectorsSize;
// 3. 估算索引结构的大小
const indexStructureSize = this.calculateIndexStructureSize();
totalSize += indexStructureSize;
console.log(
`VectorDatabase: Storage size breakdown - Documents: ${documentsSize}, Vectors: ${vectorsSize}, Index: ${indexStructureSize}, Total: ${totalSize} bytes`,
);
} catch (error) {
console.warn('VectorDatabase: Failed to calculate storage size:', error);
// 返回一个基于文档数量的估算值
totalSize = this.documents.size * 1024; // 每个文档估算1KB
}
return totalSize;
}
/**
* Calculate document mappings size
*/
private calculateDocumentMappingsSize(): number {
let size = 0;
// Calculate documents Map size
for (const [label, document] of this.documents.entries()) {
// label (number): 8 bytes
size += 8;
// document object
size += this.calculateObjectSize(document);
}
// Calculate tabDocuments Map size
for (const [tabId, labels] of this.tabDocuments.entries()) {
// tabId (number): 8 bytes
size += 8;
// Set of labels: 8 bytes per label + Set overhead
size += labels.size * 8 + 32; // 32 bytes Set overhead
}
return size;
}
/**
* Calculate vectors data size
*/
private calculateVectorsSize(): number {
const documentCount = this.documents.size;
const dimension = this.config.dimension;
// Each vector: dimension * 4 bytes (Float32)
const vectorSize = dimension * 4;
return documentCount * vectorSize;
}
/**
* Estimate index structure size
*/
private calculateIndexStructureSize(): number {
const documentCount = this.documents.size;
if (documentCount === 0) return 0;
// HNSW index size estimation
// Based on papers and actual testing, HNSW index size is about 20-40% of vector data
const vectorsSize = this.calculateVectorsSize();
const indexOverhead = Math.floor(vectorsSize * 0.3); // 30% overhead
// Additional graph structure overhead
const graphOverhead = documentCount * 64; // About 64 bytes graph structure overhead per node
return indexOverhead + graphOverhead;
}
/**
* Calculate object size (rough estimation)
*/
private calculateObjectSize(obj: any): number {
let size = 0;
try {
const jsonString = JSON.stringify(obj);
// UTF-8 encoding, most characters 1 byte, Chinese etc 3 bytes, average 2 bytes
size = jsonString.length * 2;
} catch (error) {
// If JSON serialization fails, use default estimation
size = 512; // Default 512 bytes
}
return size;
}
/**
* Clear entire database
*/
public async clear(): Promise<void> {
console.log('VectorDatabase: Starting complete database clear...');
try {
// Clear in-memory data structures
this.documents.clear();
this.tabDocuments.clear();
this.nextLabel = 0;
// Clear HNSW index file (in hnswlib-index database)
if (this.isInitialized && this.index) {
try {
console.log('VectorDatabase: Clearing HNSW index file from IndexedDB...');
// 1. First try to physically delete index file (using EmscriptenFileSystemManager)
try {
if (
globalHnswlib &&
globalHnswlib.EmscriptenFileSystemManager.checkFileExists(this.config.indexFileName)
) {
console.log(
`VectorDatabase: Deleting physical index file: ${this.config.indexFileName}`,
);
globalHnswlib.EmscriptenFileSystemManager.deleteFile(this.config.indexFileName);
await this.syncFileSystem('write'); // Ensure deletion is synced to persistent storage
console.log(
`VectorDatabase: Physical index file ${this.config.indexFileName} deleted successfully`,
);
} else {
console.log(
`VectorDatabase: Physical index file ${this.config.indexFileName} does not exist or already deleted`,
);
}
} catch (fileError) {
console.warn(
`VectorDatabase: Failed to delete physical index file ${this.config.indexFileName}:`,
fileError,
);
// Continue with other cleanup operations, don't block the process
}
// 2. Delete index file from IndexedDB
await this.index.deleteIndex(this.config.indexFileName);
console.log('VectorDatabase: HNSW index file cleared from IndexedDB');
// 3. Reinitialize empty index
console.log('VectorDatabase: Reinitializing empty HNSW index...');
this.index.initIndex(
this.config.maxElements,
this.config.M,
this.config.efConstruction,
200,
);
this.index.setEfSearch(this.config.efSearch);
// 4. Force save empty index
await this.forceSaveIndex();
} catch (indexError) {
console.warn('VectorDatabase: Failed to clear HNSW index file:', indexError);
// Continue with other cleanup operations
}
}
// Clear document mappings from IndexedDB (in VectorDatabaseStorage database)
try {
console.log('VectorDatabase: Clearing document mappings from IndexedDB...');
await IndexedDBHelper.deleteData(this.config.indexFileName);
console.log('VectorDatabase: Document mappings cleared from IndexedDB');
} catch (idbError) {
console.warn(
'VectorDatabase: Failed to clear document mappings from IndexedDB, trying chrome.storage fallback:',
idbError,
);
// Clear backup data from chrome.storage
try {
const storageKey = `hnswlib_document_mappings_${this.config.indexFileName}`;
await chrome.storage.local.remove([storageKey]);
console.log('VectorDatabase: Chrome storage fallback cleared');
} catch (storageError) {
console.warn('VectorDatabase: Failed to clear chrome.storage fallback:', storageError);
}
}
// Save empty document mappings to ensure consistency
await this.saveDocumentMappings();
console.log('VectorDatabase: Complete database clear finished successfully');
} catch (error) {
console.error('VectorDatabase: Failed to clear database:', error);
throw error;
}
}
/**
* Force save index and sync filesystem
*/
private async forceSaveIndex(): Promise<void> {
try {
await this.index.writeIndex(this.config.indexFileName);
await this.syncFileSystem('write'); // Force sync
} catch (error) {
console.error('VectorDatabase: Failed to force save index:', error);
}
}
/**
* Check and perform auto cleanup
*/
private async checkAndPerformAutoCleanup(): Promise<void> {
try {
const currentCount = this.documents.size;
const maxElements = this.config.maxElements;
console.log(
`VectorDatabase: Auto cleanup check - current: ${currentCount}, max: ${maxElements}`,
);
// Check if maximum element count is exceeded
if (currentCount >= maxElements) {
console.log('VectorDatabase: Document count reached limit, performing cleanup...');
await this.performLRUCleanup(Math.floor(maxElements * 0.2)); // Clean up 20% of data
}
// Check if there's expired data
if (this.config.maxRetentionDays && this.config.maxRetentionDays > 0) {
await this.performTimeBasedCleanup();
}
} catch (error) {
console.error('VectorDatabase: Auto cleanup failed:', error);
}
}
/**
* Perform LRU-based cleanup (delete oldest documents)
*/
private async performLRUCleanup(cleanupCount: number): Promise<void> {
try {
console.log(
`VectorDatabase: Starting LRU cleanup, removing ${cleanupCount} oldest documents`,
);
// Get all documents and sort by timestamp
const allDocuments = Array.from(this.documents.entries());
allDocuments.sort((a, b) => a[1].timestamp - b[1].timestamp);
// Select documents to delete
const documentsToDelete = allDocuments.slice(0, cleanupCount);
for (const [label, _document] of documentsToDelete) {
await this.removeDocumentByLabel(label);
}
// Save updated index and mappings
await this.saveIndex();
await this.saveDocumentMappings();
console.log(
`VectorDatabase: LRU cleanup completed, removed ${documentsToDelete.length} documents`,
);
} catch (error) {
console.error('VectorDatabase: LRU cleanup failed:', error);
}
}
/**
* Perform time-based cleanup (delete expired documents)
*/
private async performTimeBasedCleanup(): Promise<void> {
try {
const maxRetentionMs = this.config.maxRetentionDays! * 24 * 60 * 60 * 1000;
const cutoffTime = Date.now() - maxRetentionMs;
console.log(
`VectorDatabase: Starting time-based cleanup, removing documents older than ${this.config.maxRetentionDays} days`,
);
const documentsToDelete: number[] = [];
for (const [label, document] of this.documents.entries()) {
if (document.timestamp < cutoffTime) {
documentsToDelete.push(label);
}
}
for (const label of documentsToDelete) {
await this.removeDocumentByLabel(label);
}
// Save updated index and mappings
if (documentsToDelete.length > 0) {
await this.saveIndex();
await this.saveDocumentMappings();
}
console.log(
`VectorDatabase: Time-based cleanup completed, removed ${documentsToDelete.length} expired documents`,
);
} catch (error) {
console.error('VectorDatabase: Time-based cleanup failed:', error);
}
}
/**
* Remove single document by label
*/
private async removeDocumentByLabel(label: number): Promise<void> {
try {
const document = this.documents.get(label);
if (!document) {
console.warn(`VectorDatabase: Document with label ${label} not found`);
return;
}
// Remove vector from HNSW index
if (this.index) {
try {
this.index.markDelete(label);
} catch (indexError) {
console.warn(
`VectorDatabase: Failed to mark delete in index for label ${label}:`,
indexError,
);
}
}
// Remove from memory mapping
this.documents.delete(label);
// Remove from tab mapping
const tabId = document.tabId;
if (this.tabDocuments.has(tabId)) {
this.tabDocuments.get(tabId)!.delete(label);
// If tab has no other documents, delete entire tab mapping
if (this.tabDocuments.get(tabId)!.size === 0) {
this.tabDocuments.delete(tabId);
}
}
console.log(`VectorDatabase: Removed document with label ${label} from tab ${tabId}`);
} catch (error) {
console.error(`VectorDatabase: Failed to remove document with label ${label}:`, error);
}
}
// 私有辅助方法
private generateDocumentId(tabId: number, chunkIndex: number): string {
return `tab_${tabId}_chunk_${chunkIndex}_${Date.now()}`;
}
private findDocumentByLabel(label: number): VectorDocument | null {
return this.documents.get(label) || null;
}
private async syncFileSystem(direction: 'read' | 'write'): Promise<void> {
try {
if (!globalHnswlib) {
return;
}
// If sync operation is already in progress, wait for it to complete
if (syncInProgress && pendingSyncPromise) {
console.log(`VectorDatabase: Sync already in progress, waiting...`);
await pendingSyncPromise;
return;
}
// Mark sync start
syncInProgress = true;
// Create sync Promise with timeout mechanism
pendingSyncPromise = new Promise<void>((resolve, reject) => {
const timeout = setTimeout(() => {
console.warn(`VectorDatabase: Filesystem sync (${direction}) timeout`);
syncInProgress = false;
pendingSyncPromise = null;
reject(new Error('Sync timeout'));
}, 5000); // 5 second timeout
try {
globalHnswlib.EmscriptenFileSystemManager.syncFS(direction === 'read', () => {
clearTimeout(timeout);
console.log(`VectorDatabase: Filesystem sync (${direction}) completed`);
syncInProgress = false;
pendingSyncPromise = null;
resolve();
});
} catch (error) {
clearTimeout(timeout);
console.warn(`VectorDatabase: Failed to sync filesystem (${direction}):`, error);
syncInProgress = false;
pendingSyncPromise = null;
reject(error);
}
});
await pendingSyncPromise;
} catch (error) {
console.warn(`VectorDatabase: Failed to sync filesystem (${direction}):`, error);
syncInProgress = false;
pendingSyncPromise = null;
}
}
private async saveIndex(): Promise<void> {
try {
await this.index.writeIndex(this.config.indexFileName);
// Reduce sync frequency, only sync when necessary
if (this.documents.size % 10 === 0) {
// Sync every 10 documents
await this.syncFileSystem('write');
}
} catch (error) {
console.error('VectorDatabase: Failed to save index:', error);
}
}
private async saveDocumentMappings(): Promise<void> {
try {
// Save document mappings to IndexedDB
const mappingData = {
documents: Array.from(this.documents.entries()),
tabDocuments: Array.from(this.tabDocuments.entries()).map(([tabId, labels]) => [
tabId,
Array.from(labels),
]),
nextLabel: this.nextLabel,
};
try {
// Use IndexedDB to save data, supports larger storage capacity
await IndexedDBHelper.saveData(this.config.indexFileName, mappingData);
console.log('VectorDatabase: Document mappings saved to IndexedDB');
} catch (idbError) {
console.warn(
'VectorDatabase: Failed to save to IndexedDB, falling back to chrome.storage:',
idbError,
);
// Fall back to chrome.storage.local
try {
const storageKey = `hnswlib_document_mappings_${this.config.indexFileName}`;
await chrome.storage.local.set({ [storageKey]: mappingData });
console.log('VectorDatabase: Document mappings saved to chrome.storage.local (fallback)');
} catch (storageError) {
console.error(
'VectorDatabase: Failed to save to both IndexedDB and chrome.storage:',
storageError,
);
}
}
} catch (error) {
console.error('VectorDatabase: Failed to save document mappings:', error);
}
}
public async loadDocumentMappings(): Promise<void> {
try {
// Load document mappings from IndexedDB
if (!globalHnswlib) {
return;
}
let mappingData = null;
try {
// First try to read from IndexedDB
mappingData = await IndexedDBHelper.loadData(this.config.indexFileName);
if (mappingData) {
console.log(`VectorDatabase: Loaded document mappings from IndexedDB`);
}
} catch (idbError) {
console.warn(
'VectorDatabase: Failed to read from IndexedDB, trying chrome.storage:',
idbError,
);
}
// If IndexedDB has no data, try reading from chrome.storage.local (backward compatibility)
if (!mappingData) {
try {
const storageKey = `hnswlib_document_mappings_${this.config.indexFileName}`;
const result = await chrome.storage.local.get([storageKey]);
mappingData = result[storageKey];
if (mappingData) {
console.log(
`VectorDatabase: Loaded document mappings from chrome.storage.local (fallback)`,
);
// Migrate to IndexedDB
try {
await IndexedDBHelper.saveData(this.config.indexFileName, mappingData);
console.log('VectorDatabase: Migrated data from chrome.storage to IndexedDB');
} catch (migrationError) {
console.warn('VectorDatabase: Failed to migrate data to IndexedDB:', migrationError);
}
}
} catch (storageError) {
console.warn('VectorDatabase: Failed to read from chrome.storage.local:', storageError);
}
}
if (mappingData) {
// Restore document mappings
this.documents.clear();
for (const [label, doc] of mappingData.documents) {
this.documents.set(label, doc);
}
// Restore tab mappings
this.tabDocuments.clear();
for (const [tabId, labels] of mappingData.tabDocuments) {
this.tabDocuments.set(tabId, new Set(labels));
}
// Restore nextLabel - use saved value or calculate max label + 1
if (mappingData.nextLabel !== undefined) {
this.nextLabel = mappingData.nextLabel;
} else if (this.documents.size > 0) {
// If no saved nextLabel, calculate max label + 1
const maxLabel = Math.max(...Array.from(this.documents.keys()));
this.nextLabel = maxLabel + 1;
} else {
this.nextLabel = 0;
}
console.log(
`VectorDatabase: Loaded ${this.documents.size} document mappings, next label: ${this.nextLabel}`,
);
} else {
console.log('VectorDatabase: No existing document mappings found');
}
} catch (error) {
console.error('VectorDatabase: Failed to load document mappings:', error);
}
}
}
// Global VectorDatabase singleton
let globalVectorDatabase: VectorDatabase | null = null;
let currentDimension: number | null = null;
/**
* Get global VectorDatabase singleton instance
* If dimension changes, will recreate instance to ensure compatibility
*/
export async function getGlobalVectorDatabase(
config?: Partial<VectorDatabaseConfig>,
): Promise<VectorDatabase> {
const newDimension = config?.dimension || 384;
// If dimension changes, need to recreate vector database
if (globalVectorDatabase && currentDimension !== null && currentDimension !== newDimension) {
console.log(
`VectorDatabase: Dimension changed from ${currentDimension} to ${newDimension}, recreating instance`,
);
// Clean up old instance - this will clean up index files and document mappings
try {
await globalVectorDatabase.clear();
console.log('VectorDatabase: Successfully cleared old instance for dimension change');
} catch (error) {
console.warn('VectorDatabase: Error during cleanup:', error);
}
globalVectorDatabase = null;
currentDimension = null;
}
if (!globalVectorDatabase) {
globalVectorDatabase = new VectorDatabase(config);
currentDimension = newDimension;
console.log(
`VectorDatabase: Created global singleton instance with dimension ${currentDimension}`,
);
}
return globalVectorDatabase;
}
/**
* Synchronous version of getting global VectorDatabase instance (for backward compatibility)
* Note: If dimension change is needed, recommend using async version
*/
export function getGlobalVectorDatabaseSync(
config?: Partial<VectorDatabaseConfig>,
): VectorDatabase {
const newDimension = config?.dimension || 384;
// If dimension changes, log warning but don't clean up (avoid race conditions)
if (globalVectorDatabase && currentDimension !== null && currentDimension !== newDimension) {
console.warn(
`VectorDatabase: Dimension mismatch detected (${currentDimension} vs ${newDimension}). Consider using async version for proper cleanup.`,
);
}
if (!globalVectorDatabase) {
globalVectorDatabase = new VectorDatabase(config);
currentDimension = newDimension;
console.log(
`VectorDatabase: Created global singleton instance with dimension ${currentDimension}`,
);
}
return globalVectorDatabase;
}
/**
* Reset global VectorDatabase instance (mainly for testing or model switching)
*/
export async function resetGlobalVectorDatabase(): Promise<void> {
console.log('VectorDatabase: Starting global instance reset...');
if (globalVectorDatabase) {
try {
console.log('VectorDatabase: Clearing existing global instance...');
await globalVectorDatabase.clear();
console.log('VectorDatabase: Global instance cleared successfully');
} catch (error) {
console.warn('VectorDatabase: Failed to clear during reset:', error);
}
}
// Additional cleanup: ensure all possible IndexedDB data is cleared
try {
console.log('VectorDatabase: Performing comprehensive IndexedDB cleanup...');
// Clear all data in VectorDatabaseStorage database
await IndexedDBHelper.clearAllData();
// Clear index files from hnswlib-index database
try {
console.log('VectorDatabase: Clearing HNSW index files from IndexedDB...');
// Try to clean up possible existing index files
const possibleIndexFiles = ['tab_content_index.dat', 'content_index.dat', 'vector_index.dat'];
// If global hnswlib instance exists, try to delete known index files
if (typeof globalHnswlib !== 'undefined' && globalHnswlib) {
for (const fileName of possibleIndexFiles) {
try {
// 1. First try to physically delete index file (using EmscriptenFileSystemManager)
try {
if (globalHnswlib.EmscriptenFileSystemManager.checkFileExists(fileName)) {
console.log(`VectorDatabase: Deleting physical index file: ${fileName}`);
globalHnswlib.EmscriptenFileSystemManager.deleteFile(fileName);
console.log(`VectorDatabase: Physical index file ${fileName} deleted successfully`);
}
} catch (fileError) {
console.log(
`VectorDatabase: Physical index file ${fileName} not found or failed to delete:`,
fileError,
);
}
// 2. Delete index file from IndexedDB
const tempIndex = new globalHnswlib.HierarchicalNSW('cosine', 384);
await tempIndex.deleteIndex(fileName);
console.log(`VectorDatabase: Deleted IndexedDB index file: ${fileName}`);
} catch (deleteError) {
// File might not exist, this is normal
console.log(`VectorDatabase: Index file ${fileName} not found or already deleted`);
}
}
// 3. Force sync filesystem to ensure deletion takes effect
try {
await new Promise<void>((resolve) => {
const timeout = setTimeout(() => {
console.warn('VectorDatabase: Filesystem sync timeout during cleanup');
resolve(); // Don't block the process
}, 3000);
globalHnswlib.EmscriptenFileSystemManager.syncFS(false, () => {
clearTimeout(timeout);
console.log('VectorDatabase: Filesystem sync completed during cleanup');
resolve();
});
});
} catch (syncError) {
console.warn('VectorDatabase: Failed to sync filesystem during cleanup:', syncError);
}
}
} catch (hnswError) {
console.warn('VectorDatabase: Failed to clear HNSW index files:', hnswError);
}
// Clear possible chrome.storage backup data (only clear vector database related data, preserve user preferences)
const possibleKeys = [
'hnswlib_document_mappings_tab_content_index.dat',
'hnswlib_document_mappings_content_index.dat',
'hnswlib_document_mappings_vector_index.dat',
// Note: Don't clear selectedModel and selectedVersion, these are user preference settings
// Note: Don't clear modelState, this contains model state info and should be handled by model management logic
];
if (possibleKeys.length > 0) {
try {
await chrome.storage.local.remove(possibleKeys);
console.log('VectorDatabase: Chrome storage backup data cleared');
} catch (storageError) {
console.warn('VectorDatabase: Failed to clear chrome.storage backup:', storageError);
}
}
console.log('VectorDatabase: Comprehensive cleanup completed');
} catch (cleanupError) {
console.warn('VectorDatabase: Comprehensive cleanup failed:', cleanupError);
}
globalVectorDatabase = null;
currentDimension = null;
console.log('VectorDatabase: Global singleton instance reset completed');
}
/**
* Specifically for data cleanup during model switching
* Clear all IndexedDB data, including HNSW index files and document mappings
*/
export async function clearAllVectorData(): Promise<void> {
console.log('VectorDatabase: Starting comprehensive vector data cleanup for model switch...');
try {
// 1. Clear global instance
if (globalVectorDatabase) {
try {
await globalVectorDatabase.clear();
} catch (error) {
console.warn('VectorDatabase: Failed to clear global instance:', error);
}
}
// 2. Clear VectorDatabaseStorage database
try {
console.log('VectorDatabase: Clearing VectorDatabaseStorage database...');
await IndexedDBHelper.clearAllData();
} catch (error) {
console.warn('VectorDatabase: Failed to clear VectorDatabaseStorage:', error);
}
// 3. Clear hnswlib-index database and physical files
try {
console.log('VectorDatabase: Clearing hnswlib-index database and physical files...');
// 3.1 First try to physically delete index files (using EmscriptenFileSystemManager)
if (typeof globalHnswlib !== 'undefined' && globalHnswlib) {
const possibleIndexFiles = [
'tab_content_index.dat',
'content_index.dat',
'vector_index.dat',
];
for (const fileName of possibleIndexFiles) {
try {
if (globalHnswlib.EmscriptenFileSystemManager.checkFileExists(fileName)) {
console.log(`VectorDatabase: Deleting physical index file: ${fileName}`);
globalHnswlib.EmscriptenFileSystemManager.deleteFile(fileName);
console.log(`VectorDatabase: Physical index file ${fileName} deleted successfully`);
}
} catch (fileError) {
console.log(
`VectorDatabase: Physical index file ${fileName} not found or failed to delete:`,
fileError,
);
}
}
// Force sync filesystem
try {
await new Promise<void>((resolve) => {
const timeout = setTimeout(() => {
console.warn('VectorDatabase: Filesystem sync timeout during model switch cleanup');
resolve();
}, 3000);
globalHnswlib.EmscriptenFileSystemManager.syncFS(false, () => {
clearTimeout(timeout);
console.log('VectorDatabase: Filesystem sync completed during model switch cleanup');
resolve();
});
});
} catch (syncError) {
console.warn(
'VectorDatabase: Failed to sync filesystem during model switch cleanup:',
syncError,
);
}
}
// 3.2 Delete entire hnswlib-index database
await new Promise<void>((resolve) => {
const deleteRequest = indexedDB.deleteDatabase('/hnswlib-index');
deleteRequest.onsuccess = () => {
console.log('VectorDatabase: Successfully deleted /hnswlib-index database');
resolve();
};
deleteRequest.onerror = () => {
console.warn(
'VectorDatabase: Failed to delete /hnswlib-index database:',
deleteRequest.error,
);
resolve(); // Don't block the process
};
deleteRequest.onblocked = () => {
console.warn('VectorDatabase: Deletion of /hnswlib-index database was blocked');
resolve(); // Don't block the process
};
});
} catch (error) {
console.warn(
'VectorDatabase: Failed to clear hnswlib-index database and physical files:',
error,
);
}
// 4. Clear backup data from chrome.storage
try {
const storageKeys = [
'hnswlib_document_mappings_tab_content_index.dat',
'hnswlib_document_mappings_content_index.dat',
'hnswlib_document_mappings_vector_index.dat',
];
await chrome.storage.local.remove(storageKeys);
console.log('VectorDatabase: Chrome storage backup data cleared');
} catch (error) {
console.warn('VectorDatabase: Failed to clear chrome.storage backup:', error);
}
// 5. Reset global state
globalVectorDatabase = null;
currentDimension = null;
console.log('VectorDatabase: Comprehensive vector data cleanup completed successfully');
} catch (error) {
console.error('VectorDatabase: Comprehensive vector data cleanup failed:', error);
throw error;
}
}
```
--------------------------------------------------------------------------------
/app/chrome-extension/utils/semantic-similarity-engine.ts:
--------------------------------------------------------------------------------
```typescript
import { AutoTokenizer, env as TransformersEnv } from '@xenova/transformers';
import type { Tensor as TransformersTensor, PreTrainedTokenizer } from '@xenova/transformers';
import LRUCache from './lru-cache';
import { SIMDMathEngine } from './simd-math-engine';
import { OffscreenManager } from './offscreen-manager';
import { STORAGE_KEYS } from '@/common/constants';
import { OFFSCREEN_MESSAGE_TYPES } from '@/common/message-types';
import { ModelCacheManager } from './model-cache-manager';
/**
* Get cached model data, prioritizing cache reads and handling redirected URLs.
* @param {string} modelUrl Stable, permanent URL of the model
* @returns {Promise<ArrayBuffer>} Model data as ArrayBuffer
*/
async function getCachedModelData(modelUrl: string): Promise<ArrayBuffer> {
const cacheManager = ModelCacheManager.getInstance();
// 1. 尝试从缓存获取数据
const cachedData = await cacheManager.getCachedModelData(modelUrl);
if (cachedData) {
return cachedData;
}
console.log('Model not found in cache or expired. Fetching from network...');
try {
// 2. 从网络获取数据
const response = await fetch(modelUrl);
if (!response.ok) {
throw new Error(`Failed to fetch model: ${response.status} ${response.statusText}`);
}
// 3. 获取数据并存储到缓存
const arrayBuffer = await response.arrayBuffer();
await cacheManager.storeModelData(modelUrl, arrayBuffer);
console.log(
`Model fetched from network and successfully cached (${(arrayBuffer.byteLength / 1024 / 1024).toFixed(2)}MB).`,
);
return arrayBuffer;
} catch (error) {
console.error(`Error fetching or caching model:`, error);
// 如果获取失败，清理可能不完整的缓存条目
await cacheManager.deleteCacheEntry(modelUrl);
throw error;
}
}
/**
* Clear all model cache entries
*/
export async function clearModelCache(): Promise<void> {
try {
const cacheManager = ModelCacheManager.getInstance();
await cacheManager.clearAllCache();
} catch (error) {
console.error('Failed to clear model cache:', error);
throw error;
}
}
/**
* Get cache statistics
*/
export async function getCacheStats(): Promise<{
totalSize: number;
totalSizeMB: number;
entryCount: number;
entries: Array<{
url: string;
size: number;
sizeMB: number;
timestamp: number;
age: string;
expired: boolean;
}>;
}> {
try {
const cacheManager = ModelCacheManager.getInstance();
return await cacheManager.getCacheStats();
} catch (error) {
console.error('Failed to get cache stats:', error);
throw error;
}
}
/**
* Manually trigger cache cleanup
*/
export async function cleanupModelCache(): Promise<void> {
try {
const cacheManager = ModelCacheManager.getInstance();
await cacheManager.manualCleanup();
} catch (error) {
console.error('Failed to cleanup cache:', error);
throw error;
}
}
/**
* Check if the default model is cached and available
* @returns Promise<boolean> True if default model is cached and valid
*/
export async function isDefaultModelCached(): Promise<boolean> {
try {
// Get the default model configuration
const result = await chrome.storage.local.get([STORAGE_KEYS.SEMANTIC_MODEL]);
const defaultModel =
(result[STORAGE_KEYS.SEMANTIC_MODEL] as ModelPreset) || 'multilingual-e5-small';
// Build the model URL
const modelInfo = PREDEFINED_MODELS[defaultModel];
const modelIdentifier = modelInfo.modelIdentifier;
const onnxModelFile = 'model.onnx'; // Default ONNX file name
const modelIdParts = modelIdentifier.split('/');
const modelNameForUrl = modelIdParts.length > 1 ? modelIdentifier : `Xenova/${modelIdentifier}`;
const onnxModelUrl = `https://huggingface.co/${modelNameForUrl}/resolve/main/onnx/${onnxModelFile}`;
// Check if this model is cached
const cacheManager = ModelCacheManager.getInstance();
return await cacheManager.isModelCached(onnxModelUrl);
} catch (error) {
console.error('Error checking if default model is cached:', error);
return false;
}
}
/**
* Check if any model cache exists (for conditional initialization)
* @returns Promise<boolean> True if any valid model cache exists
*/
export async function hasAnyModelCache(): Promise<boolean> {
try {
const cacheManager = ModelCacheManager.getInstance();
return await cacheManager.hasAnyValidCache();
} catch (error) {
console.error('Error checking for any model cache:', error);
return false;
}
}
// Predefined model configurations - 2025 curated recommended models, using quantized versions to reduce file size
export const PREDEFINED_MODELS = {
// Multilingual model - default recommendation
'multilingual-e5-small': {
modelIdentifier: 'Xenova/multilingual-e5-small',
dimension: 384,
description: 'Multilingual E5 Small - Lightweight multilingual model supporting 100+ languages',
language: 'multilingual',
performance: 'excellent',
size: '116MB', // Quantized version
latency: '20ms',
multilingualFeatures: {
languageSupport: '100+',
crossLanguageRetrieval: 'good',
chineseEnglishMixed: 'good',
},
modelSpecificConfig: {
requiresTokenTypeIds: false, // E5 model doesn't require token_type_ids
},
},
'multilingual-e5-base': {
modelIdentifier: 'Xenova/multilingual-e5-base',
dimension: 768,
description: 'Multilingual E5 base - Medium-scale multilingual model supporting 100+ languages',
language: 'multilingual',
performance: 'excellent',
size: '279MB', // Quantized version
latency: '30ms',
multilingualFeatures: {
languageSupport: '100+',
crossLanguageRetrieval: 'excellent',
chineseEnglishMixed: 'excellent',
},
modelSpecificConfig: {
requiresTokenTypeIds: false, // E5 model doesn't require token_type_ids
},
},
} as const;
export type ModelPreset = keyof typeof PREDEFINED_MODELS;
/**
* Get model information
*/
export function getModelInfo(preset: ModelPreset) {
return PREDEFINED_MODELS[preset];
}
/**
* List all available models
*/
export function listAvailableModels() {
return Object.entries(PREDEFINED_MODELS).map(([key, value]) => ({
preset: key as ModelPreset,
...value,
}));
}
/**
* Recommend model based on language - only uses multilingual-e5 series models
*/
export function recommendModelForLanguage(
_language: 'en' | 'zh' | 'multilingual' = 'multilingual',
scenario: 'speed' | 'balanced' | 'quality' = 'balanced',
): ModelPreset {
// All languages use multilingual models
if (scenario === 'quality') {
return 'multilingual-e5-base'; // High quality choice
}
return 'multilingual-e5-small'; // Default lightweight choice
}
/**
* Intelligently recommend model based on device performance and usage scenario - only uses multilingual-e5 series models
*/
export function recommendModelForDevice(
_language: 'en' | 'zh' | 'multilingual' = 'multilingual',
deviceMemory: number = 4, // GB
networkSpeed: 'slow' | 'fast' = 'fast',
prioritizeSpeed: boolean = false,
): ModelPreset {
// Low memory devices or slow network, prioritize small models
if (deviceMemory < 4 || networkSpeed === 'slow' || prioritizeSpeed) {
return 'multilingual-e5-small'; // Lightweight choice
}
// High performance devices can use better models
if (deviceMemory >= 8 && !prioritizeSpeed) {
return 'multilingual-e5-base'; // High performance choice
}
// Default balanced choice
return 'multilingual-e5-small';
}
/**
* Get model size information (only supports quantized version)
*/
export function getModelSizeInfo(
preset: ModelPreset,
_version: 'full' | 'quantized' | 'compressed' = 'quantized',
) {
const model = PREDEFINED_MODELS[preset];
return {
size: model.size,
recommended: 'quantized',
description: `${model.description} (Size: ${model.size})`,
};
}
/**
* Compare performance and size of multiple models
*/
export function compareModels(presets: ModelPreset[]) {
return presets.map((preset) => {
const model = PREDEFINED_MODELS[preset];
return {
preset,
name: model.description.split(' - ')[0],
language: model.language,
performance: model.performance,
dimension: model.dimension,
latency: model.latency,
size: model.size,
features: (model as any).multilingualFeatures || {},
maxLength: (model as any).maxLength || 512,
recommendedFor: getRecommendationContext(preset),
};
});
}
/**
* Get recommended use cases for model
*/
function getRecommendationContext(preset: ModelPreset): string[] {
const contexts: string[] = [];
const model = PREDEFINED_MODELS[preset];
// All models are multilingual
contexts.push('Multilingual document processing');
if (model.performance === 'excellent') contexts.push('High accuracy requirements');
if (model.latency.includes('20ms')) contexts.push('Fast response');
// Add scenarios based on model size
const sizeInMB = parseInt(model.size.replace('MB', ''));
if (sizeInMB < 300) {
contexts.push('Mobile devices');
contexts.push('Lightweight deployment');
}
if (preset === 'multilingual-e5-small') {
contexts.push('Lightweight deployment');
} else if (preset === 'multilingual-e5-base') {
contexts.push('High accuracy requirements');
}
return contexts;
}
/**
* Get ONNX model filename (only supports quantized version)
*/
export function getOnnxFileNameForVersion(
_version: 'full' | 'quantized' | 'compressed' = 'quantized',
): string {
// Only return quantized version filename
return 'model_quantized.onnx';
}
/**
* Get model identifier (only supports quantized version)
*/
export function getModelIdentifierWithVersion(
preset: ModelPreset,
_version: 'full' | 'quantized' | 'compressed' = 'quantized',
): string {
const model = PREDEFINED_MODELS[preset];
return model.modelIdentifier;
}
/**
* Get size comparison of all available models
*/
export function getAllModelSizes() {
const models = Object.entries(PREDEFINED_MODELS).map(([preset, config]) => {
return {
preset: preset as ModelPreset,
name: config.description.split(' - ')[0],
language: config.language,
size: config.size,
performance: config.performance,
latency: config.latency,
};
});
// Sort by size
return models.sort((a, b) => {
const sizeA = parseInt(a.size.replace('MB', ''));
const sizeB = parseInt(b.size.replace('MB', ''));
return sizeA - sizeB;
});
}
// Define necessary types
interface ModelConfig {
modelIdentifier: string;
localModelPathPrefix?: string; // Base path for local models (relative to public)
onnxModelFile?: string; // ONNX model filename
maxLength?: number;
cacheSize?: number;
numThreads?: number;
executionProviders?: string[];
useLocalFiles?: boolean;
workerPath?: string; // Worker script path (relative to extension root)
concurrentLimit?: number; // Worker task concurrency limit
forceOffscreen?: boolean; // Force offscreen mode (for testing)
modelPreset?: ModelPreset; // Predefined model selection
dimension?: number; // Vector dimension (auto-obtained from preset model)
modelVersion?: 'full' | 'quantized' | 'compressed'; // Model version selection
requiresTokenTypeIds?: boolean; // Whether model requires token_type_ids input
}
interface WorkerMessagePayload {
modelPath?: string;
modelData?: ArrayBuffer;
numThreads?: number;
executionProviders?: string[];
input_ids?: number[];
attention_mask?: number[];
token_type_ids?: number[];
dims?: {
input_ids: number[];
attention_mask: number[];
token_type_ids?: number[];
};
}
interface WorkerResponsePayload {
data?: Float32Array | number[]; // Tensor data as Float32Array or number array
dims?: number[]; // Tensor dimensions
message?: string; // For error or status messages
}
interface WorkerStats {
inferenceTime?: number;
totalInferences?: number;
averageInferenceTime?: number;
memoryAllocations?: number;
batchSize?: number;
}
// Memory pool manager
class EmbeddingMemoryPool {
private pools: Map<number, Float32Array[]> = new Map();
private maxPoolSize: number = 10;
private stats = { allocated: 0, reused: 0, released: 0 };
getEmbedding(size: number): Float32Array {
const pool = this.pools.get(size);
if (pool && pool.length > 0) {
this.stats.reused++;
return pool.pop()!;
}
this.stats.allocated++;
return new Float32Array(size);
}
releaseEmbedding(embedding: Float32Array): void {
const size = embedding.length;
if (!this.pools.has(size)) {
this.pools.set(size, []);
}
const pool = this.pools.get(size)!;
if (pool.length < this.maxPoolSize) {
// Clear array for reuse
embedding.fill(0);
pool.push(embedding);
this.stats.released++;
}
}
getStats() {
return { ...this.stats };
}
clear(): void {
this.pools.clear();
this.stats = { allocated: 0, reused: 0, released: 0 };
}
}
interface PendingMessage {
resolve: (value: WorkerResponsePayload | PromiseLike<WorkerResponsePayload>) => void;
reject: (reason?: any) => void;
type: string;
}
interface TokenizedOutput {
// Simulates part of transformers.js tokenizer output
input_ids: TransformersTensor;
attention_mask: TransformersTensor;
token_type_ids?: TransformersTensor;
}
/**
* SemanticSimilarityEngine proxy class
* Used by ContentIndexer and other components to reuse engine instance in offscreen, avoiding duplicate model downloads
*/
export class SemanticSimilarityEngineProxy {
private _isInitialized = false;
private config: Partial<ModelConfig>;
private offscreenManager: OffscreenManager;
private _isEnsuring = false; // Flag to prevent concurrent ensureOffscreenEngineInitialized calls
constructor(config: Partial<ModelConfig> = {}) {
this.config = config;
this.offscreenManager = OffscreenManager.getInstance();
console.log('SemanticSimilarityEngineProxy: Proxy created with config:', {
modelPreset: config.modelPreset,
modelVersion: config.modelVersion,
dimension: config.dimension,
});
}
async initialize(): Promise<void> {
try {
console.log('SemanticSimilarityEngineProxy: Starting proxy initialization...');
// Ensure offscreen document exists
console.log('SemanticSimilarityEngineProxy: Ensuring offscreen document exists...');
await this.offscreenManager.ensureOffscreenDocument();
console.log('SemanticSimilarityEngineProxy: Offscreen document ready');
// Ensure engine in offscreen is initialized
console.log('SemanticSimilarityEngineProxy: Ensuring offscreen engine is initialized...');
await this.ensureOffscreenEngineInitialized();
this._isInitialized = true;
console.log(
'SemanticSimilarityEngineProxy: Proxy initialized, delegating to offscreen engine',
);
} catch (error) {
console.error('SemanticSimilarityEngineProxy: Initialization failed:', error);
throw new Error(
`Failed to initialize proxy: ${error instanceof Error ? error.message : 'Unknown error'}`,
);
}
}
/**
* Check engine status in offscreen
*/
private async checkOffscreenEngineStatus(): Promise<{
isInitialized: boolean;
currentConfig: any;
}> {
try {
const response = await chrome.runtime.sendMessage({
target: 'offscreen',
type: OFFSCREEN_MESSAGE_TYPES.SIMILARITY_ENGINE_STATUS,
});
if (response && response.success) {
return {
isInitialized: response.isInitialized || false,
currentConfig: response.currentConfig || null,
};
}
} catch (error) {
console.warn('SemanticSimilarityEngineProxy: Failed to check engine status:', error);
}
return { isInitialized: false, currentConfig: null };
}
/**
* Ensure engine in offscreen is initialized (with concurrency protection)
*/
private async ensureOffscreenEngineInitialized(): Promise<void> {
// Prevent concurrent initialization attempts
if (this._isEnsuring) {
console.log('SemanticSimilarityEngineProxy: Already ensuring initialization, waiting...');
// Wait a bit and check again
await new Promise((resolve) => setTimeout(resolve, 100));
return;
}
try {
this._isEnsuring = true;
const status = await this.checkOffscreenEngineStatus();
if (!status.isInitialized) {
console.log(
'SemanticSimilarityEngineProxy: Engine not initialized in offscreen, initializing...',
);
// Reinitialize engine
const response = await chrome.runtime.sendMessage({
target: 'offscreen',
type: OFFSCREEN_MESSAGE_TYPES.SIMILARITY_ENGINE_INIT,
config: this.config,
});
if (!response || !response.success) {
throw new Error(response?.error || 'Failed to initialize engine in offscreen document');
}
console.log('SemanticSimilarityEngineProxy: Engine reinitialized successfully');
}
} finally {
this._isEnsuring = false;
}
}
/**
* Send message to offscreen document with retry mechanism and auto-reinitialization
*/
private async sendMessageToOffscreen(message: any, maxRetries: number = 3): Promise<any> {
// 确保offscreen document存在
await this.offscreenManager.ensureOffscreenDocument();
let lastError: Error | null = null;
for (let attempt = 1; attempt <= maxRetries; attempt++) {
try {
console.log(
`SemanticSimilarityEngineProxy: Sending message (attempt ${attempt}/${maxRetries}):`,
message.type,
);
const response = await chrome.runtime.sendMessage(message);
if (!response) {
throw new Error('No response received from offscreen document');
}
// If engine not initialized error received, try to reinitialize
if (!response.success && response.error && response.error.includes('not initialized')) {
console.log(
'SemanticSimilarityEngineProxy: Engine not initialized, attempting to reinitialize...',
);
await this.ensureOffscreenEngineInitialized();
// Resend original message
const retryResponse = await chrome.runtime.sendMessage(message);
if (retryResponse && retryResponse.success) {
return retryResponse;
}
}
return response;
} catch (error) {
lastError = error as Error;
console.warn(
`SemanticSimilarityEngineProxy: Message failed (attempt ${attempt}/${maxRetries}):`,
error,
);
// If engine not initialized error, try to reinitialize
if (error instanceof Error && error.message.includes('not initialized')) {
try {
console.log(
'SemanticSimilarityEngineProxy: Attempting to reinitialize engine due to error...',
);
await this.ensureOffscreenEngineInitialized();
// Resend original message
const retryResponse = await chrome.runtime.sendMessage(message);
if (retryResponse && retryResponse.success) {
return retryResponse;
}
} catch (reinitError) {
console.warn(
'SemanticSimilarityEngineProxy: Failed to reinitialize engine:',
reinitError,
);
}
}
if (attempt < maxRetries) {
// Wait before retry
await new Promise((resolve) => setTimeout(resolve, 100 * attempt));
// Re-ensure offscreen document exists
try {
await this.offscreenManager.ensureOffscreenDocument();
} catch (offscreenError) {
console.warn(
'SemanticSimilarityEngineProxy: Failed to ensure offscreen document:',
offscreenError,
);
}
}
}
}
throw new Error(
`Failed to communicate with offscreen document after ${maxRetries} attempts. Last error: ${lastError?.message}`,
);
}
async getEmbedding(text: string, options: Record<string, any> = {}): Promise<Float32Array> {
if (!this._isInitialized) {
await this.initialize();
}
// Check and ensure engine is initialized before each call
await this.ensureOffscreenEngineInitialized();
const response = await this.sendMessageToOffscreen({
target: 'offscreen',
type: OFFSCREEN_MESSAGE_TYPES.SIMILARITY_ENGINE_COMPUTE,
text: text,
options: options,
});
if (!response || !response.success) {
throw new Error(response?.error || 'Failed to get embedding from offscreen document');
}
if (!response.embedding || !Array.isArray(response.embedding)) {
throw new Error('Invalid embedding data received from offscreen document');
}
return new Float32Array(response.embedding);
}
async getEmbeddingsBatch(
texts: string[],
options: Record<string, any> = {},
): Promise<Float32Array[]> {
if (!this._isInitialized) {
await this.initialize();
}
if (!texts || texts.length === 0) return [];
// Check and ensure engine is initialized before each call
await this.ensureOffscreenEngineInitialized();
const response = await this.sendMessageToOffscreen({
target: 'offscreen',
type: OFFSCREEN_MESSAGE_TYPES.SIMILARITY_ENGINE_BATCH_COMPUTE,
texts: texts,
options: options,
});
if (!response || !response.success) {
throw new Error(response?.error || 'Failed to get embeddings batch from offscreen document');
}
return response.embeddings.map((emb: number[]) => new Float32Array(emb));
}
async computeSimilarity(
text1: string,
text2: string,
options: Record<string, any> = {},
): Promise<number> {
const [embedding1, embedding2] = await this.getEmbeddingsBatch([text1, text2], options);
return this.cosineSimilarity(embedding1, embedding2);
}
async computeSimilarityBatch(
pairs: { text1: string; text2: string }[],
options: Record<string, any> = {},
): Promise<number[]> {
if (!this._isInitialized) {
await this.initialize();
}
// Check and ensure engine is initialized before each call
await this.ensureOffscreenEngineInitialized();
const response = await this.sendMessageToOffscreen({
target: 'offscreen',
type: OFFSCREEN_MESSAGE_TYPES.SIMILARITY_ENGINE_BATCH_COMPUTE,
pairs: pairs,
options: options,
});
if (!response || !response.success) {
throw new Error(
response?.error || 'Failed to compute similarity batch from offscreen document',
);
}
return response.similarities;
}
private cosineSimilarity(a: Float32Array, b: Float32Array): number {
if (a.length !== b.length) {
throw new Error(`Vector dimensions don't match: ${a.length} vs ${b.length}`);
}
let dotProduct = 0;
let normA = 0;
let normB = 0;
for (let i = 0; i < a.length; i++) {
dotProduct += a[i] * b[i];
normA += a[i] * a[i];
normB += b[i] * b[i];
}
const magnitude = Math.sqrt(normA) * Math.sqrt(normB);
return magnitude === 0 ? 0 : dotProduct / magnitude;
}
get isInitialized(): boolean {
return this._isInitialized;
}
async dispose(): Promise<void> {
// Proxy class doesn't need to clean up resources, actual resources are managed by offscreen
this._isInitialized = false;
console.log('SemanticSimilarityEngineProxy: Proxy disposed');
}
}
export class SemanticSimilarityEngine {
private worker: Worker | null = null;
private tokenizer: PreTrainedTokenizer | null = null;
public isInitialized = false;
private isInitializing = false;
private initPromise: Promise<void> | null = null;
private nextTokenId = 0;
private pendingMessages = new Map<number, PendingMessage>();
private useOffscreen = false; // Whether to use offscreen mode
public readonly config: Required<ModelConfig>;
private embeddingCache: LRUCache<string, Float32Array>;
// Added: tokenization cache
private tokenizationCache: LRUCache<string, TokenizedOutput>;
// Added: memory pool manager
private memoryPool: EmbeddingMemoryPool;
// Added: SIMD math engine
private simdMath: SIMDMathEngine | null = null;
private useSIMD = false;
public cacheStats = {
embedding: { hits: 0, misses: 0, size: 0 },
tokenization: { hits: 0, misses: 0, size: 0 },
};
public performanceStats = {
totalEmbeddingComputations: 0,
totalEmbeddingTime: 0,
averageEmbeddingTime: 0,
totalTokenizationTime: 0,
averageTokenizationTime: 0,
totalSimilarityComputations: 0,
totalSimilarityTime: 0,
averageSimilarityTime: 0,
workerStats: null as WorkerStats | null,
};
private runningWorkerTasks = 0;
private workerTaskQueue: (() => void)[] = [];
/**
* Detect if current runtime environment supports Worker
*/
private isWorkerSupported(): boolean {
try {
// Check if in Service Worker environment (background script)
if (typeof importScripts === 'function') {
return false;
}
// Check if Worker constructor is available
return typeof Worker !== 'undefined';
} catch {
return false;
}
}
/**
* Detect if in offscreen document environment
*/
private isInOffscreenDocument(): boolean {
try {
// In offscreen document, window.location.pathname is usually '/offscreen.html'
return (
typeof window !== 'undefined' &&
window.location &&
window.location.pathname.includes('offscreen')
);
} catch {
return false;
}
}
/**
* Ensure offscreen document exists
*/
private async ensureOffscreenDocument(): Promise<void> {
return OffscreenManager.getInstance().ensureOffscreenDocument();
}
// Helper function to safely convert tensor data to number array
private convertTensorDataToNumbers(data: any): number[] {
if (data instanceof BigInt64Array) {
return Array.from(data, (val: bigint) => Number(val));
} else if (data instanceof Int32Array) {
return Array.from(data);
} else {
return Array.from(data);
}
}
constructor(options: Partial<ModelConfig> = {}) {
console.log('SemanticSimilarityEngine: Constructor called with options:', {
useLocalFiles: options.useLocalFiles,
modelIdentifier: options.modelIdentifier,
forceOffscreen: options.forceOffscreen,
modelPreset: options.modelPreset,
modelVersion: options.modelVersion,
});
// Handle model presets
let modelConfig = { ...options };
if (options.modelPreset && PREDEFINED_MODELS[options.modelPreset]) {
const preset = PREDEFINED_MODELS[options.modelPreset];
const modelVersion = options.modelVersion || 'quantized'; // Default to quantized version
const baseModelIdentifier = preset.modelIdentifier; // Use base identifier without version suffix
const onnxFileName = getOnnxFileNameForVersion(modelVersion); // Get ONNX filename based on version
// Get model-specific configuration
const modelSpecificConfig = (preset as any).modelSpecificConfig || {};
modelConfig = {
...options,
modelIdentifier: baseModelIdentifier, // Use base identifier
onnxModelFile: onnxFileName, // Set corresponding version ONNX filename
dimension: preset.dimension,
modelVersion: modelVersion,
requiresTokenTypeIds: modelSpecificConfig.requiresTokenTypeIds !== false, // Default to true unless explicitly set to false
};
console.log(
`SemanticSimilarityEngine: Using model preset "${options.modelPreset}" with version "${modelVersion}":`,
preset,
);
console.log(`SemanticSimilarityEngine: Base model identifier: ${baseModelIdentifier}`);
console.log(`SemanticSimilarityEngine: ONNX file for version: ${onnxFileName}`);
console.log(
`SemanticSimilarityEngine: Requires token_type_ids: ${modelConfig.requiresTokenTypeIds}`,
);
}
// Set default configuration - using 2025 recommended default model
this.config = {
...modelConfig,
modelIdentifier: modelConfig.modelIdentifier || 'Xenova/bge-small-en-v1.5',
localModelPathPrefix: modelConfig.localModelPathPrefix || 'models/',
onnxModelFile: modelConfig.onnxModelFile || 'model.onnx',
maxLength: modelConfig.maxLength || 256,
cacheSize: modelConfig.cacheSize || 500,
numThreads:
modelConfig.numThreads ||
(typeof navigator !== 'undefined' && navigator.hardwareConcurrency
? Math.max(1, Math.floor(navigator.hardwareConcurrency / 2))
: 2),
executionProviders:
modelConfig.executionProviders ||
(typeof WebAssembly === 'object' &&
WebAssembly.validate(new Uint8Array([0, 97, 115, 109, 1, 0, 0, 0]))
? ['wasm']
: ['webgl']),
useLocalFiles: (() => {
console.log(
'SemanticSimilarityEngine: DEBUG - modelConfig.useLocalFiles:',
modelConfig.useLocalFiles,
);
console.log(
'SemanticSimilarityEngine: DEBUG - modelConfig.useLocalFiles !== undefined:',
modelConfig.useLocalFiles !== undefined,
);
const result = modelConfig.useLocalFiles !== undefined ? modelConfig.useLocalFiles : true;
console.log('SemanticSimilarityEngine: DEBUG - final useLocalFiles value:', result);
return result;
})(),
workerPath: modelConfig.workerPath || 'js/similarity.worker.js', // Will be overridden by WXT's `new URL`
concurrentLimit:
modelConfig.concurrentLimit ||
Math.max(
1,
modelConfig.numThreads ||
(typeof navigator !== 'undefined' && navigator.hardwareConcurrency
? Math.max(1, Math.floor(navigator.hardwareConcurrency / 2))
: 2),
),
forceOffscreen: modelConfig.forceOffscreen || false,
modelPreset: modelConfig.modelPreset || 'bge-small-en-v1.5',
dimension: modelConfig.dimension || 384,
modelVersion: modelConfig.modelVersion || 'quantized',
requiresTokenTypeIds: modelConfig.requiresTokenTypeIds !== false, // Default to true
} as Required<ModelConfig>;
console.log('SemanticSimilarityEngine: Final config:', {
useLocalFiles: this.config.useLocalFiles,
modelIdentifier: this.config.modelIdentifier,
forceOffscreen: this.config.forceOffscreen,
});
this.embeddingCache = new LRUCache<string, Float32Array>(this.config.cacheSize);
this.tokenizationCache = new LRUCache<string, TokenizedOutput>(
Math.min(this.config.cacheSize, 200),
);
this.memoryPool = new EmbeddingMemoryPool();
this.simdMath = new SIMDMathEngine();
}
private _sendMessageToWorker(
type: string,
payload?: WorkerMessagePayload,
transferList?: Transferable[],
): Promise<WorkerResponsePayload> {
return new Promise((resolve, reject) => {
if (!this.worker) {
reject(new Error('Worker is not initialized.'));
return;
}
const id = this.nextTokenId++;
this.pendingMessages.set(id, { resolve, reject, type });
// Use transferable objects if provided for zero-copy transfer
if (transferList && transferList.length > 0) {
this.worker.postMessage({ id, type, payload }, transferList);
} else {
this.worker.postMessage({ id, type, payload });
}
});
}
private _setupWorker(): void {
console.log('SemanticSimilarityEngine: Setting up worker...');
// 方式1: Chrome extension URL (推荐，生产环境最可靠)
try {
const workerUrl = chrome.runtime.getURL('workers/similarity.worker.js');
console.log(`SemanticSimilarityEngine: Trying chrome.runtime.getURL ${workerUrl}`);
this.worker = new Worker(workerUrl);
console.log(`SemanticSimilarityEngine: Method 1 successful with path`);
} catch (error) {
console.warn('Method (chrome.runtime.getURL) failed:', error);
}
if (!this.worker) {
throw new Error('Worker creation failed');
}
this.worker.onmessage = (
event: MessageEvent<{
id: number;
type: string;
status: string;
payload: WorkerResponsePayload;
stats?: WorkerStats;
}>,
) => {
const { id, status, payload, stats } = event.data;
const promiseCallbacks = this.pendingMessages.get(id);
if (!promiseCallbacks) return;
this.pendingMessages.delete(id);
// 更新 Worker 统计信息
if (stats) {
this.performanceStats.workerStats = stats;
}
if (status === 'success') {
promiseCallbacks.resolve(payload);
} else {
const error = new Error(
payload?.message || `Worker error for task ${promiseCallbacks.type}`,
);
(error as any).name = (payload as any)?.name || 'WorkerError';
(error as any).stack = (payload as any)?.stack || undefined;
console.error(
`Error from worker (task ${id}, type ${promiseCallbacks.type}):`,
error,
event.data,
);
promiseCallbacks.reject(error);
}
};
this.worker.onerror = (error: ErrorEvent) => {
console.error('==== Unhandled error in SemanticSimilarityEngine Worker ====');
console.error('Event Message:', error.message);
console.error('Event Filename:', error.filename);
console.error('Event Lineno:', error.lineno);
console.error('Event Colno:', error.colno);
if (error.error) {
// 检查 event.error 是否存在
console.error('Actual Error Name:', error.error.name);
console.error('Actual Error Message:', error.error.message);
console.error('Actual Error Stack:', error.error.stack);
} else {
console.error('Actual Error object (event.error) is not available. Error details:', {
message: error.message,
filename: error.filename,
lineno: error.lineno,
colno: error.colno,
});
}
console.error('==========================================================');
this.pendingMessages.forEach((callbacks) => {
callbacks.reject(new Error(`Worker terminated or unhandled error: ${error.message}`));
});
this.pendingMessages.clear();
this.isInitialized = false;
this.isInitializing = false;
};
}
public async initialize(): Promise<void> {
if (this.isInitialized) return Promise.resolve();
if (this.isInitializing && this.initPromise) return this.initPromise;
this.isInitializing = true;
this.initPromise = this._doInitialize().finally(() => {
this.isInitializing = false;
// this.warmupModel();
});
return this.initPromise;
}
/**
* 带进度回调的初始化方法
*/
public async initializeWithProgress(
onProgress?: (progress: { status: string; progress: number; message?: string }) => void,
): Promise<void> {
if (this.isInitialized) return Promise.resolve();
if (this.isInitializing && this.initPromise) return this.initPromise;
this.isInitializing = true;
this.initPromise = this._doInitializeWithProgress(onProgress).finally(() => {
this.isInitializing = false;
// this.warmupModel();
});
return this.initPromise;
}
/**
* 带进度回调的内部初始化方法
*/
private async _doInitializeWithProgress(
onProgress?: (progress: { status: string; progress: number; message?: string }) => void,
): Promise<void> {
console.log('SemanticSimilarityEngine: Initializing with progress tracking...');
const startTime = performance.now();
// 进度报告辅助函数
const reportProgress = (status: string, progress: number, message?: string) => {
if (onProgress) {
onProgress({ status, progress, message });
}
};
try {
reportProgress('initializing', 5, 'Starting initialization...');
// 检测环境并决定使用哪种模式
const workerSupported = this.isWorkerSupported();
const inOffscreenDocument = this.isInOffscreenDocument();
// 🛠️ 防止死循环：如果已经在 offscreen document 中，强制使用直接 Worker 模式
if (inOffscreenDocument) {
this.useOffscreen = false;
console.log(
'SemanticSimilarityEngine: Running in offscreen document, using direct Worker mode to prevent recursion',
);
} else {
this.useOffscreen = this.config.forceOffscreen || !workerSupported;
}
console.log(
`SemanticSimilarityEngine: Worker supported: ${workerSupported}, In offscreen: ${inOffscreenDocument}, Using offscreen: ${this.useOffscreen}`,
);
reportProgress('initializing', 10, 'Environment detection complete');
if (this.useOffscreen) {
// 使用offscreen模式 - 委托给offscreen document，它会处理自己的进度
reportProgress('initializing', 15, 'Setting up offscreen document...');
await this.ensureOffscreenDocument();
// 发送初始化消息到offscreen document
console.log('SemanticSimilarityEngine: Sending config to offscreen:', {
useLocalFiles: this.config.useLocalFiles,
modelIdentifier: this.config.modelIdentifier,
localModelPathPrefix: this.config.localModelPathPrefix,
});
// 确保配置对象被正确序列化，显式设置所有属性
const configToSend = {
modelIdentifier: this.config.modelIdentifier,
localModelPathPrefix: this.config.localModelPathPrefix,
onnxModelFile: this.config.onnxModelFile,
maxLength: this.config.maxLength,
cacheSize: this.config.cacheSize,
numThreads: this.config.numThreads,
executionProviders: this.config.executionProviders,
useLocalFiles: Boolean(this.config.useLocalFiles), // 强制转换为布尔值
workerPath: this.config.workerPath,
concurrentLimit: this.config.concurrentLimit,
forceOffscreen: this.config.forceOffscreen,
modelPreset: this.config.modelPreset,
modelVersion: this.config.modelVersion,
dimension: this.config.dimension,
};
// 使用 JSON 序列化确保数据完整性
const serializedConfig = JSON.parse(JSON.stringify(configToSend));
reportProgress('initializing', 20, 'Delegating to offscreen document...');
const response = await chrome.runtime.sendMessage({
target: 'offscreen',
type: OFFSCREEN_MESSAGE_TYPES.SIMILARITY_ENGINE_INIT,
config: serializedConfig,
});
if (!response || !response.success) {
throw new Error(response?.error || 'Failed to initialize engine in offscreen document');
}
reportProgress('ready', 100, 'Initialized via offscreen document');
console.log('SemanticSimilarityEngine: Initialized via offscreen document');
} else {
// 使用直接Worker模式 - 这里我们可以提供真实的进度跟踪
await this._initializeDirectWorkerWithProgress(reportProgress);
}
this.isInitialized = true;
console.log(
`SemanticSimilarityEngine: Initialization complete in ${(performance.now() - startTime).toFixed(2)}ms`,
);
} catch (error) {
console.error('SemanticSimilarityEngine: Initialization failed.', error);
const errorMessage = error instanceof Error ? error.message : 'Unknown error';
reportProgress('error', 0, `Initialization failed: ${errorMessage}`);
if (this.worker) this.worker.terminate();
this.worker = null;
this.isInitialized = false;
this.isInitializing = false;
this.initPromise = null;
// 创建一个更详细的错误对象
const enhancedError = new Error(errorMessage);
enhancedError.name = 'ModelInitializationError';
throw enhancedError;
}
}
private async _doInitialize(): Promise<void> {
console.log('SemanticSimilarityEngine: Initializing...');
const startTime = performance.now();
try {
// 检测环境并决定使用哪种模式
const workerSupported = this.isWorkerSupported();
const inOffscreenDocument = this.isInOffscreenDocument();
// 🛠️ 防止死循环：如果已经在 offscreen document 中，强制使用直接 Worker 模式
if (inOffscreenDocument) {
this.useOffscreen = false;
console.log(
'SemanticSimilarityEngine: Running in offscreen document, using direct Worker mode to prevent recursion',
);
} else {
this.useOffscreen = this.config.forceOffscreen || !workerSupported;
}
console.log(
`SemanticSimilarityEngine: Worker supported: ${workerSupported}, In offscreen: ${inOffscreenDocument}, Using offscreen: ${this.useOffscreen}`,
);
if (this.useOffscreen) {
// 使用offscreen模式
await this.ensureOffscreenDocument();
// 发送初始化消息到offscreen document
console.log('SemanticSimilarityEngine: Sending config to offscreen:', {
useLocalFiles: this.config.useLocalFiles,
modelIdentifier: this.config.modelIdentifier,
localModelPathPrefix: this.config.localModelPathPrefix,
});
// 确保配置对象被正确序列化，显式设置所有属性
const configToSend = {
modelIdentifier: this.config.modelIdentifier,
localModelPathPrefix: this.config.localModelPathPrefix,
onnxModelFile: this.config.onnxModelFile,
maxLength: this.config.maxLength,
cacheSize: this.config.cacheSize,
numThreads: this.config.numThreads,
executionProviders: this.config.executionProviders,
useLocalFiles: Boolean(this.config.useLocalFiles), // 强制转换为布尔值
workerPath: this.config.workerPath,
concurrentLimit: this.config.concurrentLimit,
forceOffscreen: this.config.forceOffscreen,
modelPreset: this.config.modelPreset,
modelVersion: this.config.modelVersion,
dimension: this.config.dimension,
};
console.log(
'SemanticSimilarityEngine: DEBUG - configToSend.useLocalFiles:',
configToSend.useLocalFiles,
);
console.log(
'SemanticSimilarityEngine: DEBUG - typeof configToSend.useLocalFiles:',
typeof configToSend.useLocalFiles,
);
console.log('SemanticSimilarityEngine: Explicit config to send:', configToSend);
console.log(
'SemanticSimilarityEngine: DEBUG - this.config.useLocalFiles value:',
this.config.useLocalFiles,
);
console.log(
'SemanticSimilarityEngine: DEBUG - typeof this.config.useLocalFiles:',
typeof this.config.useLocalFiles,
);
// 使用 JSON 序列化确保数据完整性
const serializedConfig = JSON.parse(JSON.stringify(configToSend));
console.log(
'SemanticSimilarityEngine: DEBUG - serializedConfig.useLocalFiles:',
serializedConfig.useLocalFiles,
);
const response = await chrome.runtime.sendMessage({
target: 'offscreen',
type: OFFSCREEN_MESSAGE_TYPES.SIMILARITY_ENGINE_INIT,
config: serializedConfig, // 使用原始配置，不强制修改 useLocalFiles
});
if (!response || !response.success) {
throw new Error(response?.error || 'Failed to initialize engine in offscreen document');
}
console.log('SemanticSimilarityEngine: Initialized via offscreen document');
} else {
// 使用直接Worker模式
this._setupWorker();
TransformersEnv.allowRemoteModels = !this.config.useLocalFiles;
TransformersEnv.allowLocalModels = this.config.useLocalFiles;
console.log(`SemanticSimilarityEngine: TransformersEnv config:`, {
allowRemoteModels: TransformersEnv.allowRemoteModels,
allowLocalModels: TransformersEnv.allowLocalModels,
useLocalFiles: this.config.useLocalFiles,
});
if (TransformersEnv.backends?.onnx?.wasm) {
// 检查路径是否存在
TransformersEnv.backends.onnx.wasm.numThreads = this.config.numThreads;
}
let tokenizerIdentifier = this.config.modelIdentifier;
if (this.config.useLocalFiles) {
// 对于WXT，public目录下的资源在运行时位于根路径
// 直接使用模型标识符，transformers.js 会自动添加 /models/ 前缀
tokenizerIdentifier = this.config.modelIdentifier;
}
console.log(
`SemanticSimilarityEngine: Loading tokenizer from ${tokenizerIdentifier} (local_files_only: ${this.config.useLocalFiles})`,
);
const tokenizerConfig: any = {
quantized: false,
local_files_only: this.config.useLocalFiles,
};
// 对于不需要token_type_ids的模型，在tokenizer配置中明确设置
if (!this.config.requiresTokenTypeIds) {
tokenizerConfig.return_token_type_ids = false;
}
console.log(`SemanticSimilarityEngine: Full tokenizer config:`, {
tokenizerIdentifier,
localModelPathPrefix: this.config.localModelPathPrefix,
modelIdentifier: this.config.modelIdentifier,
useLocalFiles: this.config.useLocalFiles,
local_files_only: this.config.useLocalFiles,
requiresTokenTypeIds: this.config.requiresTokenTypeIds,
tokenizerConfig,
});
this.tokenizer = await AutoTokenizer.from_pretrained(tokenizerIdentifier, tokenizerConfig);
console.log('SemanticSimilarityEngine: Tokenizer loaded.');
if (this.config.useLocalFiles) {
// Local files mode - use URL path as before
const onnxModelPathForWorker = chrome.runtime.getURL(
`models/${this.config.modelIdentifier}/${this.config.onnxModelFile}`,
);
console.log(
`SemanticSimilarityEngine: Instructing worker to load local ONNX model from ${onnxModelPathForWorker}`,
);
await this._sendMessageToWorker('init', {
modelPath: onnxModelPathForWorker,
numThreads: this.config.numThreads,
executionProviders: this.config.executionProviders,
});
} else {
// Remote files mode - use cached model data
const modelIdParts = this.config.modelIdentifier.split('/');
const modelNameForUrl =
modelIdParts.length > 1
? this.config.modelIdentifier
: `Xenova/${this.config.modelIdentifier}`;
const onnxModelUrl = `https://huggingface.co/${modelNameForUrl}/resolve/main/onnx/${this.config.onnxModelFile}`;
if (!this.config.modelIdentifier.includes('/')) {
console.warn(
`Warning: modelIdentifier "${this.config.modelIdentifier}" might not be a full HuggingFace path. Assuming Xenova prefix for remote URL.`,
);
}
console.log(`SemanticSimilarityEngine: Getting cached model data from ${onnxModelUrl}`);
// Get model data from cache (may download if not cached)
const modelData = await getCachedModelData(onnxModelUrl);
console.log(
`SemanticSimilarityEngine: Sending cached model data to worker (${modelData.byteLength} bytes)`,
);
// Send ArrayBuffer to worker with transferable objects for zero-copy
await this._sendMessageToWorker(
'init',
{
modelData: modelData,
numThreads: this.config.numThreads,
executionProviders: this.config.executionProviders,
},
[modelData],
);
}
console.log('SemanticSimilarityEngine: Worker reported model initialized.');
// 尝试初始化 SIMD 加速
try {
console.log('SemanticSimilarityEngine: Checking SIMD support...');
const simdSupported = await SIMDMathEngine.checkSIMDSupport();
if (simdSupported) {
console.log('SemanticSimilarityEngine: SIMD supported, initializing...');
await this.simdMath!.initialize();
this.useSIMD = true;
console.log('SemanticSimilarityEngine: ✅ SIMD acceleration enabled');
} else {
console.log(
'SemanticSimilarityEngine: ❌ SIMD not supported, using JavaScript fallback',
);
console.log('SemanticSimilarityEngine: To enable SIMD, please use:');
console.log(' - Chrome 91+ (May 2021)');
console.log(' - Firefox 89+ (June 2021)');
console.log(' - Safari 16.4+ (March 2023)');
console.log(' - Edge 91+ (May 2021)');
this.useSIMD = false;
}
} catch (simdError) {
console.warn(
'SemanticSimilarityEngine: SIMD initialization failed, using JavaScript fallback:',
simdError,
);
this.useSIMD = false;
}
}
this.isInitialized = true;
console.log(
`SemanticSimilarityEngine: Initialization complete in ${(performance.now() - startTime).toFixed(2)}ms`,
);
} catch (error) {
console.error('SemanticSimilarityEngine: Initialization failed.', error);
if (this.worker) this.worker.terminate();
this.worker = null;
this.isInitialized = false;
this.isInitializing = false;
this.initPromise = null;
// 创建一个更详细的错误对象
const errorMessage = error instanceof Error ? error.message : 'Unknown error';
const enhancedError = new Error(errorMessage);
enhancedError.name = 'ModelInitializationError';
throw enhancedError;
}
}
/**
* 直接Worker模式的初始化，支持进度回调
*/
private async _initializeDirectWorkerWithProgress(
reportProgress: (status: string, progress: number, message?: string) => void,
): Promise<void> {
// 使用直接Worker模式
reportProgress('initializing', 25, 'Setting up worker...');
this._setupWorker();
TransformersEnv.allowRemoteModels = !this.config.useLocalFiles;
TransformersEnv.allowLocalModels = this.config.useLocalFiles;
console.log(`SemanticSimilarityEngine: TransformersEnv config:`, {
allowRemoteModels: TransformersEnv.allowRemoteModels,
allowLocalModels: TransformersEnv.allowLocalModels,
useLocalFiles: this.config.useLocalFiles,
});
if (TransformersEnv.backends?.onnx?.wasm) {
TransformersEnv.backends.onnx.wasm.numThreads = this.config.numThreads;
}
let tokenizerIdentifier = this.config.modelIdentifier;
if (this.config.useLocalFiles) {
tokenizerIdentifier = this.config.modelIdentifier;
}
reportProgress('downloading', 40, 'Loading tokenizer...');
console.log(
`SemanticSimilarityEngine: Loading tokenizer from ${tokenizerIdentifier} (local_files_only: ${this.config.useLocalFiles})`,
);
// 使用 transformers.js 2.17+ 的进度回调功能
const tokenizerProgressCallback = (progress: any) => {
if (progress.status === 'downloading') {
const progressPercent = Math.min(40 + (progress.progress || 0) * 0.3, 70);
reportProgress(
'downloading',
progressPercent,
`Downloading tokenizer: ${progress.file || ''}`,
);
}
};
const tokenizerConfig: any = {
quantized: false,
local_files_only: this.config.useLocalFiles,
};
// 对于不需要token_type_ids的模型，在tokenizer配置中明确设置
if (!this.config.requiresTokenTypeIds) {
tokenizerConfig.return_token_type_ids = false;
}
try {
if (!this.config.useLocalFiles) {
tokenizerConfig.progress_callback = tokenizerProgressCallback;
}
this.tokenizer = await AutoTokenizer.from_pretrained(tokenizerIdentifier, tokenizerConfig);
} catch (error) {
// 如果进度回调不支持，回退到标准方式
console.log(
'SemanticSimilarityEngine: Progress callback not supported, using standard loading',
);
delete tokenizerConfig.progress_callback;
this.tokenizer = await AutoTokenizer.from_pretrained(tokenizerIdentifier, tokenizerConfig);
}
reportProgress('downloading', 70, 'Tokenizer loaded, setting up ONNX model...');
console.log('SemanticSimilarityEngine: Tokenizer loaded.');
if (this.config.useLocalFiles) {
// Local files mode - use URL path as before
const onnxModelPathForWorker = chrome.runtime.getURL(
`models/${this.config.modelIdentifier}/${this.config.onnxModelFile}`,
);
reportProgress('downloading', 80, 'Loading local ONNX model...');
console.log(
`SemanticSimilarityEngine: Instructing worker to load local ONNX model from ${onnxModelPathForWorker}`,
);
await this._sendMessageToWorker('init', {
modelPath: onnxModelPathForWorker,
numThreads: this.config.numThreads,
executionProviders: this.config.executionProviders,
});
} else {
// Remote files mode - use cached model data
const modelIdParts = this.config.modelIdentifier.split('/');
const modelNameForUrl =
modelIdParts.length > 1
? this.config.modelIdentifier
: `Xenova/${this.config.modelIdentifier}`;
const onnxModelUrl = `https://huggingface.co/${modelNameForUrl}/resolve/main/onnx/${this.config.onnxModelFile}`;
if (!this.config.modelIdentifier.includes('/')) {
console.warn(
`Warning: modelIdentifier "${this.config.modelIdentifier}" might not be a full HuggingFace path. Assuming Xenova prefix for remote URL.`,
);
}
reportProgress('downloading', 80, 'Loading cached ONNX model...');
console.log(`SemanticSimilarityEngine: Getting cached model data from ${onnxModelUrl}`);
// Get model data from cache (may download if not cached)
const modelData = await getCachedModelData(onnxModelUrl);
console.log(
`SemanticSimilarityEngine: Sending cached model data to worker (${modelData.byteLength} bytes)`,
);
// Send ArrayBuffer to worker with transferable objects for zero-copy
await this._sendMessageToWorker(
'init',
{
modelData: modelData,
numThreads: this.config.numThreads,
executionProviders: this.config.executionProviders,
},
[modelData],
);
}
console.log('SemanticSimilarityEngine: Worker reported model initialized.');
reportProgress('initializing', 90, 'Setting up SIMD acceleration...');
// 尝试初始化 SIMD 加速
try {
console.log('SemanticSimilarityEngine: Checking SIMD support...');
const simdSupported = await SIMDMathEngine.checkSIMDSupport();
if (simdSupported) {
console.log('SemanticSimilarityEngine: SIMD supported, initializing...');
await this.simdMath!.initialize();
this.useSIMD = true;
console.log('SemanticSimilarityEngine: ✅ SIMD acceleration enabled');
} else {
console.log('SemanticSimilarityEngine: ❌ SIMD not supported, using JavaScript fallback');
this.useSIMD = false;
}
} catch (simdError) {
console.warn(
'SemanticSimilarityEngine: SIMD initialization failed, using JavaScript fallback:',
simdError,
);
this.useSIMD = false;
}
reportProgress('ready', 100, 'Initialization complete');
}
public async warmupModel(): Promise<void> {
if (!this.isInitialized && !this.isInitializing) {
await this.initialize();
} else if (this.isInitializing && this.initPromise) {
await this.initPromise;
}
if (!this.isInitialized) throw new Error('Engine not initialized after warmup attempt.');
console.log('SemanticSimilarityEngine: Warming up model...');
// 更有代表性的预热文本，包含不同长度和语言
const warmupTexts = [
// 短文本
'Hello',
'你好',
'Test',
// 中等长度文本
'Hello world, this is a test.',
'你好世界，这是一个测试。',
'The quick brown fox jumps over the lazy dog.',
// 长文本
'This is a longer text that contains multiple sentences. It helps warm up the model for various text lengths.',
'这是一个包含多个句子的较长文本。它有助于为各种文本长度预热模型。',
];
try {
// 渐进式预热：先单个，再批量
console.log('SemanticSimilarityEngine: Phase 1 - Individual warmup...');
for (const text of warmupTexts.slice(0, 4)) {
await this.getEmbedding(text);
}
console.log('SemanticSimilarityEngine: Phase 2 - Batch warmup...');
await this.getEmbeddingsBatch(warmupTexts.slice(4));
// 保留预热结果，不清空缓存
console.log('SemanticSimilarityEngine: Model warmup complete. Cache preserved.');
console.log(`Embedding cache: ${this.cacheStats.embedding.size} items`);
console.log(`Tokenization cache: ${this.cacheStats.tokenization.size} items`);
} catch (error) {
console.warn('SemanticSimilarityEngine: Warmup failed. This might not be critical.', error);
}
}
private async _tokenizeText(text: string | string[]): Promise<TokenizedOutput> {
if (!this.tokenizer) throw new Error('Tokenizer not initialized.');
// 对于单个文本，尝试使用缓存
if (typeof text === 'string') {
const cacheKey = `tokenize:${text}`;
const cached = this.tokenizationCache.get(cacheKey);
if (cached) {
this.cacheStats.tokenization.hits++;
this.cacheStats.tokenization.size = this.tokenizationCache.size;
return cached;
}
this.cacheStats.tokenization.misses++;
const startTime = performance.now();
const tokenizerOptions: any = {
padding: true,
truncation: true,
max_length: this.config.maxLength,
return_tensors: 'np',
};
// 对于不需要token_type_ids的模型，明确设置return_token_type_ids为false
if (!this.config.requiresTokenTypeIds) {
tokenizerOptions.return_token_type_ids = false;
}
const result = (await this.tokenizer(text, tokenizerOptions)) as TokenizedOutput;
// 更新性能统计
this.performanceStats.totalTokenizationTime += performance.now() - startTime;
this.performanceStats.averageTokenizationTime =
this.performanceStats.totalTokenizationTime /
(this.cacheStats.tokenization.hits + this.cacheStats.tokenization.misses);
// 缓存结果
this.tokenizationCache.set(cacheKey, result);
this.cacheStats.tokenization.size = this.tokenizationCache.size;
return result;
}
// 对于批量文本，直接处理（批量处理通常不重复）
const startTime = performance.now();
const tokenizerOptions: any = {
padding: true,
truncation: true,
max_length: this.config.maxLength,
return_tensors: 'np',
};
// 对于不需要token_type_ids的模型，明确设置return_token_type_ids为false
if (!this.config.requiresTokenTypeIds) {
tokenizerOptions.return_token_type_ids = false;
}
const result = (await this.tokenizer(text, tokenizerOptions)) as TokenizedOutput;
this.performanceStats.totalTokenizationTime += performance.now() - startTime;
return result;
}
private _extractEmbeddingFromWorkerOutput(
workerOutput: WorkerResponsePayload,
attentionMaskArray: number[],
): Float32Array {
if (!workerOutput.data || !workerOutput.dims)
throw new Error('Invalid worker output for embedding extraction.');
// 优化：直接使用 Float32Array，避免不必要的转换
const lastHiddenStateData =
workerOutput.data instanceof Float32Array
? workerOutput.data
: new Float32Array(workerOutput.data);
const dims = workerOutput.dims;
const seqLength = dims[1];
const hiddenSize = dims[2];
// 使用内存池获取 embedding 数组
const embedding = this.memoryPool.getEmbedding(hiddenSize);
let validTokens = 0;
for (let i = 0; i < seqLength; i++) {
if (attentionMaskArray[i] === 1) {
const offset = i * hiddenSize;
for (let j = 0; j < hiddenSize; j++) {
embedding[j] += lastHiddenStateData[offset + j];
}
validTokens++;
}
}
if (validTokens > 0) {
for (let i = 0; i < hiddenSize; i++) {
embedding[i] /= validTokens;
}
}
return this.normalizeVector(embedding);
}
private _extractBatchEmbeddingsFromWorkerOutput(
workerOutput: WorkerResponsePayload,
attentionMasksBatch: number[][],
): Float32Array[] {
if (!workerOutput.data || !workerOutput.dims)
throw new Error('Invalid worker output for batch embedding extraction.');
// 优化：直接使用 Float32Array，避免不必要的转换
const lastHiddenStateData =
workerOutput.data instanceof Float32Array
? workerOutput.data
: new Float32Array(workerOutput.data);
const dims = workerOutput.dims;
const batchSize = dims[0];
const seqLength = dims[1];
const hiddenSize = dims[2];
const embeddings: Float32Array[] = [];
for (let b = 0; b < batchSize; b++) {
// 使用内存池获取 embedding 数组
const embedding = this.memoryPool.getEmbedding(hiddenSize);
let validTokens = 0;
const currentAttentionMask = attentionMasksBatch[b];
for (let i = 0; i < seqLength; i++) {
if (currentAttentionMask[i] === 1) {
const offset = (b * seqLength + i) * hiddenSize;
for (let j = 0; j < hiddenSize; j++) {
embedding[j] += lastHiddenStateData[offset + j];
}
validTokens++;
}
}
if (validTokens > 0) {
for (let i = 0; i < hiddenSize; i++) {
embedding[i] /= validTokens;
}
}
embeddings.push(this.normalizeVector(embedding));
}
return embeddings;
}
public async getEmbedding(
text: string,
options: Record<string, any> = {},
): Promise<Float32Array> {
if (!this.isInitialized) await this.initialize();
const cacheKey = this.getCacheKey(text, options);
const cached = this.embeddingCache.get(cacheKey);
if (cached) {
this.cacheStats.embedding.hits++;
this.cacheStats.embedding.size = this.embeddingCache.size;
return cached;
}
this.cacheStats.embedding.misses++;
// 如果使用offscreen模式，委托给offscreen document
if (this.useOffscreen) {
const response = await chrome.runtime.sendMessage({
target: 'offscreen',
type: OFFSCREEN_MESSAGE_TYPES.SIMILARITY_ENGINE_COMPUTE,
text: text,
options: options,
});
if (!response || !response.success) {
throw new Error(response?.error || 'Failed to get embedding from offscreen document');
}
// 验证响应数据
if (!response.embedding || !Array.isArray(response.embedding)) {
throw new Error('Invalid embedding data received from offscreen document');
}
console.log('SemanticSimilarityEngine: Received embedding from offscreen:', {
length: response.embedding.length,
type: typeof response.embedding,
isArray: Array.isArray(response.embedding),
firstFewValues: response.embedding.slice(0, 5),
});
const embedding = new Float32Array(response.embedding);
// 验证转换后的数据
console.log('SemanticSimilarityEngine: Converted embedding:', {
length: embedding.length,
type: typeof embedding,
constructor: embedding.constructor.name,
isFloat32Array: embedding instanceof Float32Array,
firstFewValues: Array.from(embedding.slice(0, 5)),
});
this.embeddingCache.set(cacheKey, embedding);
this.cacheStats.embedding.size = this.embeddingCache.size;
// 更新性能统计
this.performanceStats.totalEmbeddingComputations++;
return embedding;
}
if (this.runningWorkerTasks >= this.config.concurrentLimit) {
await this.waitForWorkerSlot();
}
this.runningWorkerTasks++;
const startTime = performance.now();
try {
const tokenized = await this._tokenizeText(text);
const inputIdsData = this.convertTensorDataToNumbers(tokenized.input_ids.data);
const attentionMaskData = this.convertTensorDataToNumbers(tokenized.attention_mask.data);
const tokenTypeIdsData = tokenized.token_type_ids
? this.convertTensorDataToNumbers(tokenized.token_type_ids.data)
: undefined;
const workerPayload: WorkerMessagePayload = {
input_ids: inputIdsData,
attention_mask: attentionMaskData,
token_type_ids: tokenTypeIdsData,
dims: {
input_ids: tokenized.input_ids.dims,
attention_mask: tokenized.attention_mask.dims,
token_type_ids: tokenized.token_type_ids?.dims,
},
};
const workerOutput = await this._sendMessageToWorker('infer', workerPayload);
const embedding = this._extractEmbeddingFromWorkerOutput(workerOutput, attentionMaskData);
this.embeddingCache.set(cacheKey, embedding);
this.cacheStats.embedding.size = this.embeddingCache.size;
this.performanceStats.totalEmbeddingComputations++;
this.performanceStats.totalEmbeddingTime += performance.now() - startTime;
this.performanceStats.averageEmbeddingTime =
this.performanceStats.totalEmbeddingTime / this.performanceStats.totalEmbeddingComputations;
return embedding;
} finally {
this.runningWorkerTasks--;
this.processWorkerQueue();
}
}
public async getEmbeddingsBatch(
texts: string[],
options: Record<string, any> = {},
): Promise<Float32Array[]> {
if (!this.isInitialized) await this.initialize();
if (!texts || texts.length === 0) return [];
// 如果使用offscreen模式，委托给offscreen document
if (this.useOffscreen) {
// 先检查缓存
const results: (Float32Array | undefined)[] = new Array(texts.length).fill(undefined);
const uncachedTexts: string[] = [];
const uncachedIndices: number[] = [];
texts.forEach((text, index) => {
const cacheKey = this.getCacheKey(text, options);
const cached = this.embeddingCache.get(cacheKey);
if (cached) {
results[index] = cached;
this.cacheStats.embedding.hits++;
} else {
uncachedTexts.push(text);
uncachedIndices.push(index);
this.cacheStats.embedding.misses++;
}
});
// 如果所有都在缓存中，直接返回
if (uncachedTexts.length === 0) {
return results as Float32Array[];
}
// 只请求未缓存的文本
const response = await chrome.runtime.sendMessage({
target: 'offscreen',
type: OFFSCREEN_MESSAGE_TYPES.SIMILARITY_ENGINE_BATCH_COMPUTE,
texts: uncachedTexts,
options: options,
});
if (!response || !response.success) {
throw new Error(
response?.error || 'Failed to get embeddings batch from offscreen document',
);
}
// 将结果放回对应位置并缓存
response.embeddings.forEach((embeddingArray: number[], batchIndex: number) => {
const embedding = new Float32Array(embeddingArray);
const originalIndex = uncachedIndices[batchIndex];
const originalText = uncachedTexts[batchIndex];
results[originalIndex] = embedding;
// 缓存结果
const cacheKey = this.getCacheKey(originalText, options);
this.embeddingCache.set(cacheKey, embedding);
});
this.cacheStats.embedding.size = this.embeddingCache.size;
this.performanceStats.totalEmbeddingComputations += uncachedTexts.length;
return results as Float32Array[];
}
const results: (Float32Array | undefined)[] = new Array(texts.length).fill(undefined);
const uncachedTextsMap = new Map<string, number[]>();
const textsToTokenize: string[] = [];
texts.forEach((text, index) => {
const cacheKey = this.getCacheKey(text, options);
const cached = this.embeddingCache.get(cacheKey);
if (cached) {
results[index] = cached;
this.cacheStats.embedding.hits++;
} else {
if (!uncachedTextsMap.has(text)) {
uncachedTextsMap.set(text, []);
textsToTokenize.push(text);
}
uncachedTextsMap.get(text)!.push(index);
this.cacheStats.embedding.misses++;
}
});
this.cacheStats.embedding.size = this.embeddingCache.size;
if (textsToTokenize.length === 0) return results as Float32Array[];
if (this.runningWorkerTasks >= this.config.concurrentLimit) {
await this.waitForWorkerSlot();
}
this.runningWorkerTasks++;
const startTime = performance.now();
try {
const tokenizedBatch = await this._tokenizeText(textsToTokenize);
const workerPayload: WorkerMessagePayload = {
input_ids: this.convertTensorDataToNumbers(tokenizedBatch.input_ids.data),
attention_mask: this.convertTensorDataToNumbers(tokenizedBatch.attention_mask.data),
token_type_ids: tokenizedBatch.token_type_ids
? this.convertTensorDataToNumbers(tokenizedBatch.token_type_ids.data)
: undefined,
dims: {
input_ids: tokenizedBatch.input_ids.dims,
attention_mask: tokenizedBatch.attention_mask.dims,
token_type_ids: tokenizedBatch.token_type_ids?.dims,
},
};
// 使用真正的批处理推理
const workerOutput = await this._sendMessageToWorker('batchInfer', workerPayload);
const attentionMasksForBatch: number[][] = [];
const batchSize = tokenizedBatch.input_ids.dims[0];
const seqLength = tokenizedBatch.input_ids.dims[1];
const rawAttentionMaskData = this.convertTensorDataToNumbers(
tokenizedBatch.attention_mask.data,
);
for (let i = 0; i < batchSize; ++i) {
attentionMasksForBatch.push(rawAttentionMaskData.slice(i * seqLength, (i + 1) * seqLength));
}
const batchEmbeddings = this._extractBatchEmbeddingsFromWorkerOutput(
workerOutput,
attentionMasksForBatch,
);
batchEmbeddings.forEach((embedding, batchIdx) => {
const originalText = textsToTokenize[batchIdx];
const cacheKey = this.getCacheKey(originalText, options);
this.embeddingCache.set(cacheKey, embedding);
const originalResultIndices = uncachedTextsMap.get(originalText)!;
originalResultIndices.forEach((idx) => {
results[idx] = embedding;
});
});
this.cacheStats.embedding.size = this.embeddingCache.size;
this.performanceStats.totalEmbeddingComputations += textsToTokenize.length;
this.performanceStats.totalEmbeddingTime += performance.now() - startTime;
this.performanceStats.averageEmbeddingTime =
this.performanceStats.totalEmbeddingTime / this.performanceStats.totalEmbeddingComputations;
return results as Float32Array[];
} finally {
this.runningWorkerTasks--;
this.processWorkerQueue();
}
}
public async computeSimilarity(
text1: string,
text2: string,
options: Record<string, any> = {},
): Promise<number> {
if (!this.isInitialized) await this.initialize();
this.validateInput(text1, text2);
const simStartTime = performance.now();
const [embedding1, embedding2] = await Promise.all([
this.getEmbedding(text1, options),
this.getEmbedding(text2, options),
]);
const similarity = this.cosineSimilarity(embedding1, embedding2);
console.log('computeSimilarity:', similarity);
this.performanceStats.totalSimilarityComputations++;
this.performanceStats.totalSimilarityTime += performance.now() - simStartTime;
this.performanceStats.averageSimilarityTime =
this.performanceStats.totalSimilarityTime / this.performanceStats.totalSimilarityComputations;
return similarity;
}
public async computeSimilarityBatch(
pairs: { text1: string; text2: string }[],
options: Record<string, any> = {},
): Promise<number[]> {
if (!this.isInitialized) await this.initialize();
if (!pairs || pairs.length === 0) return [];
// 如果使用offscreen模式，委托给offscreen document
if (this.useOffscreen) {
const response = await chrome.runtime.sendMessage({
target: 'offscreen',
type: OFFSCREEN_MESSAGE_TYPES.SIMILARITY_ENGINE_BATCH_COMPUTE,
pairs: pairs,
options: options,
});
if (!response || !response.success) {
throw new Error(response?.error || 'Failed to compute similarities in offscreen document');
}
return response.similarities;
}
// 直接模式的原有逻辑
const simStartTime = performance.now();
const uniqueTextsSet = new Set<string>();
pairs.forEach((pair) => {
this.validateInput(pair.text1, pair.text2);
uniqueTextsSet.add(pair.text1);
uniqueTextsSet.add(pair.text2);
});
const uniqueTextsArray = Array.from(uniqueTextsSet);
const embeddingsArray = await this.getEmbeddingsBatch(uniqueTextsArray, options);
const embeddingMap = new Map<string, Float32Array>();
uniqueTextsArray.forEach((text, index) => {
embeddingMap.set(text, embeddingsArray[index]);
});
const similarities = pairs.map((pair) => {
const emb1 = embeddingMap.get(pair.text1);
const emb2 = embeddingMap.get(pair.text2);
if (!emb1 || !emb2) {
console.warn('Embeddings not found for pair:', pair);
return 0;
}
return this.cosineSimilarity(emb1, emb2);
});
this.performanceStats.totalSimilarityComputations += pairs.length;
this.performanceStats.totalSimilarityTime += performance.now() - simStartTime;
this.performanceStats.averageSimilarityTime =
this.performanceStats.totalSimilarityTime / this.performanceStats.totalSimilarityComputations;
return similarities;
}
public async computeSimilarityMatrix(
texts1: string[],
texts2: string[],
options: Record<string, any> = {},
): Promise<number[][]> {
if (!this.isInitialized) await this.initialize();
if (!texts1 || !texts2 || texts1.length === 0 || texts2.length === 0) return [];
const simStartTime = performance.now();
const allTextsSet = new Set<string>([...texts1, ...texts2]);
texts1.forEach((t) => this.validateInput(t, 'valid_dummy'));
texts2.forEach((t) => this.validateInput(t, 'valid_dummy'));
const allTextsArray = Array.from(allTextsSet);
const embeddingsArray = await this.getEmbeddingsBatch(allTextsArray, options);
const embeddingMap = new Map<string, Float32Array>();
allTextsArray.forEach((text, index) => {
embeddingMap.set(text, embeddingsArray[index]);
});
// 使用 SIMD 优化的矩阵计算（如果可用）
if (this.useSIMD && this.simdMath) {
try {
const embeddings1 = texts1.map((text) => embeddingMap.get(text)!).filter(Boolean);
const embeddings2 = texts2.map((text) => embeddingMap.get(text)!).filter(Boolean);
if (embeddings1.length === texts1.length && embeddings2.length === texts2.length) {
const matrix = await this.simdMath.similarityMatrix(embeddings1, embeddings2);
this.performanceStats.totalSimilarityComputations += texts1.length * texts2.length;
this.performanceStats.totalSimilarityTime += performance.now() - simStartTime;
this.performanceStats.averageSimilarityTime =
this.performanceStats.totalSimilarityTime /
this.performanceStats.totalSimilarityComputations;
return matrix;
}
} catch (error) {
console.warn('SIMD matrix computation failed, falling back to JavaScript:', error);
}
}
// JavaScript 回退版本
const matrix: number[][] = [];
for (const textA of texts1) {
const row: number[] = [];
const embA = embeddingMap.get(textA);
if (!embA) {
console.warn(`Embedding not found for text1: "${textA}"`);
texts2.forEach(() => row.push(0));
matrix.push(row);
continue;
}
for (const textB of texts2) {
const embB = embeddingMap.get(textB);
if (!embB) {
console.warn(`Embedding not found for text2: "${textB}"`);
row.push(0);
continue;
}
row.push(this.cosineSimilarity(embA, embB));
}
matrix.push(row);
}
this.performanceStats.totalSimilarityComputations += texts1.length * texts2.length;
this.performanceStats.totalSimilarityTime += performance.now() - simStartTime;
this.performanceStats.averageSimilarityTime =
this.performanceStats.totalSimilarityTime / this.performanceStats.totalSimilarityComputations;
return matrix;
}
public cosineSimilarity(vecA: Float32Array, vecB: Float32Array): number {
if (!vecA || !vecB || vecA.length !== vecB.length) {
console.warn('Cosine similarity: Invalid vectors provided.', vecA, vecB);
return 0;
}
// 使用 SIMD 优化版本（如果可用）
if (this.useSIMD && this.simdMath) {
try {
// SIMD 版本是异步的，但为了保持接口兼容性，我们需要同步版本
// 这里我们回退到 JavaScript 版本，或者可以考虑重构为异步
return this.cosineSimilarityJS(vecA, vecB);
} catch (error) {
console.warn('SIMD cosine similarity failed, falling back to JavaScript:', error);
return this.cosineSimilarityJS(vecA, vecB);
}
}
return this.cosineSimilarityJS(vecA, vecB);
}
private cosineSimilarityJS(vecA: Float32Array, vecB: Float32Array): number {
let dotProduct = 0;
let normA = 0;
let normB = 0;
for (let i = 0; i < vecA.length; i++) {
dotProduct += vecA[i] * vecB[i];
normA += vecA[i] * vecA[i];
normB += vecB[i] * vecB[i];
}
const magnitude = Math.sqrt(normA) * Math.sqrt(normB);
return magnitude === 0 ? 0 : dotProduct / magnitude;
}
// 新增：异步 SIMD 优化的余弦相似度
public async cosineSimilaritySIMD(vecA: Float32Array, vecB: Float32Array): Promise<number> {
if (!vecA || !vecB || vecA.length !== vecB.length) {
console.warn('Cosine similarity: Invalid vectors provided.', vecA, vecB);
return 0;
}
if (this.useSIMD && this.simdMath) {
try {
return await this.simdMath.cosineSimilarity(vecA, vecB);
} catch (error) {
console.warn('SIMD cosine similarity failed, falling back to JavaScript:', error);
}
}
return this.cosineSimilarityJS(vecA, vecB);
}
public normalizeVector(vector: Float32Array): Float32Array {
let norm = 0;
for (let i = 0; i < vector.length; i++) norm += vector[i] * vector[i];
norm = Math.sqrt(norm);
if (norm === 0) return vector;
const normalized = new Float32Array(vector.length);
for (let i = 0; i < vector.length; i++) normalized[i] = vector[i] / norm;
return normalized;
}
public validateInput(text1: string, text2: string | 'valid_dummy'): void {
if (typeof text1 !== 'string' || (text2 !== 'valid_dummy' && typeof text2 !== 'string')) {
throw new Error('输入必须是字符串');
}
if (text1.trim().length === 0 || (text2 !== 'valid_dummy' && text2.trim().length === 0)) {
throw new Error('输入文本不能为空');
}
const roughCharLimit = this.config.maxLength * 5;
if (
text1.length > roughCharLimit ||
(text2 !== 'valid_dummy' && text2.length > roughCharLimit)
) {
console.warn('输入文本可能过长，将由分词器截断。');
}
}
private getCacheKey(text: string, _options: Record<string, any> = {}): string {
return text; // Options currently not used to vary embedding, simplify key
}
public getPerformanceStats(): Record<string, any> {
return {
...this.performanceStats,
cacheStats: {
...this.cacheStats,
embedding: {
...this.cacheStats.embedding,
hitRate:
this.cacheStats.embedding.hits + this.cacheStats.embedding.misses > 0
? this.cacheStats.embedding.hits /
(this.cacheStats.embedding.hits + this.cacheStats.embedding.misses)
: 0,
},
tokenization: {
...this.cacheStats.tokenization,
hitRate:
this.cacheStats.tokenization.hits + this.cacheStats.tokenization.misses > 0
? this.cacheStats.tokenization.hits /
(this.cacheStats.tokenization.hits + this.cacheStats.tokenization.misses)
: 0,
},
},
memoryPool: this.memoryPool.getStats(),
memoryUsage: this.getMemoryUsage(),
isInitialized: this.isInitialized,
isInitializing: this.isInitializing,
config: this.config,
pendingWorkerTasks: this.workerTaskQueue.length,
runningWorkerTasks: this.runningWorkerTasks,
};
}
private async waitForWorkerSlot(): Promise<void> {
return new Promise((resolve) => {
this.workerTaskQueue.push(resolve);
});
}
private processWorkerQueue(): void {
if (this.workerTaskQueue.length > 0 && this.runningWorkerTasks < this.config.concurrentLimit) {
const resolve = this.workerTaskQueue.shift();
if (resolve) resolve();
}
}
// 新增：获取 Worker 统计信息
public async getWorkerStats(): Promise<WorkerStats | null> {
if (!this.worker || !this.isInitialized) return null;
try {
const response = await this._sendMessageToWorker('getStats');
return response as WorkerStats;
} catch (error) {
console.warn('Failed to get worker stats:', error);
return null;
}
}
// 新增：清理 Worker 缓冲区
public async clearWorkerBuffers(): Promise<void> {
if (!this.worker || !this.isInitialized) return;
try {
await this._sendMessageToWorker('clearBuffers');
console.log('SemanticSimilarityEngine: Worker buffers cleared.');
} catch (error) {
console.warn('Failed to clear worker buffers:', error);
}
}
// 新增：清理所有缓存
public clearAllCaches(): void {
this.embeddingCache.clear();
this.tokenizationCache.clear();
this.cacheStats = {
embedding: { hits: 0, misses: 0, size: 0 },
tokenization: { hits: 0, misses: 0, size: 0 },
};
console.log('SemanticSimilarityEngine: All caches cleared.');
}
// 新增：获取内存使用情况
public getMemoryUsage(): {
embeddingCacheUsage: number;
tokenizationCacheUsage: number;
totalCacheUsage: number;
} {
const embeddingStats = this.embeddingCache.getStats();
const tokenizationStats = this.tokenizationCache.getStats();
return {
embeddingCacheUsage: embeddingStats.usage,
tokenizationCacheUsage: tokenizationStats.usage,
totalCacheUsage: (embeddingStats.usage + tokenizationStats.usage) / 2,
};
}
public async dispose(): Promise<void> {
console.log('SemanticSimilarityEngine: Disposing...');
// 清理 Worker 缓冲区
await this.clearWorkerBuffers();
if (this.worker) {
this.worker.terminate();
this.worker = null;
}
// 清理 SIMD 引擎
if (this.simdMath) {
this.simdMath.dispose();
this.simdMath = null;
}
this.tokenizer = null;
this.embeddingCache.clear();
this.tokenizationCache.clear();
this.memoryPool.clear();
this.pendingMessages.clear();
this.workerTaskQueue = [];
this.isInitialized = false;
this.isInitializing = false;
this.initPromise = null;
this.useSIMD = false;
console.log('SemanticSimilarityEngine: Disposed.');
}
}
```