diff --git a/README.md b/README.md index e461b0d..9d19665 100644 --- a/README.md +++ b/README.md @@ -8,10 +8,10 @@ A local context retrieval library that enables semantic search over your documen ## Features -- ๐Ÿ“„ **Multi-format Support**: Markdown, JSON, Text ๆ–‡ๆกฃ่‡ชๅŠจๅŠ ่ฝฝไธŽๅ‘้‡ๅŒ– -- ๐Ÿ” **Hybrid Retrieval**: ๅ‘้‡่ฏญไน‰ + FTS ๅ…จๆ–‡ๆฃ€็ดขๅŒ่ทฏๅฌๅ›ž๏ผŒRRF ่žๅˆๆŽ’ๅบ -- ๐Ÿ” **Two-stage Reranking**: KeywordReranker ็ฒพๆŽ’๏ผŒๅ…ณ้”ฎ่ฏๅ‘ฝไธญไผ˜ๅ…ˆ -- ๐ŸŒ **Query Expansion**: ็”จๆˆท่‡ชๅฎšไน‰ๅŒไน‰่ฏ่กจ๏ผŒCNโ†”EN ่ทจ่ฏญ่จ€ๅฌๅ›žๅขžๅผบ +- **Multi-format Loading**: Automatic parsing and vectorization of Markdown, JSON, and plain text files +- **Hybrid Search**: Combines semantic vectors with full-text search using RRF fusion for better recall +- **Two-stage Ranking**: Coarse vector search followed by keyword-based reranking for precision +- **Query Expansion**: Extends queries with user-defined synonym maps for cross-language and domain-specific matching ## Quick Start @@ -23,8 +23,8 @@ npm install @antv/context ```typescript import { Context } from '@antv/context'; -// Standard creation โ€” specify vectorsDir -const ctx = await Context.create({ vectorsDir: './vectors' }); +// Create context (vectorsDir is optional, defaults to .context/vectors) +const ctx = await Context.create(); // Load documents into a specific library with automatic vectorization await ctx.load('g2', './g2-docs/**/*.md'); @@ -45,7 +45,7 @@ await ctx.close(); | Parameter | Type | Default | Description | |-----------|------|---------|-------------| -| `vectorsDir` | `string` | โ€” | **Required**. Directory to store vector files | +| `vectorsDir` | `string` | `.context/vectors` | Directory to store vector files | | `basePath` | `string` | `process.cwd()` | Base path for resolving document IDs. Set for cross-machine consistent IDs. | | `onProgress` | `(phase, detail) => void` | โ€” | Progress callback for `load()` phases: `'load'` โ†’ `'embed'` โ†’ `'insert'`. | | `queryExpansion` | `QueryExpansionOptions | false` | `false` (no-op) | Query expansion with user-provided synonym map. `false` disables. Without `synonyms`, expansion is a no-op. | @@ -57,7 +57,7 @@ await ctx.close(); ```typescript const ctx = await Context.create({ - vectorsDir: './vectors', + vectorsDir: '.context/vectors', // Boost title matches 3x over content matches ftsFieldWeights: { content: 1, title: 3 }, // More "winner-takes-all" ranking @@ -69,7 +69,7 @@ const ctx = await Context.create({ ```typescript const ctx = await Context.create({ - vectorsDir: './vectors', + vectorsDir: '.context/vectors', // Define your own CNโ†”EN synonym bridges (no built-in defaults) queryExpansion: { synonyms: { @@ -82,7 +82,7 @@ const ctx = await Context.create({ // Disable query expansion entirely const ctxNoExpand = await Context.create({ - vectorsDir: './vectors', + vectorsDir: '.context/vectors', queryExpansion: false, }); ``` @@ -107,7 +107,7 @@ Load phases emit progress via the `onProgress` callback: ```typescript const ctx = await Context.create({ - vectorsDir: './vectors', + vectorsDir: '.context/vectors', onProgress: (phase, detail) => { console.log(`${phase}: ${detail.loaded}/${detail.total}`); }, @@ -138,9 +138,8 @@ Each result includes: | `id` | `string` | Document ID | | `content` | `string` | Document content | | `score` | `number` | Similarity score (0โ€“1) | -| `scoreMode` | `'vector' | 'hybrid' | 'reranked'` | How the score was computed | | `meta` | `Record` | Front-matter metadata (if present) | -| `sourceFilePath` | `string` | Original file path relative to `basePath` | +| `path` | `string` | Original file path relative to `basePath` | ### `ctx.close()` @@ -167,13 +166,13 @@ await ctx.close(); +----+-----+ +----+-----+ +----+-----+ +----+-----+ | | | | +--------------+--------------+ | - v v - +-----------------+ +-----------------+ - | FileLoader | | QueryExpander | + | | + +---------v-------+ +-------v----------+ + | FileLoader | | QueryExpander | +--------+--------+ | (SynonymExpander)| - | +--------+--------+ + | +--------+---------+ +--------v--------+ | - | EmbedBatch | v + | EmbedBatch | | +--------+--------+ +--------v--------+ | | Embedder | +--------v--------+ +--------+--------+ @@ -182,40 +181,31 @@ await ctx.close(); | Vectorize | +--------+--------+ | - +--------+-----------+ - | +-----------v-----------+ | | - +-------v-------+ +-------v-------+ + +-------v--------+ +-------v-------+ | FTS Text Path | | Vector Path | - |(ftsFieldWeights| | | - +-------+-------+ +-------+-------+ + | | | | + +-------+--------+ +-------+-------+ | | +-----------+-----------+ | +-----------v-----------+ - | RRF Fusion | - | (rankConstant) | + | RRF Fusion | + | (rankConstant) | +-----------+-----------+ | - +-----------v-----------+ + +-----------v------------+ | KeywordReranker | | (optional, 2nd stage) | - +-----------+-----------+ + +-----------+------------+ | + v Query Result +------------------------------------------------------------------------+ ``` -### Module Structure - -- **Public API**: `Context`, `QueryOptions`, `QueryResult`, `Document`, `Loader`, `MarkdownLoader`, `JsonLoader`, `TextLoader`, `pathToId` -- **Reranking**: `KeywordReranker`, `createReranker`, `Reranker`, `RerankCandidate`, `RerankResult`, `RerankOptions` -- **Query Expansion**: `SynonymExpander`, `NoopExpander`, `QueryExpander`, `QueryExpansionOptions` -- **Advanced API**: `Embedder`, `TransformersEmbedder`, `EmbedderManager`, `IZvecStore`, `ActualZvecStore`, `Store` - - ## License MIT \ No newline at end of file diff --git a/src/context.ts b/src/context.ts index cb4e5a8..943241a 100644 --- a/src/context.ts +++ b/src/context.ts @@ -5,246 +5,130 @@ import { ContextOptions, QueryOptions, QueryResult, - Document, - LoadPhase, - LoadProgress, + LoadedDoc, } from './types'; -import { resolveEmbedder } from './embedder'; -import type { Embedder, EmbedderInfo } from './embedder'; -import { Loader, MarkdownLoader, JsonLoader, TextLoader } from './loaders'; -import { Store } from './storage/store'; -import type { ZvecDoc } from './storage/zvec-store'; -import { pathToId } from './loaders/util'; +import { Embedder } from './embedder'; + +type EmbedderInfo = { dimensions: number }; +import { getLoader } from './loaders'; +import { Store } from './storage'; +import type { ZvecDoc } from './storage'; import { - createReranker, - SynonymExpander, - NoopExpander, - safeParseMeta, + safeJsonParse, computeContentHash, - selectSampleFiles, + loadSampleText, + detectTokenizer, + pathToId, } from './utils'; -import type { Reranker, RerankCandidate, QueryExpander } from './utils'; - -// --------------------------------------------------------------------------- -// Context class -// --------------------------------------------------------------------------- +import { expand } from './expander'; +import { applyRerank } from './reranker'; export class Context { - private readonly basePath: string; + private readonly options: ContextOptions; private readonly embedder: Embedder; - private readonly _embedderInfo: EmbedderInfo; + readonly embedderInfo: EmbedderInfo; private readonly store: Store; - private readonly loaders: Loader[]; - private readonly reranker: Reranker | null; - private readonly queryExpander: QueryExpander; - private readonly _onProgress?: (phase: LoadPhase, detail: LoadProgress) => void; private constructor(options: ContextOptions, embedder: Embedder, embedderInfo: EmbedderInfo) { - this.basePath = options.basePath ?? process.cwd(); + this.options = { + basePath: process.cwd(), + vectorsDir: '.context/vectors', + ...options, + }; this.embedder = embedder; - this._embedderInfo = embedderInfo; - this.store = new Store(options.vectorsDir, embedder, options); - this.loaders = [new MarkdownLoader(), new JsonLoader(), new TextLoader()]; - this.reranker = createReranker(options.rerankWeights); - this.queryExpander = - options.queryExpansion === false - ? new NoopExpander() - : new SynonymExpander( - options.queryExpansion && typeof options.queryExpansion === 'object' - ? options.queryExpansion.synonyms - : undefined, - ); - this._onProgress = options.onProgress; + this.embedderInfo = embedderInfo; + this.store = new Store(this.options.vectorsDir!, embedder, this.options); } static async create(options: ContextOptions): Promise { - const result = await resolveEmbedder(); - const embedder = result.embedder; - const embedderInfo = result.info; - - // Ensure vectors directory exists - if (!fs.existsSync(options.vectorsDir)) { - fs.mkdirSync(options.vectorsDir, { recursive: true }); - } - - const ctx = new Context(options, embedder, embedderInfo); + const embedder = new Embedder(); + await embedder.embed('probe'); - return ctx; - } + const vectorsDir = options.vectorsDir ?? '.context/vectors'; - /** - * Diagnostic information about the active embedder. - */ - get embedderInfo(): EmbedderInfo { - return this._embedderInfo; - } + if (!fs.existsSync(vectorsDir)) { + fs.mkdirSync(vectorsDir, { recursive: true }); + } - private getLoader(filePath: string): Loader | undefined { - return this.loaders.find((loader) => loader.canHandle(filePath)); + return new Context(options, embedder, { dimensions: embedder.dimensions }); } - /** - * Load documents into a library with automatic vectorization. - * - * Deduplication uses zvec's native document catalog as the single source - * of truth โ€” already-loaded documents whose content hasn't changed - * (same contentHash) are skipped. No separate registry file is needed. - * - * @param library Library name for organizing documents. - * @param pattern Glob pattern(s) matching files to load. - */ async load(library: string, pattern: string | string[]): Promise { const patterns = Array.isArray(pattern) ? pattern : [pattern]; const files = await glob(patterns, { absolute: true }); - // Sample multiple files for FTS tokenizer auto-detection. - let sampleText: string | undefined; - if (files.length > 0) { - try { - const sampleFiles = selectSampleFiles(files, 5); - const samples = await Promise.allSettled( - sampleFiles.map((f) => fs.promises.readFile(f, 'utf-8')), - ); - const validSamples = samples - .filter((r) => r.status === 'fulfilled') - .map((r) => (r as PromiseFulfilledResult).value); - if (validSamples.length > 0) { - sampleText = validSamples.join('\n'); - } - } catch { - // Sample failure is non-fatal โ€” fall back to default tokenizer - } - } - - await this.store.create(library, sampleText); - - // Internal type that extends Document with load-phase metadata. - interface LoadedDoc extends Document { - id: string; - contentHash: string; - sourceFilePath: string; - } + const sampleText = await loadSampleText(files); + const tokenizerName = sampleText ? detectTokenizer(sampleText) : 'jieba'; + this.store.acquireZvec(library, tokenizerName); - // Phase 1: Load all files concurrently and collect candidates. - const loadSettled = await Promise.allSettled( + const docs: LoadedDoc[] = await Promise.all( files.map(async (filePath) => { - const loader = this.getLoader(filePath); - if (!loader) return null; + const loader = getLoader(filePath); + if (!loader) throw new Error(`No loader for ${filePath}`); const doc = await loader.load(filePath); - const relativePath = path.relative(this.basePath, filePath); - const docId = pathToId(relativePath); - const contentHash = computeContentHash(doc.content); - - return { ...doc, id: docId, contentHash, sourceFilePath: relativePath }; + const relativePath = path.relative(this.options.basePath!, filePath); + return { + ...doc, + id: pathToId(relativePath), + contentHash: computeContentHash(doc.content), + path: relativePath, + }; }), ); - const candidates: LoadedDoc[] = []; - let failCount = 0; - for (const r of loadSettled) { - if (r.status === 'fulfilled' && r.value !== null) { - candidates.push(r.value); - } else if (r.status === 'rejected') { - failCount++; - } - } - if (failCount > 0) { - console.warn( - `[context] ${failCount}/${files.length} file(s) failed to load in library "${library}" and were skipped.`, - ); - } - - if (candidates.length === 0) return; - - // Phase 1b: Dedup via zvec โ€” batch-fetch stored contentHashes. - // Zvec is the single source of truth; no separate registry needed. - const candidateIds = candidates.map((d) => d.id); - const existing = await this.store.fetchDocs(library, candidateIds, ['contentHash']); + const docIds = docs.map((d) => d.id); + const existing = await this.store.fetchDocs(library, docIds, ['contentHash']); - const docsToEmbed: LoadedDoc[] = []; - for (const doc of candidates) { - const stored = existing[doc.id]; - if (!stored) { - // New document โ€” embed and insert - docsToEmbed.push(doc); - } else if (stored.fields.contentHash !== doc.contentHash) { - // Content changed โ€” re-embed and update - docsToEmbed.push(doc); - } - // else: contentHash matches โ†’ skip (already up to date) - } + const docsToEmbed = docs.filter( + (doc) => !existing[doc.id] || existing[doc.id].fields.contentHash !== doc.contentHash, + ); - // Progress: load phase complete - if (this._onProgress) { - this._onProgress('load', { loaded: docsToEmbed.length, total: files.length }); + if (this.options.onProgress) { + this.options.onProgress('load', { loaded: docsToEmbed.length, total: files.length }); } if (docsToEmbed.length === 0) return; - // Phase 2: Batch embed all items for better performance const contents = docsToEmbed.map((doc) => doc.content); const vectors = await this.embedder.embedBatch(contents); - // Progress: embed phase complete - if (this._onProgress) { - this._onProgress('embed', { loaded: vectors.length, total: docsToEmbed.length }); + if (this.options.onProgress) { + this.options.onProgress('embed', { loaded: vectors.length, total: docsToEmbed.length }); } - // Phase 3: Batch insert into store (upsert โ€” handles both new and changed docs) const zvecDocs: ZvecDoc[] = docsToEmbed.map((doc, index) => ({ id: doc.id, vector: vectors[index], fields: { content: doc.content, - meta: doc.meta && Object.keys(doc.meta).length > 0 ? JSON.stringify(doc.meta) : '', - sourceFilePath: doc.sourceFilePath ?? '', + meta: doc.meta ? JSON.stringify(doc.meta) : '', + path: doc.path, contentHash: doc.contentHash, }, })); await this.store.addDoc(library, zvecDocs); - // Progress: insert phase complete - if (this._onProgress) { - this._onProgress('insert', { loaded: zvecDocs.length, total: docsToEmbed.length }); + if (this.options.onProgress) { + this.options.onProgress('insert', { loaded: zvecDocs.length, total: docsToEmbed.length }); } } - /** - * Query documents by semantic similarity. - * - * Default mode is `'hybrid'` which combines vector similarity + FTS text - * matching via RRF fusion for better recall. Use `mode: 'vector'` for - * pure semantic search when FTS is not needed. - * - * @param text Query text. - * @param options Query options (library, topK, mode). - * @returns Ranked results with content and score. - */ async query(text: string, options: QueryOptions): Promise { - const library = options.library; + const opt = { topK: 5, mode: 'hybrid' as const, rerank: { rerankFactor: 3, minCandidates: 10 }, ...options }; + + const { library, topK, mode, rerank } = opt; // Expand the query with synonyms / cross-language bridging terms. - // The expanded text is used for both embedding and FTS, so a single - // expanded query can match CN and EN documents simultaneously. - const expandedText = this.queryExpander.expand(text); + const expandedText = expand(text, this.options.queryExpansion); + const vector = await this.embedder.embed(expandedText); - const topK = options.topK ?? 5; - const mode = options.mode ?? 'hybrid'; - const rerankEnabled = options.rerank !== false; - // Reranking pipeline: pull extra candidates from the coarse search so the - // reranker has a larger pool to select from. Without reranking, search - // exactly topK for efficiency. - const rerankFactor = - (options.rerank && typeof options.rerank === 'object' - ? options.rerank.rerankFactor - : undefined) ?? 3; - const minCandidates = - (options.rerank && typeof options.rerank === 'object' - ? options.rerank.minCandidates - : undefined) ?? 10; + const rerankConfig = typeof rerank === 'object' ? rerank : null; + const rerankEnabled = rerankConfig !== null; + const rerankFactor = rerankConfig?.rerankFactor ?? 3; + const minCandidates = rerankConfig?.minCandidates ?? 10; const searchTopK = rerankEnabled ? Math.max(topK * rerankFactor, minCandidates) : topK; const searchResults = await this.store.queryDoc(library, { @@ -255,58 +139,29 @@ export class Context { filter: options.filter, }); - if (searchResults.length === 0) return []; - - const allResults: QueryResult[] = searchResults.map((result) => { + const queryResults: QueryResult[] = searchResults.map((result) => { const content = String(result.fields?.content ?? ''); const metaStr = result.fields?.meta as string | undefined; - const meta = safeParseMeta(metaStr); + const meta = safeJsonParse(metaStr) as Record | undefined; return { id: result.id, content, score: result.score, - scoreMode: mode === 'hybrid' ? ('hybrid' as const) : ('vector' as const), meta, - sourceFilePath: result.fields?.sourceFilePath as string | undefined, - embedderKind: this._embedderInfo.kind, + path: result.fields?.path as string | undefined, }; }); - // Stage 2: Rerank the candidate pool for precision (when enabled). - // Pull extra candidates, re-score each against the query, then keep topK. - if (rerankEnabled && allResults.length > topK) { - const candidates: RerankCandidate[] = allResults.map((r) => ({ - id: r.id, - content: r.content, - score: r.score, - })); - - const reranked = await this.reranker!.rerank(text, candidates); - - // Merge reranked scores back into results - const scoreMap = new Map(reranked.map((r) => [r.id, r.score])); - for (const result of allResults) { - const newScore = scoreMap.get(result.id); - if (newScore !== undefined) { - result.score = newScore; - result.scoreMode = 'reranked'; - } - } + if (rerankEnabled) { + await applyRerank(this.options.rerankWeights, text, queryResults, topK); } - // Final sort by (possibly reranked) score and return topK - allResults.sort((a, b) => b.score - a.score); - return allResults.slice(0, topK); + queryResults.sort((a, b) => b.score - a.score); + return queryResults.slice(0, topK); } - /** - * Close all stores and release resources. - * - * Call this when you are done using the Context instance (e.g. at process - * exit or before re-creating a new instance). - */ async close(): Promise { - await this.store.closeAll(); + this.store.close(); } -} +} \ No newline at end of file diff --git a/src/embedder/embedder.ts b/src/embedder/embedder.ts new file mode 100644 index 0000000..a5cbc0f --- /dev/null +++ b/src/embedder/embedder.ts @@ -0,0 +1,38 @@ +/** + * Embedder โ€” local sentence-transformers model via @huggingface/transformers. + */ + +const DEFAULT_MODEL_ID = 'onnx-community/bge-small-zh-v1.5-ONNX'; +const DEFAULT_DIMENSIONS = 512; + +type TransformersPipeline = (texts: string[], options: Record) => Promise<{ tolist(): number[][] }>; + +export class Embedder { + private static pipeline: TransformersPipeline | null = null; + + readonly dimensions = DEFAULT_DIMENSIONS; + + private static async getPipeline(): Promise { + if (Embedder.pipeline) return Embedder.pipeline; + + const mod = await import('@huggingface/transformers'); + const hfEndpoint = process.env.HF_ENDPOINT; + if (hfEndpoint && (mod as any).env) { + (mod as any).env.remoteHost = hfEndpoint; + } + + const pipe = await mod.pipeline('feature-extraction', DEFAULT_MODEL_ID) as TransformersPipeline; + Embedder.pipeline = pipe; + return Embedder.pipeline; + } + + async embed(text: string): Promise { + return (await this.embedBatch([text]))[0]; + } + + async embedBatch(texts: string[]): Promise { + const pipe = await Embedder.getPipeline(); + const outputs = await pipe(texts, { pooling: 'mean', normalize: true }); + return outputs.tolist(); + } +} diff --git a/src/embedder/index.ts b/src/embedder/index.ts index 9d3ba7b..5b49951 100644 --- a/src/embedder/index.ts +++ b/src/embedder/index.ts @@ -1,29 +1 @@ -/** - * embedder โ€” aggregate entry point for all embedding modules. - * - * Re-exports from split files: - * types.ts โ†’ Embedder interface - * transformers.ts โ†’ TransformersEmbedder - * manager.ts โ†’ EmbedderManager, getEmbedder, resetEmbedder - * resolve.ts โ†’ resolveEmbedder - * - * Language detection & tokenizer utilities have moved to src/utils/tokenizer.ts. - * They are still available via the embedder module for backward compatibility. - */ - -// Types -export type { Embedder } from './types'; - -// Language detection & CJK utilities โ€” re-exported from utils for backward compatibility -export { isCJK, detectLanguage, tokenizerForLanguage, detectTokenizer } from '../utils/tokenizer'; -export type { LanguageHint } from '../utils/tokenizer'; - -// TransformersEmbedder โ€” production-quality model embedder -export { TransformersEmbedder } from './transformers'; - -// EmbedderManager & global convenience functions -export { EmbedderManager, getEmbedder, resetEmbedder } from './manager'; - -// Embedder resolution -export { resolveEmbedder } from './resolve'; -export type { EmbedderInfo, EmbedderKind } from './resolve'; \ No newline at end of file +export { Embedder } from './embedder'; \ No newline at end of file diff --git a/src/embedder/manager.ts b/src/embedder/manager.ts deleted file mode 100644 index d66212a..0000000 --- a/src/embedder/manager.ts +++ /dev/null @@ -1,134 +0,0 @@ -/** - * EmbedderManager โ€” encapsulates embedder resolution and lifecycle. - * - * Each instance maintains its own embedder cache and (optionally) its own - * transformers module loader, enabling true instance-level isolation for - * multi-tenant servers and test environments. - */ - -import { Embedder } from './types'; -import { - TransformersEmbedder, - loadTransformersModule, - resetTransformersModule, - createTransformersLoader, -} from './transformers'; -import type { TransformersLoader } from './transformers'; - -// --------------------------------------------------------------------------- -// EmbedderManager -// --------------------------------------------------------------------------- - -export interface EmbedderManagerOptions { - /** - * Custom transformers module loader for instance-level isolation. - * When omitted, uses the shared global loader (backward-compatible). - */ - transformersLoader?: TransformersLoader; -} - -export class EmbedderManager { - private _defaultEmbedder: Embedder | null = null; - private readonly _loader: TransformersLoader; - - constructor(options?: EmbedderManagerOptions) { - if (options?.transformersLoader) { - this._loader = options.transformersLoader; - } else { - // Default: use global shared loader (backward-compatible) - this._loader = { - load: loadTransformersModule, - reset: resetTransformersModule, - }; - } - } - - /** - * Create an EmbedderManager with a fully isolated transformers loader. - * - * The returned manager has its own module cache and failure state, - * completely independent of other managers and the global state. - */ - static createIsolated(): EmbedderManager { - const loader = createTransformersLoader(); - return new EmbedderManager({ transformersLoader: loader }); - } - - /** - * Return a shared Embedder instance (async). - * - * Requires @huggingface/transformers to be installed and the model to be - * loadable. Throws an error when either condition is not met. - */ - async getEmbedder(): Promise { - if (this._defaultEmbedder) return this._defaultEmbedder; - - const t = await this._loader.load(); - if (!t) { - throw new Error( - '@huggingface/transformers is not installed. Semantic search requires a model-based embedder.\n' + - ' Install it with:\n' + - ' npm install @huggingface/transformers\n' + - ' Then download the model:\n' + - ' node scripts/download-model.mjs\n' + - ' Or set mirror for China:\n' + - ' export HF_ENDPOINT=https://hf-mirror.com', - ); - } - - try { - const probe = new TransformersEmbedder(() => this._loader.load()); - await probe.embed('probe'); - this._defaultEmbedder = probe; - } catch (err) { - throw new Error( - `Failed to load embedding model (bge-small-zh-v1.5): ${(err as Error).message?.split('\n')[0] ?? 'unknown'}\n` + - ' To fix model download:\n' + - ' 1. Set mirror: export HF_ENDPOINT=https://hf-mirror.com\n' + - ' 2. Manual download: node scripts/download-model.mjs', - ); - } - return this._defaultEmbedder; - } - - /** - * Force-reset all cached state (useful for tests). - * - * Only resets the owned loader if this manager was created via - * `createIsolated()`. Global-loader managers delegate to the shared reset. - */ - reset(): void { - this._defaultEmbedder = null; - this._loader.reset(); - } -} - -// --------------------------------------------------------------------------- -// Global convenience functions (backward-compatible, prefer EmbedderManager) -// --------------------------------------------------------------------------- - -/** Global manager instance for backward-compatible convenience functions. */ -const _globalManager = new EmbedderManager(); - -/** - * Return a shared Embedder instance (async) via the global manager. - * - * @deprecated Prefer `new EmbedderManager().getEmbedder()` or pass an - * explicit embedder to `Context.create({ embedder })` to avoid hidden - * global state. This function remains for backward compatibility. - * - */ -export async function getEmbedder(): Promise { - return _globalManager.getEmbedder(); -} - -/** - * Force-reset the global cached default embedder (useful for tests). - * - * @deprecated Prefer `new EmbedderManager().reset()` on a manager instance - * to avoid affecting other Context instances. This function remains for - * backward compatibility and test cleanup. - */ -export function resetEmbedder(): void { - _globalManager.reset(); -} diff --git a/src/embedder/resolve.ts b/src/embedder/resolve.ts deleted file mode 100644 index d0076af..0000000 --- a/src/embedder/resolve.ts +++ /dev/null @@ -1,81 +0,0 @@ -/** - * Embedder resolution โ€” resolves the TransformersEmbedder from the - * configured model. Throws an error when the model cannot be loaded - * or @huggingface/transformers is not installed. - */ - -import { Embedder } from './types'; -import { TransformersEmbedder, loadTransformersModule } from './transformers'; - -// --------------------------------------------------------------------------- -// Embedder type info -// --------------------------------------------------------------------------- - -/** Describes the kind of embedder being used. */ -export type EmbedderKind = 'transformers'; - -/** Diagnostic information about the resolved embedder. */ -export interface EmbedderInfo { - /** Which embedder implementation is active. */ - kind: EmbedderKind; - /** Vector dimensions of the active embedder. */ - dimensions: number; - /** Model ID (only set for TransformersEmbedder). */ - modelId?: string; -} - -// --------------------------------------------------------------------------- -// Resolve result -// --------------------------------------------------------------------------- - -/** The result of embedder resolution โ€” embedder instance + diagnostic info. */ -export interface ResolveResult { - embedder: Embedder; - info: EmbedderInfo; -} - -// --------------------------------------------------------------------------- -// Public helpers -// --------------------------------------------------------------------------- - -/** - * Resolve which Embedder to use based on the model option. - * - * Requires @huggingface/transformers to be installed and the model to be - * loadable. Throws a descriptive error when either condition is not met. - */ -export async function resolveEmbedder(model?: string): Promise { - const t = await loadTransformersModule(); - - if (!t) { - throw new Error( - '@huggingface/transformers is not installed. Semantic search requires a model-based embedder.\n' + - ' Install it with:\n' + - ' npm install @huggingface/transformers\n' + - ' Then download the model:\n' + - ' node scripts/download-model.mjs\n' + - ' Or set mirror for China:\n' + - ' export HF_ENDPOINT=https://hf-mirror.com' - ); - } - - try { - const embedder = new TransformersEmbedder(undefined, { modelId: model }); - await embedder.embed('probe'); - return { - embedder, - info: { - kind: 'transformers', - dimensions: embedder.dimensions, - modelId: model ?? 'onnx-community/bge-small-zh-v1.5-ONNX', - }, - }; - } catch (err) { - throw new Error( - `Failed to load embedding model (${model ?? 'bge-small-zh-v1.5'}): ${(err as Error).message?.split('\n')[0] ?? 'unknown'}\n` + - ' To fix model download:\n' + - ' 1. Set mirror: export HF_ENDPOINT=https://hf-mirror.com\n' + - ' 2. Manual download: node scripts/download-model.mjs' - ); - } -} diff --git a/src/embedder/transformers.ts b/src/embedder/transformers.ts deleted file mode 100644 index f70fd50..0000000 --- a/src/embedder/transformers.ts +++ /dev/null @@ -1,225 +0,0 @@ -/** - * TransformersEmbedder โ€” local sentence-transformers model via @huggingface/transformers. - * - * Constructor is cheap โ€“ the model is loaded lazily on the first embed() call. - * Uses the provided EmbedderManager for module loading (or the global one). - */ - -import { Embedder } from './types'; - -// --------------------------------------------------------------------------- -// Constants -// --------------------------------------------------------------------------- - -const DEFAULT_MODEL_ID = 'onnx-community/bge-small-zh-v1.5-ONNX'; -const DEFAULT_TRANSFORMERS_DIMS = 512; - -/** Cooldown period (ms) before retrying Transformers module load after failure. */ -const TRANSFORMERS_RETRY_COOLDOWN_MS = 60_000; // 1 minute - -// --------------------------------------------------------------------------- -// Minimal type interfaces for @huggingface/transformers -// --------------------------------------------------------------------------- - -/** Transformers module type โ€” minimal interface we rely on. */ -interface TransformersModule { - pipeline(task: string, modelId: string): Promise; - env?: { remoteHost?: string }; -} - -/** Transformers pipeline type โ€” minimal interface we rely on. */ -interface TransformersPipeline { - (texts: string[], options: Record): Promise; -} - -/** Transformers output type โ€” minimal interface we rely on. */ -interface TransformersOutput { - tolist(): number[][]; -} - -// --------------------------------------------------------------------------- -// Module loading -// --------------------------------------------------------------------------- - -/** - * Internal module cache โ€” shared across all instances by default. - * - * Node.js `import()` already caches ESM modules, so this additional cache - * avoids repeated async overhead and failure-tracking logic. The cache is - * module-level because the underlying native/WASM module can only be loaded - * once per process; multiple Context instances correctly share it. - * - * For test isolation, use `createTransformersLoader()` to get an independent - * loader with its own cache, or call `resetTransformersModule()` between tests. - */ -let _transformersModule: TransformersModule | undefined; -let _transformersLoadFailed = false; -let _transformersLoadFailedAt = 0; - -/** - * A self-contained transformers module loader with its own cache. - * - * Use this when you need instance-level isolation (e.g. tests, multi-tenant - * servers). Each loader maintains independent failure tracking and retry state. - */ -export interface TransformersLoader { - load(): Promise; - reset(): void; -} - -/** - * Create an independent transformers module loader. - * - * The returned loader has its own cache and failure state, completely - * isolated from the global `loadTransformersModule()` / `resetTransformersModule()`. - */ -export function createTransformersLoader(): TransformersLoader { - let localModule: TransformersModule | undefined; - let localFailed = false; - let localFailedAt = 0; - - return { - async load(): Promise { - if (localModule) return localModule; - - if (localFailed) { - const elapsed = Date.now() - localFailedAt; - if (elapsed < TRANSFORMERS_RETRY_COOLDOWN_MS) return undefined; - localFailed = false; - } - - try { - const mod = await import('@huggingface/transformers'); - localModule = mod as TransformersModule; - } catch { - localFailed = true; - localFailedAt = Date.now(); - return undefined; - } - - const hfEndpoint = process.env.HF_ENDPOINT; - if (hfEndpoint && localModule?.env) { - localModule.env.remoteHost = hfEndpoint; - } - - return localModule; - }, - - reset(): void { - localModule = undefined; - localFailed = false; - localFailedAt = 0; - }, - }; -} - -/** - * Load the @huggingface/transformers module with TTL-based retry. - * - * Uses the shared module-level cache. For instance-level isolation, - * use `createTransformersLoader()` instead. - */ -export async function loadTransformersModule(): Promise { - if (_transformersModule) return _transformersModule; - - // TTL-based retry: allow re-attempting after cooldown period - if (_transformersLoadFailed) { - const elapsed = Date.now() - _transformersLoadFailedAt; - if (elapsed < TRANSFORMERS_RETRY_COOLDOWN_MS) { - return undefined; - } - // Cooldown expired โ€” reset failure flag and try again - _transformersLoadFailed = false; - } - - try { - // @huggingface/transformers v4 is an ESM-first package (type: "module"). - // Dynamic import() loads the proper ESM bundle where `pipeline` is a - // real async function. - const mod = await import('@huggingface/transformers'); - _transformersModule = mod as TransformersModule; - } catch { - _transformersLoadFailed = true; - _transformersLoadFailedAt = Date.now(); - return undefined; - } - - // Apply HF_ENDPOINT mirror if set (e.g. https://hf-mirror.com for China). - // @huggingface/transformers v4 does NOT read HF_ENDPOINT automatically; - // it hardcodes "https://huggingface.co/" as env.remoteHost. - const hfEndpoint = process.env.HF_ENDPOINT; - if (hfEndpoint && _transformersModule?.env) { - _transformersModule.env.remoteHost = hfEndpoint; - } - - return _transformersModule; -} - -/** Reset the shared transformers module cache (for tests). */ -export function resetTransformersModule(): void { - _transformersModule = undefined; - _transformersLoadFailed = false; - _transformersLoadFailedAt = 0; -} - -// --------------------------------------------------------------------------- -// TransformersEmbedder -// --------------------------------------------------------------------------- - -export class TransformersEmbedder implements Embedder { - readonly dimensions: number; - private _modelId: string; - private _pipeline: TransformersPipeline | null = null; - private _loadPromise: Promise | null = null; - private _loadTransformers: () => Promise; - - constructor( - loadTransformers?: () => Promise, - options?: { modelId?: string; dimensions?: number } - ) { - this._loadTransformers = loadTransformers ?? loadTransformersModule; - this._modelId = options?.modelId ?? DEFAULT_MODEL_ID; - this.dimensions = options?.dimensions ?? DEFAULT_TRANSFORMERS_DIMS; - } - - private async _getPipeline(): Promise { - if (this._pipeline) { - return this._pipeline; - } - - if (!this._loadPromise) { - this._loadPromise = (async () => { - const t = await this._loadTransformers(); - if (!t) { - throw new Error( - '@huggingface/transformers is not installed. Install it with:\n' + - ' npm install @huggingface/transformers\n' + - ' Or set HF_ENDPOINT for mirror download.' - ); - } - this._pipeline = await t.pipeline('feature-extraction', this._modelId); - return this._pipeline; - })().catch((err) => { - // Reset so subsequent calls can retry instead of returning a - // forever-rejected cached promise. - this._loadPromise = null; - throw err; - }); - } - - return this._loadPromise; - } - - async embed(text: string): Promise { - return (await this.embedBatch([text]))[0]; - } - - async embedBatch(texts: string[]): Promise { - const pipe = await this._getPipeline(); - const outputs = await pipe(texts, { - pooling: 'mean', - normalize: true - }); - return outputs.tolist(); - } -} diff --git a/src/embedder/types.ts b/src/embedder/types.ts deleted file mode 100644 index 4d7e711..0000000 --- a/src/embedder/types.ts +++ /dev/null @@ -1,9 +0,0 @@ -/** - * Embedder interface โ€” text-to-vector conversion contract. - */ - -export interface Embedder { - readonly dimensions: number; - embed(text: string): Promise; - embedBatch(texts: string[]): Promise; -} diff --git a/src/expander/expand.ts b/src/expander/expand.ts new file mode 100644 index 0000000..2455dc7 --- /dev/null +++ b/src/expander/expand.ts @@ -0,0 +1,54 @@ +/** + * Synonym query expander โ€” pure function. + */ + +import type { QueryExpansionOptions } from '../types'; +import { containsCJK } from '../utils'; + +const WORD_BOUNDARY_RE = /[\s,๏ผŒใ€‚.!๏ผ?๏ผŸ;๏ผ›:๏ผš\(\)\[\]{}""''\"\'\-_\/\\|@#$%^&*+=<>~`]/; + +function containsTerm(text: string, term: string): boolean { + if (!text.includes(term)) return false; + + // CJK terms use substring match + if (containsCJK(term)) return true; + + // Latin terms use word boundary match + let searchFrom = 0; + while (searchFrom <= text.length - term.length) { + const idx = text.indexOf(term, searchFrom); + if (idx === -1) return false; + + const prevOk = idx === 0 || WORD_BOUNDARY_RE.test(text[idx - 1]); + const afterIdx = idx + term.length; + const afterOk = afterIdx >= text.length || WORD_BOUNDARY_RE.test(text[afterIdx]); + + if (prevOk && afterOk) return true; + searchFrom = idx + 1; + } + return false; +} + +export function expand(query: string, queryExpansion?: QueryExpansionOptions | false): string { + if (queryExpansion === false) return query; + const synonyms = queryExpansion?.synonyms; + if (!synonyms || Object.keys(synonyms).length === 0) return query; + + const queryLower = query.toLowerCase().trim(); + const additions = new Set(); + + for (const [term, syns] of Object.entries(synonyms as Record)) { + const termLower = term.toLowerCase(); + + if (containsTerm(queryLower, termLower)) { + for (const syn of syns) { + if (!containsTerm(queryLower, syn.toLowerCase())) { + additions.add(syn); + } + } + } + } + + if (additions.size === 0) return query; + return `${query} ${[...additions].join(' ')}`; +} \ No newline at end of file diff --git a/src/expander/index.ts b/src/expander/index.ts new file mode 100644 index 0000000..dd613e7 --- /dev/null +++ b/src/expander/index.ts @@ -0,0 +1 @@ +export { expand } from './expand'; \ No newline at end of file diff --git a/src/index.ts b/src/index.ts index 57866cf..2e2f5f6 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,66 +1,2 @@ -// --------------------------------------------------------------------------- -// Public API โ€” main entry points for typical usage -// --------------------------------------------------------------------------- - export { Context } from './context'; -export * from './types'; - -// Loaders โ€” file format handlers (commonly extended by users) -export { - Loader, - MarkdownLoader, - JsonLoader, - TextLoader, -} from './loaders'; -export { pathToId } from './loaders/util'; - -// Reranker โ€” two-stage retrieval precision scoring -export { KeywordReranker, createReranker } from './utils/reranker'; -export type { Reranker, RerankCandidate, RerankResult } from './utils/reranker'; - -// Query expansion โ€” synonym bridging for cross-language recall -export { SynonymExpander, NoopExpander } from './utils/expander'; -export type { QueryExpander } from './utils/expander'; - -// --------------------------------------------------------------------------- -// Advanced API โ€” for extending or customizing internals -// --------------------------------------------------------------------------- - -// Embedder โ€” custom embedding strategies -export { - TransformersEmbedder, - EmbedderManager, - getEmbedder, - resetEmbedder, - isCJK, - detectLanguage, - tokenizerForLanguage, - detectTokenizer, -} from './embedder'; -export type { Embedder, LanguageHint, EmbedderInfo, EmbedderKind } from './embedder'; - -// Zvec store โ€” custom vector storage backends -export { - ActualZvecStore, - createZvecStore, - openZvecStore, - openZvecStoreSync, - isZvecAvailable, - buildZvecSchema, -} from './storage/zvec-store'; -export type { - IZvecStore, - ZvecDoc, - ZvecQueryResult, - ZvecSearchParams, - ZvecHybridParams, - ZvecFieldSchema, - ZvecStoreConfig, - ActualZvecStoreOptions, -} from './storage/zvec-store'; - -// Store โ€” zvec store lifecycle management -export { Store } from './storage/store'; -export type { StoreQueryParams } from './storage/store'; - - +export * from './types'; \ No newline at end of file diff --git a/src/loaders/index.ts b/src/loaders/index.ts index 799617e..7fa09cb 100644 --- a/src/loaders/index.ts +++ b/src/loaders/index.ts @@ -1,5 +1,10 @@ -export { Loader } from './base'; -export { MarkdownLoader } from './markdown'; -export { JsonLoader } from './json'; -export { TextLoader } from './text'; -export { pathToId } from './util'; \ No newline at end of file +import { Loader } from './base'; +import { MarkdownLoader } from './markdown'; +import { JsonLoader } from './json'; +import { TextLoader } from './text'; + +const loaders = [new MarkdownLoader(), new JsonLoader(), new TextLoader()]; + +export function getLoader(filePath: string): Loader | undefined { + return loaders.find((l) => l.canHandle(filePath)); +} \ No newline at end of file diff --git a/src/loaders/util.ts b/src/loaders/util.ts deleted file mode 100644 index 1c17ced..0000000 --- a/src/loaders/util.ts +++ /dev/null @@ -1,33 +0,0 @@ -import * as path from 'path'; -import { computeContentHash } from '../utils/hash'; - -/** - * Convert a file path to a safe, collision-resistant ID for zvec. - * - * zvec doc IDs have a 64-character limit and only allow alphanumeric + underscore. - * Strategy: hash the full path to guarantee uniqueness and fit within the limit, - * while appending a short readable suffix for debugging. - * - * Cross-platform consistency: paths are normalized to forward slashes before - * hashing, so the same relative path produces the same ID on Windows, macOS, - * and Linux. - * - * Example: "/long/path/getting-started.md" โ†’ "f3a1b2c4__getting_started" - */ -export function pathToId(filePath: string): string { - // Normalize separators to forward slash for cross-platform consistency. - // Also collapse redundant segments (e.g. "a/../b" โ†’ "b") via path.normalize. - const normalized = path.normalize(filePath).replace(/\\/g, '/').replace(/^\/+/, ''); - - // Generate a compact hash of the full path (16 hex chars = 64-bit) - const hash = computeContentHash(normalized); - - // Derive a short readable suffix from the filename (without extension) - const basename = normalized.split('/').pop() ?? ''; - const suffix = basename - .replace(/\.[a-zA-Z0-9]+$/, '') // remove extension - .replace(/[^a-zA-Z0-9]/g, '_') // sanitize - .slice(0, 20); // truncate to keep total under 64 - - return `${hash}__${suffix}`; -} diff --git a/src/reranker/helpers.ts b/src/reranker/helpers.ts new file mode 100644 index 0000000..48db010 --- /dev/null +++ b/src/reranker/helpers.ts @@ -0,0 +1,38 @@ +import { containsCJK } from '../utils'; + +const WORD_BOUNDARY = /[\s\n.,;:!?๏ผŒใ€‚๏ผ๏ผŸใ€๏ผš๏ผ›"'(๏ผˆใ€ใ€Š\-_]/; + +export function tokenizeQuery(query: string): string[] { + const tokens = query.split(/[\s,๏ผŒใ€‚.!๏ผ๏ผŸ?ใ€๏ผš:๏ผ›;]+/).filter(Boolean); + const result: string[] = []; + for (const r of tokens) { + result.push(r); + if (containsCJK(r)) { + for (let i = 0; i + 2 <= r.length; i++) result.push(r.slice(i, i + 2)); + } + } + return [...new Set(result)]; +} + +export function countOccurrences(text: string, sub: string): number { + let count = 0, pos = 0; + while ((pos = text.indexOf(sub, pos)) !== -1) { count++; pos += sub.length; } + return count; +} + +export function countTermMatches(text: string, term: string): number { + let count = 0, pos = 0; + while ((pos = text.indexOf(term, pos)) !== -1) { + if (isWordBoundary(text, term, pos)) count++; + pos += term.length; + } + return count; +} + +export function isWordBoundary(text: string, term: string, pos?: number): boolean { + const idx = pos ?? text.indexOf(term); + if (idx === -1) return false; + const before = idx === 0 || WORD_BOUNDARY.test(text[idx - 1]); + const after = idx + term.length >= text.length || WORD_BOUNDARY.test(text[idx + term.length]); + return before && after; +} \ No newline at end of file diff --git a/src/reranker/index.ts b/src/reranker/index.ts new file mode 100644 index 0000000..f1e4c66 --- /dev/null +++ b/src/reranker/index.ts @@ -0,0 +1,35 @@ +import type { RerankCandidate, RerankOptions } from './types'; +import type { QueryResult } from '../types'; +import { rerank } from './reranker'; + +export type { RerankOptions }; + +/** + * Apply reranking to query results. + */ +export async function applyRerank( + rerankOptions: RerankOptions | undefined, + query: string, + results: QueryResult[], + topK: number, +): Promise { + if (results.length <= topK) return results; + + const candidates: RerankCandidate[] = results.map((r) => ({ + id: r.id, + content: r.content, + score: r.score, + })); + + const reranked = await rerank(query, candidates, rerankOptions); + + const scoreMap = new Map(reranked.map((r) => [r.id, r.score])); + for (const result of results) { + const newScore = scoreMap.get(result.id); + if (newScore !== undefined) { + result.score = newScore; + } + } + + return results; +} \ No newline at end of file diff --git a/src/reranker/reranker.ts b/src/reranker/reranker.ts new file mode 100644 index 0000000..35c0179 --- /dev/null +++ b/src/reranker/reranker.ts @@ -0,0 +1,81 @@ +import type { RerankCandidate, RerankResult, RerankOptions } from './types'; +import { tokenizeQuery, countOccurrences, countTermMatches, isWordBoundary } from './helpers'; + +const DEFAULTS = { + phraseWeight: 3.0, + phraseRepeatBonus: 0.5, + termWeight: 1.0, + termRepeatBonus: 0.2, + substringWeight: 0.3, + headingTermBonus: 2.0, + headingPhraseBonus: 2.5, + originalScoreCarry: 0.1, +} as const; + +/** + * Rerank candidates by keyword / phrase overlap with the query. + * Scores are normalised to [0, 1] via min-max scaling. + */ +export async function rerank( + query: string, + candidates: RerankCandidate[], + options?: RerankOptions, +): Promise { + if (candidates.length === 0) return []; + + const w = { + phraseWeight: options?.phraseWeight ?? DEFAULTS.phraseWeight, + phraseRepeatBonus: options?.phraseRepeatBonus ?? DEFAULTS.phraseRepeatBonus, + termWeight: options?.termWeight ?? DEFAULTS.termWeight, + termRepeatBonus: options?.termRepeatBonus ?? DEFAULTS.termRepeatBonus, + substringWeight: options?.substringWeight ?? DEFAULTS.substringWeight, + headingTermBonus: options?.headingTermBonus ?? DEFAULTS.headingTermBonus, + headingPhraseBonus: options?.headingPhraseBonus ?? DEFAULTS.headingPhraseBonus, + originalScoreCarry: options?.originalScoreCarry ?? DEFAULTS.originalScoreCarry, + }; + + const lowerQuery = query.toLowerCase(); + const queryTerms = tokenizeQuery(lowerQuery); + const queryPhrase = lowerQuery.trim(); + + const scored = candidates.map((c) => { + const content = c.content.toLowerCase(); + let score = 0; + + // 1. Exact phrase match + if (content.includes(queryPhrase)) { + score += w.phraseWeight + (countOccurrences(content, queryPhrase) - 1) * w.phraseRepeatBonus; + } + + // 2. Per-term matching + for (const term of queryTerms) { + if (content.includes(term)) { + score += isWordBoundary(content, term) ? w.termWeight + (countTermMatches(content, term) - 1) * w.termRepeatBonus : w.substringWeight; + } + } + + // 3. Heading path bonus + if (c.headingPath) { + const heading = c.headingPath.toLowerCase(); + for (const term of queryTerms) { + if (heading.includes(term)) score += w.headingTermBonus; + } + if (queryPhrase.length > 2 && heading.includes(queryPhrase)) { + score += w.headingPhraseBonus; + } + } + + // 4. Carry over original score + score += c.score * w.originalScoreCarry; + + return { id: c.id, score }; + }); + + // Min-max normalise to [0, 1] + const scores = scored.map((s) => s.score); + const min = Math.min(...scores); + const max = Math.max(...scores); + const range = max - min || 1; + + return scored.map((s) => ({ id: s.id, score: (s.score - min) / range })); +} \ No newline at end of file diff --git a/src/reranker/types.ts b/src/reranker/types.ts new file mode 100644 index 0000000..a7abde8 --- /dev/null +++ b/src/reranker/types.ts @@ -0,0 +1,43 @@ +/** + * Reranker โ€” second-stage precision scoring for search results. + */ + +// --------------------------------------------------------------------------- +// Types +// --------------------------------------------------------------------------- + +/** A candidate document for reranking. */ +export interface RerankCandidate { + id: string; + content: string; + /** Original score from the coarse search stage. */ + score: number; + /** Heading path as a string (e.g. "Line Chart > Tooltip"). */ + headingPath?: string; +} + +/** A reranked result. */ +export interface RerankResult { + id: string; + /** Final score after reranking (higher is better). */ + score: number; +} + +/** Reranker type */ +export type Reranker = { + rerank(query: string, candidates: RerankCandidate[]): Promise; +}; + +/** Configuration for reranking. */ +export interface RerankOptions { + rerankFactor?: number; + minCandidates?: number; + phraseWeight?: number; + phraseRepeatBonus?: number; + termWeight?: number; + termRepeatBonus?: number; + substringWeight?: number; + headingTermBonus?: number; + headingPhraseBonus?: number; + originalScoreCarry?: number; +} \ No newline at end of file diff --git a/src/storage/index.ts b/src/storage/index.ts new file mode 100644 index 0000000..57f2cc3 --- /dev/null +++ b/src/storage/index.ts @@ -0,0 +1,2 @@ +export { Store } from './store'; +export type { StoreQueryParams, ZvecDoc, ZvecQueryResult, ZvecFieldSchema, ZvecStoreConfig } from './types'; \ No newline at end of file diff --git a/src/storage/schema.ts b/src/storage/schema.ts new file mode 100644 index 0000000..cb3f816 --- /dev/null +++ b/src/storage/schema.ts @@ -0,0 +1,51 @@ +/** + * Zvec schema builder. + */ + +import type { ZvecFieldSchema } from './types'; + +import { ZVecCollectionSchema, ZVecDataType, ZVecIndexType, ZVecMetricType } from '@zvec/zvec'; + +const FIELD_TYPES = { + STRING: ZVecDataType.STRING, + INT64: ZVecDataType.INT64, + FLOAT: ZVecDataType.FLOAT, + VECTOR_FP32: ZVecDataType.VECTOR_FP32, +}; + +const INDEX_TYPES = { + FTS: ZVecIndexType.FTS, + INVERT: ZVecIndexType.INVERT, + HNSW: ZVecIndexType.HNSW, +}; + +export function buildZvecSchema(dims: number, tokenizerName: string = 'jieba'): ZVecCollectionSchema { + const fields: ZvecFieldSchema[] = [ + { name: 'content', dataType: 'STRING', indexType: 'FTS', indexOptions: { tokenizerName } }, + { name: 'meta', dataType: 'STRING' }, + { name: 'path', dataType: 'STRING' }, + { name: 'contentHash', dataType: 'STRING' }, + ]; + + return new ZVecCollectionSchema({ + name: 'context_docs', + vectors: { + name: 'embedding', + dataType: ZVecDataType.VECTOR_FP32, + dimension: dims, + indexParams: { + indexType: ZVecIndexType.HNSW, + metricType: ZVecMetricType.COSINE, + m: 32, + efConstruction: 200, + }, + }, + fields: fields.map((f: ZvecFieldSchema) => ({ + name: f.name, + dataType: FIELD_TYPES[f.dataType], + ...(f.indexType && f.indexType !== 'NONE' && INDEX_TYPES[f.indexType] + ? { indexParams: { indexType: INDEX_TYPES[f.indexType], ...f.indexOptions } as never } + : {}), + })) as never, + }) as ZVecCollectionSchema; +} \ No newline at end of file diff --git a/src/storage/store.ts b/src/storage/store.ts index ab5a7cf..18a0889 100644 --- a/src/storage/store.ts +++ b/src/storage/store.ts @@ -1,274 +1,169 @@ import * as fs from 'fs'; import * as path from 'path'; -import { - createZvecStore, - openZvecStoreSync, -} from './zvec-store'; -import type { IZvecStore, ZvecStoreConfig, ActualZvecStoreOptions } from './zvec-store'; -import type { ZvecDoc, ZvecQueryResult } from './types'; -import type { Embedder } from '../embedder'; -import { detectTokenizer } from '../utils/tokenizer'; -import type { ContextOptions } from '../types'; - -// --------------------------------------------------------------------------- -// Default schema constants -// --------------------------------------------------------------------------- - -const DEFAULT_VECTOR_FIELD = 'embedding'; -const DEFAULT_FTS_FIELDS = ['content']; -const DEFAULT_RANK_CONSTANT = 60; - -function resolveTokenizer(sampleText?: string): string { - // Always auto-detect based on sample text content. - // Falls back to 'jieba' as safe default for mixed-language content - // when no sample has been loaded yet. - if (sampleText && sampleText.trim().length > 0) { - return detectTokenizer(sampleText); - } - return 'jieba'; -} - -function contextStoreConfig(dims: number, sampleText?: string): ZvecStoreConfig { - const ftsFields = DEFAULT_FTS_FIELDS; - const tokenizerName = resolveTokenizer(sampleText); - - return { - collectionName: 'context_docs', - vectorField: DEFAULT_VECTOR_FIELD, - vectorDims: dims, - ftsFields, - fields: [ - { - name: 'content', - dataType: 'STRING', - indexType: ftsFields.includes('content') ? 'FTS' : 'NONE', - indexOptions: { tokenizerName }, - }, - { name: 'meta', dataType: 'STRING' }, - { name: 'sourceFilePath', dataType: 'STRING' }, - { name: 'contentHash', dataType: 'STRING' }, - ], - }; -} - -function storeOpenOptions(options?: ContextOptions): ActualZvecStoreOptions { - return { - vectorField: DEFAULT_VECTOR_FIELD, - ftsFields: options?.ftsFields ?? DEFAULT_FTS_FIELDS, - rankConstant: options?.rankConstant ?? DEFAULT_RANK_CONSTANT, - }; -} +import type { ZvecDoc, ZvecQueryResult, StoreQueryParams } from './types'; -// --------------------------------------------------------------------------- -// Query params type for Store.queryDoc() -// --------------------------------------------------------------------------- - -export interface StoreQueryParams { - /** Search mode: 'hybrid' combines vector + FTS, 'vector' is pure semantic. */ - mode: 'hybrid' | 'vector'; - /** Query text for the FTS path (hybrid mode only). */ - queryText?: string; - /** Query vector (pre-computed by the embedder). */ - queryVector: number[]; - /** Number of results to return. */ - topK: number; - /** Optional field-level filter expression. */ - filter?: string; -} +import { ZVecCreateAndOpen, ZVecOpen, ZVecIndexType, ZVecCollection } from '@zvec/zvec'; +import { buildZvecSchema } from './schema'; +import { Embedder } from '../embedder'; +import type { ContextOptions } from '../types'; -// --------------------------------------------------------------------------- -// Store โ€” manages zvec store lifecycle per library -// --------------------------------------------------------------------------- +const VECTOR_FIELD = 'embedding'; +const FTS_FIELDS = ['content']; /** - * Store manages creation, caching, and querying of zvec store instances. + * Store - zvec storage for multiple libraries. * - * Each library gets its own `.zvec` file on disk. The public API is - * intentionally minimal โ€” three methods cover all usage: - * - * - `create(library)` โ€” create / open a store for a library - * - `addDoc(docs)` โ€” batch-insert documents into a store - * - `queryDoc(params)` โ€” search a store (vector or hybrid) + * Usage: + * const store = new Store('./data', embedder, options); + * store.acquireZvec('lib', 'jieba'); // get or create + * store.addDoc('lib', docs); // insert docs + * store.queryDoc('lib', params); // query + * store.close(); // close all */ export class Store { - private readonly vectorsDir: string; - private readonly embedder: Embedder; - private readonly contextOptions?: ContextOptions; - private readonly stores: Map = new Map(); - /** In-flight creation promises โ€” prevents duplicate stores from concurrent calls. */ - private readonly pending: Map> = new Map(); + private vectorsDir: string; + private embedder: Embedder; + private options?: ContextOptions; + private zvecs: Map = new Map(); constructor(vectorsDir: string, embedder: Embedder, options?: ContextOptions) { this.vectorsDir = vectorsDir; this.embedder = embedder; - this.contextOptions = options; + this.options = options; } - // โ”€โ”€ Public API โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - - /** - * Create (or re-open) a zvec store for a library. - * - * - If the store is already cached, returns it immediately. - * - If a `.zvec` file exists on disk, opens it. - * - Otherwise, creates a new store with the configured schema. - * - * Uses a Promise-based lock so concurrent calls for the same library - * share a single creation attempt instead of racing. - * - * @param library Library name. - * @param sampleText Optional document sample for auto-detecting FTS tokenizer - * when `tokenizer` is `'auto'`. Only used for new stores. - */ - async create(library: string, sampleText?: string): Promise { - const cached = this.stores.get(library); + /** Get or create zvec instance. tokenizerName configures tokenizer on first creation. */ + acquireZvec(library: string, tokenizerName?: string): ZVecCollection { + const cached = this.zvecs.get(library); if (cached) return cached; - const pendingPromise = this.pending.get(library); - if (pendingPromise) return pendingPromise; + const filePath = path.join(this.vectorsDir, `${library}.zvec`); + let collection: ZVecCollection; - const promise = this._doCreate(library, sampleText); - this.pending.set(library, promise); - try { - const store = await promise; - this.stores.set(library, store); - return store; - } finally { - this.pending.delete(library); + if (fs.existsSync(filePath)) { + collection = ZVecOpen(filePath); + } else { + const schema = buildZvecSchema(this.embedder.dimensions, tokenizerName); + collection = ZVecCreateAndOpen(filePath, schema); } - } - /** - * Batch-insert documents into a library's store (upsert semantics). - * - * The store must have been created via `create()` first. Uses upsert - * internally so already-existing IDs are updated in-place. - * - * @param library Library name whose store to insert into. - * @param docs Documents to insert (id, vector, fields). - */ - async addDoc(library: string, docs: ZvecDoc[]): Promise { - const store = this.stores.get(library); - if (!store) { - throw new Error(`Store for library "${library}" has not been created. Call create() first.`); - } - await store.insert(docs); + this.zvecs.set(library, collection); + return collection; } - /** - * Batch-fetch stored documents by ID, returning only the requested fields. - * - * This is the single-source-of-truth for dedup โ€” zvec owns the document - * catalog so there is no separate registry to keep in sync. - * - * @param library Library name to query. - * @param ids Document IDs to look up. - * @param outputFields Fields to return (default: all). - * @returns Map of found IDs โ†’ their fields. Missing IDs are absent. - */ - async fetchDocs( - library: string, - ids: string[], - outputFields?: string[], - ): Promise> { - const store = this.stores.get(library) ?? this._tryOpenFromDisk(library); - if (!store) return {}; - return store.fetch(ids, outputFields); + /** Insert docs (upsert semantics). */ + addDoc(library: string, docs: ZvecDoc[]): void { + const collection = this.acquireZvec(library); + if (docs.length === 0) return; + + const opts = this.getStoreOptions(); + const records = docs.map((d) => ({ + id: d.id, + vectors: { [opts.vectorField]: d.vector }, + fields: d.fields, + })); + collection.upsertSync(records); } - /** - * Query a library's store for similar documents. - * - * - `mode: 'hybrid'` combines vector + FTS text matching via RRF fusion. - * - `mode: 'vector'` is pure semantic ANN search. - * - * The store must have been created via `create()` first. Returns - * empty array if the library has no store. - * - * @param library Library name to query. - * @param params Search parameters (mode, vectors, topK, filter). - * @returns Ranked query results with scores and fields. - */ - async queryDoc(library: string, params: StoreQueryParams): Promise { - const store = this.stores.get(library) ?? this._tryOpenFromDisk(library); - if (!store) return []; + /** Fetch docs by IDs. */ + fetchDocs(library: string, ids: string[], outputFields?: string[]): Record { + const collection = this.acquireZvec(library); - if (params.mode === 'hybrid') { - return store.searchHybrid({ - queryText: params.queryText ?? '', - queryVector: params.queryVector, - topK: params.topK, - filter: params.filter, - }); - } + if (ids.length === 0) return {}; - return store.search({ - vector: params.queryVector, - topK: params.topK, - filter: params.filter, + const raw = collection.fetchSync({ + ids, + outputFields, + includeVector: false, }); + + const result: Record = {}; + for (const [id, doc] of Object.entries(raw)) { + result[id] = { + id, + score: 0, + fields: doc.fields ?? {}, + }; + } + return result; } - // โ”€โ”€ Lifecycle helpers (used by Context internally) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + /** Query docs. Supports hybrid (vector + fulltext) or vector-only mode. */ + queryDoc(library: string, params: StoreQueryParams): ZvecQueryResult[] { + const collection = this.acquireZvec(library); - /** - * Close and remove a single library's store from cache. - */ - async close(library: string): Promise { - const store = this.stores.get(library); - if (store) { - await store.close(); - this.stores.delete(library); + const opts = this.getStoreOptions(); + + if (params.mode === 'hybrid') { + return this.hybridSearch(collection, opts, params); } + + return this.vectorSearch(collection, opts, params); } - /** - * Close all cached stores and release resources. - * - * Called by `Context.close()` at process exit. - */ - async closeAll(): Promise { - for (const [, store] of this.stores) { - await store.close(); + /** Close all zvec instances. */ + close(): void { + for (const [, collection] of this.zvecs) { + collection.closeSync(); } - this.stores.clear(); + this.zvecs.clear(); } - // โ”€โ”€ Private helpers โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - - private _getStorePath(library: string): string { - return path.join(this.vectorsDir, `${library}.zvec`); + private getStoreOptions() { + return { + vectorField: VECTOR_FIELD, + ftsFields: this.options?.ftsFields ?? FTS_FIELDS, + rankConstant: this.options?.rankConstant ?? 60, + }; } - private async _doCreate(library: string, sampleText?: string): Promise { - const filePath = this._getStorePath(library); - - if (fs.existsSync(filePath)) { - return openZvecStoreSync(filePath, storeOpenOptions(this.contextOptions)); - } - - return await createZvecStore( - filePath, - contextStoreConfig(this.embedder.dimensions, sampleText), - ); + private vectorSearch(collection: ZVecCollection, opts: { vectorField: string }, params: StoreQueryParams): ZvecQueryResult[] { + const query = { + fieldName: opts.vectorField, + vector: params.queryVector, + topk: params.topK, + ...(params.filter ? { filter: params.filter } : {}), + }; + + const rawResults = collection.querySync(query); + return rawResults.map((r) => ({ + id: r.id, + score: r.score, + fields: r.fields ?? {}, + })); } - /** - * Try to lazily open an existing store from disk if not already cached. - * Used by queryDoc() to auto-open stores on first query. - */ - private _tryOpenFromDisk(library: string): IZvecStore | undefined { - const filePath = this._getStorePath(library); - if (fs.existsSync(filePath)) { - try { - const store = openZvecStoreSync(filePath, storeOpenOptions(this.contextOptions)); - this.stores.set(library, store); - return store; - } catch { - // @zvec/zvec not available โ€” cannot open - return undefined; - } + private hybridSearch(collection: ZVecCollection, opts: { vectorField: string; ftsFields: string[]; rankConstant: number }, params: StoreQueryParams): ZvecQueryResult[] { + if (!opts.ftsFields?.length) { + return this.vectorSearch(collection, opts, params); } - return undefined; + + const ftsQueries = opts.ftsFields.map((fieldName) => ({ + fieldName, + fts: { matchString: params.queryText ?? '' }, + numCandidates: params.topK * 2, + params: { indexType: ZVecIndexType.FTS, defaultOperator: 'OR' as const }, + })); + + const query = { + queries: [ + { + fieldName: opts.vectorField, + vector: params.queryVector, + numCandidates: params.topK * 2, + }, + ...ftsQueries, + ], + topk: params.topK, + rerank: { type: 'rrf' as const, rankConstant: opts.rankConstant }, + ...(params.filter ? { filter: params.filter } : {}), + }; + + const rawResults = collection.multiQuerySync(query); + return rawResults.map((r) => ({ + id: r.id, + score: r.score, + fields: r.fields ?? {}, + })); } -} +} \ No newline at end of file diff --git a/src/storage/types.ts b/src/storage/types.ts index 0387fc8..d1344a2 100644 --- a/src/storage/types.ts +++ b/src/storage/types.ts @@ -1,5 +1,5 @@ /** - * ZvecStore public types and interfaces. + * ZvecStore internal types. */ export interface ZvecDoc { @@ -17,39 +17,16 @@ export interface ZvecQueryResult { export interface ZvecSearchParams { vector: number[]; topK: number; - /** Optional field-level filter (exact match expression like `library = 'g2'`). */ filter?: string; } export interface ZvecHybridParams { - /** Query text for FTS path (no embedding needed). */ queryText: string; - /** Query vector for ANN path (must be pre-computed). */ queryVector: number[]; topK: number; filter?: string; } -export interface IZvecStore { - insert(docs: ZvecDoc[]): Promise; - /** - * Batch-fetch documents by ID, returning only the requested fields. - * - * This is the single-source-of-truth for dedup โ€” zvec owns the document - * catalog so there is no separate registry to keep in sync. - */ - fetch(ids: string[], outputFields?: string[]): Promise>; - /** Pure ANN vector search. */ - search(params: ZvecSearchParams): Promise; - /** Hybrid FTS + Vector with native RRF fusion (when available). */ - searchHybrid(params: ZvecHybridParams): Promise; - /** Synchronous ANN vector search. */ - searchSync(params: ZvecSearchParams): ZvecQueryResult[]; - /** Synchronous hybrid FTS + Vector search. */ - searchHybridSync(params: ZvecHybridParams): ZvecQueryResult[]; - close(): Promise; -} - /** Field schema entry for building a zvec collection. */ export interface ZvecFieldSchema { name: string; @@ -83,3 +60,11 @@ export interface ActualZvecStoreOptions { */ rankConstant?: number; } + +export interface StoreQueryParams { + mode: 'hybrid' | 'vector'; + queryText?: string; + queryVector: number[]; + topK: number; + filter?: string; +} diff --git a/src/storage/zvec-store.ts b/src/storage/zvec-store.ts deleted file mode 100644 index 299817c..0000000 --- a/src/storage/zvec-store.ts +++ /dev/null @@ -1,426 +0,0 @@ -/** - * zvec-store โ€” zvec storage module: types, implementation, and factories. - * - * Combines the former split files into a single module: - * types.ts โ†’ public types & interfaces - * actual-store.ts โ†’ ActualZvecStore, factories, schema builder - */ - -// --------------------------------------------------------------------------- -// Public types (re-exported from types.ts) -// --------------------------------------------------------------------------- - -export type { - ZvecDoc, - ZvecQueryResult, - ZvecSearchParams, - ZvecHybridParams, - IZvecStore, - ZvecFieldSchema, - ZvecStoreConfig, - ActualZvecStoreOptions, -} from './types'; - -// --------------------------------------------------------------------------- -// Internal types โ€” zvec SDK & collection interfaces -// --------------------------------------------------------------------------- - -import type { - ZvecDoc, - ZvecQueryResult, - ZvecSearchParams, - ZvecHybridParams, - IZvecStore, - ZvecStoreConfig, - ActualZvecStoreOptions, - ZvecFieldSchema, -} from './types'; - -let _zvecModule: unknown = undefined; -let _zvecLoadFailed = false; - -/** Zvec SDK type โ€” minimal interface we rely on. */ -interface ZvecSDK { - ZVecCollectionSchema: new (config: unknown) => unknown; - ZVecDataType: Record; - ZVecIndexType: Record; - ZVecMetricType: Record; - ZVecCreateAndOpen(path: string, schema: unknown): ZvecCollection; - ZVecOpen(path: string, options?: Record): ZvecCollection; -} - -/** Zvec collection type โ€” minimal interface we rely on. */ -interface ZvecCollection { - insertSync(records: unknown[]): void; - upsertSync(records: unknown[]): void; - fetchSync(params: { ids: string | string[]; outputFields?: string[]; includeVector?: boolean }): Record; - querySync(params: unknown): ZvecQueryRaw[]; - multiQuerySync(params: unknown): ZvecQueryRaw[]; - closeSync(): void; -} - -interface ZvecQueryRaw { - id: string; - score: number; - fields?: Record; -} - -function loadZvecSync(): ZvecSDK | undefined { - if (_zvecModule) return _zvecModule as ZvecSDK; - if (_zvecLoadFailed) return undefined; - try { - // eslint-disable-next-line @typescript-eslint/no-require-imports - _zvecModule = require('@zvec/zvec'); - } catch { - _zvecLoadFailed = true; - return undefined; - } - return _zvecModule as ZvecSDK; -} - -function requireZvecSync(): ZvecSDK { - const z = loadZvecSync(); - if (!z) { - throw new Error( - '@zvec/zvec is not installed. Install it with:\n' + - ' pnpm add @zvec/zvec' - ); - } - return z; -} - -// --------------------------------------------------------------------------- -// Schema builder -// --------------------------------------------------------------------------- - -/** - * Build a ZVecCollectionSchema from a generic ZvecStoreConfig. - * - * This is a helper for callers that want to create zvec collections without - * dealing with the low-level zvec SDK schema API directly. - */ -export function buildZvecSchema(z: ZvecSDK, config: ZvecStoreConfig): unknown { - const { ZVecCollectionSchema, ZVecDataType, ZVecIndexType, ZVecMetricType } = z; - - const vectorSchema = { - name: config.vectorField, - dataType: ZVecDataType.VECTOR_FP32, - dimension: config.vectorDims, - indexParams: { - indexType: ZVecIndexType.HNSW, - metricType: ZVecMetricType.COSINE, - m: 32, - efConstruction: 200 - } - }; - - const fieldTypeMap: Record = { - STRING: ZVecDataType.STRING, - INT64: ZVecDataType.INT64, - FLOAT: ZVecDataType.FLOAT, - VECTOR_FP32: ZVecDataType.VECTOR_FP32, - }; - - const indexTypeMap: Record = { - FTS: ZVecIndexType.FTS, - INVERT: ZVecIndexType.INVERT, - HNSW: ZVecIndexType.HNSW, - }; - - const fieldSchemas = config.fields.map((f: ZvecFieldSchema) => { - const schema: Record = { - name: f.name, - dataType: fieldTypeMap[f.dataType], - }; - - if (f.indexType && f.indexType !== 'NONE' && indexTypeMap[f.indexType]) { - schema.indexParams = { - indexType: indexTypeMap[f.indexType], - ...(f.indexOptions ?? {}), - }; - } - - return schema; - }); - - return new ZVecCollectionSchema({ - name: config.collectionName, - vectors: vectorSchema, - fields: fieldSchemas, - }); -} - -// --------------------------------------------------------------------------- -// ActualZvecStore implementation -// --------------------------------------------------------------------------- - -const DEFAULT_OPEN_OPTIONS: ActualZvecStoreOptions = { - vectorField: 'embedding', - ftsFields: [], - rankConstant: 60, -}; - -/** - * ZvecStore backed by the native @zvec/zvec library. - * - * Provides true hybrid search: zvec's FTS path accepts raw text (no embedding), - * the vector path uses pre-computed embeddings, and `multiQuerySync` with RRF - * fuses both in the engine. - */ -export class ActualZvecStore implements IZvecStore { - private _collection: ZvecCollection; - private _closed = false; - private _vectorField: string; - private _ftsFields: string[]; - private _rankConstant: number; - - constructor(collection: ZvecCollection, options: ActualZvecStoreOptions) { - this._collection = collection; - this._vectorField = options.vectorField; - this._ftsFields = options.ftsFields; - this._rankConstant = options.rankConstant ?? 60; - } - - /** - * Create a new zvec collection with a generic schema config. - */ - static async create(path: string, config: ZvecStoreConfig): Promise { - const z = requireZvecSync(); - const schema = buildZvecSchema(z, config); - const collection = z.ZVecCreateAndOpen(path, schema); - return new ActualZvecStore(collection, { - vectorField: config.vectorField, - ftsFields: config.ftsFields, - }); - } - - /** Synchronous version of `create`. */ - static createSync(path: string, config: ZvecStoreConfig): ActualZvecStore { - const z = requireZvecSync(); - const schema = buildZvecSchema(z, config); - const collection = z.ZVecCreateAndOpen(path, schema); - return new ActualZvecStore(collection, { - vectorField: config.vectorField, - ftsFields: config.ftsFields, - }); - } - - static async open( - path: string, - options?: ActualZvecStoreOptions - ): Promise { - const z = requireZvecSync(); - const collection = z.ZVecOpen(path); - return new ActualZvecStore(collection, options ?? DEFAULT_OPEN_OPTIONS); - } - - /** Synchronous version of `open`. Uses read-only to allow concurrent readers. */ - static openSync( - path: string, - options?: ActualZvecStoreOptions - ): ActualZvecStore { - const z = requireZvecSync(); - const collection = z.ZVecOpen(path, { readOnly: true }); - return new ActualZvecStore(collection, options ?? DEFAULT_OPEN_OPTIONS); - } - - async insert(docs: ZvecDoc[]): Promise { - if (this._closed) throw new Error('Store is closed'); - if (docs.length === 0) return; - - const records = docs.map((d) => ({ - id: d.id, - vectors: { [this._vectorField]: d.vector }, - fields: d.fields - })); - // upsertSync handles both new docs and content-changed updates in one call. - this._collection.upsertSync(records); - } - - async fetch(ids: string[], outputFields?: string[]): Promise> { - if (this._closed) throw new Error('Store is closed'); - if (ids.length === 0) return {}; - - const raw = this._collection.fetchSync({ - ids, - outputFields, - includeVector: false, - }); - - const result: Record = {}; - for (const [id, doc] of Object.entries(raw)) { - result[id] = { - id, - score: 0, - fields: (doc as ZvecQueryRaw).fields ?? {}, - }; - } - return result; - } - - async search(params: ZvecSearchParams): Promise { - if (this._closed) throw new Error('Store is closed'); - - const rawResults = this._collection.querySync( - this._buildVectorQuery(params) - ); - - return rawResults.map((r: ZvecQueryRaw) => ({ - id: r.id, - score: r.score, - fields: r.fields ?? {} - })); - } - - async searchHybrid(params: ZvecHybridParams): Promise { - if (this._closed) throw new Error('Store is closed'); - - const rawResults = this._collection.multiQuerySync( - this._buildHybridQuery(params) - ); - - return rawResults.map((r: ZvecQueryRaw) => ({ - id: r.id, - score: r.score, - fields: r.fields ?? {} - })); - } - - searchSync(params: ZvecSearchParams): ZvecQueryResult[] { - if (this._closed) throw new Error('Store is closed'); - - const rawResults = this._collection.querySync( - this._buildVectorQuery(params) - ); - - return rawResults.map((r: ZvecQueryRaw) => ({ - id: r.id, - score: r.score, - fields: r.fields ?? {} - })); - } - - searchHybridSync(params: ZvecHybridParams): ZvecQueryResult[] { - if (this._closed) throw new Error('Store is closed'); - - const rawResults = this._collection.multiQuerySync( - this._buildHybridQuery(params) - ); - - return rawResults.map((r: ZvecQueryRaw) => ({ - id: r.id, - score: r.score, - fields: r.fields ?? {} - })); - } - - /** Build a query params object for vector search, omitting filter when undefined. */ - private _buildVectorQuery(params: ZvecSearchParams): Record { - const q: Record = { - fieldName: this._vectorField, - vector: params.vector, - topk: params.topK, - }; - if (params.filter) { - q.filter = params.filter; - } - return q; - } - - /** Build multi-field FTS query paths for hybrid search from configured ftsFields. */ - private _buildFtsQueries(queryText: string, topK: number): Record[] { - const z = loadZvecSync(); - const ftsParams = z - ? { indexType: z.ZVecIndexType.FTS, defaultOperator: 'OR' as const } - : undefined; - - return this._ftsFields.map((fieldName) => ({ - fieldName, - fts: { matchString: queryText }, - numCandidates: topK * 2, - params: ftsParams, - })); - } - - /** Build multiQuery params for hybrid search, omitting filter when undefined. */ - private _buildHybridQuery(params: ZvecHybridParams): Record { - const q: Record = { - queries: [ - { - fieldName: this._vectorField, - vector: params.queryVector, - numCandidates: params.topK * 2 - }, - ...this._buildFtsQueries(params.queryText, params.topK), - ], - topk: params.topK, - rerank: { type: 'rrf', rankConstant: this._rankConstant } - }; - if (params.filter) { - q.filter = params.filter; - } - return q; - } - - async close(): Promise { - if (this._closed) return; - this._closed = true; - try { - this._collection.closeSync(); - } catch { - // best-effort - } - } -} - -// --------------------------------------------------------------------------- -// Convenience factories -// --------------------------------------------------------------------------- - -/** - * Create a zvec store with the given config. - * - * Requires @zvec/zvec to be installed. Throws an error when it is not available. - */ -export async function createZvecStore( - path: string, - config: ZvecStoreConfig, -): Promise { - return ActualZvecStore.create(path, config); -} - -/** - * Asynchronously open a zvec store. - * - * Requires @zvec/zvec to be installed. - */ -export async function openZvecStore( - path: string, - options?: ActualZvecStoreOptions -): Promise { - return openZvecStoreSync(path, options); -} - -/** - * Synchronously open a zvec store. - * - * Requires @zvec/zvec to be installed. Throws an error when it is not available. - */ -export function openZvecStoreSync( - path: string, - options?: ActualZvecStoreOptions -): IZvecStore { - const z = loadZvecSync(); - if (z) { - return ActualZvecStore.openSync(path, options); - } - throw new Error( - '@zvec/zvec is not installed. Install it with:\n' + - ' pnpm add @zvec/zvec' - ); -} - -/** Synchronous check: is @zvec/zvec available? */ -export function isZvecAvailable(): boolean { - return loadZvecSync() !== undefined; -} diff --git a/src/types.ts b/src/types.ts index 34d99fe..2ab749b 100644 --- a/src/types.ts +++ b/src/types.ts @@ -1,5 +1,4 @@ -import type { RerankOptions } from './utils/reranker'; -import type { EmbedderKind } from './embedder/resolve'; +import type { RerankOptions } from './reranker'; /** * Document structure @@ -11,6 +10,18 @@ export interface Document { meta?: Record; } +/** + * Document with computed fields for loading. + */ +export interface LoadedDoc extends Document { + /** Document ID */ + id: string; + /** Hash of document content */ + contentHash: string; + /** Source file path relative to base path */ + path: string; +} + /** Configuration for query expansion. */ export interface QueryExpansionOptions { /** @@ -31,8 +42,12 @@ export interface QueryExpansionOptions { * Context initialization options */ export interface ContextOptions { - /** Directory to store vector files */ - vectorsDir: string; + /** + * Directory to store vector files. + * + * Defaults to `'.context/vectors'`. + */ + vectorsDir?: string; /** * Base path for resolving document IDs. * @@ -179,15 +194,6 @@ export interface QueryResult { content: string; /** Similarity score (normalized to [0, 1] range) */ score: number; - /** - * How the score was computed โ€” helps users interpret and compare - * scores across different query modes. - * - * - `'vector'`: cosine similarity from pure vector search (0~1) - * - `'hybrid'`: RRF fusion score from vector + FTS text path - * - `'reranked'`: keyword reranker score, min-max normalized to [0, 1] - */ - scoreMode?: 'vector' | 'hybrid' | 'reranked'; /** Document metadata */ meta?: Record; /** Original file path relative to `basePath` โ€” allows users to trace @@ -195,13 +201,7 @@ export interface QueryResult { * * Populated during `load()` and stored as a zvec field. */ - sourceFilePath?: string; - /** - * The embedder kind used to produce the vectors for this result's store. - * - * - `'transformers'`: high-quality bilingual model (bge-small-zh-v1.5) - */ - embedderKind?: EmbedderKind; + path?: string; } /** diff --git a/src/utils/common.ts b/src/utils/common.ts index c72dd3e..1a45206 100644 --- a/src/utils/common.ts +++ b/src/utils/common.ts @@ -1,43 +1,11 @@ -// --------------------------------------------------------------------------- -// JSON / meta helpers -// --------------------------------------------------------------------------- - /** - * Safely parse a meta JSON string. Returns undefined on invalid JSON. + * Safely parse JSON string. Returns undefined on invalid JSON. */ -export function safeParseMeta(metaStr: string | undefined): Record | undefined { - if (!metaStr) return undefined; +export function safeJsonParse(str: string | undefined): unknown { + if (!str) return undefined; try { - return JSON.parse(metaStr); + return JSON.parse(str); } catch { return undefined; } -} - -// --------------------------------------------------------------------------- -// Sampling -// --------------------------------------------------------------------------- - -/** - * Select a representative sample of files from a list for tokenizer detection. - * - * Picks files spread across the list (first, middle, last, and evenly-spaced) - * to avoid bias when the file order doesn't reflect content distribution. - * Returns at most `maxCount` file paths. - */ -export function selectSampleFiles(files: string[], maxCount: number): string[] { - if (files.length <= maxCount) return files; - - const result: string[] = []; - result.push(files[0]); - const step = Math.floor((files.length - 1) / (maxCount - 1)); - for (let i = step; i < files.length - 1; i += step) { - if (result.length < maxCount) { - result.push(files[i]); - } - } - if (result[result.length - 1] !== files[files.length - 1] && result.length < maxCount) { - result.push(files[files.length - 1]); - } - return result; -} +} \ No newline at end of file diff --git a/src/utils/doc.ts b/src/utils/doc.ts new file mode 100644 index 0000000..910df3f --- /dev/null +++ b/src/utils/doc.ts @@ -0,0 +1,16 @@ +import * as path from 'path'; +import { computeContentHash } from './hash'; + +/** + * Convert a file path to a safe, collision-resistant ID for zvec. + */ +export function pathToId(filePath: string): string { + const normalized = path.normalize(filePath).replace(/\\/g, '/').replace(/^\/+/, ''); + const hash = computeContentHash(normalized); + const basename = normalized.split('/').pop() ?? ''; + const suffix = basename + .replace(/\.[a-zA-Z0-9]+$/, '') + .replace(/[^a-zA-Z0-9]/g, '_') + .slice(0, 20); + return `${hash}__${suffix}`; +} \ No newline at end of file diff --git a/src/utils/expander.ts b/src/utils/expander.ts deleted file mode 100644 index 608fbb9..0000000 --- a/src/utils/expander.ts +++ /dev/null @@ -1,128 +0,0 @@ -/** - * Query Expander โ€” augments user queries with synonyms and cross-language - * bridging terms to improve recall across different naming conventions. - * - * The expanded query text is used for both embedding and FTS, so a single - * expanded query can match documents written in either Chinese or English, - * or documents using different terminology for the same concept. - * - * Synonyms are entirely user-provided โ€” no built-in defaults. Pass your - * own map to `SynonymExpander` to define domain-specific term bridges. - */ - -// --------------------------------------------------------------------------- -// Types -// --------------------------------------------------------------------------- - -/** Query expander interface โ€” custom expanders can implement this. */ -export interface QueryExpander { - /** - * Expand a query string with additional terms. - * - * Returns the augmented query text. The original query terms are always - * preserved; the expander only appends additional terms. - */ - expand(query: string): string; -} - -// --------------------------------------------------------------------------- -// SynonymExpander -// --------------------------------------------------------------------------- - -/** - * Characters that delimit word boundaries in Latin/script-based text. - * Does NOT include CJK characters โ€” CJK text has no inter-word spacing, - * so substring matching is the correct strategy for CJK terms. - */ -const WORD_BOUNDARY_RE = /[\s,๏ผŒใ€‚.!๏ผ?๏ผŸ;๏ผ›:๏ผš\(\)\[\]{}""''\"\'\-_\/\\|@#$%^&*+=<>~`]/; - -/** Unicode range for CJK Unified Ideographs (basic block). */ -const CJK_CHAR_RE = /[\u4e00-\u9fff\u3400-\u4dbf\uf900-\ufaff]/; - -/** - * Check whether a string contains any CJK character. - */ -function containsCJK(text: string): boolean { - return CJK_CHAR_RE.test(text); -} - -/** - * Expands query text by injecting synonym terms. - * - * For each word/phrase in the query that matches a synonym entry, the - * corresponding alternative terms are appended to the query. This bridges - * CNโ†”EN terminology gaps and handles different naming conventions. - * - * Matching strategy: - * - CJK terms: substring match (CJK text has no word boundaries) - * - Latin terms: strict word-boundary match (avoids partial matches like - * "line" matching inside "inline") - */ -export class SynonymExpander implements QueryExpander { - private readonly synonyms: Record; - - constructor(synonyms?: Record) { - this.synonyms = synonyms ?? {}; - } - - expand(query: string): string { - const queryLower = query.toLowerCase().trim(); - const additions = new Set(); - - for (const [term, syns] of Object.entries(this.synonyms)) { - const termLower = term.toLowerCase(); - - if (this.containsTerm(queryLower, termLower)) { - for (const syn of syns) { - if (!this.containsTerm(queryLower, syn.toLowerCase())) { - additions.add(syn); - } - } - } - } - - if (additions.size === 0) return query; - - return `${query} ${[...additions].join(' ')}`; - } - - /** - * Check if a term appears in the text with appropriate boundary rules. - * - * - CJK terms use substring matching (no word boundaries in CJK text). - * - Latin terms use strict word-boundary detection to avoid false positives. - */ - private containsTerm(text: string, term: string): boolean { - if (!text.includes(term)) return false; - - if (containsCJK(term)) { - return true; - } - - let searchFrom = 0; - while (searchFrom <= text.length - term.length) { - const idx = text.indexOf(term, searchFrom); - if (idx === -1) return false; - - const prevOk = idx === 0 || WORD_BOUNDARY_RE.test(text[idx - 1]); - const afterIdx = idx + term.length; - const afterOk = afterIdx >= text.length || WORD_BOUNDARY_RE.test(text[afterIdx]); - - if (prevOk && afterOk) return true; - - searchFrom = idx + 1; - } - - return false; - } -} - -// --------------------------------------------------------------------------- -// NoopExpander โ€” pass-through for when expansion is not desired -// --------------------------------------------------------------------------- - -export class NoopExpander implements QueryExpander { - expand(query: string): string { - return query; - } -} diff --git a/src/utils/index.ts b/src/utils/index.ts index bf5b472..0da8493 100644 --- a/src/utils/index.ts +++ b/src/utils/index.ts @@ -1,32 +1,7 @@ -/** - * utils โ€” aggregate entry point for all utility modules. - * - * Pure functions and lightweight classes that don't belong to - * a specific domain module (embedder, loaders, storage). - */ - -// Common helpers โ€” JSON parsing, sampling -export { - safeParseMeta, - selectSampleFiles, -} from './common'; - -// Hashing โ€” SHA-256 content hash for dedup & ID generation +export { pathToId } from './doc'; export { computeContentHash } from './hash'; - -// Tokenizer selection โ€” language detection & FTS tokenizer auto-configuration -export { - isCJK, - detectLanguage, - tokenizerForLanguage, - detectTokenizer, -} from './tokenizer'; +export { safeJsonParse } from './common'; +export { loadSampleText } from './sample'; +export { isCJK, detectLanguage, tokenizerForLanguage, detectTokenizer } from './tokenizer'; export type { LanguageHint } from './tokenizer'; - -// Query expansion โ€” synonym-based query augmentation -export { SynonymExpander, NoopExpander } from './expander'; -export type { QueryExpander } from './expander'; - -// Reranking โ€” second-stage keyword precision scoring -export { KeywordReranker, createReranker } from './reranker'; -export type { Reranker, RerankCandidate, RerankResult, RerankOptions } from './reranker'; +export { containsCJK } from './str'; \ No newline at end of file diff --git a/src/utils/reranker.ts b/src/utils/reranker.ts deleted file mode 100644 index 7ebf146..0000000 --- a/src/utils/reranker.ts +++ /dev/null @@ -1,224 +0,0 @@ -/** - * Reranker โ€” second-stage precision scoring for search results. - * - * The pipeline: - * 1. Coarse search (vector / hybrid) โ†’ topK ร— rerankFactor candidates - * 2. Reranker scores each candidate against the query - * 3. Final sort by reranked score โ†’ topK results - */ - -// --------------------------------------------------------------------------- -// Types -// --------------------------------------------------------------------------- - -/** A candidate document for reranking. */ -export interface RerankCandidate { - id: string; - content: string; - /** Original score from the coarse search stage. */ - score: number; - /** Heading path as a string (e.g. "Line Chart > Tooltip"). */ - headingPath?: string; -} - -/** A reranked result. */ -export interface RerankResult { - id: string; - /** Final score after reranking (higher is better). */ - score: number; -} - -/** Reranker interface โ€” custom rerankers can implement this. */ -export interface Reranker { - rerank(query: string, candidates: RerankCandidate[]): Promise; -} - -/** Configuration for reranking. */ -export interface RerankOptions { - rerankFactor?: number; - minCandidates?: number; - phraseWeight?: number; - phraseRepeatBonus?: number; - termWeight?: number; - termRepeatBonus?: number; - substringWeight?: number; - headingTermBonus?: number; - headingPhraseBonus?: number; - originalScoreCarry?: number; -} - -// --------------------------------------------------------------------------- -// Defaults -// --------------------------------------------------------------------------- - -const DEFAULT_WEIGHTS = { - phraseWeight: 3.0, - phraseRepeatBonus: 0.5, - termWeight: 1.0, - termRepeatBonus: 0.2, - substringWeight: 0.3, - headingTermBonus: 2.0, - headingPhraseBonus: 2.5, - originalScoreCarry: 0.1, -} as const; - -// --------------------------------------------------------------------------- -// KeywordReranker โ€” lexical overlap scoring -// --------------------------------------------------------------------------- - -/** - * Reranker that scores candidates by keyword / phrase overlap with the query. - * - * All scoring weights are configurable via constructor options. - * Scores are normalised to [0, 1] via min-max scaling. - */ -export class KeywordReranker implements Reranker { - private readonly weights: Record; - - constructor(options?: RerankOptions) { - this.weights = { - phraseWeight: options?.phraseWeight ?? DEFAULT_WEIGHTS.phraseWeight, - phraseRepeatBonus: options?.phraseRepeatBonus ?? DEFAULT_WEIGHTS.phraseRepeatBonus, - termWeight: options?.termWeight ?? DEFAULT_WEIGHTS.termWeight, - termRepeatBonus: options?.termRepeatBonus ?? DEFAULT_WEIGHTS.termRepeatBonus, - substringWeight: options?.substringWeight ?? DEFAULT_WEIGHTS.substringWeight, - headingTermBonus: options?.headingTermBonus ?? DEFAULT_WEIGHTS.headingTermBonus, - headingPhraseBonus: options?.headingPhraseBonus ?? DEFAULT_WEIGHTS.headingPhraseBonus, - originalScoreCarry: options?.originalScoreCarry ?? DEFAULT_WEIGHTS.originalScoreCarry, - }; - } - - async rerank(query: string, candidates: RerankCandidate[]): Promise { - if (candidates.length === 0) return []; - - const queryLower = query.toLowerCase(); - const queryTerms = tokenizeQuery(queryLower); - const queryPhrase = queryLower.trim(); - const w = this.weights; - - const scored = candidates.map((c) => { - const contentLower = c.content.toLowerCase(); - let score = 0; - - // 1. Exact phrase match โ€” strongest signal - if (contentLower.includes(queryPhrase)) { - score += w.phraseWeight; - const phraseCount = countOccurrences(contentLower, queryPhrase); - score += (phraseCount - 1) * w.phraseRepeatBonus; - } - - // 2. Per-term matching - for (const term of queryTerms) { - if (contentLower.includes(term)) { - if (isWordBoundary(contentLower, term)) { - score += w.termWeight; - const termCount = countTermMatches(contentLower, term); - score += (termCount - 1) * w.termRepeatBonus; - } else { - score += w.substringWeight; - } - } - } - - // 3. Heading path bonus - if (c.headingPath) { - const headingLower = c.headingPath.toLowerCase(); - for (const term of queryTerms) { - if (headingLower.includes(term)) { - score += w.headingTermBonus; - } - } - if (queryPhrase.length > 2 && headingLower.includes(queryPhrase)) { - score += w.headingPhraseBonus; - } - } - - // 4. Carry over a fraction of the original vector score - score += c.score * w.originalScoreCarry; - - return { id: c.id, score }; - }); - - // Min-max normalise to [0, 1] - const scores = scored.map((s) => s.score); - const min = Math.min(...scores); - const max = Math.max(...scores); - const range = max - min || 1; - - return scored.map((s) => ({ - id: s.id, - score: (s.score - min) / range, - })); - } -} - -// --------------------------------------------------------------------------- -// Helpers -// --------------------------------------------------------------------------- - -/** Split a query into meaningful tokens. */ -function tokenizeQuery(query: string): string[] { - const raw = query - .split(/[\s,๏ผŒใ€‚.!๏ผ๏ผŸ?ใ€๏ผš:๏ผ›;]+/) - .filter(Boolean); - - const tokens: string[] = []; - for (const r of raw) { - tokens.push(r); - if (/[ไธ€-้ฟฟ]{3,}/.test(r)) { - for (let i = 0; i + 2 <= r.length; i++) { - tokens.push(r.slice(i, i + 2)); - } - } - } - return [...new Set(tokens)]; -} - -/** Count non-overlapping occurrences of a substring. */ -function countOccurrences(text: string, sub: string): number { - let count = 0; - let pos = 0; - while ((pos = text.indexOf(sub, pos)) !== -1) { - count++; - pos += sub.length; - } - return count; -} - -/** Count whole-word matches of a term. */ -function countTermMatches(text: string, term: string): number { - let count = 0; - let pos = 0; - while ((pos = text.indexOf(term, pos)) !== -1) { - if (isWordBoundary(text, term, pos)) { - count++; - } - pos += term.length; - } - return count; -} - -/** Check if a substring occurrence is at a word boundary. */ -function isWordBoundary( - text: string, - term: string, - pos?: number, -): boolean { - const idx = pos ?? text.indexOf(term); - if (idx === -1) return false; - const before = idx === 0 || /[\s\n.,;:!?๏ผŒใ€‚๏ผ๏ผŸใ€๏ผš๏ผ›"'(๏ผˆใ€ใ€Š\-_]/.test(text[idx - 1]); - const afterIdx = idx + term.length; - const after = afterIdx >= text.length || /[\s\n.,;:!?๏ผŒใ€‚๏ผ๏ผŸใ€๏ผš๏ผ›"')๏ผ‰ใ€‘ใ€‹\-_]/.test(text[afterIdx]); - return before && after; -} - -// --------------------------------------------------------------------------- -// Factory -// --------------------------------------------------------------------------- - -/** - * Create a reranker with optional weight configuration. - */ -export function createReranker(options?: RerankOptions): Reranker { - return new KeywordReranker(options); -} diff --git a/src/utils/sample.ts b/src/utils/sample.ts new file mode 100644 index 0000000..d773bb7 --- /dev/null +++ b/src/utils/sample.ts @@ -0,0 +1,45 @@ +import * as fs from 'fs'; + +/** + * Select a representative sample of files for tokenizer detection. + */ +function selectSampleFiles(files: string[], maxCount: number): string[] { + if (files.length <= maxCount) return files; + + const result: string[] = []; + result.push(files[0]); + const step = Math.floor((files.length - 1) / (maxCount - 1)); + for (let i = step; i < files.length - 1; i += step) { + if (result.length < maxCount) { + result.push(files[i]); + } + } + if (result[result.length - 1] !== files[files.length - 1] && result.length < maxCount) { + result.push(files[files.length - 1]); + } + return result; +} + +/** + * Load sample text from a list of files for tokenizer detection. + * Returns undefined if no valid samples could be loaded. + */ +export async function loadSampleText(files: string[], sampleCount = 5): Promise { + if (files.length === 0) return undefined; + + try { + const sampleFiles = selectSampleFiles(files, sampleCount); + const samples = await Promise.allSettled( + sampleFiles.map((f) => fs.promises.readFile(f, 'utf-8')), + ); + const validSamples = samples + .filter((r): r is PromiseFulfilledResult => r.status === 'fulfilled') + .map((r) => r.value); + if (validSamples.length > 0) { + return validSamples.join('\n'); + } + } catch { + // Sample failure is non-fatal + } + return undefined; +} \ No newline at end of file diff --git a/src/utils/str.ts b/src/utils/str.ts new file mode 100644 index 0000000..8a05a2c --- /dev/null +++ b/src/utils/str.ts @@ -0,0 +1,7 @@ +/** + * String utilities. + */ + +export function containsCJK(text: string): boolean { + return /[ไธ€-้ฟฟใ€-ไถฟ๏ค€-๏ซฟ]/.test(text); +} \ No newline at end of file diff --git a/test/reranker.test.ts b/test/reranker.test.ts index 977bbab..e8567a5 100644 --- a/test/reranker.test.ts +++ b/test/reranker.test.ts @@ -1,93 +1,91 @@ import { describe, it, expect } from 'vitest'; -import { KeywordReranker, createReranker } from '../src/utils/reranker'; -import type { RerankCandidate } from '../src/utils/reranker'; +import { rerank } from '../src/reranker/reranker'; +import type { RerankCandidate } from '../src/reranker/types'; -describe('KeywordReranker', () => { - const reranker = new KeywordReranker(); - - describe('rerank', () => { - it('should return empty array for empty candidates', async () => { - const results = await reranker.rerank('test', []); - expect(results).toEqual([]); - }); +describe('rerank', () => { + it('should return empty array for empty candidates', async () => { + const results = await rerank('test', []); + expect(results).toEqual([]); + }); - it('should boost candidates with exact phrase match', async () => { - const candidates: RerankCandidate[] = [ - { id: '1', content: 'tooltip configuration settings', score: 0.5 }, - { id: '2', content: 'chart axis configuration', score: 0.5 }, - { id: '3', content: 'random unrelated content', score: 0.5 }, - ]; + it('should boost candidates with exact phrase match', async () => { + const candidates: RerankCandidate[] = [ + { id: '1', content: 'tooltip configuration settings', score: 0.5 }, + { id: '2', content: 'chart axis configuration', score: 0.5 }, + { id: '3', content: 'random unrelated content', score: 0.5 }, + ]; - const results = await reranker.rerank('tooltip configuration', candidates); - // Candidate 1 should score highest (exact phrase match) - const sorted = [...results].sort((a, b) => b.score - a.score); - expect(sorted[0].id).toBe('1'); - }); + const results = await rerank('tooltip configuration', candidates); + const sorted = [...results].sort((a, b) => b.score - a.score); + expect(sorted[0].id).toBe('1'); + }); - it('should boost candidates with term-level matches', async () => { - const candidates: RerankCandidate[] = [ - { id: '1', content: 'the tooltip provides hover information', score: 0.3 }, - { id: '2', content: 'a simple chart example', score: 0.8 }, - ]; + it('should boost candidates with term-level matches', async () => { + const candidates: RerankCandidate[] = [ + { id: '1', content: 'the tooltip provides hover information', score: 0.3 }, + { id: '2', content: 'a simple chart example', score: 0.8 }, + { id: '3', content: 'chart axis labels configuration', score: 0.2 }, + ]; - const results = await reranker.rerank('tooltip', candidates); - const sorted = [...results].sort((a, b) => b.score - a.score); - expect(sorted[0].id).toBe('1'); - }); + const results = await rerank('tooltip', candidates); + const sorted = [...results].sort((a, b) => b.score - a.score); + expect(sorted[0].id).toBe('1'); + }); - it('should boost heading path matches', async () => { - const candidates: RerankCandidate[] = [ - { id: '1', content: 'general information', score: 0.5, headingPath: 'Line Chart > Tooltip' }, - { id: '2', content: 'chart general information', score: 0.5 }, - ]; + it('should boost candidates with substring matches', async () => { + const candidates: RerankCandidate[] = [ + { id: '1', content: 'some chart API', score: 0.5 }, + { id: '2', content: 'line API', score: 0.5 }, + ]; - const results = await reranker.rerank('tooltip', candidates); - const sorted = [...results].sort((a, b) => b.score - a.score); - expect(sorted[0].id).toBe('1'); - }); + const results = await rerank('chart', candidates); + const sorted = [...results].sort((a, b) => b.score - a.score); + expect(sorted[0].id).toBe('1'); + }); - it('should carry over original score as minor factor', async () => { - const candidates: RerankCandidate[] = [ - { id: '1', content: 'random unrelated', score: 0.9 }, - { id: '2', content: 'random other', score: 0.1 }, - ]; + it('should respect custom weights', async () => { + const candidates: RerankCandidate[] = [ + { id: '1', content: 'tooltip', score: 0.5 }, + { id: '2', content: 'chart', score: 0.5 }, + ]; - const results = await reranker.rerank('unrelated query', candidates); - // Even with no keyword matches, higher original score gives slight advantage - // (but reranker score normalization means the difference is small) - expect(results.length).toBe(2); - }); + const results = await rerank('tooltip', candidates, { phraseWeight: 10, termWeight: 0 }); + const sorted = [...results].sort((a, b) => b.score - a.score); + expect(sorted[0].id).toBe('1'); + }); - it('should normalize scores to [0, 1] range', async () => { - const candidates: RerankCandidate[] = [ - { id: '1', content: 'exact match test', score: 0.5 }, - { id: '2', content: 'partial test', score: 0.3 }, - { id: '3', content: 'no match', score: 0.1 }, - ]; + it('should handle heading path bonus', async () => { + const candidates: RerankCandidate[] = [ + { id: '1', content: 'some content', score: 0.5, headingPath: 'Chart > Tooltip' }, + { id: '2', content: 'tooltip content', score: 0.5, headingPath: 'Other Section' }, + ]; - const results = await reranker.rerank('test', candidates); - for (const r of results) { - expect(r.score).toBeGreaterThanOrEqual(0); - expect(r.score).toBeLessThanOrEqual(1); - } - }); + const results = await rerank('tooltip', candidates); + const sorted = [...results].sort((a, b) => b.score - a.score); + expect(sorted[0].id).toBe('1'); + }); - it('should handle CJK queries', async () => { - const candidates: RerankCandidate[] = [ - { id: '1', content: 'ๆŠ˜็บฟๅ›พ้…็ฝฎๆ–นๆณ•่ฏฆ่งฃ', score: 0.5 }, - { id: '2', content: 'ๆŸฑ็Šถๅ›พไฝฟ็”จ่ฏดๆ˜Ž', score: 0.5 }, - ]; + it('should normalize scores to [0, 1]', async () => { + const candidates: RerankCandidate[] = [ + { id: '1', content: 'test content', score: 0.1 }, + { id: '2', content: 'another test content', score: 0.9 }, + ]; - const results = await reranker.rerank('ๆŠ˜็บฟๅ›พ', candidates); - const sorted = [...results].sort((a, b) => b.score - a.score); - expect(sorted[0].id).toBe('1'); - }); + const results = await rerank('test', candidates); + const scores = results.map((r) => r.score); + expect(Math.min(...scores)).toBe(0); + expect(Math.max(...scores)).toBe(1); }); -}); -describe('createReranker', () => { - it('should return a KeywordReranker instance', () => { - const reranker = createReranker(); - expect(reranker).toBeInstanceOf(KeywordReranker); + it('should handle Chinese queries', async () => { + const candidates: RerankCandidate[] = [ + { id: '1', content: 'ๆŠ˜็บฟๅ›พ้…็ฝฎ', score: 0.5 }, + { id: '2', content: 'ๆŸฑ็Šถๅ›พ้…็ฝฎ', score: 0.5 }, + { id: '3', content: 'ๅ…ถไป–ๅ†…ๅฎน', score: 0.5 }, + ]; + + const results = await rerank('ๆŠ˜็บฟๅ›พ', candidates); + const sorted = [...results].sort((a, b) => b.score - a.score); + expect(sorted[0].id).toBe('1'); }); -}); +}); \ No newline at end of file diff --git a/test/synonym-expander.test.ts b/test/synonym-expander.test.ts index a5a7b80..ae0e75c 100644 --- a/test/synonym-expander.test.ts +++ b/test/synonym-expander.test.ts @@ -1,84 +1,60 @@ import { describe, it, expect } from 'vitest'; -import { SynonymExpander, NoopExpander } from '../src/utils/expander'; +import { expand } from '../src/expander'; -describe('SynonymExpander', () => { - describe('with user-provided synonyms', () => { - const expander = new SynonymExpander({ +describe('expand', () => { + const queryExpansion = { + synonyms: { 'ๆŠ˜็บฟๅ›พ': ['line chart', 'ๆŠ˜็บฟ'], 'tooltip': ['ๆ็คบๆก†', 'ๆ็คบ', 'hover'], 'animation': ['ๅŠจๆ•ˆ', 'animate', 'transition'], 'config': ['้…็ฝฎ', 'configuration', '่ฎพ็ฝฎ'], - }); + }, + }; - it('should expand CN chart type to EN equivalent', () => { - const result = expander.expand('ๆŠ˜็บฟๅ›พ'); - expect(result).toContain('line chart'); - }); - - it('should expand EN term to CN equivalent', () => { - const result = expander.expand('tooltip'); - expect(result).toContain('ๆ็คบๆก†'); - }); - - it('should not duplicate terms already in query', () => { - const result = expander.expand('tooltip ๆ็คบๆก†'); - const promptCount = result.split('ๆ็คบๆก†').length - 1; - expect(promptCount).toBe(1); - }); - - it('should expand multiple terms in one query', () => { - const result = expander.expand('tooltip config'); - expect(result).toContain('ๆ็คบๆก†'); - expect(result).toContain('้…็ฝฎ'); - }); - - it('should preserve original query text', () => { - const result = expander.expand('animation settings'); - expect(result.startsWith('animation settings')).toBe(true); - }); - - it('should return original query when no synonyms match', () => { - const result = expander.expand('random unrelated terms'); - expect(result).toBe('random unrelated terms'); - }); + it('should expand CN chart type to EN equivalent', () => { + const result = expand('ๆŠ˜็บฟๅ›พ', queryExpansion); + expect(result).toContain('line chart'); + }); - it('should handle empty query', () => { - const result = expander.expand(''); - expect(result).toBe(''); - }); + it('should expand EN term to CN equivalent', () => { + const result = expand('tooltip', queryExpansion); + expect(result).toContain('ๆ็คบๆก†'); }); - describe('with no synonyms (empty map)', () => { - const expanderEmpty = new SynonymExpander({}); - const expanderDefault = new SynonymExpander(); + it('should not duplicate terms already in query', () => { + const result = expand('tooltip ๆ็คบๆก†', queryExpansion); + const promptCount = result.split('ๆ็คบๆก†').length - 1; + expect(promptCount).toBe(1); + }); - it('should return original query unchanged with empty map', () => { - const result = expanderEmpty.expand('tooltip configuration'); - expect(result).toBe('tooltip configuration'); - }); + it('should expand multiple terms in one query', () => { + const result = expand('tooltip config', queryExpansion); + expect(result).toContain('ๆ็คบๆก†'); + expect(result).toContain('้…็ฝฎ'); + }); - it('should return original query unchanged with no arguments', () => { - const result = expanderDefault.expand('tooltip configuration'); - expect(result).toBe('tooltip configuration'); - }); + it('should preserve original query text', () => { + const result = expand('animation settings', queryExpansion); + expect(result.startsWith('animation settings')).toBe(true); + }); - it('should handle empty query', () => { - const result = expanderDefault.expand(''); - expect(result).toBe(''); - }); + it('should return original query when no synonyms match', () => { + const result = expand('random unrelated terms', queryExpansion); + expect(result).toBe('random unrelated terms'); }); -}); -describe('NoopExpander', () => { - const expander = new NoopExpander(); + it('should handle empty query', () => { + const result = expand('', queryExpansion); + expect(result).toBe(''); + }); - it('should return query unchanged', () => { - const result = expander.expand('tooltip configuration'); + it('should return original query with no synonyms', () => { + const result = expand('tooltip configuration'); expect(result).toBe('tooltip configuration'); }); - it('should handle empty query', () => { - const result = expander.expand(''); - expect(result).toBe(''); + it('should return original query with empty synonyms', () => { + const result = expand('tooltip configuration', {}); + expect(result).toBe('tooltip configuration'); }); -}); +}); \ No newline at end of file diff --git a/test/zvec-store.test.ts b/test/zvec-store.test.ts deleted file mode 100644 index 1370cd1..0000000 --- a/test/zvec-store.test.ts +++ /dev/null @@ -1,86 +0,0 @@ -import { describe, it, expect, beforeEach, afterEach } from 'vitest'; -import * as fs from 'fs'; -import * as path from 'path'; -import { - createZvecStore, - isZvecAvailable, -} from '../src/storage/zvec-store'; -import type { IZvecStore, ZvecStoreConfig } from '../src/storage/zvec-store'; - -const TEST_DIR = path.join(__dirname, '.test-tmp'); -const STORE_PATH = path.join(TEST_DIR, 'test.zvec'); - -function makeConfig(dims: number = 4): ZvecStoreConfig { - return { - collectionName: 'test_collection', - vectorField: 'embedding', - vectorDims: dims, - ftsFields: ['content'], - fields: [ - { name: 'content', dataType: 'STRING', indexType: 'FTS', indexOptions: { tokenizerName: 'jieba' } }, - ], - }; -} - -describe('ActualZvecStore (native)', () => { - let store: IZvecStore; - - beforeEach(async () => { - if (fs.existsSync(TEST_DIR)) { - fs.rmSync(TEST_DIR, { recursive: true, force: true }); - } - }); - - afterEach(async () => { - if (store) { - try { await store.close(); } catch { /* ok */ } - } - await new Promise((r) => setTimeout(r, 200)); - if (fs.existsSync(TEST_DIR)) { - try { - fs.rmSync(TEST_DIR, { recursive: true, force: true }); - } catch { /* may be locked */ } - } - }); - - it('should create ActualZvecStore when zvec is available', async () => { - if (!isZvecAvailable()) { - console.warn('Skipping: @zvec/zvec not available'); - return; - } - - const config = makeConfig(4); - store = await createZvecStore(STORE_PATH, config); - - await store.insert([ - { id: 'doc1', vector: [0.1, 0.2, 0.3, 0.4], fields: { content: 'hello' } }, - { id: 'doc2', vector: [0.4, 0.3, 0.2, 0.1], fields: { content: 'world' } }, - ]); - - const results = store.searchSync({ vector: [0.1, 0.2, 0.3, 0.4], topK: 2 }); - expect(results.length).toBe(2); - expect(results[0].id).toBe('doc1'); - }); - - it('should support hybrid search with Full Text Search', async () => { - if (!isZvecAvailable()) { - console.warn('Skipping: @zvec/zvec not available'); - return; - } - - const config = makeConfig(4); - store = await createZvecStore(STORE_PATH, config); - - await store.insert([ - { id: 't1', vector: [1, 0, 0, 0], fields: { content: 'sankey diagram visualization' } }, - { id: 't2', vector: [0, 1, 0, 0], fields: { content: 'bar chart example' } }, - ]); - - const results = store.searchHybridSync({ - queryText: 'sankey', - queryVector: [0.9, 0.1, 0, 0], - topK: 2, - }); - expect(results.length).toBeGreaterThan(0); - }); -});