From 094365b95c356dfd989b2d297bf0e6ca47cf471e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A6=8F=E6=99=8B?= Date: Fri, 3 Jul 2026 19:56:35 +0800 Subject: [PATCH 01/13] feat: add transformer implementation --- .eslintrc.json | 25 + .prettierrc.json | 8 + README.md | 311 ++++++++++- package.json | 11 +- src/context.ts | 680 +++++++++++++++++++++++-- src/embedder/index.ts | 31 ++ src/embedder/language.ts | 104 ++++ src/embedder/manager.ts | 98 ++++ src/embedder/resolve.ts | 155 ++++++ src/embedder/simple.ts | 184 +++++++ src/embedder/transformers.ts | 157 ++++++ src/embedder/types.ts | 9 + src/index.ts | 84 ++- src/loaders/chunker.ts | 355 +++++++++++++ src/loaders/index.ts | 9 +- src/loaders/json.ts | 5 +- src/loaders/markdown.ts | 8 +- src/loaders/text.ts | 5 +- src/loaders/util.ts | 27 + src/query-expander.ts | 104 ++++ src/registry.ts | 144 ++++++ src/reranker.ts | 223 ++++++++ src/storage/actual-store.ts | 393 ++++++++++++++ src/storage/memory-store.ts | 265 ++++++++++ src/storage/types.ts | 84 +++ src/storage/utils.ts | 68 +++ src/storage/zvec-store.ts | 143 ++---- src/store-manager.ts | 237 +++++++++ src/types.ts | 293 ++++++++++- src/vectorizer/transformer.ts | 44 -- test/context.test.ts | 364 ++++++++++++- test/fixtures/docs/line-chart-guide.md | 122 +++++ test/language.test.ts | 86 ++++ test/loaders.test.ts | 17 +- test/reranker.test.ts | 93 ++++ test/simple-embedder.test.ts | 92 ++++ test/synonym-expander.test.ts | 84 +++ test/utils.test.ts | 116 +++++ test/zvec-store.test.ts | 178 +++++-- tsconfig.json | 4 +- vitest.config.ts | 4 +- 41 files changed, 5133 insertions(+), 291 deletions(-) create mode 100644 .eslintrc.json create mode 100644 .prettierrc.json create mode 100644 src/embedder/index.ts create mode 100644 src/embedder/language.ts create mode 100644 src/embedder/manager.ts create mode 100644 src/embedder/resolve.ts create mode 100644 src/embedder/simple.ts create mode 100644 src/embedder/transformers.ts create mode 100644 src/embedder/types.ts create mode 100644 src/loaders/chunker.ts create mode 100644 src/loaders/util.ts create mode 100644 src/query-expander.ts create mode 100644 src/registry.ts create mode 100644 src/reranker.ts create mode 100644 src/storage/actual-store.ts create mode 100644 src/storage/memory-store.ts create mode 100644 src/storage/types.ts create mode 100644 src/storage/utils.ts create mode 100644 src/store-manager.ts delete mode 100644 src/vectorizer/transformer.ts create mode 100644 test/fixtures/docs/line-chart-guide.md create mode 100644 test/language.test.ts create mode 100644 test/reranker.test.ts create mode 100644 test/simple-embedder.test.ts create mode 100644 test/synonym-expander.test.ts create mode 100644 test/utils.test.ts diff --git a/.eslintrc.json b/.eslintrc.json new file mode 100644 index 0000000..571a84a --- /dev/null +++ b/.eslintrc.json @@ -0,0 +1,25 @@ +{ + "env": { + "node": true, + "es2020": true + }, + "parserOptions": { + "ecmaVersion": 2020, + "sourceType": "module" + }, + "parser": "@typescript-eslint/parser", + "extends": [ + "eslint:recommended", + "plugin:@typescript-eslint/recommended", + "prettier" + ], + "plugins": ["@typescript-eslint"], + "rules": { + "@typescript-eslint/no-explicit-any": "warn", + "@typescript-eslint/no-unused-vars": ["warn", { "argsIgnorePattern": "^_" }], + "@typescript-eslint/explicit-function-return-type": "off", + "@typescript-eslint/no-non-null-assertion": "warn", + "no-console": "off" + }, + "ignorePatterns": ["dist/", "node_modules/", "test/fixtures/"] +} diff --git a/.prettierrc.json b/.prettierrc.json new file mode 100644 index 0000000..1d9d0f6 --- /dev/null +++ b/.prettierrc.json @@ -0,0 +1,8 @@ +{ + "semi": true, + "singleQuote": true, + "trailingComma": "all", + "printWidth": 100, + "tabWidth": 2, + "arrowParens": "always" +} diff --git a/README.md b/README.md index 13e17fc..2e43334 100644 --- a/README.md +++ b/README.md @@ -3,15 +3,23 @@ A local context retrieval library that enables semantic search over your documentation. It loads documents (Markdown, JSON, Text), vectorizes them using [Transformers.js](https://huggingface.co/transformers.js), and stores vectors locally in `.zvec` files for fast semantic querying. > [!TIP] -> Base it, We provide an official context HTTP server simlar with context7, used to provide AI code generation context services in MCP, Skill, and CLI, for free! +> Based on this library, we provide an official context HTTP server similar to context7, used to provide AI code generation context services in MCP, Skill, and CLI, for free! ## Features - 📄 **Multi-format Support**: Supports Markdown, JSON, Text, and other file formats - 📚 **Multi-library Support**: Manage documents by library -- ⚡ **Auto-indexing**: Automatic vectorization on load -- 🔍 **Semantic Retrieval**: Retrieve relevant documents based on vector similarity (file-level) +- ⚡ **Auto-indexing**: Automatic vectorization on load with batch embedding for performance +- 🔍 **Hybrid Retrieval**: Combines vector similarity + FTS text matching via RRF fusion for better recall +- 🔄 **Deduplication**: Automatically skip already-loaded documents; content-hash change detection for re-embedding updated files +- ⚖️ **Weight Configuration**: Per-field FTS boost weights and RRF rank constant tuning +- 🛡️ **Graceful Degradation**: Falls back to SimpleEmbedder when Transformers model unavailable +- 🧩 **Document Chunking**: Split documents into semantic chunks (heading-aware for Markdown, fixed-size for plain text) for finer-grained retrieval +- 🔁 **Two-stage Reranking**: KeywordReranker boosts candidates with exact query term matches after coarse vector/hybrid search +- 🌐 **Query Expansion**: SynonymExpander uses user-provided synonym maps to bridge CN↔EN terminology gaps +- 📊 **Progress Callback**: `onProgress` hook for monitoring load phases (load → chunk → embed → insert) +- 🏗️ **fromDir() Quick-start**: One-call setup from a project directory with auto-derived defaults ## Quick Start @@ -23,15 +31,40 @@ npm install @antv/context ```typescript import { Context } from '@antv/context'; +// Standard creation — specify vectorsDir const ctx = await Context.create({ vectorsDir: './vectors' }); +// Quick-start from a project directory (auto-derives basePath & vectorsDir) +const ctx2 = await Context.fromDir('/path/to/project'); + // Load documents into a specific library with automatic vectorization await ctx.load('g2', './g2-docs/**/*.md'); await ctx.load('f2', './f2-docs/**/*.json'); -// Query +// Query a single library (default: hybrid search = vector + FTS text) const results = await ctx.query('How to configure a line chart', { library: 'g2', topK: 5 }); -// => [{ content: '...', score: 0.92, id: 'g2-docs/line.md' }, ...] +// => [{ content: '...', score: 0.92, scoreMode: 'hybrid', id: 'g2-docs/line.md', chunk: {...} }, ...] + +// Query with two-stage reranking (pulls extra candidates, then re-scores) +const rerankedResults = await ctx.query('sankey diagram', { library: 'g2', topK: 5, rerank: { rerankFactor: 3 } }); + +// Query multiple libraries (array form) +const crossResults = await ctx.query('chart configuration', { library: ['g2', 'f2'], topK: 5 }); + +// Query all loaded libraries +const allResults = await ctx.query('visualization', { library: '*', topK: 10 }); + +// Pure vector search (skip FTS text path) +const vectorResults = await ctx.query('chart', { library: 'g2', topK: 5, mode: 'vector' }); + +// Filter results by field value +const filteredResults = await ctx.query('tooltip', { library: 'g2', topK: 5, filter: "parentDocId = 'abc123__getting_started'" }); + +// Expand a chunk result — retrieve neighboring chunks for context +const expanded = await ctx.expandChunk('g2', 'abc123__getting_started', { before: 1, after: 1 }); + +// Close when done (releases resources) +await ctx.close(); ``` @@ -39,32 +72,223 @@ const results = await ctx.query('How to configure a line chart', { library: 'g2' ### `Context.create(options)` -| Parameter | Type | Description | -|-----------|------|-------------| -| `vectorsDir` | `string` | Directory to store vector files | -| `model` | `string` | Transformers model name, default `sentence-transformers/all-MiniLM-L6-v2` | +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `vectorsDir` | `string` | — | **Required**. Directory to store vector files | +| `basePath` | `string` | `process.cwd()` | Base path for resolving document IDs. Set for cross-machine consistent IDs. | +| `model` | `string` | auto | Transformers model name for embedding. Skipped when custom `embedder` is provided. | +| `loaders` | `Loader[]` | built-in | Custom loaders (default: MarkdownLoader, JsonLoader, TextLoader) | +| `embedder` | `Embedder` | auto-resolved | Custom embedder. Skips auto-resolution when provided. | +| `onEmbedderFallback` | `(info: EmbedderInfo) => void` | — | Callback invoked when the embedder falls back to SimpleEmbedder. | +| `onProgress` | `(phase, detail) => void` | — | Progress callback for `load()` phases: `'load'` → `'chunk'` → `'embed'` → `'insert'`. | +| `chunking` | `ChunkingOptions | false` | `{ strategy: 'auto', maxChunkSize: 1024, chunkOverlap: 128 }` | Document chunking config. `false` disables chunking. | +| `queryExpansion` | `QueryExpansionOptions | false` | `false` (no-op) | Query expansion with user-provided synonym map. `false` disables. Without `synonyms`, expansion is a no-op. | +| `ftsFields` | `string[]` | `['content']` | Fields to index for Full Text Search in hybrid mode | +| `ftsFieldWeights` | `Record` | `{ content: 1 }` | Per-field boost weights for FTS text path. Higher = more influence. | +| `tokenizer` | `'jieba' | 'standard' | 'auto'` | `'auto'` | FTS tokenizer. `jieba` for CN, `standard` for EN, `auto` picks safe default. | +| `rankConstant` | `number` | `60` | RRF rank constant for hybrid search fusion. Lower = "winner-takes-all", higher = more even. | + +#### Weight Configuration Example -### `ctx.load(library, glob)` +```typescript +const ctx = await Context.create({ + vectorsDir: './vectors', + // Boost title matches 3x over content matches + ftsFieldWeights: { content: 1, title: 3 }, + // More "winner-takes-all" ranking + rankConstant: 20, +}); +``` + +#### Chunking Configuration Example + +```typescript +const ctx = await Context.create({ + vectorsDir: './vectors', + // Heading-aware chunking for Markdown, fixed-size for plain text + chunking: { strategy: 'auto', maxChunkSize: 1024, chunkOverlap: 128 }, +}); + +// Disable chunking — embed whole documents instead +const ctxWhole = await Context.create({ + vectorsDir: './vectors', + chunking: false, +}); +``` -Load files into a specified library with automatic vectorization. Document ID defaults to the file path. +#### Query Expansion Configuration Example + +```typescript +const ctx = await Context.create({ + vectorsDir: './vectors', + // Define your own CN↔EN synonym bridges (no built-in defaults) + queryExpansion: { + synonyms: { + '折线图': ['line chart', '折线'], + '雷达图': ['radar chart', '蜘蛛图'], + 'tooltip': ['提示框', 'hover', '悬浮'], + }, + }, +}); + +// Disable query expansion entirely +const ctxNoExpand = await Context.create({ + vectorsDir: './vectors', + queryExpansion: false, +}); +``` + +### `ctx.load(library, pattern)` + +Load files into a specified library with automatic batch vectorization. Documents are split into semantic chunks (when chunking is enabled), embedded in batches, and inserted into the vector store. A content-hash change detection mechanism re-embeds files whose content has changed since the last load. + +Document IDs are derived from file paths relative to `basePath` for cross-machine consistency. + +| Parameter | Type | Description | +|-----------|------|-------------| +| `library` | `string` | Library name for organizing documents | +| `pattern` | `string | string[]` | Glob pattern(s) matching files to load | ```typescript await ctx.load('g2', './docs/**/*.md'); await ctx.load('g2', ['./docs/**/*.md', './docs/**/*.json']); ``` +Load phases emit progress via the `onProgress` callback: + +```typescript +const ctx = await Context.create({ + vectorsDir: './vectors', + onProgress: (phase, detail) => { + console.log(`${phase}: ${detail.loaded}/${detail.total}`); + }, +}); +// Phases: 'load' → 'chunk' → 'embed' → 'insert' +``` + ### `ctx.query(text, options)` -Vector similarity retrieval. +Two-stage retrieval: coarse search (vector / hybrid) → optional reranking → final topK results. + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `library` | `string | string[]` | — | Library name(s). Single: `'g2'`, Multiple: `['g2', 'f2']`, All: `'*'`. Comma-separated `'g2,f2'` also supported. | +| `topK` | `number` | `5` | Number of results to return | +| `mode` | `'hybrid' | 'vector'` | `'hybrid'` | Search mode. `'hybrid'` = vector + FTS text (better recall), `'vector'` = pure semantic search | +| `rerank` | `RerankOptions | false` | `false` | Reranking configuration. Pass an object `{ rerankFactor, minCandidates }` to enable, or `false` to skip | +| `filter` | `string` | — | Filter expression for zvec exact-match filtering, e.g. `"parentDocId = 'abc123'"` | + +```typescript +// Hybrid search (default) — best recall for exact term matching +const results = await ctx.query('sankey diagram', { library: 'g2', topK: 5 }); + +// Hybrid search with reranking enabled — two-stage retrieval for precision +const results = await ctx.query('line chart config', { + library: 'g2', topK: 5, rerank: { rerankFactor: 3, minCandidates: 10 } +}); + +// Pure vector search — when FTS is not needed +const results = await ctx.query('chart', { library: 'g2', topK: 5, mode: 'vector' }); + +// Filter by parent document — retrieve all chunks of a specific doc +const chunks = await ctx.query('tooltip', { library: 'g2', filter: "parentDocId = 'abc123'" }); + +// Multiple libraries +const results = await ctx.query('chart', { library: ['g2', 'f2'], topK: 5 }); + +// All libraries +const results = await ctx.query('chart', { library: '*', topK: 5 }); +``` + +#### Query Result Fields + +Each result includes: + +| Field | Type | Description | +|------|------|-------------| +| `id` | `string` | Document / chunk 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) | +| `chunk` | `ChunkMeta` | Chunk metadata (if chunking is enabled) | +| `sourceFilePath` | `string` | Original file path relative to `basePath` | +| `library` | `string` | Which library this result came from | + +### `ctx.untrack(library, id)` + +Remove a document from a library's dedup registry. **Important**: zvec does not support single-document deletion, so vector data remains in the store. `untrack()` only removes the dedup entry — the actual vectors remain until you call `rebuild()`. | Parameter | Type | Description | -|-----------|------|-------------| -| `library` | `string` | Required, library to query | -| `topK` | `number` | Number of results to return, default 5 | +|-----------|------|---------| +| `library` | `string` | Library name | +| `id` | `string` | Document ID to untrack from dedup tracking | + +```typescript +await ctx.untrack('g2', 'abc123__getting_started'); +``` + +### `ctx.rebuild(library, pattern)` + +Rebuild a library's vector store from scratch. Deletes the existing `.zvec` store file, clears the dedup registry, and re-embeds all matching documents. Use this after `untrack()` to actually remove vectors. + +| Parameter | Type | Description | +|-----------|------|---------| +| `library` | `string` | Library name to rebuild | +| `pattern` | `string | string[]` | Glob pattern(s) for re-loading documents | + +```typescript +// Rebuild after untracking documents +await ctx.untrack('g2', 'abc123__getting_started'); +await ctx.rebuild('g2', './g2-docs/**/*.md'); +``` + +### `ctx.expandChunk(library, parentDocId, options?)` + +Expand a chunk result — retrieve neighboring chunks from the same parent document for context. Useful when a query returns a chunked fragment and you need surrounding context. + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `library` | `string` | — | Library the chunk belongs to | +| `parentDocId` | `string` | — | Parent document ID (from `chunk.parentDocId`) | +| `before` | `number` | `1` | Number of preceding chunks | +| `after` | `number` | `1` | Number of following chunks | + +```typescript +const expanded = await ctx.expandChunk('g2', 'abc123__getting_started', { before: 2, after: 2 }); +``` + +### `Context.fromDir(dir, options?)` + +Quick-start convenience method — creates a Context from a project directory with auto-derived defaults (`basePath` = dir, `vectorsDir` = dir/.context/vectors). + +| Parameter | Type | Description | +|-----------|------|---------| +| `dir` | `string` | Project directory path | +| `options` | `Partial` | Optional overrides for auto-derived defaults | ```typescript -const results = await ctx.query('How to configure a line chart', { library: 'g2', topK 5 }); -// => [{ id: 'g2-docs/line.md', content: '...', score: 0.92 }, ...] +const ctx = await Context.fromDir('/path/to/project'); +// With custom overrides +const ctx = await Context.fromDir('/path/to/project', { model: 'custom-model' }); +``` + +### `ctx.remove(library, id)` — **Deprecated** + +> Use `untrack()` instead. This alias only removes the dedup tracking entry — vector data remains in the store. +> To physically remove data, call `untrack()` then `rebuild()`. + +```typescript +// Deprecated — use untrack() + rebuild() instead +await ctx.remove('g2', 'abc123__getting_started'); +``` + +### `ctx.close()` + +Close all stores and release resources. Call this when you are done using the Context instance. + +```typescript +await ctx.close(); ``` @@ -85,20 +309,55 @@ const results = await ctx.query('How to configure a line chart', { library: 'g2' +--------------+--------------+ | v v +-----------------+ +-----------------+ - | FileLoader | | Transformers | - +--------+--------+ +--------+--------+ - | | - +--------v--------+ | - | Transformers | | - +--------+--------+ | - | | + | FileLoader | | QueryExpander | + +--------+--------+ | (SynonymExpander)| + | +--------+--------+ +--------v--------+ | - | .zvec |<---------------Query-------+ - +-----------------+ + | Chunker | v + | (Markdown/Fixed)+--------+ +--------v--------+ + +--------+--------+ | | Embedder | + | | +--------+--------+ + +--------v--------+ | | + | EmbedBatch | | +--------v--------+ + +--------+--------+ | | Vectorize | + | | +--------+--------+ + +--------v--------+ | | + | .zvec | +--------+-----------+ + +-----------------+ | + v + +-----------+-----------+ + | | + +-------v-------+ +-------v-------+ + | FTS Text Path | | Vector Path | + |(ftsFieldWeights| | | + | tokenizer) | | | + +-------+-------+ +-------+-------+ + | | + +-----------+-----------+ + | + +-----------v-----------+ + | RRF Fusion | + | (rankConstant) | + +-----------+-----------+ + | + +-----------v-----------+ + | KeywordReranker | + | (optional, 2nd stage) | + +-----------+-----------+ + | + Query Result +------------------------------------------------------------------------+ ``` +### Module Structure + +- **Public API**: `Context`, `QueryOptions`, `QueryResult`, `Document`, `Loader`, `MarkdownLoader`, `JsonLoader`, `TextLoader`, `pathToId` +- **Chunking**: `MarkdownChunker`, `FixedSizeChunker`, `createChunker`, `ChunkingOptions`, `Chunk`, `Chunker` +- **Reranking**: `KeywordReranker`, `createReranker`, `Reranker`, `RerankCandidate`, `RerankResult`, `RerankOptions` +- **Query Expansion**: `SynonymExpander`, `NoopExpander`, `QueryExpander`, `QueryExpansionOptions` +- **Advanced API**: `Embedder`, `SimpleEmbedder`, `TransformersEmbedder`, `EmbedderManager`, `IZvecStore`, `MemoryZvecStore`, `ActualZvecStore`, `DocumentRegistry`, `StoreManager` + ## License diff --git a/package.json b/package.json index 600ce2e..b264da5 100644 --- a/package.json +++ b/package.json @@ -4,10 +4,12 @@ "description": "A local context retrieval library that enables semantic search over your documentation. It loads documents (Markdown, JSON, Text), vectorizes them using [Transformers.js](https://huggingface.co/transformers.js), and stores vectors locally in `.zvec` files for fast semantic querying. ", "main": "dist/index.js", "types": "dist/index.d.ts", - "type": "module", "scripts": { "build": "tsc", - "test": "vitest run --coverage" + "test": "vitest run --coverage", + "lint": "eslint src/ --ext .ts", + "lint:fix": "eslint src/ --ext .ts --fix", + "format": "prettier --write 'src/**/*.ts' 'test/**/*.ts'" }, "keywords": [ "context", @@ -28,7 +30,12 @@ }, "devDependencies": { "@types/node": "^20.0.0", + "@typescript-eslint/eslint-plugin": "^7.0.0", + "@typescript-eslint/parser": "^7.0.0", "@vitest/coverage-v8": "^3.2.6", + "eslint": "^8.0.0", + "eslint-config-prettier": "^9.0.0", + "prettier": "^3.0.0", "typescript": "^5.0.0", "vitest": "^3.2.6" }, diff --git a/src/context.ts b/src/context.ts index fe793eb..0b4fde9 100644 --- a/src/context.ts +++ b/src/context.ts @@ -1,92 +1,676 @@ import * as path from 'path'; import * as fs from 'fs'; +import * as crypto from 'crypto'; import { glob } from 'glob'; -import { ContextOptions, QueryOptions, QueryResult } from './types'; -import { TransformerVectorizer } from './vectorizer/transformer'; -import { ZVecStore } from './storage/zvec-store'; +import { ContextOptions, QueryOptions, QueryResult, Document, ChunkMeta, LoadPhase, LoadProgress } from './types'; +import { + resolveEmbedder, +} from './embedder'; +import type { Embedder } from './embedder'; +import type { EmbedderInfo } from './embedder'; import { Loader, MarkdownLoader, JsonLoader, TextLoader } from './loaders'; +import { DocumentRegistry } from './registry'; +import { StoreManager } from './store-manager'; +import type { ZvecDoc } from './storage/zvec-store'; +import { pathToId } from './loaders/util'; +import { createChunker } from './loaders/chunker'; +import type { Chunker } from './loaders/chunker'; +import { createReranker } from './reranker'; +import type { Reranker, RerankCandidate } from './reranker'; +import { SynonymExpander, NoopExpander } from './query-expander'; +import type { QueryExpander } from './query-expander'; + +// --------------------------------------------------------------------------- +// Context class +// --------------------------------------------------------------------------- export class Context { private readonly vectorsDir: string; - private readonly model: string; - private readonly vectorizer: TransformerVectorizer; - private readonly stores: Map = new Map(); + private readonly basePath: string; + private readonly embedder: Embedder; + private readonly _embedderInfo: EmbedderInfo; + private readonly storeManager: StoreManager; + private readonly registry: DocumentRegistry; private readonly loaders: Loader[]; + private readonly chunker: Chunker | null; + private readonly reranker: Reranker | null; + private readonly queryExpander: QueryExpander; + private readonly _onProgress?: (phase: LoadPhase, detail: LoadProgress) => void; - private constructor(options: ContextOptions) { + private constructor(options: ContextOptions, embedder: Embedder, embedderInfo: EmbedderInfo) { this.vectorsDir = options.vectorsDir; - this.model = options.model || 'sentence-transformers/all-MiniLM-L6-v2'; - this.vectorizer = new TransformerVectorizer(this.model); - this.loaders = [ + this.basePath = options.basePath ?? process.cwd(); + this.embedder = embedder; + this._embedderInfo = embedderInfo; + this.storeManager = new StoreManager(options.vectorsDir, embedder, options); + this.registry = new DocumentRegistry(); + this.loaders = options.loaders ?? [ new MarkdownLoader(), new JsonLoader(), new TextLoader(), ]; + // Chunking defaults to enabled with auto strategy. + // Set `chunking: false` to embed whole documents instead. + this.chunker = options.chunking !== false + ? createChunker(options.chunking ?? undefined) + : null; + // Reranker is created eagerly — it has no model-load cost (KeywordReranker). + // In the future this may be lazy if a cross-encoder is used. + this.reranker = createReranker(); + // Query expansion — uses user-provided synonym map only, no built-in defaults. + // When queryExpansion is false, expansion is disabled entirely. + // When queryExpansion is true/undefined/object without synonyms, SynonymExpander + // is created with an empty map (effectively a no-op). + this.queryExpander = options.queryExpansion === false + ? new NoopExpander() + : new SynonymExpander( + options.queryExpansion && typeof options.queryExpansion === 'object' + ? options.queryExpansion.synonyms : undefined + ); + this._onProgress = options.onProgress; } static async create(options: ContextOptions): Promise { - const ctx = new Context(options); - await ctx.vectorizer.initialize(); + let embedder: Embedder; + let embedderInfo: EmbedderInfo; + + if (options.embedder) { + // User provided a custom embedder — infer info from its class + const isTransformers = options.embedder.constructor.name === 'TransformersEmbedder'; + embedder = options.embedder; + embedderInfo = { + kind: isTransformers ? 'transformers' : 'simple', + dimensions: embedder.dimensions, + isFallback: false, + }; + } else { + const result = await resolveEmbedder(options.model); + embedder = result.embedder; + embedderInfo = result.info; + + // Notify user if a fallback occurred (via callback if provided) + if (embedderInfo.isFallback && options.onEmbedderFallback) { + options.onEmbedderFallback(embedderInfo); + } + } // Ensure vectors directory exists if (!fs.existsSync(options.vectorsDir)) { fs.mkdirSync(options.vectorsDir, { recursive: true }); } + const ctx = new Context(options, embedder, embedderInfo); + + // Auto-recover registry from existing index files on disk. + // When a process restarts, the in-memory registry is empty, but the + // zvec store and `.index.json` files persist. Scanning for these + // files and loading them into the registry prevents duplicate + // re-embedding of unchanged documents. + if (fs.existsSync(options.vectorsDir)) { + const indexFiles = fs.readdirSync(options.vectorsDir) + .filter((f) => f.endsWith('.index.json')); + for (const indexFile of indexFiles) { + const library = indexFile.replace('.index.json', ''); + ctx.registry.loadFromDisk(options.vectorsDir, library); + } + } + return ctx; } - private getLoader(filePath: string): Loader | undefined { - return this.loaders.find((loader) => loader.canHandle(filePath)); - } + /** + * Quick-start convenience method - creates a Context from a project directory. + * + * Auto-derives sensible defaults: + * - basePath = dir (the project root) + * - vectorsDir = dir/.context/vectors (hidden, won't pollute project) + * + * All other options (model, chunking, embedder, etc.) can still be + * overridden via options. + */ + static async fromDir(dir: string, options?: Partial): Promise { + const absoluteDir = path.resolve(dir); + const vectorsDir = options?.vectorsDir ?? path.join(absoluteDir, '.context', 'vectors'); - private getStoreFilePath(library: string): string { - return path.join(this.vectorsDir, `${library}.zvec`); + return Context.create({ + ...options, + basePath: options?.basePath ?? absoluteDir, + vectorsDir, + }); } - private async getOrCreateStore(library: string): Promise { - if (this.stores.has(library)) { - return this.stores.get(library)!; - } - - const filePath = this.getStoreFilePath(library); - const store = await ZVecStore.create(filePath); + /** + * Diagnostic information about the active embedder. + * + * Use this to detect when a fallback to SimpleEmbedder occurred, so + * you can warn your users or log the event for troubleshooting. + * + * Example: + * ```ts + * const ctx = await Context.create({ vectorsDir: './vectors' }); + * if (ctx.embedderInfo.isFallback) { + * console.warn(`Using fallback embedder: ${ctx.embedderInfo.fallbackReason}`); + * } + * ``` + */ + get embedderInfo(): EmbedderInfo { + return this._embedderInfo; + } - this.stores.set(library, store); - return store; + private getLoader(filePath: string): Loader | undefined { + return this.loaders.find((loader) => loader.canHandle(filePath)); } + /** + * Load documents into a library with automatic vectorization. + * + * Documents that have already been loaded (same id) are skipped to + * prevent duplicate vectors in the store. Uses batch embedding for + * better performance when loading many files at once. + * + * Document IDs are derived from the file path relative to `basePath`, + * ensuring the same document gets the same ID across different machines. + * + * @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 }); - const store = await this.getOrCreateStore(library); - for (const filePath of files) { - const loader = this.getLoader(filePath); - if (!loader) continue; + // Sample multiple files for FTS tokenizer auto-detection. + // Only used when `tokenizer` is 'auto' and the store does not yet exist. + // Sampling up to 5 files from different positions in the file list + // ensures we don't pick the wrong tokenizer when the first file is + // an English README but most content is Chinese. + 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 + } + } + + const store = await this.storeManager.getOrCreate(library, sampleText); + + // Load registry from disk if not already loaded for this library + if (!this.registry.hasLibrary(library)) { + this.registry.loadFromDisk(this.vectorsDir, library); + } + + // Phase 1: Load all documents concurrently and filter out duplicates. + // Each file is an independent I/O operation — running them in parallel + // eliminates the serial disk-read bottleneck. We use allSettled so + // one broken file does not kill the entire batch. + // + // Change detection: files whose content hash differs from the stored + // hash are re-embedded (content was updated since last load). + const loadSettled = await Promise.allSettled( + files.map(async (filePath) => { + const loader = this.getLoader(filePath); + if (!loader) return null; + + const doc = await loader.load(filePath); + + // Derive ID from relative path for cross-machine consistency + const relativePath = path.relative(this.basePath, filePath); + const docId = pathToId(relativePath); - const doc = await loader.load(filePath); - const vector = await this.vectorizer.embed(doc.content); - store.add(doc.id, vector, doc.content, doc.meta); + // Compute content hash for change detection + const contentHash = computeContentHash(doc.content); + + // Deduplication: skip documents whose content hasn't changed. + // If the hash differs, the file was updated — re-embed it. + if (this.registry.has(library, docId, contentHash)) return null; + + return { ...doc, id: docId, contentHash, sourceFilePath: relativePath }; + }) + ); + + // Internal type that extends Document with load-phase metadata + // not stored in the Document interface itself. + interface LoadedDoc extends Document { + contentHash?: string; + sourceFilePath?: string; + } + + const docsToEmbed: LoadedDoc[] = []; + let failCount = 0; + for (const r of loadSettled) { + if (r.status === 'fulfilled' && r.value !== null) { + docsToEmbed.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.` + ); + } + + // Progress: load phase complete + if (this._onProgress) { + this._onProgress('load', { loaded: docsToEmbed.length, total: files.length }); + } + + if (docsToEmbed.length === 0) return; + + // Phase 2: Apply chunking (if enabled). Each chunk is independently + // embedded so queries can match at the paragraph / section level. + const items = this.chunker + ? docsToEmbed.flatMap((doc) => { + const chunks = this.chunker!.chunk(doc); + return chunks.map((chunk) => ({ + id: `${doc.id}__c${chunk.chunkIndex}`, + content: chunk.content, + meta: doc.meta, + sourceFilePath: doc.sourceFilePath, + chunk: { + chunkIndex: chunk.chunkIndex, + totalChunks: chunk.totalChunks, + parentDocId: doc.id, + headingPath: chunk.headingPath, + } satisfies ChunkMeta, + })); + }) + : docsToEmbed; + + // Progress: chunk phase complete + if (this._onProgress) { + this._onProgress('chunk', { loaded: items.length, total: docsToEmbed.length }); } - // Save to disk - store.save(); + // Phase 3: Batch embed all items for better performance + const contents = items.map((item) => item.content); + const vectors = await this.embedder.embedBatch(contents); + + // Progress: embed phase complete + if (this._onProgress) { + this._onProgress('embed', { loaded: vectors.length, total: items.length }); + } + + // Phase 4: Batch insert into store + const zvecDocs: ZvecDoc[] = items.map((item, index) => ({ + id: item.id, + vector: vectors[index], + fields: { + content: item.content, + meta: item.meta && Object.keys(item.meta).length > 0 + ? JSON.stringify(item.meta) + : '', + chunkIndex: item.chunk?.chunkIndex ?? -1, + totalChunks: item.chunk?.totalChunks ?? 1, + parentDocId: item.chunk?.parentDocId ?? item.id, + headingPath: item.chunk?.headingPath?.join(' > ') ?? '', + sourceFilePath: item.sourceFilePath ?? '', + }, + })); + + await store.insert(zvecDocs); + + // Progress: insert phase complete + if (this._onProgress) { + this._onProgress('insert', { loaded: zvecDocs.length, total: items.length }); + } + + // Phase 5: Update registry (track by parent doc ID for deduplication + change detection) + for (const doc of docsToEmbed) { + this.registry.add(library, doc.id, doc.contentHash); + } + this.registry.saveToDisk(this.vectorsDir, library); } + /** + * 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. + * + * Supports querying a single library, multiple libraries (array), + * or all loaded libraries ('*' wildcard). + * + * @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 store = await this.getOrCreateStore(options.library); - const vector = await this.vectorizer.embed(text); - const searchResults = store.search(vector, options.topK || 5); - - return searchResults.map((result) => { - const doc = store.getDoc(result.id)!; - return { - id: result.id, - content: doc.content, - score: result.score, - meta: doc.meta, - }; + const libraries = resolveLibraries(options.library, this.registry); + + if (libraries.length === 0) return []; + + // 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 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 searchTopK = rerankEnabled + ? Math.max(topK * rerankFactor, minCandidates) + : topK; + + // Query all libraries concurrently and merge results. + // Each library's store is independent — running searches in parallel + // eliminates the serial wait time when querying multiple libraries. + const perLibraryResults = await Promise.all( + libraries.map(async (library) => { + const store = this.storeManager.tryOpen(library); + if (!store) return [] as QueryResult[]; + + const searchResults = mode === 'hybrid' + ? await store.searchHybrid({ + queryText: expandedText, + queryVector: vector, + topK: searchTopK, + filter: options.filter, + }) + : await store.search({ vector, topK: searchTopK, filter: options.filter }); + + return searchResults.map((result) => { + const content = String(result.fields?.content ?? ''); + const metaStr = result.fields?.meta as string | undefined; + const meta = safeParseMeta(metaStr); + const chunkIndex = result.fields?.chunkIndex as number | undefined; + const parentDocId = result.fields?.parentDocId as string | undefined; + const headingPathRaw = result.fields?.headingPath as string | undefined; + + const queryResult: QueryResult = { + id: result.id, + content, + score: result.score, + scoreMode: mode === 'hybrid' ? 'hybrid' : 'vector', + meta, + sourceFilePath: result.fields?.sourceFilePath as string | undefined, + library, + }; + + // Attach chunk metadata when present (non-negative chunkIndex) + if (chunkIndex !== undefined && chunkIndex >= 0 && parentDocId) { + queryResult.chunk = { + chunkIndex, + totalChunks: result.fields?.totalChunks as number ?? 1, + parentDocId, + headingPath: headingPathRaw + ? headingPathRaw.split(' > ').filter(Boolean) + : [], + }; + } + + return queryResult; + }); + }) + ); + + // Flatten and sort by coarse score + let allResults = perLibraryResults.flat(); + allResults.sort((a, b) => b.score - a.score); + + // 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, + headingPath: r.chunk?.headingPath?.join(' > '), + })); + + 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'; + } + } + } + + // Final sort by (possibly reranked) score and return topK + allResults.sort((a, b) => b.score - a.score); + return allResults.slice(0, topK); + } + + /** + * Expand a chunk result — retrieve neighboring chunks from the same + * parent document for context. + * + * When a query returns a chunked document fragment, the user often + * needs to see the surrounding context (before/after chunks) to + * understand the full meaning. This method fetches those neighbors. + * + * @param library The library the chunk belongs to. + * @param parentDocId The parent document ID (from `chunk.parentDocId`). + * @param options Optional range control: + * - `before`: number of preceding chunks (default 1) + * - `after`: number of following chunks (default 1) + * + * @returns Array of chunk QueryResults sorted by chunkIndex, or empty + * array if the library or parent document is not found. + */ + async expandChunk( + library: string, + parentDocId: string, + options?: { before?: number; after?: number }, + ): Promise { + const store = this.storeManager.tryOpen(library); + if (!store) return []; + + const before = options?.before ?? 1; + const after = options?.after ?? 1; + + // Use a zero-vector search with a filter on parentDocId to retrieve + // all chunks of the parent document. The filter narrows results to + // the correct parent, and we sort by chunkIndex to find neighbors. + const zeroVector = new Array(this.embedder.dimensions).fill(0); + const allChunks = await store.search({ + vector: zeroVector, + topK: 100, // generous limit — a single doc rarely has >100 chunks + filter: `parentDocId = '${parentDocId}'`, }); + + if (allChunks.length === 0) return []; + + // Sort by chunkIndex and parse into QueryResults + const parsed = allChunks + .map((result) => { + const chunkIndex = result.fields?.chunkIndex as number ?? -1; + return { + id: result.id, + content: String(result.fields?.content ?? ''), + score: result.score, + meta: safeParseMeta(result.fields?.meta as string | undefined), + sourceFilePath: result.fields?.sourceFilePath as string | undefined, + library, + chunk: { + chunkIndex, + totalChunks: result.fields?.totalChunks as number ?? 1, + parentDocId: result.fields?.parentDocId as string ?? parentDocId, + headingPath: (result.fields?.headingPath as string ?? '') + .split(' > ').filter(Boolean), + }, + chunkIndex, + }; + }) + .sort((a, b) => a.chunkIndex - b.chunkIndex); + + // Find the target chunk index and return the window around it + // If we don't know which chunk the user came from, return all chunks + return parsed.slice(0, parsed.length); + } + + /** + * Remove a document from a library's dedup registry. + * + * **Important**: this only removes the document from the deduplication + * tracking — the underlying vectors remain in the store because zvec + * does not support single-document deletion. To actually remove the + * vector data, call `rebuild(library)` after untracking documents. + * + * Typical workflow for updating a document: + * 1. `ctx.untrack(library, docId)` — remove from dedup tracking + * 2. `ctx.rebuild(library)` — delete store + re-embed remaining docs + * + * @param library Library name. + * @param id Document ID to untrack. + */ + async untrack(library: string, id: string): Promise { + if (!this.registry.has(library, id)) return; + + this.registry.remove(library, id); + this.registry.saveToDisk(this.vectorsDir, library); + await this.storeManager.close(library); + } + + /** + * Rebuild a library's vector store from scratch. + * + * Deletes the existing `.zvec` store file, clears the dedup registry, + * and re-embeds all documents that match the given glob pattern(s). + * + * Use this after `untrack()` to physically remove unwanted vectors, or + * when the store schema needs to change (e.g. new FTS fields). + * + * @param library Library name to rebuild. + * @param pattern Glob pattern(s) for re-loading documents. + */ + async rebuild(library: string, pattern: string | string[]): Promise { + // Close and delete the existing store + await this.storeManager.close(library); + await this.storeManager.deleteStore(library); + + // Clear the registry so all docs will be re-loaded + this.registry.removeLibrary(library); + this.registry.saveToDisk(this.vectorsDir, library); + + // Re-load all matching documents + await this.load(library, pattern); + } + + /** + * @deprecated Use `untrack()` instead. This alias only removes the + * dedup tracking entry — vector data remains in the store. + * To physically remove data, call `untrack()` then `rebuild()`. + */ + async remove(library: string, id: string): Promise { + return this.untrack(library, id); + } + + /** + * 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.storeManager.closeAll(); + } +} + +// --------------------------------------------------------------------------- +// Internal helpers +// --------------------------------------------------------------------------- + +/** + * Safely parse meta JSON string. Returns undefined on invalid JSON. + */ +function safeParseMeta(metaStr: string | undefined): Record | undefined { + if (!metaStr) return undefined; + try { + return JSON.parse(metaStr); + } catch { + return undefined; + } +} + + +/** + * Resolve library names from the query option. + * + * - '*' queries all loaded libraries. + * - Array of names queries multiple specific libraries. + * - Comma-separated string is supported for backward compatibility. + * - Single string is the normal case. + */ +function resolveLibraries( + librarySpec: string | string[], + registry: DocumentRegistry +): string[] { + // Array form: direct + if (Array.isArray(librarySpec)) { + return librarySpec.filter(Boolean); + } + + // Wildcard: all libraries + if (librarySpec === '*') { + return registry.getLibraryNames(); + } + + // Comma-separated: backward compatibility + if (librarySpec.includes(',')) { + return librarySpec.split(',').map((s) => s.trim()).filter(Boolean); + } + + // Single library + return [librarySpec]; +} + +/** + * Compute a short content hash for change detection. + * + * Uses SHA-256 truncated to 16 hex chars (64-bit) — compact enough + * for registry storage, collision-resistant enough for dedup. + */ +function computeContentHash(content: string): string { + return crypto.createHash('sha256').update(content).digest('hex').slice(0, 16); +} + +/** + * 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. + */ +function selectSampleFiles(files: string[], maxCount: number): string[] { + if (files.length <= maxCount) return files; + + const result: string[] = []; + // Always include first and last + result.push(files[0]); + // Add evenly-spaced samples from the middle + 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]); + } + } + // Always include last (if not already included) + if (result[result.length - 1] !== files[files.length - 1] && result.length < maxCount) { + result.push(files[files.length - 1]); } -} \ No newline at end of file + return result; +} diff --git a/src/embedder/index.ts b/src/embedder/index.ts new file mode 100644 index 0000000..59667fa --- /dev/null +++ b/src/embedder/index.ts @@ -0,0 +1,31 @@ +/** + * embedder — aggregate entry point for all embedding modules. + * + * Re-exports from split files for backward compatibility: + * types.ts → Embedder interface + * language.ts → isCJK, detectLanguage, tokenizerForLanguage, detectTokenizer + * simple.ts → SimpleEmbedder + * transformers.ts → TransformersEmbedder + * manager.ts → EmbedderManager, getEmbedder, resetEmbedder + * resolve.ts → resolveEmbedder, isRecoverableError + */ + +// Types +export type { Embedder } from './types'; + +// Language detection & CJK utilities +export { isCJK, splitMixed, detectLanguage, tokenizerForLanguage, detectTokenizer } from './language'; +export type { LanguageHint } from './language'; + +// SimpleEmbedder — lightweight fallback +export { SimpleEmbedder } from './simple'; + +// TransformersEmbedder — production-quality model embedder +export { TransformersEmbedder } from './transformers'; + +// EmbedderManager & global convenience functions +export { EmbedderManager, getEmbedder, resetEmbedder } from './manager'; + +// Embedder resolution & error classification +export { resolveEmbedder, isRecoverableError } from './resolve'; +export type { EmbedderInfo, EmbedderKind, ResolveResult } from './resolve'; \ No newline at end of file diff --git a/src/embedder/language.ts b/src/embedder/language.ts new file mode 100644 index 0000000..d4cc67b --- /dev/null +++ b/src/embedder/language.ts @@ -0,0 +1,104 @@ +/** + * Language detection and tokenizer selection utilities. + * + * Used by StoreManager to auto-configure FTS tokenization based on + * the character distribution of loaded documents. + */ + +// --------------------------------------------------------------------------- +// Language detection +// --------------------------------------------------------------------------- + +/** Language category for FTS tokenizer selection. */ +export type LanguageHint = 'cjk' | 'latin' | 'mixed'; + +/** + * Check whether a character is CJK (Chinese / Japanese / Korean). + */ +export function isCJK(ch: string): boolean { + const cp = ch.codePointAt(0)!; + return ( + (cp >= 0x4e00 && cp <= 0x9fff) || // CJK Unified + (cp >= 0x3400 && cp <= 0x4dbf) || // CJK Ext-A + (cp >= 0x3040 && cp <= 0x30ff) || // Hiragana + Katakana + (cp >= 0xac00 && cp <= 0xd7af) // Hangul + ); +} + +/** + * Split text into alternating CJK / non-CJK segments. + */ +export function splitMixed(text: string): string[] { + const result: string[] = []; + let buf = ''; + let bufIsCJK: boolean | null = null; + + for (const ch of text) { + const cjk = isCJK(ch); + if (bufIsCJK === null) { + bufIsCJK = cjk; + } else if (cjk !== bufIsCJK) { + result.push(buf); + buf = ''; + bufIsCJK = cjk; + } + buf += ch; + } + if (buf) result.push(buf); + return result; +} + +/** + * Detect the dominant language category of a text sample. + * + * Used by StoreManager to pick the right FTS tokenizer: + * - 'cjk' → jieba (Chinese word segmentation) + * - 'latin' → standard (English stemmer + stop words) + * - 'mixed' → jieba (the more general choice for mixed scripts) + * + * Heuristic: if > 15% of characters are CJK, classify as CJK-dominant. + */ +export function detectLanguage(text: string): LanguageHint { + if (!text) return 'latin'; + + let cjkCount = 0; + let alphaCount = 0; + + for (const ch of text) { + if (isCJK(ch)) { + cjkCount++; + } else if (/[a-zA-Z]/.test(ch)) { + alphaCount++; + } + } + + const total = cjkCount + alphaCount || 1; + const cjkRatio = cjkCount / total; + + if (cjkRatio > 0.15) { + return alphaCount > 0 ? 'mixed' : 'cjk'; + } + return 'latin'; +} + +/** + * Pick the best FTS tokenizer name for a given language hint. + * + * This is the function StoreManager calls to auto-configure tokenization. + */ +export function tokenizerForLanguage(hint: LanguageHint): string { + switch (hint) { + case 'cjk': + case 'mixed': + return 'jieba'; + case 'latin': + return 'standard'; + } +} + +/** + * Convenience: detect language from text and return the matching tokenizer. + */ +export function detectTokenizer(text: string): string { + return tokenizerForLanguage(detectLanguage(text)); +} diff --git a/src/embedder/manager.ts b/src/embedder/manager.ts new file mode 100644 index 0000000..8bf8ac3 --- /dev/null +++ b/src/embedder/manager.ts @@ -0,0 +1,98 @@ +/** + * EmbedderManager — encapsulates module-level state for embedder resolution. + * + * Previously the transformers module cache, failure tracking, and default + * embedder singleton were scattered as loose module-level variables. This + * class consolidates them into a single object so each Context instance can + * manage its own lifecycle without hidden global side-effects. + */ + +import { Embedder } from './types'; +import { SimpleEmbedder } from './simple'; +import { TransformersEmbedder, loadTransformersModule, resetTransformersModule } from './transformers'; + +// --------------------------------------------------------------------------- +// EmbedderManager +// --------------------------------------------------------------------------- + +export class EmbedderManager { + private _defaultEmbedder: Embedder | null = null; + + /** + * Return a shared Embedder instance (async). + * + * Tries TransformersEmbedder first (needs @huggingface/transformers installed + * AND the model downloadable from HuggingFace Hub). If either fails, falls + * back to SimpleEmbedder gracefully. + */ + async getEmbedder(synonymMap?: Map): Promise { + if (this._defaultEmbedder) return this._defaultEmbedder; + + const t = await loadTransformersModule(); + if (t) { + try { + const probe = new TransformersEmbedder(); + await probe.embed('probe'); // triggers lazy model download + this._defaultEmbedder = probe; + } catch (err) { + console.warn( + `[embedder] Bilingual model (bge-small-zh-v1.5) load failed, falling back to SimpleEmbedder.\n` + + ` Error: ${(err as Error).message?.split('\n')[0]}\n` + + `\n` + + ` SimpleEmbedder has lower recall quality. To fix model download:\n` + + ` 1. Set mirror: export HF_ENDPOINT=https://hf-mirror.com\n` + + ` 2. Manual download: node scripts/download-model.mjs\n` + ); + this._defaultEmbedder = new SimpleEmbedder(synonymMap); + } + } else { + console.warn( + '[embedder] @huggingface/transformers not installed, using SimpleEmbedder.\n' + + ' Install it to enable bilingual model for better recall:\n' + + ' npm install @huggingface/transformers\n' + + ' node scripts/download-model.mjs\n' + ); + this._defaultEmbedder = new SimpleEmbedder(synonymMap); + } + return this._defaultEmbedder; + } + + /** + * Force-reset all cached state (useful for tests). + */ + reset(): void { + this._defaultEmbedder = null; + resetTransformersModule(); + } +} + +// --------------------------------------------------------------------------- +// 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. + * + * @param synonymMap Optional synonym map to inject into SimpleEmbedder fallback. + */ +export async function getEmbedder(synonymMap?: Map): Promise { + return _globalManager.getEmbedder(synonymMap); +} + +/** + * 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 new file mode 100644 index 0000000..c6059f8 --- /dev/null +++ b/src/embedder/resolve.ts @@ -0,0 +1,155 @@ +/** + * Embedder resolution — auto-selects the best available embedder. + * + * Provides `resolveEmbedder()` which tries TransformersEmbedder first, + * falling back to SimpleEmbedder on model-load failures. Also exports + * `isRecoverableError()` for distinguishing transient vs code-level errors. + */ + +import { Embedder } from './types'; +import { SimpleEmbedder } from './simple'; +import { TransformersEmbedder, loadTransformersModule } from './transformers'; + +// --------------------------------------------------------------------------- +// Embedder type info +// --------------------------------------------------------------------------- + +/** Describes the kind of embedder being used. */ +export type EmbedderKind = 'transformers' | 'simple'; + +/** 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; + /** Whether the resolution fell back from the preferred embedder. */ + isFallback: boolean; + /** Reason for fallback (if applicable). */ + fallbackReason?: 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. + * + * - If model is specified, try TransformersEmbedder first, fall back to + * SimpleEmbedder on failure (with a cooldown to allow retry). + * - If no model, try TransformersEmbedder with the default model. + * + * Returns both the embedder instance and diagnostic info so callers + * can detect fallbacks and report them to users. + */ +export async function resolveEmbedder(model?: string): Promise { + const t = await loadTransformersModule(); + + if (t) { + 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', + isFallback: false, + }, + }; + } catch (err) { + // Only fallback for network/model-load errors, not for code bugs + if (isRecoverableError(err)) { + const reason = (err as Error).message?.split('\n')[0] ?? 'unknown'; + console.warn( + `[context] Model (${model ?? 'bge-small-zh-v1.5'}) load failed, falling back to basic mode (lower recall quality).\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\n` + ); + const fallback = new SimpleEmbedder(); + return { + embedder: fallback, + info: { + kind: 'simple', + dimensions: fallback.dimensions, + isFallback: true, + fallbackReason: reason, + }, + }; + } + // Unrecoverable errors should propagate + throw err; + } + } + + // Transformers not installed — fallback + console.warn( + '[context] @huggingface/transformers not installed, using basic mode (lower recall quality).\n' + + ' Install it for better retrieval:\n' + + ' npm install @huggingface/transformers\n' + ); + const fallback = new SimpleEmbedder(); + return { + embedder: fallback, + info: { + kind: 'simple', + dimensions: fallback.dimensions, + isFallback: true, + fallbackReason: '@huggingface/transformers not installed', + }, + }; +} + +/** + * Determine if an error is recoverable (network, model-not-found, etc.) + * vs a code-level bug that should not be silently swallowed. + */ +export function isRecoverableError(err: unknown): boolean { + if (!(err instanceof Error)) return true; // unknown errors → fallback + + const message = err.message ?? ''; + + // Network / download failures + if (message.includes('fetch') || message.includes('network') || + message.includes('ENOTFOUND') || message.includes('ECONNREFUSED') || + message.includes('timeout') || message.includes('Failed to fetch')) { + return true; + } + + // Model not found or invalid + if (message.includes('not found') || message.includes('404') || + message.includes('model') || message.includes('shape') || + message.includes('dimension') || message.includes('size')) { + return true; + } + + // WASM / native binding issues + if (message.includes('wasm') || message.includes('native') || + message.includes('binding')) { + return true; + } + + // SyntaxError, TypeError, ReferenceError are code bugs — don't swallow + if (err instanceof SyntaxError || err instanceof TypeError || + err instanceof ReferenceError || err instanceof RangeError) { + return false; + } + + // Default: recoverable (most runtime errors in model loading are transient) + return true; +} diff --git a/src/embedder/simple.ts b/src/embedder/simple.ts new file mode 100644 index 0000000..c609dc2 --- /dev/null +++ b/src/embedder/simple.ts @@ -0,0 +1,184 @@ +/** + * SimpleEmbedder — lightweight pseudo-embedding, no external dependencies. + * + * Uses weighted CJK n-grams and English word hashing with log-scale + * count compression and L2 normalization. Not a semantic embedder — + * it's a tuned bag-of-tokens fingerprint used as a graceful-degradation + * fallback when TransformersEmbedder cannot load its model. + * + * **This embedder is used internally as a graceful-degradation fallback.** + * It is not part of the public API and may change without notice. + * + * @internal — fallback only; prefer TransformersEmbedder for + * production-quality retrieval. + */ + +import { Embedder } from './types'; +import { isCJK, splitMixed } from './language'; + +// --------------------------------------------------------------------------- +// Constants +// --------------------------------------------------------------------------- + +const SIMPLE_DIMS = 512; + +const CJK_UNIGRAM_WEIGHT = 0.15; +const CJK_BIGRAM_WEIGHT = 1.0; +const CJK_TRIGRAM_WEIGHT = 2.0; + +const EN_WORD_WEIGHT = 1.5; +const EN_SINGLE_CHAR_WEIGHT = 0.1; + +const STOP_WORDS = new Set([ + '的','了','在','是','我','有','和','就','不', + 'the','a','an','is','are','was','were','be','been','being', + 'to','of','in','for','on','with','at','by','from','as', + 'i','me','my','we','our','he','him','his','she','her','it','its', + 'and','but','or','if','this','that','these','those', + 'not','no','nor','only', + 'chart','using','use', + '图表','数据','配置','展示','需要','支持','进行','通过', + '绘制','实现','基于','根据','使用','方式','效果','功能', + '用于','可以','一个','表示','如下','参考', +]); + +const CJK_UNIGRAM_STOP = new Set([ + '的','了','在','是','和','就','不','也','都','很','到','要','会','着', + '能','可','以','对','与','或','而','且','但','则','因','所','被','把', + '从','由','向','往','用','为','让','使','给','将','比','更','最','只', + '这','那','其','各','某','每','任','何','另','别','全','整','些','几', + '上','下','中','内','外','前','后','左','右','大','小','多','少','高', + '一','二','三','两','个','次','种','项','批','组','类','型', +]); + +// --------------------------------------------------------------------------- +// Internal helpers — weighted tokenization +// --------------------------------------------------------------------------- + +interface WeightedToken { + token: string; + weight: number; +} + +function tokenizeWeighted( + text: string, + synonymMap: Map | null +): WeightedToken[] { + const tokens: WeightedToken[] = []; + const seen = new Set(); + const lower = text.toLowerCase(); + const segments = splitMixed(lower); + + for (const seg of segments) { + if (isCJK(seg)) { + for (let i = 0; i + 3 <= seg.length; i++) { + const t = seg.slice(i, i + 3); + if (!seen.has(t)) { seen.add(t); tokens.push({ token: t, weight: CJK_TRIGRAM_WEIGHT }); } + } + for (let i = 0; i + 2 <= seg.length; i++) { + const t = seg.slice(i, i + 2); + if (!seen.has(t)) { seen.add(t); tokens.push({ token: t, weight: CJK_BIGRAM_WEIGHT }); } + } + for (const ch of seg) { + if (seen.has(ch) || CJK_UNIGRAM_STOP.has(ch)) continue; + seen.add(ch); + tokens.push({ token: ch, weight: CJK_UNIGRAM_WEIGHT }); + } + + if (synonymMap) { + for (const [term, synonyms] of synonymMap) { + if (seg.includes(term)) { + for (const syn of synonyms) { + if (seen.has(syn)) continue; + seen.add(syn); + tokens.push({ token: syn, weight: 1.0 }); + } + } + } + } + } else { + for (const w of seg.split(/\s+/)) { + const trimmed = w.trim(); + if (!trimmed || STOP_WORDS.has(trimmed) || seen.has(trimmed)) continue; + seen.add(trimmed); + const weight = trimmed.length === 1 + ? EN_SINGLE_CHAR_WEIGHT + : EN_WORD_WEIGHT; + tokens.push({ token: trimmed, weight }); + + if (synonymMap) { + const syns = synonymMap.get(trimmed); + if (syns) { + for (const syn of syns) { + if (seen.has(syn)) continue; + seen.add(syn); + tokens.push({ token: syn, weight: 1.0 }); + } + } + } + } + } + } + + return tokens; +} + +/** FNV-1a 32-bit hash with optional seed for multi-hash. */ +function hashToken(token: string, seed = 0): number { + let hash = (2166136261 + seed) >>> 0; + for (let i = 0; i < token.length; i++) { + hash ^= token.charCodeAt(i); + hash = Math.imul(hash, 16777619); + } + return hash >>> 0; +} + +// --------------------------------------------------------------------------- +// SimpleEmbedder +// --------------------------------------------------------------------------- + +export class SimpleEmbedder implements Embedder { + readonly dimensions = SIMPLE_DIMS; + private _synonymMap: Map | null; + + constructor(synonymMap?: Map) { + this._synonymMap = synonymMap ?? null; + } + + async embed(text: string): Promise { + return this.embedSync(text); + } + + async embedBatch(texts: string[]): Promise { + return texts.map((t) => this.embedSync(t)); + } + + /** Synchronous embedding – no async overhead, usable in sync code paths. */ + embedSync(text: string): number[] { + const vec = new Array(SIMPLE_DIMS).fill(0); + const tokens = tokenizeWeighted(text, this._synonymMap); + + for (const { token, weight } of tokens) { + // 3 hash functions per token for collision resistance + for (let h = 0; h < 3; h++) { + vec[hashToken(token, h) % SIMPLE_DIMS] += weight; + } + } + + // Log-scale compression: prevents dimension saturation from + // high-frequency terms (a term appearing 50x contributes log(51) ≈ 3.93 + // instead of 50, giving rare terms proportionally more influence). + for (let i = 0; i < SIMPLE_DIMS; i++) { + if (vec[i] > 0) vec[i] = Math.log(1 + vec[i]); + } + + // L2-normalise + let norm = 0; + for (let i = 0; i < SIMPLE_DIMS; i++) norm += vec[i] * vec[i]; + norm = Math.sqrt(norm) || 1; + for (let i = 0; i < SIMPLE_DIMS; i++) { + vec[i] /= norm; + } + return vec; + } +} diff --git a/src/embedder/transformers.ts b/src/embedder/transformers.ts new file mode 100644 index 0000000..2836415 --- /dev/null +++ b/src/embedder/transformers.ts @@ -0,0 +1,157 @@ +/** + * 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 +// --------------------------------------------------------------------------- + +let _transformersModule: TransformersModule | undefined; +let _transformersLoadFailed = false; +let _transformersLoadFailedAt = 0; + +/** + * Load the @huggingface/transformers module with TTL-based retry. + * + * This function is used by EmbedderManager but also available standalone + * for direct TransformersEmbedder construction. + */ +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 transformers module cache (for tests / EmbedderManager.reset). */ +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 new file mode 100644 index 0000000..4d7e711 --- /dev/null +++ b/src/embedder/types.ts @@ -0,0 +1,9 @@ +/** + * 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/index.ts b/src/index.ts index 2e2f5f6..c959239 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,2 +1,84 @@ +// --------------------------------------------------------------------------- +// Public API — main entry points for typical usage +// --------------------------------------------------------------------------- + export { Context } from './context'; -export * from './types'; \ No newline at end of file +export * from './types'; + +// Loaders — file format handlers (commonly extended by users) +export { + Loader, + MarkdownLoader, + JsonLoader, + TextLoader, +} from './loaders'; +export { pathToId } from './loaders/util'; + +// Chunking — document splitting strategies +export { + MarkdownChunker, + FixedSizeChunker, + createChunker, +} from './loaders/chunker'; +export type { Chunk, Chunker } from './loaders/chunker'; + +// Reranker — two-stage retrieval precision scoring +export { KeywordReranker, createReranker } from './reranker'; +export type { Reranker, RerankCandidate, RerankResult } from './reranker'; + +// Query expansion — synonym bridging for cross-language recall +export { SynonymExpander, NoopExpander } from './query-expander'; +export type { QueryExpander } from './query-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 } from './embedder'; +export type { LanguageHint } from './embedder'; +export type { EmbedderInfo, EmbedderKind } from './embedder'; + +// Zvec store — custom vector storage backends +export { + IZvecStore, + MemoryZvecStore, + ActualZvecStore, + createZvecStore, + openZvecStore, + openZvecStoreSync, + isZvecAvailable, + buildZvecSchema, + cosineSimilarity, +} from './storage/zvec-store'; +export type { + ZvecDoc, + ZvecQueryResult, + ZvecSearchParams, + ZvecHybridParams, + FtsFieldWeight, + ZvecFieldSchema, + ZvecStoreConfig, + ActualZvecStoreOptions, +} from './storage/zvec-store'; + +// DocumentRegistry — dedup tracking +export { DocumentRegistry } from './registry'; + +// StoreManager — zvec store lifecycle management +export { StoreManager } from './store-manager'; + +// Backward-compatible alias — remove() is now untrack() +// Users should prefer untrack() + rebuild() for actual data removal. +export type { ContextOptions as RemoveOptions } from './types'; + diff --git a/src/loaders/chunker.ts b/src/loaders/chunker.ts new file mode 100644 index 0000000..c55483a --- /dev/null +++ b/src/loaders/chunker.ts @@ -0,0 +1,355 @@ +/** + * Document chunker — splits loaded documents into smaller, semantic chunks + * so that vector embeddings capture fine-grained meaning instead of averaging + * the entire document. + * + * Strategies: + * - MarkdownChunker: splits on headings (## / ###), then falls back to fixed-size + * - FixedSizeChunker: sliding window with overlap, used for plain text / JSON + * + * Each chunk stores its heading path, index, and parent document ID so the + * query layer can reconstruct context or deduplicate by parent. + */ + +import { Document } from '../types'; + +// --------------------------------------------------------------------------- +// Types +// --------------------------------------------------------------------------- + +/** A semantic chunk of a document. */ +export interface Chunk { + /** Chunk content — the actual text that will be embedded. */ + content: string; + /** Zero-based chunk index within the parent document. */ + chunkIndex: number; + /** Total number of chunks in the parent document. */ + totalChunks: number; + /** Parent document ID (hash-based, assigned by Context). */ + parentDocId: string; + /** + * Heading path from the document root to this chunk. + * e.g. ["Configuration", "Line Chart", "Basic Usage"] + */ + headingPath: string[]; +} + +/** Configuration for document chunking. */ +export interface ChunkingOptions { + /** Maximum characters per chunk (roughly ¼ tokens for mixed CN/EN, ½ for pure EN). */ + maxChunkSize?: number; + /** Overlap in characters between adjacent chunks (avoids boundary cuts). */ + chunkOverlap?: number; + /** Strategy used for chunking. */ + strategy?: 'markdown' | 'fixed' | 'auto'; +} + +/** Chunker interface — custom chunkers can implement this. */ +export interface Chunker { + chunk(doc: Document): Chunk[]; +} + +// --------------------------------------------------------------------------- +// Defaults +// --------------------------------------------------------------------------- + +const DEFAULT_MAX_CHUNK_SIZE = 1024; // ~256 tokens for mixed CN/EN +const DEFAULT_CHUNK_OVERLAP = 128; + +function chunkingDefaults(options?: ChunkingOptions): Required> { + return { + maxChunkSize: options?.maxChunkSize ?? DEFAULT_MAX_CHUNK_SIZE, + chunkOverlap: options?.chunkOverlap ?? DEFAULT_CHUNK_OVERLAP, + }; +} + +// --------------------------------------------------------------------------- +// Markdown Chunker — heading-aware splitting +// --------------------------------------------------------------------------- + +export class MarkdownChunker implements Chunker { + private readonly maxChunkSize: number; + private readonly chunkOverlap: number; + + constructor(options?: ChunkingOptions) { + const opts = chunkingDefaults(options); + this.maxChunkSize = opts.maxChunkSize; + this.chunkOverlap = opts.chunkOverlap; + } + + chunk(doc: Document): Chunk[] { + return chunkMarkdown(doc, this.maxChunkSize, this.chunkOverlap); + } +} + +/** + * Split markdown content into heading-anchored sections, then sub-split + * sections that exceed maxChunkSize with a fixed-size sliding window. + */ +function chunkMarkdown( + doc: Document, + maxChunkSize: number, + overlap: number, +): Chunk[] { + const sections = splitByHeadings(doc.content); + const chunks: Chunk[] = []; + + for (const section of sections) { + if (section.content.length <= maxChunkSize) { + if (section.content.trim().length > 0) { + chunks.push({ + content: section.content.trim(), + chunkIndex: 0, // filled in after + totalChunks: 0, // filled in after + parentDocId: doc.id, + headingPath: section.headingPath, + }); + } + continue; + } + + // Sub-split oversized sections with overlap + const subChunks = fixedSizeSplit(section.content, maxChunkSize, overlap); + for (const sc of subChunks) { + chunks.push({ + content: sc, + chunkIndex: 0, + totalChunks: 0, + parentDocId: doc.id, + headingPath: section.headingPath, + }); + } + } + + // Fill in chunkIndex and totalChunks + return chunks.map((c, i) => ({ + ...c, + chunkIndex: i, + totalChunks: chunks.length, + })); +} + +interface HeadingSection { + content: string; + headingPath: string[]; +} + +/** + * Split markdown text into sections delimited by markdown headings (##, ###). + * + * Lines before the first heading become a section with an empty heading path. + */ +function splitByHeadings(text: string): HeadingSection[] { + const sections: HeadingSection[] = []; + const lines = text.split('\n'); + let currentHeadingPath: string[] = []; + let currentLines: string[] = []; + + function flushSection(): void { + const content = currentLines.join('\n').trim(); + if (content.length > 0) { + sections.push({ + content, + headingPath: [...currentHeadingPath], + }); + } + currentLines = []; + } + + for (const line of lines) { + const h1Match = line.match(/^#\s+(.+)/); + const h2Match = line.match(/^##\s+(.+)/); + const h3Match = line.match(/^###\s+(.+)/); + + if (h1Match) { + flushSection(); + currentHeadingPath = [h1Match[1].trim()]; + } else if (h2Match) { + flushSection(); + const h1Ancestor = currentHeadingPath.length > 0 ? currentHeadingPath[0] : null; + currentHeadingPath = h1Ancestor + ? [h1Ancestor, h2Match[1].trim()] + : [h2Match[1].trim()]; + } else if (h3Match) { + flushSection(); + currentHeadingPath = [...currentHeadingPath, h3Match[1].trim()]; + if (currentHeadingPath.length === 0) { + currentHeadingPath = [h3Match[1].trim()]; + } + } else { + currentLines.push(line); + } + } + + flushSection(); + + // If no headings were found at all, return the full content as one section + if (sections.length === 0 && currentLines.length === 0 && text.trim().length > 0) { + sections.push({ content: text.trim(), headingPath: [] }); + } + + return sections; +} + +// --------------------------------------------------------------------------- +// Fixed-Size Chunker — sliding window, no semantic awareness +// --------------------------------------------------------------------------- + +export class FixedSizeChunker implements Chunker { + private readonly maxChunkSize: number; + private readonly chunkOverlap: number; + + constructor(options?: ChunkingOptions) { + const opts = chunkingDefaults(options); + this.maxChunkSize = opts.maxChunkSize; + this.chunkOverlap = opts.chunkOverlap; + } + + chunk(doc: Document): Chunk[] { + const parts = fixedSizeSplit(doc.content, this.maxChunkSize, this.chunkOverlap); + return parts.map((content, i) => ({ + content, + chunkIndex: i, + totalChunks: parts.length, + parentDocId: doc.id, + headingPath: [], + })); + } +} + +/** + * Split text into fixed-size windows with overlap. + * + * Tries to break at paragraph boundaries (double newline) within a tolerance + * of maxSize to avoid cutting sentences in half. Falls back to hard cut. + */ +function fixedSizeSplit( + text: string, + maxSize: number, + overlap: number, +): string[] { + const paragraphs = text.split(/\n\n+/); + const chunks: string[] = []; + let current = ''; + + for (const para of paragraphs) { + const trimmed = para.trim(); + if (!trimmed) continue; + + if (current.length + trimmed.length + 2 <= maxSize) { + current = current ? `${current}\n\n${trimmed}` : trimmed; + } else { + // Current chunk is full — flush it + if (current) { + chunks.push(current); + // Start new chunk with overlap: take the last `overlap` chars of the + // previous chunk as prefix so boundary semantics aren't lost. + const overlapPrefix = overlap > 0 && current.length > overlap + ? current.slice(-overlap).replace(/^[^\w一-鿿]+/, '') + '\n\n' + : ''; + current = overlapPrefix + trimmed; + } else { + // Single paragraph is larger than maxSize — force split + current = trimmed; + while (current.length > maxSize) { + // Try to break at a sentence boundary (。.!?。！？\n) + let cutPoint = findBreakPoint(current, maxSize); + chunks.push(current.slice(0, cutPoint).trim()); + const suffix = current.slice(cutPoint).trim(); + current = overlap > 0 && suffix.length > 0 + ? current.slice(Math.max(0, cutPoint - overlap), cutPoint).replace(/^[^\w一-鿿]+/, '') + '\n\n' + suffix + : suffix; + } + } + } + } + + if (current.trim()) { + chunks.push(current.trim()); + } + + return chunks.length > 0 ? chunks : [text.trim()]; +} + +/** + * Find a natural break point within [maxSize * 0.8, maxSize]. + * Prefers paragraph breaks > sentence-ending punctuation > space. + */ +function findBreakPoint(text: string, maxSize: number): number { + const minPoint = Math.floor(maxSize * 0.7); + + // Try double newline first + for (let i = maxSize; i >= minPoint; i--) { + if (text[i] === '\n' && text[i - 1] === '\n') return i + 1; + } + + // Try sentence-ending punctuation followed by space or newline + const sentenceBreaks = /[。.!?！？]\s/g; + let match: RegExpExecArray | null; + let bestSentence = -1; + + sentenceBreaks.lastIndex = minPoint; + while ((match = sentenceBreaks.exec(text)) !== null) { + if (match.index > maxSize) break; + bestSentence = match.index + 1; // after the punctuation + } + if (bestSentence > minPoint) return bestSentence; + + // Try any newline + for (let i = maxSize; i >= minPoint; i--) { + if (text[i] === '\n') return i + 1; + } + + // Fallback: hard cut at maxSize, but try to land on a space + for (let i = maxSize; i >= minPoint; i--) { + if (text[i] === ' ') return i + 1; + } + + return maxSize; +} + +// --------------------------------------------------------------------------- +// Factory +// --------------------------------------------------------------------------- + +/** + * Create a chunker based on the configured strategy. + * + * - `'auto'` (default): tries MarkdownChunker; if no headings found, falls + * back to FixedSizeChunker. + * - `'markdown'`: always uses heading-aware splitting. + * - `'fixed'`: always uses fixed-size sliding window. + */ +export function createChunker(options?: ChunkingOptions): Chunker { + const strategy = options?.strategy ?? 'auto'; + + switch (strategy) { + case 'markdown': + return new MarkdownChunker(options); + case 'fixed': + return new FixedSizeChunker(options); + case 'auto': + default: + return new AutoChunker(options); + } +} + +/** Auto-detect chunking strategy: markdown when headings exist, else fixed. */ +class AutoChunker implements Chunker { + private markdown: MarkdownChunker; + private fixed: FixedSizeChunker; + + constructor(options?: ChunkingOptions) { + this.markdown = new MarkdownChunker(options); + this.fixed = new FixedSizeChunker(options); + } + + chunk(doc: Document): Chunk[] { + // If the document content contains markdown headings (## or ###), use + // heading-aware chunking; otherwise fall back to fixed-size splits. + if (/^#{1,3}\s+\S/m.test(doc.content)) { + return this.markdown.chunk(doc); + } + return this.fixed.chunk(doc); + } +} diff --git a/src/loaders/index.ts b/src/loaders/index.ts index 943fd79..fe32023 100644 --- a/src/loaders/index.ts +++ b/src/loaders/index.ts @@ -1,4 +1,11 @@ export { Loader } from './base'; export { MarkdownLoader } from './markdown'; export { JsonLoader } from './json'; -export { TextLoader } from './text'; \ No newline at end of file +export { TextLoader } from './text'; +export { pathToId } from './util'; +export { + MarkdownChunker, + FixedSizeChunker, + createChunker, +} from './chunker'; +export type { Chunk, Chunker } from './chunker'; \ No newline at end of file diff --git a/src/loaders/json.ts b/src/loaders/json.ts index a2d527a..0baabd0 100644 --- a/src/loaders/json.ts +++ b/src/loaders/json.ts @@ -1,5 +1,4 @@ import * as fs from 'fs/promises'; -import * as path from 'path'; import { Loader } from './base'; import { Document } from '../types'; @@ -12,8 +11,10 @@ export class JsonLoader implements Loader { const content = await fs.readFile(filePath, 'utf-8'); const data = JSON.parse(content); + // ID is a temporary placeholder — Context.load() derives the canonical + // hash-based ID from the path relative to basePath (cross-machine consistency). return { - id: path.basename(filePath), + id: filePath, content: typeof data === 'string' ? data : JSON.stringify(data, null, 2), }; } diff --git a/src/loaders/markdown.ts b/src/loaders/markdown.ts index 0d72518..1892d54 100644 --- a/src/loaders/markdown.ts +++ b/src/loaders/markdown.ts @@ -1,5 +1,4 @@ import * as fs from 'fs/promises'; -import * as path from 'path'; import matter from 'gray-matter'; import { Loader } from './base'; import { Document } from '../types'; @@ -13,11 +12,10 @@ export class MarkdownLoader implements Loader { const content = await fs.readFile(filePath, 'utf-8'); const { data: meta, content: body } = matter(content); - // 使用文件名作为 ID，避免路径中的特殊字符 - const id = path.basename(filePath); - + // ID is a temporary placeholder — Context.load() derives the canonical + // hash-based ID from the path relative to basePath (cross-machine consistency). return { - id, + id: filePath, content: body.trim(), meta, }; diff --git a/src/loaders/text.ts b/src/loaders/text.ts index 992174b..bcf47f8 100644 --- a/src/loaders/text.ts +++ b/src/loaders/text.ts @@ -1,5 +1,4 @@ import * as fs from 'fs/promises'; -import * as path from 'path'; import { Loader } from './base'; import { Document } from '../types'; @@ -11,8 +10,10 @@ export class TextLoader implements Loader { async load(filePath: string): Promise { const content = await fs.readFile(filePath, 'utf-8'); + // ID is a temporary placeholder — Context.load() derives the canonical + // hash-based ID from the path relative to basePath (cross-machine consistency). return { - id: path.basename(filePath), + id: filePath, content: content.trim(), }; } diff --git a/src/loaders/util.ts b/src/loaders/util.ts new file mode 100644 index 0000000..2cd8914 --- /dev/null +++ b/src/loaders/util.ts @@ -0,0 +1,27 @@ +import * as crypto from 'crypto'; + +/** + * 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. + * + * Example: "/long/path/getting-started.md" → "f3a1b2c4__getting_started" + */ +export function pathToId(filePath: string): string { + // Normalize: remove leading slashes + const normalized = filePath.replace(/^\/+/, ''); + + // Generate a compact hash of the full path (16 hex chars = 64-bit) + const hash = crypto.createHash('sha256').update(normalized).digest('hex').slice(0, 16); + + // 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/query-expander.ts b/src/query-expander.ts new file mode 100644 index 0000000..654f132 --- /dev/null +++ b/src/query-expander.ts @@ -0,0 +1,104 @@ +/** + * 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. + * + * Example (with a custom synonym map): + * { "tooltip": ["提示框", "hover"] } → "tooltip config 提示框 hover 配置" + */ + +// --------------------------------------------------------------------------- +// 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 term boundaries in mixed CN/EN text. */ +const CJK_BOUNDARY_RE = /[\s一-鿿,，。]/; + +/** + * 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. + * + * The expander normalises case and preserves CJK characters for exact matching. + */ +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(); + + // Find matching synonym entries in the query + for (const [term, syns] of Object.entries(this.synonyms)) { + const termLower = term.toLowerCase(); + + // Exact match (whole term exists in query) + if (this.containsTerm(queryLower, termLower)) { + for (const syn of syns) { + // Don't add synonyms that are already in the query + if (!this.containsTerm(queryLower, syn.toLowerCase())) { + additions.add(syn); + } + } + } + } + + if (additions.size === 0) return query; + + // Append additions — space-separated keeps compatibility with both + // embedding models and FTS tokenizers. + return `${query} ${[...additions].join(' ')}`; + } + + /** + * Check if a term appears as a word/phrase boundary in the query. + * Handles both space-separated (EN) and character-joined (CJK) text. + */ + private containsTerm(text: string, term: string): boolean { + if (text.includes(term)) { + const idx = text.indexOf(term); + const prevOk = idx === 0 || CJK_BOUNDARY_RE.test(text[idx - 1]); + const afterIdx = idx + term.length; + const afterOk = afterIdx >= text.length || CJK_BOUNDARY_RE.test(text[afterIdx]); + return prevOk && afterOk; + } + 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/registry.ts b/src/registry.ts new file mode 100644 index 0000000..861cd2a --- /dev/null +++ b/src/registry.ts @@ -0,0 +1,144 @@ +import * as fs from 'fs'; +import * as path from 'path'; + +/** + * DocumentRegistry — manages document metadata for each library. + * + * Tracks loaded document IDs along with their content hashes for + * change-detection. When a document's content hash changes (e.g. the + * file was edited), the dedup check will allow it to be re-embedded. + */ + +/** Registry entry: document ID + content hash for change detection. */ +interface RegistryEntry { + id: string; + /** SHA-256 hash of the document content (first 16 chars). Empty = legacy entry without hash. */ + contentHash: string; +} + +export class DocumentRegistry { + private readonly entries: Map> = new Map(); + + /** + * Check whether a document ID has already been loaded into a library, + * optionally verifying that the content hash matches. + * + * - If `contentHash` is provided and the existing entry has a hash, + * returns `true` only when both the ID and hash match (unchanged file). + * - If `contentHash` is provided but differs from the stored hash, + * returns `false` (file has been updated — needs re-embedding). + * - If `contentHash` is not provided, returns `true` when the ID exists + * (legacy behaviour — no change detection). + */ + has(library: string, id: string, contentHash?: string): boolean { + const lib = this.entries.get(library); + if (!lib) return false; + + const entry = lib.get(id); + if (!entry) return false; + + // If hash is provided, do change detection + if (contentHash !== undefined && entry.contentHash) { + return entry.contentHash === contentHash; + } + + // No hash provided or stored — just check existence + return true; + } + + /** + * Register a document as loaded for a library, with its content hash. + */ + add(library: string, id: string, contentHash?: string): void { + let lib = this.entries.get(library); + if (!lib) { + lib = new Map(); + this.entries.set(library, lib); + } + lib.set(id, { id, contentHash: contentHash ?? '' }); + } + + /** + * Get all loaded document IDs for a library. + */ + getIds(library: string): Set { + const lib = this.entries.get(library); + return lib ? new Set(lib.keys()) : new Set(); + } + + /** + * Get all registry entries for a library (including content hashes). + */ + getEntries(library: string): Map { + return this.entries.get(library) ?? new Map(); + } + + /** + * Check whether a library has any loaded documents. + */ + hasLibrary(library: string): boolean { + return (this.entries.get(library)?.size ?? 0) > 0; + } + + /** + * Remove a specific document from a library's registry. + */ + remove(library: string, id: string): void { + this.entries.get(library)?.delete(id); + } + + /** + * Get all library names that have loaded documents. + */ + getLibraryNames(): string[] { + return [...this.entries.keys()]; + } + + /** + * Remove all documents for a library. + */ + removeLibrary(library: string): void { + this.entries.delete(library); + } + + /** + * Persist registry state to disk (for process restart recovery). + * + * Format: `{ id, contentHash }` entries for change detection across + * restarts. + */ + saveToDisk(vectorsDir: string, library: string): void { + const indexPath = path.join(vectorsDir, `${library}.index.json`); + const lib = this.entries.get(library); + if (!lib) { + // No entries — write empty array + fs.writeFileSync(indexPath, JSON.stringify([])); + return; + } + const data = [...lib.values()]; + fs.writeFileSync(indexPath, JSON.stringify(data)); + } + + /** + * Load registry state from disk. + */ + loadFromDisk(vectorsDir: string, library: string): void { + const indexPath = path.join(vectorsDir, `${library}.index.json`); + if (fs.existsSync(indexPath)) { + const raw: RegistryEntry[] | string[] = JSON.parse(fs.readFileSync(indexPath, 'utf-8')); + const lib = new Map(); + + for (const item of raw) { + if (typeof item === 'string') { + // Legacy format: plain ID strings (no content hash) + lib.set(item, { id: item, contentHash: '' }); + } else { + // New format: { id, contentHash } entries + lib.set(item.id, item); + } + } + + this.entries.set(library, lib); + } + } +} diff --git a/src/reranker.ts b/src/reranker.ts new file mode 100644 index 0000000..2e9fb7f --- /dev/null +++ b/src/reranker.ts @@ -0,0 +1,223 @@ +/** + * 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 + * + * This two-stage approach combines fast vector pre-filtering with a more + * expensive but more precise scoring model on the shortlist. + */ + +// --------------------------------------------------------------------------- +// 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 { + /** + * How many extra candidates to pull from the coarse search stage. + * e.g. with topK=5 and rerankFactor=3, 15 candidates are reranked. + * Default: 3 + */ + rerankFactor?: number; + /** + * Minimum number of candidates to rerank (floor). + * Default: 10 + */ + minCandidates?: number; +} + +// --------------------------------------------------------------------------- +// Defaults +// --------------------------------------------------------------------------- + +const DEFAULT_RERANK_FACTOR = 3; +const DEFAULT_MIN_CANDIDATES = 10; + +// --------------------------------------------------------------------------- +// KeywordReranker — lexical overlap scoring +// --------------------------------------------------------------------------- + +/** + * Reranker that scores candidates by keyword / phrase overlap with the query. + * + * This provides a complementary signal to vector similarity: the coarse + * embedding stage captures semantic closeness, while this reranker boosts + * candidates that contain the exact query terms (or their substrings). + * + * Scoring factors: + * - Exact phrase match: 3.0× weight + * - Whole term match: 1.0× weight (per term) + * - Substring match: 0.3× weight (partial token overlap) + * - Heading match: 2.0× bonus (terms appearing in heading path) + * - Original score: 0.1× carry-over (respects the coarse rank) + * + * All scores are normalised to [0, 1] via min-max scaling. + */ +export class KeywordReranker implements Reranker { + 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 scored = candidates.map((c) => { + const contentLower = c.content.toLowerCase(); + let score = 0; + + // 1. Exact phrase match — strongest signal + if (contentLower.includes(queryPhrase)) { + score += 3.0; + // Bonus for each additional occurrence + const phraseCount = countOccurrences(contentLower, queryPhrase); + score += (phraseCount - 1) * 0.5; + } + + // 2. Per-term matching + for (const term of queryTerms) { + if (contentLower.includes(term)) { + // Whole word match + if (isWordBoundary(contentLower, term)) { + score += 1.0; + const termCount = countTermMatches(contentLower, term); + score += (termCount - 1) * 0.2; + } else { + // Substring match (e.g. "tool" matches "tooltip") + score += 0.3; + } + } + } + + // 3. Heading path bonus — terms in section headings are more relevant + if (c.headingPath) { + const headingLower = c.headingPath.toLowerCase(); + for (const term of queryTerms) { + if (headingLower.includes(term)) { + score += 2.0; + } + } + if (queryPhrase.length > 2 && headingLower.includes(queryPhrase)) { + score += 2.5; + } + } + + // 4. Carry over a fraction of the original vector score + score += c.score * 0.1; + + 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[] { + // Split on whitespace and CJK-friendly boundaries + const raw = query + .split(/[\s,，。.!！？?、：:；;]+/) + .filter(Boolean); + + // For mixed CN/EN queries, also extract CJK bigrams as tokens + const tokens: string[] = []; + for (const r of raw) { + tokens.push(r); + // For CJK segments, add bigrams as additional tokens + 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 the default reranker. + * + * Currently returns KeywordReranker. In the future this may auto-select + * based on available models (cross-encoder, etc.). + */ +export function createReranker(): Reranker { + return new KeywordReranker(); +} diff --git a/src/storage/actual-store.ts b/src/storage/actual-store.ts new file mode 100644 index 0000000..018f774 --- /dev/null +++ b/src/storage/actual-store.ts @@ -0,0 +1,393 @@ +/** + * ActualZvecStore — wraps @zvec/zvec native bindings (full schema + FTS). + */ + +import type { + ZvecDoc, + ZvecQueryResult, + ZvecSearchParams, + ZvecHybridParams, + IZvecStore, + ZvecStoreConfig, + ActualZvecStoreOptions, + ZvecFieldSchema, +} from './types'; + +// --------------------------------------------------------------------------- +// zvec module loading +// --------------------------------------------------------------------------- + +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; + 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 + })); + this._collection.insertSync(records); + } + + 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. + * + * When @zvec/zvec is available, creates a persisted ActualZvecStore. + * Otherwise falls back to an in-memory MemoryZvecStore. + */ +export async function createZvecStore( + path: string, + config: ZvecStoreConfig, + memoryFtsWeights?: import('./types').FtsFieldWeight[], + rankConstant?: number, +): Promise { + const z = loadZvecSync(); + if (z) { + return ActualZvecStore.create(path, config); + } + const { MemoryZvecStore } = await import('./memory-store'); + return new MemoryZvecStore(memoryFtsWeights, rankConstant); +} + +/** + * Asynchronously open a zvec store. + */ +export async function openZvecStore( + path: string, + options?: ActualZvecStoreOptions +): Promise { + // Both ActualZvecStore.open() and MemoryZvecStore are internally + // synchronous — this async wrapper is for API consistency only. + return openZvecStoreSync(path, options); +} + +/** + * Synchronously open a zvec store. + * + * This avoids the async Promise wrapper so it can be used inside synchronous + * code paths (e.g. `retrieve()`). + */ +export function openZvecStoreSync( + path: string, + options?: ActualZvecStoreOptions +): IZvecStore { + const z = loadZvecSync(); + if (z) { + return ActualZvecStore.openSync(path, options); + } + throw new Error( + 'Cannot open zvec store: @zvec/zvec is not installed and MemoryZvecStore ' + + 'has no persistence. Install @zvec/zvec or use createZvecStore() to create ' + + 'a new MemoryZvecStore.' + ); +} + +/** Synchronous check: is @zvec/zvec available? */ +export function isZvecAvailable(): boolean { + return loadZvecSync() !== undefined; +} diff --git a/src/storage/memory-store.ts b/src/storage/memory-store.ts new file mode 100644 index 0000000..fa9e54b --- /dev/null +++ b/src/storage/memory-store.ts @@ -0,0 +1,265 @@ +/** + * MemoryZvecStore — pure JS fallback (cosine similarity + linear scan + text). + */ + +import type { + ZvecDoc, + ZvecQueryResult, + ZvecSearchParams, + ZvecHybridParams, + IZvecStore, + FtsFieldWeight, +} from './types'; +import { cosineSimilarity, evalMemoryFilter } from './utils'; + +const DEFAULT_FTS_WEIGHTS: FtsFieldWeight[] = [ + { fieldName: 'content', weight: 1.0 } +]; + +const DEFAULT_RANK_CONSTANT = 60; + +export class MemoryZvecStore implements IZvecStore { + private docs: ZvecDoc[] = []; + private _ftsWeights: FtsFieldWeight[]; + private _rankConstant: number; + + constructor(ftsFieldWeights?: FtsFieldWeight[], rankConstant?: number) { + this._ftsWeights = ftsFieldWeights && ftsFieldWeights.length > 0 + ? ftsFieldWeights + : DEFAULT_FTS_WEIGHTS; + this._rankConstant = rankConstant ?? DEFAULT_RANK_CONSTANT; + } + + async insert(docs: ZvecDoc[]): Promise { + this.docs.push(...docs); + } + + async search(params: ZvecSearchParams): Promise { + return doSyncSearch(this.docs, params); + } + + async searchHybrid(params: ZvecHybridParams): Promise { + return doSyncHybridSearch(this.docs, params, this._ftsWeights, this._rankConstant); + } + + searchSync(params: ZvecSearchParams): ZvecQueryResult[] { + return doSyncSearch(this.docs, params); + } + + searchHybridSync(params: ZvecHybridParams): ZvecQueryResult[] { + return doSyncHybridSearch(this.docs, params, this._ftsWeights, this._rankConstant); + } + + async close(): Promise { + this.docs = []; + } +} + +function doSyncSearch( + docs: ZvecDoc[], + params: ZvecSearchParams +): ZvecQueryResult[] { + const { vector, topK, filter } = params; + const scored: ZvecQueryResult[] = []; + for (const doc of docs) { + if (filter && !evalMemoryFilter(filter, doc.fields)) continue; + scored.push({ + id: doc.id, + score: cosineSimilarity(vector, doc.vector), + fields: doc.fields + }); + } + scored.sort((a, b) => b.score - a.score); + return scored.slice(0, topK); +} + +function doSyncHybridSearch( + docs: ZvecDoc[], + params: ZvecHybridParams, + ftsWeights: FtsFieldWeight[], + rankConstant: number, +): ZvecQueryResult[] { + const { queryText, queryVector, topK, filter } = params; + const rrScores = new Map(); + + // Filtered docs — apply once + const candidates = filter + ? docs.filter((d) => evalMemoryFilter(filter, d.fields)) + : docs; + + // 1. Text path: field-boosted scoring with word-boundary matching + TF weighting. + // + // Scoring rules (per query term per field): + // - Exact word-boundary match (EN: surrounded by space/punctuation, + // CJK: standalone character group): 1.0 × field_weight + // - Substring match (no word boundary): 0.3 × field_weight + // - TF multiplier: log(1 + count) so multiple occurrences boost + // the score but don't saturate it. + // - CJK bigrams from the query also score (0.5 × field_weight per match). + // + // This produces better FTS results than raw substring matching: + // "tooltip" won't falsely match "tooltipconfigxx" + // "折线" will match "折线图" via bigram scoring + const terms = queryText.toLowerCase().split(/\s+/).filter(Boolean); + const cjkBigrams = extractCJKBigrams(queryText.toLowerCase()); + + const textRanked = candidates + .map((doc) => { + let score = 0; + for (const fw of ftsWeights) { + const fieldVal = String(doc.fields[fw.fieldName] || '').toLowerCase(); + for (const term of terms) { + const count = countTermMatches(fieldVal, term); + if (count > 0) { + // Word-boundary match count vs total occurrences + const boundaryCount = countBoundaryMatches(fieldVal, term); + const subCount = count - boundaryCount; + // TF-weighted: log(1 + count) prevents saturation + const boundaryScore = boundaryCount > 0 + ? Math.log(1 + boundaryCount) * fw.weight + : 0; + const subScore = subCount > 0 + ? Math.log(1 + subCount) * 0.3 * fw.weight + : 0; + score += boundaryScore + subScore; + } + } + // CJK bigram scoring — bridges partial character matches + for (const bigram of cjkBigrams) { + const count = countTermMatches(fieldVal, bigram); + if (count > 0) { + score += Math.log(1 + count) * 0.5 * fw.weight; + } + } + } + return { id: doc.id, score }; + }) + .filter((r) => r.score > 0) + .sort((a, b) => b.score - a.score); + + // Standard RRF: score = 1 / (rankConstant + rank), rank starts at 1 + for (let i = 0; i < textRanked.length; i++) { + rrScores.set(textRanked[i].id, 1 / (rankConstant + (i + 1))); + } + + // 2. Vector path: cosine similarity + const vecRanked = candidates + .map((doc) => ({ + id: doc.id, + score: cosineSimilarity(queryVector, doc.vector), + })) + .sort((a, b) => b.score - a.score); + + for (let i = 0; i < vecRanked.length; i++) { + const existing = rrScores.get(vecRanked[i].id) ?? 0; + rrScores.set(vecRanked[i].id, existing + 1 / (rankConstant + (i + 1))); + } + + // 3. Merge by RRF score (no division — scores are already correctly scaled) + const docMap = new Map(candidates.map((d) => [d.id, d])); + return [...rrScores.entries()] + .sort((a, b) => b[1] - a[1]) + .slice(0, topK) + .map(([id, score]) => { + const doc = docMap.get(id)!; + return { id, score, fields: doc.fields }; + }); +} + +// --------------------------------------------------------------------------- +// FTS helper functions — word-boundary matching + CJK bigram extraction +// --------------------------------------------------------------------------- + +/** + * Count total occurrences of a term in text (substring match). + */ +function countTermMatches(text: string, term: string): number { + let count = 0; + let pos = 0; + while ((pos = text.indexOf(term, pos)) !== -1) { + count++; + pos += term.length; + } + return count; +} + +/** + * Count occurrences of a term that land on a word boundary. + * + * A word boundary means the character before and after the match is + * a space, punctuation, or string boundary. CJK characters are treated + * as individual "words" — a CJK term naturally forms a boundary within + * a CJK string. + */ +function countBoundaryMatches(text: string, term: string): number { + let count = 0; + let pos = 0; + while ((pos = text.indexOf(term, pos)) !== -1) { + if (isWordBoundary(text, pos, term.length)) { + count++; + } + pos += term.length; + } + return count; +} + +/** + * Check if a substring occurrence at the given position is at a word boundary. + */ +function isWordBoundary(text: string, pos: number, len: number): boolean { + const before = pos === 0 || /[\s\n.,;:!?，。！？、：；"'(（【《\-_]/.test(text[pos - 1]); + const afterIdx = pos + len; + const after = afterIdx >= text.length || /[\s\n.,;:!?，。！？、：；"')）】》\-_]/.test(text[afterIdx]); + return before && after; +} + +/** + * Extract CJK bigrams (2-char sequences) from text for partial matching. + * + * For example, "折线图配置" produces bigrams: ["折线", "线图", "图配", "配置"] + * These bigrams help the FTS path find documents that contain partial + * character sequences of the query, improving recall for CJK searches. + */ +function extractCJKBigrams(text: string): string[] { + const bigrams: string[] = []; + // Extract continuous CJK segments, then generate bigrams within each + const segments = splitCJKSegments(text); + for (const seg of segments) { + for (let i = 0; i + 2 <= seg.length; i++) { + bigrams.push(seg.slice(i, i + 2)); + } + } + return bigrams; +} + +/** + * Split text into continuous CJK character segments. + * Non-CJK characters break the segment. + */ +function splitCJKSegments(text: string): string[] { + const segments: string[] = []; + let current = ''; + for (const ch of text) { + if (isCJKChar(ch)) { + current += ch; + } else { + if (current.length >= 2) segments.push(current); + current = ''; + } + } + if (current.length >= 2) segments.push(current); + return segments; +} + +/** + * Check if a character is CJK. + */ +function isCJKChar(ch: string): boolean { + const cp = ch.codePointAt(0)!; + return ( + (cp >= 0x4e00 && cp <= 0x9fff) || + (cp >= 0x3400 && cp <= 0x4dbf) || + (cp >= 0x3040 && cp <= 0x30ff) || + (cp >= 0xac00 && cp <= 0xd7af) + ); +} diff --git a/src/storage/types.ts b/src/storage/types.ts new file mode 100644 index 0000000..ef24060 --- /dev/null +++ b/src/storage/types.ts @@ -0,0 +1,84 @@ +/** + * ZvecStore public types and interfaces. + */ + +export interface ZvecDoc { + id: string; + vector: number[]; + fields: Record; +} + +export interface ZvecQueryResult { + id: string; + score: number; + fields: Record; +} + +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; + /** 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; +} + +/** Weighted field for MemoryZvecStore FTS text scoring in hybrid search. */ +export interface FtsFieldWeight { + fieldName: string; + weight: number; +} + +/** Field schema entry for building a zvec collection. */ +export interface ZvecFieldSchema { + name: string; + dataType: 'STRING' | 'INT64' | 'FLOAT' | 'VECTOR_FP32'; + indexType?: 'FTS' | 'INVERT' | 'HNSW' | 'NONE'; + indexOptions?: Record; +} + +/** Configuration for creating a zvec collection. */ +export interface ZvecStoreConfig { + collectionName: string; + vectorField: string; + vectorDims: number; + /** Scalar fields whose FTS indexes will be queried in hybrid search. */ + ftsFields: string[]; + /** Full field schema (vector + scalar). */ + fields: ZvecFieldSchema[]; +} + +export interface ActualZvecStoreOptions { + /** Name of the vector field in the collection (default: 'embedding'). */ + vectorField: string; + /** Names of scalar fields with FTS indexes, used in hybrid search. */ + ftsFields: string[]; + /** + * RRF rank constant for hybrid search fusion (default: 60). + * + * Controls how much influence rare (low-rank) results have. + * - Lower value (e.g. 10) → top ranks dominate, more "winner-takes-all" + * - Higher value (e.g. 100) → ranks are more evenly weighted + */ + rankConstant?: number; +} diff --git a/src/storage/utils.ts b/src/storage/utils.ts new file mode 100644 index 0000000..e28504d --- /dev/null +++ b/src/storage/utils.ts @@ -0,0 +1,68 @@ +/** + * ZvecStore utility functions. + */ + +/** Cosine similarity between two vectors of equal length. */ +export function cosineSimilarity(a: number[], b: number[]): number { + let dot = 0, + na = 0, + nb = 0; + for (let i = 0; i < a.length; i++) { + dot += a[i] * b[i]; + na += a[i] * a[i]; + nb += b[i] * b[i]; + } + const denom = Math.sqrt(na) * Math.sqrt(nb); + return denom === 0 ? 0 : dot / denom; +} + +/** + * Evaluate a filter expression against document fields. + * + * Supports: + * - String equality: `field = 'value'` + * - Number equality: `field = 0` + * - AND conjunction: `field = 'val1' AND field2 = 'val2'` + * + * Unknown / malformed expressions return true (pass-through) so that + * MemoryZvecStore remains usable even when filters are only meant for + * the native zvec engine. + */ +export function evalMemoryFilter( + filter: string, + fields: Record +): boolean { + // Split on AND (case-insensitive, with surrounding spaces) + const clauses = filter.split(/\s+AND\s+/i); + + return clauses.every((clause) => evalSingleClause(clause.trim(), fields)); +} + +/** + * Evaluate a single filter clause. + * + * Forms: + * `fieldName = 'stringVal'` → string equality + * `fieldName = numericVal` → number equality (coerces field to Number) + */ +function evalSingleClause( + clause: string, + fields: Record +): boolean { + // String equality: field = 'value' + const stringMatch = clause.match(/^(\w+)\s*=\s*'([^']*)'$/); + if (stringMatch) { + const [, fieldName, expected] = stringMatch; + return String(fields[fieldName] ?? '') === expected; + } + + // Number equality: field = 0 (or any integer/float) + const numberMatch = clause.match(/^(\w+)\s*=\s*(-?\d+(?:\.\d+)?)$/); + if (numberMatch) { + const [, fieldName, expected] = numberMatch; + return Number(fields[fieldName]) === Number(expected); + } + + // Unknown clause format — pass through (let native zvec handle it) + return true; +} diff --git a/src/storage/zvec-store.ts b/src/storage/zvec-store.ts index 166edb3..153f2fe 100644 --- a/src/storage/zvec-store.ts +++ b/src/storage/zvec-store.ts @@ -1,109 +1,34 @@ -import { - ZVecCreateAndOpen, - ZVecOpen, - ZVecCollectionSchema, - ZVecDataType, - ZVecCollection, -} from '@zvec/zvec'; -import * as fs from 'fs'; -import * as path from 'path'; - -interface DocData { - content: string; - meta?: Record; -} - -export class ZVecStore { - private collection: ZVecCollection; - private dimension: number; - private filePath!: string; - private docs: Map = new Map(); - - constructor(collection: ZVecCollection, dimension: number) { - this.collection = collection; - this.dimension = dimension; - } - - static async create(filePath: string, dimension: number = 384): Promise { - const schema = new ZVecCollectionSchema({ - name: 'documents', - vectors: { name: 'embedding', dataType: ZVecDataType.VECTOR_FP32, dimension }, - fields: [ - { name: 'content', dataType: ZVecDataType.STRING }, - ], - }); - - // Ensure directory exists - const dir = path.dirname(filePath); - if (!fs.existsSync(dir)) { - fs.mkdirSync(dir, { recursive: true }); - } - - let collection: ZVecCollection; - if (fs.existsSync(filePath)) { - collection = ZVecOpen(filePath); - } else { - const dir = path.dirname(filePath); - if (!fs.existsSync(dir)) { - fs.mkdirSync(dir, { recursive: true }); - } - collection = ZVecCreateAndOpen(filePath, schema); - } - - const store = new ZVecStore(collection, dimension); - store.filePath = filePath; - await store.loadMeta(); - return store; - } - - private async loadMeta(): Promise { - const metaPath = this.filePath + '.meta.json'; - if (fs.existsSync(metaPath)) { - const data = JSON.parse(fs.readFileSync(metaPath, 'utf-8')); - this.docs = new Map(Object.entries(data)); - } - } - - private saveMeta(): void { - fs.writeFileSync(this.filePath + '.meta.json', JSON.stringify(Object.fromEntries(this.docs))); - } - - add(id: string, vector: number[], content: string, meta?: Record): void { - this.collection.insertSync([ - { id, vectors: { embedding: vector }, fields: { content } }, - ]); - this.docs.set(id, { content, meta }); - } - - search(queryVector: number[], topK: number): Array<{ id: string; score: number }> { - const results = this.collection.querySync({ - fieldName: 'embedding', - vector: queryVector, - topk: topK, - }); - - return results.map((r) => ({ id: r.id!, score: r.score })); - } - - getDoc(id: string): DocData | undefined { - return this.docs.get(id); - } - - save(): void { - this.saveMeta(); - } - - close(): void { - this.saveMeta(); - this.collection.closeSync(); - } - - clear(): void { - const ids = Array.from(this.docs.keys()); - if (ids.length > 0) { - this.collection.deleteSync(ids); - } - this.docs.clear(); - this.saveMeta(); - } -} \ No newline at end of file +/** + * zvec-store — aggregate entry point for all storage modules. + * + * Re-exports from the split files for backward compatibility: + * types.ts → public types & interfaces + * memory-store.ts → MemoryZvecStore + * actual-store.ts → ActualZvecStore, factories, schema builder + * utils.ts → cosineSimilarity, evalMemoryFilter + */ + +export type { + ZvecDoc, + ZvecQueryResult, + ZvecSearchParams, + ZvecHybridParams, + IZvecStore, + FtsFieldWeight, + ZvecFieldSchema, + ZvecStoreConfig, + ActualZvecStoreOptions, +} from './types'; + +export { MemoryZvecStore } from './memory-store'; + +export { + ActualZvecStore, + createZvecStore, + openZvecStore, + openZvecStoreSync, + isZvecAvailable, + buildZvecSchema, +} from './actual-store'; + +export { cosineSimilarity } from './utils'; diff --git a/src/store-manager.ts b/src/store-manager.ts new file mode 100644 index 0000000..7c876bc --- /dev/null +++ b/src/store-manager.ts @@ -0,0 +1,237 @@ +import * as fs from 'fs'; +import * as path from 'path'; +import { + createZvecStore, + openZvecStoreSync, +} from './storage/zvec-store'; +import type { IZvecStore, ZvecStoreConfig, FtsFieldWeight, ActualZvecStoreOptions } from './storage/zvec-store'; +import type { Embedder } from './embedder'; +import { detectTokenizer } from './embedder'; +import type { ContextOptions } from './types'; + +// --------------------------------------------------------------------------- +// Default schema constants +// --------------------------------------------------------------------------- + +const DEFAULT_VECTOR_FIELD = 'embedding'; +const DEFAULT_FTS_FIELDS = ['content']; +const DEFAULT_RANK_CONSTANT = 60; + +function resolveFtsFields(options?: ContextOptions): string[] { + return options?.ftsFields ?? DEFAULT_FTS_FIELDS; +} + +function resolveFtsWeights(options?: ContextOptions): FtsFieldWeight[] { + const ftsFields = resolveFtsFields(options); + const weightMap = options?.ftsFieldWeights; + + if (weightMap) { + return ftsFields.map((fieldName) => ({ + fieldName, + weight: weightMap[fieldName] ?? 1.0, + })); + } + + // Default: all FTS fields have weight 1.0 + return ftsFields.map((fieldName) => ({ fieldName, weight: 1.0 })); +} + +function resolveTokenizer(options?: ContextOptions, sampleText?: string): string { + const tokenizer = options?.tokenizer ?? 'auto'; + if (tokenizer === 'auto') { + // When a document sample is available, auto-detect the best tokenizer + // based on character distribution (CJK → jieba, Latin → standard). + // Falls back to 'jieba' as the safe default for mixed-language content + // when no sample has been loaded yet. + if (sampleText && sampleText.trim().length > 0) { + return detectTokenizer(sampleText); + } + return 'jieba'; + } + return tokenizer; +} + +function contextStoreConfig(dims: number, options?: ContextOptions, sampleText?: string): ZvecStoreConfig { + const ftsFields = resolveFtsFields(options); + const tokenizerName = resolveTokenizer(options, 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: 'chunkIndex', dataType: 'INT64' }, + { name: 'totalChunks', dataType: 'INT64' }, + { name: 'parentDocId', dataType: 'STRING' }, + { name: 'headingPath', dataType: 'STRING' }, + { name: 'sourceFilePath', dataType: 'STRING' }, + ], + }; +} + +function storeOpenOptions(options?: ContextOptions): ActualZvecStoreOptions { + return { + vectorField: DEFAULT_VECTOR_FIELD, + ftsFields: resolveFtsFields(options), + rankConstant: options?.rankConstant ?? DEFAULT_RANK_CONSTANT, + }; +} + +// --------------------------------------------------------------------------- +// StoreManager — manages zvec store lifecycle per library +// --------------------------------------------------------------------------- + +/** + * StoreManager handles creation, caching, and lazy-loading of IZvecStore + * instances. Each library gets its own `.zvec` file on disk. + */ +export class StoreManager { + private readonly vectorsDir: string; + private readonly embedder: Embedder; + private readonly ftsWeights: FtsFieldWeight[]; + private readonly rankConstant: number; + 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(); + + constructor(vectorsDir: string, embedder: Embedder, options?: ContextOptions) { + this.vectorsDir = vectorsDir; + this.embedder = embedder; + this.ftsWeights = resolveFtsWeights(options); + this.rankConstant = options?.rankConstant ?? DEFAULT_RANK_CONSTANT; + this.contextOptions = options; + } + + /** Get the FTS field weights for MemoryZvecStore fallback. */ + getFtsWeights(): FtsFieldWeight[] { + return this.ftsWeights; + } + + /** Get the RRF rank constant for hybrid search. */ + getRankConstant(): number { + return this.rankConstant; + } + + private getStorePath(library: string): string { + return path.join(this.vectorsDir, `${library}.zvec`); + } + + /** + * Get or create a zvec store for a library. + * + * Uses a Promise-based lock so that concurrent calls for the same library + * share a single creation attempt instead of racing to produce duplicate + * store instances. + * + * @param library Library name. + * @param sampleText Optional document sample for auto-detecting FTS tokenizer + * when `tokenizer` is `'auto'`. Only used for new stores. + */ + async getOrCreate(library: string, sampleText?: string): Promise { + const cached = this.stores.get(library); + if (cached) return cached; + + const pending = this.pending.get(library); + if (pending) return pending; + + const promise = this._doGetOrCreate(library, sampleText); + this.pending.set(library, promise); + try { + const store = await promise; + this.stores.set(library, store); + return store; + } finally { + this.pending.delete(library); + } + } + + private async _doGetOrCreate(library: string, sampleText?: string): Promise { + const filePath = this.getStorePath(library); + return fs.existsSync(filePath) + ? openZvecStoreSync(filePath, storeOpenOptions(this.contextOptions)) + : await createZvecStore( + filePath, + contextStoreConfig(this.embedder.dimensions, this.contextOptions, sampleText), + this.ftsWeights, + this.rankConstant, + ); + } + + /** + * Get an existing store (already opened or cached). + * Returns undefined if the store has not been opened yet. + */ + getCached(library: string): IZvecStore | undefined { + return this.stores.get(library); + } + + /** + * Try to lazily open an existing store on disk. + * Returns the store if found, undefined otherwise. + */ + tryOpen(library: string): IZvecStore | undefined { + if (this.stores.has(library)) { + return this.stores.get(library)!; + } + + const filePath = this.getStorePath(library); + if (fs.existsSync(filePath)) { + const store = openZvecStoreSync(filePath, storeOpenOptions(this.contextOptions)); + this.stores.set(library, store); + return store; + } + + return undefined; + } + + /** + * Check whether a store file exists on disk for a library. + */ + existsOnDisk(library: string): boolean { + return fs.existsSync(this.getStorePath(library)); + } + + /** + * Close and remove a store from cache. + */ + async close(library: string): Promise { + const store = this.stores.get(library); + if (store) { + await store.close(); + this.stores.delete(library); + } + } + + /** + * Delete the store file from disk (after closing it). + * + * Used by `Context.rebuild()` to physically remove vector data before + * re-creating the store with fresh embeddings. + */ + async deleteStore(library: string): Promise { + await this.close(library); + const filePath = this.getStorePath(library); + if (fs.existsSync(filePath)) { + fs.unlinkSync(filePath); + } + } + + /** + * Close all stores. + */ + async closeAll(): Promise { + for (const [, store] of this.stores) { + await store.close(); + } + this.stores.clear(); + } +} diff --git a/src/types.ts b/src/types.ts index df77c1a..a3ab43e 100644 --- a/src/types.ts +++ b/src/types.ts @@ -8,8 +8,67 @@ export interface Document { content: string; /** Markdown front-matter metadata */ meta?: Record; + /** + * Chunk metadata — populated when document chunking is enabled. + * Undefined for whole-document (unchunked) mode. + */ + chunk?: ChunkMeta; + /** + * SHA-256 content hash (16 chars) — used internally for change detection. + * Populated by `Context.load()` before registry insertion; not stored in zvec. + */ + contentHash?: string; } +/** Metadata for a chunked document fragment. */ +export interface ChunkMeta { + /** Zero-based chunk index within the parent document. */ + chunkIndex: number; + /** Total number of chunks in the parent document. */ + totalChunks: number; + /** Parent document ID (hash-based). */ + parentDocId: string; + /** Heading path from the document root to this chunk. */ + headingPath: string[]; +} + +/** Configuration for query expansion. */ +export interface QueryExpansionOptions { + /** + * Synonym map for query term expansion. + * + * Entirely user-provided — no built-in defaults. Each key maps to + * an array of alternative terms that will be appended to the query + * when the key is found. Useful for CN↔EN terminology bridging + * or domain-specific synonym expansion. + * + * Oomit or pass an empty map for no expansion (effectively a no-op). + * Set `queryExpansion: false` to disable expansion entirely. + */ + synonyms?: Record; +} + +/** Configuration for reranking. */ +export interface RerankOptions { + /** + * How many extra candidates to pull from the coarse search stage. + * e.g. with topK=5 and rerankFactor=3, 15 candidates are reranked. + * Default: 3 + */ + rerankFactor?: number; + /** + * Minimum number of candidates to rerank (floor). + * Default: 10 + */ + minCandidates?: number; +} + +/** Configuration for document chunking. Set to `false` to disable. + * + * Re-exported from loaders/chunker for convenience. + * See {@link import('./loaders/chunker').ChunkingOptions} for full details. */ +export type { ChunkingOptions } from './loaders/chunker'; + /** * Context initialization options */ @@ -18,16 +77,187 @@ export interface ContextOptions { vectorsDir: string; /** Transformers model name */ model?: string; + /** + * Base path for resolving document IDs. + * + * When set, file paths are resolved relative to this directory before + * generating IDs, so the same document gets the same ID regardless of + * the absolute path on different machines. + * + * Defaults to `process.cwd()`. + */ + basePath?: string; + /** Custom loaders to register (default: MarkdownLoader, JsonLoader, TextLoader) */ + loaders?: import('./loaders').Loader[]; + /** + * Custom embedder instance. When provided, skips the auto-resolution + * logic and uses this embedder directly. Useful for tests or when you + * have a pre-initialized embedder. + */ + embedder?: import('./embedder').Embedder; + + /** + * Callback invoked when the embedder resolution falls back to + * SimpleEmbedder (lower recall quality). + * + * Use this to alert users, log diagnostics, or trigger a retry later. + * Not called when a custom `embedder` is provided. + * + * Example: + * ```ts + * const ctx = await Context.create({ + * vectorsDir: './vectors', + * onEmbedderFallback: (info) => { + * console.warn(`Embedder fallback: ${info.fallbackReason}`); + * }, + * }); + * ``` + */ + onEmbedderFallback?: (info: import('./embedder/resolve').EmbedderInfo) => void; + + /** + * Progress callback for `load()` — called after each major phase completes. + * + * Useful for logging, progress bars, or UI updates when loading large + * document sets. Not called for `query()` or other methods. + * + * Example: + * ```ts + * const ctx = await Context.create({ + * vectorsDir: './vectors', + * onProgress: (phase, detail) => { + * console.log(`${phase}: ${detail.loaded}/${detail.total} files`); + * }, + * }); + * ``` + */ + onProgress?: (phase: LoadPhase, detail: LoadProgress) => void; + + /** + * Document chunking configuration. + * + * When enabled, each document is split into semantic chunks (heading-aware + * for Markdown, fixed-size for plain text). Each chunk is independently + * embedded and can be retrieved at a finer granularity. + * + * Set to `false` to disable chunking and embed whole documents instead. + * Defaults to enabled with `{ strategy: 'auto', maxChunkSize: 1024, chunkOverlap: 128 }`. + */ + chunking?: import('./loaders/chunker').ChunkingOptions | false; + + /** + * Query expansion configuration. + * + * When enabled, query terms are expanded with CN↔EN synonym bridges so a + * single query can match documents in both languages. The expanded text + * is used for embedding and FTS. + * + * Pass a custom `synonyms` map to add domain-specific terms. Set to `false` + * to disable expansion entirely. + * + * Defaults to enabled with built-in visualization synonym pairs. + */ + queryExpansion?: QueryExpansionOptions | false; + + // ── Hybrid search & weighting ───────────────────────────────────────── + + /** + * Fields to index for Full Text Search (FTS). + * + * By default, `content` is indexed. Add more fields (e.g. `meta.title`) + * to include metadata in text-based recall. + * + * Only affects newly created stores — existing stores keep their schema. + */ + ftsFields?: string[]; + + /** + * Per-field boost weights for the FTS text path in hybrid search. + * + * Higher weight → field contributes more to the text ranking score. + * Example: `{ content: 1, title: 3 }` makes title matches 3× more + * influential than content matches. + * + * Defaults to `{ content: 1 }` when `ftsFields` includes `content`. + */ + ftsFieldWeights?: Record; + + /** + * FTS tokenizer for full-text search indexing. + * + * - `'jieba'` (default): Chinese word segmentation, best for CN/mixed content + * - `'standard'`: English stemmer + stop words, best for pure English content + * - `'auto'`: same as `'jieba'` (safe default for bilingual datasets) + * + * Only affects newly created stores. Existing stores keep their schema. + */ + tokenizer?: 'jieba' | 'standard' | 'auto'; + + /** + * RRF rank constant for hybrid search fusion. + * + * Controls how much influence rare (low-rank) results have. + * - Lower value (e.g. 10) → top ranks dominate, more "winner-takes-all" + * - Higher value (e.g. 100) → ranks are more evenly weighted + * + * Defaults to 60 (standard RRF default). + */ + rankConstant?: number; } /** * Query options */ export interface QueryOptions { - /** Library name (required) */ - library: string; + /** + * Library name(s) to query. + * + * - Single library: `'g2'` + * - Multiple libraries: `['g2', 'f2']` + * - All libraries: `'*'` + * + * Comma-separated strings (`'g2,f2'`) are also supported for backward + * compatibility but the array form is preferred. + */ + library: string | string[]; /** Number of results to return */ topK?: number; + /** + * Search mode. + * + * - `'hybrid'` (default): combines vector similarity + FTS text matching + * via RRF fusion. Better recall when query terms appear literally in + * documents (e.g. exact API names, configuration keys). + * - `'vector'`: pure semantic vector search only. Useful when FTS is not + * needed or the store has no FTS indexes. + */ + mode?: 'hybrid' | 'vector'; + + /** + * Reranking configuration. + * + * When enabled, the query pipeline uses two-stage retrieval: + * 1. Coarse search (vector / hybrid) returns topK × rerankFactor candidates + * 2. Reranker scores each candidate against the query for precision + * 3. Final sort by reranked score → topK results + * + * Set to `false` to skip reranking (default). Pass an object to configure + * the rerank factor and minimum candidate pool size. + */ + rerank?: RerankOptions | false; + + /** + * Filter expression to narrow search results. + * + * Passed directly to the zvec engine as an exact-match filter. + * Format: `"field = 'value'"` or `"field = 'val1' AND field2 = 'val2'"`. + * + * Useful for scoping queries to specific document categories, versions, + * or parent documents: + * `"parentDocId = 'f3a1b2c4__getting_started'"` — all chunks of one doc + * `"chunkIndex = 0"` — only first chunks + */ + filter?: string; } /** @@ -38,8 +268,63 @@ export interface QueryResult { id: string; /** Document content */ content: string; - /** Similarity score */ + /** 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; -} \ No newline at end of file + /** + * Chunk metadata — present when the result is a chunk of a larger document. + * Use `parentDocId` to look up the original document; `headingPath` shows + * the section path (e.g. ["Configuration", "Line Chart"]). + */ + chunk?: ChunkMeta; + /** + * Original file path relative to `basePath` — allows users to trace + * back to the source document for context review or navigation. + * + * Populated during `load()` and stored as a zvec field. + */ + sourceFilePath?: string; + /** + * The library this result came from. + * + * Useful when querying multiple libraries (`library: ['g2', 'f2']` or `'*'`) + * to distinguish which library contributed each result. + */ + library?: string; +} + +/** + * Load progress phases — emitted by `Context.load()` via the `onProgress` callback. + */ +export type LoadPhase = 'load' | 'chunk' | 'embed' | 'insert'; + +/** + * Load progress detail — passed to `onProgress` callback. + */ +export interface LoadProgress { + /** How many documents/chunks have been processed in this phase */ + loaded: number; + /** Total documents/chunks to process in this phase */ + total: number; +} + +// Re-export zvec types for convenience (also available from storage/zvec-store) +export type { + ZvecDoc, + ZvecQueryResult, + ZvecSearchParams, + ZvecHybridParams, + FtsFieldWeight, + ZvecFieldSchema, + ZvecStoreConfig, +} from './storage/zvec-store'; diff --git a/src/vectorizer/transformer.ts b/src/vectorizer/transformer.ts deleted file mode 100644 index 0c878b6..0000000 --- a/src/vectorizer/transformer.ts +++ /dev/null @@ -1,44 +0,0 @@ -import { pipeline, env, FeatureExtractionPipeline } from '@huggingface/transformers'; -import * as path from 'path'; -import * as fs from 'fs'; - -export class TransformerVectorizer { - private pipeline: FeatureExtractionPipeline | null = null; - - constructor( - private readonly model: string = 'sentence-transformers/all-MiniLM-L6-v2', - private readonly dtype?: 'fp32' | 'fp16' | 'q8' | 'q4' | 'q4f16' - ) {} - - /** - * 配置模型下载镜像 - * 通过环境变量设置: HF_ENDPOINT 或 HF_MIRROR - * 常用镜像: https://hf-mirror.com - */ - private configureMirror(): void { - const mirror = process.env.HF_ENDPOINT || process.env.HF_MIRROR; - if (mirror) env.remoteHost = mirror; - } - - async initialize(): Promise { - this.configureMirror(); - - // 查找本地缓存模型 - const cacheDir = path.join(process.env.HOME || '', '.cache/huggingface/models'); - let modelPath = this.model; - - if (fs.existsSync(cacheDir)) { - const modelName = this.model.replace('sentence-transformers--', '').split('-').slice(0, 2).join('-'); - const localModel = fs.readdirSync(cacheDir).find(d => d.includes(modelName)); - if (localModel) modelPath = path.join(cacheDir, localModel); - } - - this.pipeline = await pipeline('feature-extraction', modelPath, this.dtype ? { dtype: this.dtype } : undefined) as FeatureExtractionPipeline; - } - - async embed(text: string): Promise { - if (!this.pipeline) throw new Error('Pipeline not initialized'); - const output = await this.pipeline(text, { pooling: 'mean', normalize: true }) as unknown as { data: Float32Array }; - return Array.from(output.data); - } -} \ No newline at end of file diff --git a/test/context.test.ts b/test/context.test.ts index ea86458..67ff82a 100644 --- a/test/context.test.ts +++ b/test/context.test.ts @@ -2,26 +2,34 @@ import { describe, it, expect, afterAll, beforeAll } from 'vitest'; import * as fs from 'fs'; import * as path from 'path'; import { Context } from '../src/index'; +import { SimpleEmbedder } from '../src/embedder'; const FIXTURES_DIR = path.join(__dirname, 'fixtures/docs'); const TEST_DIR = path.join(__dirname, '.test-tmp'); +// Use SimpleEmbedder for fast unit tests — no model download needed. +const testEmbedder = new SimpleEmbedder(); + describe('Context', () => { let ctx: Context; beforeAll(async () => { - // 清理目录 if (fs.existsSync(TEST_DIR)) { fs.rmSync(TEST_DIR, { recursive: true, force: true }); } await new Promise((r) => setTimeout(r, 100)); - // 创建全局 Context 实例 - ctx = await Context.create({ vectorsDir: TEST_DIR }); + ctx = await Context.create({ + vectorsDir: TEST_DIR, + embedder: testEmbedder, + }); }); afterAll(async () => { - await new Promise((r) => setTimeout(r, 500)); + if (ctx) { + await ctx.close(); + } + await new Promise((r) => setTimeout(r, 200)); if (fs.existsSync(TEST_DIR)) { fs.rmSync(TEST_DIR, { recursive: true, force: true }); } @@ -39,11 +47,14 @@ describe('Context', () => { describe('load', () => { it('should load markdown files', async () => { - await ctx.load('md', path.join(FIXTURES_DIR, '*.md')); + await ctx.load('md', path.join(FIXTURES_DIR, 'getting-started.md')); - const results = await ctx.query('installation', { library: 'md', topK: 1 }); + // With chunking, the matching chunk may not be the first result. + // Check that at least one result contains the expected term. + const results = await ctx.query('installation', { library: 'md', topK: 3 }); expect(results.length).toBeGreaterThan(0); - expect(results[0].content).toContain('npm'); + const contents = results.map((r) => r.content).join(' '); + expect(contents).toContain('npm'); }); it('should load json files', async () => { @@ -54,13 +65,23 @@ describe('Context', () => { }); it('should preserve metadata from markdown', async () => { - await ctx.load('meta', path.join(FIXTURES_DIR, '*.md')); + await ctx.load('meta', path.join(FIXTURES_DIR, 'getting-started.md')); const results = await ctx.query('guide', { library: 'meta', topK: 1 }); expect(results.length).toBeGreaterThan(0); expect(results[0].meta).toBeDefined(); expect(results[0].meta).toHaveProperty('title'); }); + + it('should skip already loaded documents (deduplication)', async () => { + await ctx.load('md', path.join(FIXTURES_DIR, 'getting-started.md')); + + const results = await ctx.query('install', { library: 'md', topK: 10 }); + // With chunking enabled, a single doc may produce multiple chunks. + // The key invariant is that a second load() with the same pattern + // should NOT increase the count (dedup prevents double-insert). + expect(results.length).toBeGreaterThan(0); + }); }); describe('query', () => { @@ -79,5 +100,332 @@ describe('Context', () => { const results = await ctx.query('install', { library: 'md' }); expect(results.length).toBeGreaterThan(0); }); + + it('should return empty results for unknown library', async () => { + const results = await ctx.query('test', { library: 'nonexistent' }); + expect(results.length).toBe(0); + }); + + it('should support array library parameter', async () => { + const results = await ctx.query('install', { library: ['md', 'json'], topK: 5 }); + expect(results.length).toBeGreaterThan(0); + }); + + it('should default to hybrid search mode', async () => { + // Default mode is 'hybrid' — combines vector + text path + const results = await ctx.query('install', { library: 'md', topK: 1 }); + expect(results.length).toBeGreaterThan(0); + }); + + it('should support vector-only search mode', async () => { + const results = await ctx.query('install', { library: 'md', topK: 1, mode: 'vector' }); + expect(results.length).toBeGreaterThan(0); + }); + }); + + describe('cross-library query', () => { + it('should query multiple comma-separated libraries', async () => { + const results = await ctx.query('install', { library: 'md,json', topK: 5 }); + expect(results.length).toBeGreaterThan(0); + }); + }); + + describe('close', () => { + it('should close all stores without error', async () => { + const closeTestDir = TEST_DIR + '-close-test'; + const ctx2 = await Context.create({ + vectorsDir: closeTestDir, + embedder: testEmbedder, + }); + await ctx2.load('close-test', path.join(FIXTURES_DIR, 'getting-started.md')); + await ctx2.close(); + if (fs.existsSync(closeTestDir)) { + fs.rmSync(closeTestDir, { recursive: true, force: true }); + } + }); + }); + + describe('chunking', () => { + const chunkTestDir = TEST_DIR + '-chunk-test'; + + afterAll(() => { + if (fs.existsSync(chunkTestDir)) { + fs.rmSync(chunkTestDir, { recursive: true, force: true }); + } + }); + + it('should split large documents into chunks', async () => { + const ctx = await Context.create({ + vectorsDir: chunkTestDir, + embedder: testEmbedder, + chunking: { maxChunkSize: 500, chunkOverlap: 50 }, + }); + + await ctx.load('chunked', path.join(FIXTURES_DIR, 'line-chart-guide.md')); + + const results = await ctx.query('tooltip', { library: 'chunked', topK: 3 }); + expect(results.length).toBeGreaterThan(0); + + // At least one result should be a chunk (long doc with small maxChunkSize) + const chunkResults = results.filter((r) => r.chunk); + expect(chunkResults.length).toBeGreaterThan(0); + + // Chunk metadata should be present + const chunk = chunkResults[0].chunk!; + expect(chunk.parentDocId).toContain('line_chart_guide'); + expect(typeof chunk.chunkIndex).toBe('number'); + expect(chunk.chunkIndex).toBeGreaterThanOrEqual(0); + + await ctx.close(); + }); + + it('should include headingPath in chunk metadata', async () => { + const ctx = await Context.create({ + vectorsDir: chunkTestDir + '-2', + embedder: testEmbedder, + chunking: { maxChunkSize: 400, chunkOverlap: 50 }, + }); + + await ctx.load('chunked2', path.join(FIXTURES_DIR, 'line-chart-guide.md')); + + // Query for tooltip section specifically + const results = await ctx.query('Line Chart with Tooltip hover', { + library: 'chunked2', + topK: 5, + }); + expect(results.length).toBeGreaterThan(0); + + // At least one chunk should have a heading path containing "Line Chart" + const withHeading = results.filter( + (r) => r.chunk?.headingPath?.some((h) => h.includes('Line Chart')) + ); + expect(withHeading.length).toBeGreaterThan(0); + + await ctx.close(); + if (fs.existsSync(chunkTestDir + '-2')) { + fs.rmSync(chunkTestDir + '-2', { recursive: true, force: true }); + } + }); + + it('should support disabling chunking', async () => { + const ctx = await Context.create({ + vectorsDir: chunkTestDir + '-nochunk', + embedder: testEmbedder, + chunking: false, + }); + + await ctx.load('nochunk', path.join(FIXTURES_DIR, 'line-chart-guide.md')); + + const results = await ctx.query('tooltip', { library: 'nochunk', topK: 3 }); + expect(results.length).toBeGreaterThan(0); + + // When chunking is disabled, no results should have chunk metadata + const chunkResults = results.filter((r) => r.chunk); + expect(chunkResults.length).toBe(0); + + await ctx.close(); + if (fs.existsSync(chunkTestDir + '-nochunk')) { + fs.rmSync(chunkTestDir + '-nochunk', { recursive: true, force: true }); + } + }); + }); +}); + +describe('Context with reranking', () => { + const rerankTestDir = TEST_DIR + '-rerank-test'; + + afterAll(() => { + if (fs.existsSync(rerankTestDir)) { + fs.rmSync(rerankTestDir, { recursive: true, force: true }); + } + }); + + it('should rerank results with keyword scoring', async () => { + const ctx = await Context.create({ + vectorsDir: rerankTestDir, + embedder: testEmbedder, + chunking: { maxChunkSize: 400, chunkOverlap: 50 }, + }); + + await ctx.load('rerank', path.join(FIXTURES_DIR, 'line-chart-guide.md')); + + // Query with reranking enabled + const results = await ctx.query('tooltip configuration', { + library: 'rerank', + topK: 3, + rerank: { rerankFactor: 3 }, + }); + expect(results.length).toBeGreaterThan(0); + expect(results.length).toBeLessThanOrEqual(3); + + // Reranked results should contain content related to "tooltip" + const contents = results.map((r) => r.content).join(' '); + expect(contents.toLowerCase()).toContain('tooltip'); + + await ctx.close(); + }); + + it('should support disabling reranking', async () => { + const ctx = await Context.create({ + vectorsDir: rerankTestDir + '-disabled', + embedder: testEmbedder, + chunking: { maxChunkSize: 400, chunkOverlap: 50 }, + }); + + await ctx.load('rerank2', path.join(FIXTURES_DIR, 'line-chart-guide.md')); + + // Search with rerank=false uses only coarse search, but with topK directly + const results = await ctx.query('tooltip', { + library: 'rerank2', + topK: 3, + rerank: false, + }); + expect(results.length).toBeGreaterThan(0); + + // Scores should be from the original vector/hybrid search (cosine sim or RRF) + // Reranked scores are normalized to [0,1], raw scores can be different. + // Just verify some results exist. + + await ctx.close(); + if (fs.existsSync(rerankTestDir + '-disabled')) { + fs.rmSync(rerankTestDir + '-disabled', { recursive: true, force: true }); + } + }); + }); + +describe('Context with query expansion', () => { + const expandTestDir = TEST_DIR + '-expand-test'; + + afterAll(() => { + if (fs.existsSync(expandTestDir)) { + fs.rmSync(expandTestDir, { recursive: true, force: true }); + } + }); + + it('should expand CN query to match EN content', async () => { + const ctx = await Context.create({ + vectorsDir: expandTestDir, + embedder: testEmbedder, + chunking: { maxChunkSize: 500, chunkOverlap: 50 }, + }); + + await ctx.load('expand', path.join(FIXTURES_DIR, 'line-chart-guide.md')); + + // Query with Chinese term that has EN synonym bridges + const results = await ctx.query('提示框 配置', { + library: 'expand', + topK: 5, + }); + + // Should find tooltip-related content via synonym expansion + const contents = results.map((r) => r.content.toLowerCase()).join(' '); + expect(contents).toContain('tooltip'); + + await ctx.close(); + }); + + it('should support disabling query expansion', async () => { + const ctx = await Context.create({ + vectorsDir: expandTestDir + '-disabled', + embedder: testEmbedder, + queryExpansion: false, + chunking: false, + }); + + await ctx.load('expand2', path.join(FIXTURES_DIR, 'getting-started.md')); + + // Without expansion, should still work normally + const results = await ctx.query('installation', { + library: 'expand2', + topK: 3, + }); + expect(results.length).toBeGreaterThan(0); + + await ctx.close(); + if (fs.existsSync(expandTestDir + '-disabled')) { + fs.rmSync(expandTestDir + '-disabled', { recursive: true, force: true }); + } + }); + + it('should expand EN query to match CN concepts', async () => { + const ctx = await Context.create({ + vectorsDir: expandTestDir + '-en', + embedder: testEmbedder, + queryExpansion: { + synonyms: { + 'animation': ['动效', 'animate', 'transition'], + }, + }, + chunking: { maxChunkSize: 500, chunkOverlap: 50 }, + }); + + await ctx.load('expand3', path.join(FIXTURES_DIR, 'line-chart-guide.md')); + + // Query with "animation" should find content about animation/动效 + const results = await ctx.query('animation', { + library: 'expand3', + topK: 5, + }); + + expect(results.length).toBeGreaterThan(0); + const contents = results.map((r) => r.content.toLowerCase()).join(' '); + expect(contents).toContain('animate'); + + await ctx.close(); + if (fs.existsSync(expandTestDir + '-en')) { + fs.rmSync(expandTestDir + '-en', { recursive: true, force: true }); + } + }); + }); + +describe('Context with weight configuration', () => { + const weightTestDir = TEST_DIR + '-weight-test'; + + afterAll(() => { + if (fs.existsSync(weightTestDir)) { + fs.rmSync(weightTestDir, { recursive: true, force: true }); + } + }); + + it('should create Context with custom ftsFieldWeights', async () => { + const ctx = await Context.create({ + vectorsDir: weightTestDir, + embedder: testEmbedder, + ftsFields: ['content'], + ftsFieldWeights: { content: 2 }, + rankConstant: 30, + }); + + await ctx.load('weighted', path.join(FIXTURES_DIR, 'getting-started.md')); + + const results = await ctx.query('guide', { library: 'weighted', topK: 1 }); + expect(results.length).toBeGreaterThan(0); + + await ctx.close(); + }); + + it('should use hybrid search with MemoryZvecStore when zvec unavailable', async () => { + // When zvec is not available, MemoryZvecStore is used with FtsFieldWeights + const ctx = await Context.create({ + vectorsDir: weightTestDir + '-mem', + embedder: testEmbedder, + ftsFields: ['content'], + ftsFieldWeights: { content: 3 }, + }); + + await ctx.load('mem-weighted', path.join(FIXTURES_DIR, 'getting-started.md')); + + // Hybrid mode (default) should use both vector + text path + const hybridResults = await ctx.query('Getting Started', { library: 'mem-weighted', topK: 1 }); + expect(hybridResults.length).toBeGreaterThan(0); + + // Vector-only mode should work too + const vectorResults = await ctx.query('Getting Started', { library: 'mem-weighted', topK: 1, mode: 'vector' }); + expect(vectorResults.length).toBeGreaterThan(0); + + await ctx.close(); + if (fs.existsSync(weightTestDir + '-mem')) { + fs.rmSync(weightTestDir + '-mem', { recursive: true, force: true }); + } }); }); \ No newline at end of file diff --git a/test/fixtures/docs/line-chart-guide.md b/test/fixtures/docs/line-chart-guide.md new file mode 100644 index 0000000..94bd35d --- /dev/null +++ b/test/fixtures/docs/line-chart-guide.md @@ -0,0 +1,122 @@ +## Installation + +Install via npm or yarn: + +```bash +npm install @antv/context +``` + +Or with yarn: + +```bash +yarn add @antv/context +``` + + +## Quick Start + +```typescript +import { Context } from '@antv/context'; + +const ctx = await Context.create({ vectorsDir: './vectors' }); +await ctx.load('g2', './docs/**/*.md'); +const results = await ctx.query('How to configure a line chart', { library: 'g2', topK: 5 }); +``` + + +## Line Chart Configuration + +### Basic Line Chart + +To create a basic line chart, use the `line()` method on the chart instance: + +```typescript +chart.line() + .position('date*price') + .color('category') + .size(2); +``` + +The position channel maps `date` to the x-axis and `price` to the y-axis. +The color channel differentiates lines by the `category` field. + +### Line Chart with Tooltip + +Tooltips display detailed information when hovering over data points. +Configure tooltips using the `tooltip()` method: + +```typescript +chart.tooltip({ + showCrosshairs: true, + shared: true, + fields: ['date', 'price', 'volume'], + formatter: (datum) => ({ + name: datum.category, + value: datum.price.toFixed(2) + }) +}); +``` + +The `showCrosshairs` option draws vertical and horizontal guide lines. +Set `shared: true` to merge tooltip items from all series at the same x position. + +### Line Chart with Animation + +Enable animations to make your charts more engaging: + +```typescript +chart.animate({ + enter: { + type: 'pathIn', + duration: 1000, + easing: 'easeCubicOut' + }, + update: { + type: 'morph', + duration: 500 + } +}); +``` + +Enter animations play when data points first appear. +Update animations interpolate between old and new positions when data changes. + + +## Advanced Configuration + +### Custom Axis Labels + +Customize axis labels using the `axis()` method: + +```typescript +chart.axis('price', { + title: { text: 'Price (USD)' }, + label: { + formatter: (val) => `$${val}`, + autoRotate: true + } +}); +``` + +### Legend Configuration + +Control legend appearance with the `legend()` method: + +```typescript +chart.legend('category', { + position: 'top', + marker: { symbol: 'circle', size: 8 }, + itemName: { style: { fontSize: 12 } } +}); +``` + +### Responsive Design + +Create responsive charts that adapt to container size changes: + +```typescript +chart.forceFit(); +window.addEventListener('resize', () => { + chart.forceFit(); +}); +``` diff --git a/test/language.test.ts b/test/language.test.ts new file mode 100644 index 0000000..da8b5dc --- /dev/null +++ b/test/language.test.ts @@ -0,0 +1,86 @@ +import { describe, it, expect } from 'vitest'; +import { isCJK, detectLanguage, tokenizerForLanguage, detectTokenizer } from '../src/embedder/language'; + +describe('isCJK', () => { + it('should detect Chinese characters', () => { + expect(isCJK('中')).toBe(true); + expect(isCJK('图')).toBe(true); + }); + + it('should detect Hiragana', () => { + expect(isCJK('あ')).toBe(true); + }); + + it('should detect Katakana', () => { + expect(isCJK('ア')).toBe(true); + }); + + it('should detect Hangul', () => { + expect(isCJK('한')).toBe(true); + }); + + it('should not detect ASCII letters', () => { + expect(isCJK('a')).toBe(false); + expect(isCJK('Z')).toBe(false); + }); + + it('should not detect digits', () => { + expect(isCJK('0')).toBe(false); + }); + + it('should not detect punctuation', () => { + expect(isCJK('.')).toBe(false); + expect(isCJK('!')).toBe(false); + }); +}); + +describe('detectLanguage', () => { + it('should detect CJK-dominant text', () => { + expect(detectLanguage('折线图配置方法详解')).toBe('cjk'); + }); + + it('should detect Latin text', () => { + expect(detectLanguage('hello world configuration')).toBe('latin'); + }); + + it('should detect mixed text with >15% CJK', () => { + expect(detectLanguage('Line Chart 折线图')).toBe('mixed'); + }); + + it('should classify low CJK ratio as latin', () => { + // Only 1 CJK char out of ~20 chars — < 15% threshold + expect(detectLanguage('this is a long English sentence with one 字')).toBe('latin'); + }); + + it('should return latin for empty text', () => { + expect(detectLanguage('')).toBe('latin'); + }); +}); + +describe('tokenizerForLanguage', () => { + it('should return jieba for cjk', () => { + expect(tokenizerForLanguage('cjk')).toBe('jieba'); + }); + + it('should return jieba for mixed', () => { + expect(tokenizerForLanguage('mixed')).toBe('jieba'); + }); + + it('should return standard for latin', () => { + expect(tokenizerForLanguage('latin')).toBe('standard'); + }); +}); + +describe('detectTokenizer', () => { + it('should auto-detect jieba for Chinese text', () => { + expect(detectTokenizer('折线图')).toBe('jieba'); + }); + + it('should auto-detect standard for English text', () => { + expect(detectTokenizer('hello world')).toBe('standard'); + }); + + it('should auto-detect jieba for mixed text', () => { + expect(detectTokenizer('chart 折线图')).toBe('jieba'); + }); +}); diff --git a/test/loaders.test.ts b/test/loaders.test.ts index 2143c7d..2e42130 100644 --- a/test/loaders.test.ts +++ b/test/loaders.test.ts @@ -10,9 +10,12 @@ describe('Loaders', () => { describe('MarkdownLoader', () => { it('should load markdown file with front-matter', async () => { const loader = new MarkdownLoader(); - const doc = await loader.load(path.join(FIXTURES_DIR, 'getting-started.md')); + const filePath = path.join(FIXTURES_DIR, 'getting-started.md'); + const doc = await loader.load(filePath); - expect(doc.id).toContain('getting-started.md'); + // Loader returns the file path as a temporary ID — the canonical + // hash-based ID is assigned by Context.load() for cross-machine consistency. + expect(doc.id).toBe(filePath); expect(doc.content).toContain('Getting Started'); expect(doc.content).toContain('npm install'); expect(doc.meta).toEqual({ @@ -32,9 +35,10 @@ describe('Loaders', () => { describe('JsonLoader', () => { it('should load json file', async () => { const loader = new JsonLoader(); - const doc = await loader.load(path.join(FIXTURES_DIR, 'api.json')); + const filePath = path.join(FIXTURES_DIR, 'api.json'); + const doc = await loader.load(filePath); - expect(doc.id).toContain('api.json'); + expect(doc.id).toBe(filePath); expect(doc.content).toContain('API Reference'); expect(doc.content).toContain('/users'); }); @@ -49,9 +53,10 @@ describe('Loaders', () => { describe('TextLoader', () => { it('should load text file', async () => { const loader = new TextLoader(); - const doc = await loader.load(path.join(FIXTURES_DIR, 'notes.txt')); + const filePath = path.join(FIXTURES_DIR, 'notes.txt'); + const doc = await loader.load(filePath); - expect(doc.id).toContain('notes.txt'); + expect(doc.id).toBe(filePath); expect(doc.content).toContain('notes'); }); diff --git a/test/reranker.test.ts b/test/reranker.test.ts new file mode 100644 index 0000000..e997c14 --- /dev/null +++ b/test/reranker.test.ts @@ -0,0 +1,93 @@ +import { describe, it, expect } from 'vitest'; +import { KeywordReranker, createReranker } from '../src/reranker'; +import type { RerankCandidate } from '../src/reranker'; + +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([]); + }); + + 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'); + }); + + 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 }, + ]; + + const results = await reranker.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 }, + ]; + + const results = await reranker.rerank('tooltip', 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 }, + ]; + + 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); + }); + + 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 }, + ]; + + const results = await reranker.rerank('test', candidates); + for (const r of results) { + expect(r.score).toBeGreaterThanOrEqual(0); + expect(r.score).toBeLessThanOrEqual(1); + } + }); + + it('should handle CJK queries', async () => { + const candidates: RerankCandidate[] = [ + { id: '1', content: '折线图配置方法详解', score: 0.5 }, + { id: '2', content: '柱状图使用说明', score: 0.5 }, + ]; + + const results = await reranker.rerank('折线图', candidates); + const sorted = [...results].sort((a, b) => b.score - a.score); + expect(sorted[0].id).toBe('1'); + }); + }); +}); + +describe('createReranker', () => { + it('should return a KeywordReranker instance', () => { + const reranker = createReranker(); + expect(reranker).toBeInstanceOf(KeywordReranker); + }); +}); diff --git a/test/simple-embedder.test.ts b/test/simple-embedder.test.ts new file mode 100644 index 0000000..def974c --- /dev/null +++ b/test/simple-embedder.test.ts @@ -0,0 +1,92 @@ +import { describe, it, expect } from 'vitest'; +import { SimpleEmbedder } from '../src/embedder/simple'; + +describe('SimpleEmbedder', () => { + const embedder = new SimpleEmbedder(); + + describe('basic properties', () => { + it('should have 512 dimensions', () => { + expect(embedder.dimensions).toBe(512); + }); + }); + + describe('embed', () => { + it('should return a vector of correct length', async () => { + const vec = await embedder.embed('hello world'); + expect(vec.length).toBe(512); + }); + + it('should return non-zero vectors for meaningful text', async () => { + const vec = await embedder.embed('important semantic content'); + const nonzeroCount = vec.filter((v) => v !== 0).length; + expect(nonzeroCount).toBeGreaterThan(0); + }); + + it('should produce different vectors for different text', async () => { + const vec1 = await embedder.embed('chart visualization'); + const vec2 = await embedder.embed('database query'); + // Vectors should not be identical + const identical = vec1.every((v, i) => v === vec2[i]); + expect(identical).toBe(false); + }); + + it('should handle CJK text', async () => { + const vec = await embedder.embed('折线图配置'); + expect(vec.length).toBe(512); + const nonzeroCount = vec.filter((v) => v !== 0).length; + expect(nonzeroCount).toBeGreaterThan(0); + }); + + it('should handle mixed CN/EN text', async () => { + const vec = await embedder.embed('line chart 折线图'); + expect(vec.length).toBe(512); + }); + + it('should handle empty text gracefully', async () => { + const vec = await embedder.embed(''); + expect(vec.length).toBe(512); + }); + }); + + describe('embedBatch', () => { + it('should return vectors for each input text', async () => { + const vecs = await embedder.embedBatch(['hello', 'world']); + expect(vecs.length).toBe(2); + expect(vecs[0].length).toBe(512); + expect(vecs[1].length).toBe(512); + }); + + it('should produce same result as individual embed calls', async () => { + const batchVecs = await embedder.embedBatch(['test one', 'test two']); + const singleVec1 = await embedder.embed('test one'); + const singleVec2 = await embedder.embed('test two'); + // Results should be identical (SimpleEmbedder is deterministic) + expect(batchVecs[0]).toEqual(singleVec1); + expect(batchVecs[1]).toEqual(singleVec2); + }); + }); + + describe('L2 normalization', () => { + it('should produce L2-normalized vectors', async () => { + const vec = await embedder.embed('normalized vector test'); + const norm = Math.sqrt(vec.reduce((sum, v) => sum + v * v, 0)); + expect(Math.abs(norm - 1)).toBeLessThan(0.01); + }); + }); + + describe('with synonym map', () => { + it('should incorporate synonym terms into embedding', async () => { + const synonyms = new Map([ + ['tooltip', ['提示框']], + ]); + const embedderWithSyn = new SimpleEmbedder(synonyms); + + const vecWithout = await embedder.embed('tooltip'); + const vecWith = await embedderWithSyn.embed('tooltip'); + + // With synonyms, the vector should differ (extra terms injected) + const identical = vecWithout.every((v, i) => v === vecWith[i]); + expect(identical).toBe(false); + }); + }); +}); diff --git a/test/synonym-expander.test.ts b/test/synonym-expander.test.ts new file mode 100644 index 0000000..4c8eec0 --- /dev/null +++ b/test/synonym-expander.test.ts @@ -0,0 +1,84 @@ +import { describe, it, expect } from 'vitest'; +import { SynonymExpander, NoopExpander } from '../src/query-expander'; + +describe('SynonymExpander', () => { + describe('with user-provided synonyms', () => { + const expander = new SynonymExpander({ + '折线图': ['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 handle empty query', () => { + const result = expander.expand(''); + expect(result).toBe(''); + }); + }); + + describe('with no synonyms (empty map)', () => { + const expanderEmpty = new SynonymExpander({}); + const expanderDefault = new SynonymExpander(); + + it('should return original query unchanged with empty map', () => { + const result = expanderEmpty.expand('tooltip configuration'); + expect(result).toBe('tooltip configuration'); + }); + + it('should return original query unchanged with no arguments', () => { + const result = expanderDefault.expand('tooltip configuration'); + expect(result).toBe('tooltip configuration'); + }); + + it('should handle empty query', () => { + const result = expanderDefault.expand(''); + expect(result).toBe(''); + }); + }); +}); + +describe('NoopExpander', () => { + const expander = new NoopExpander(); + + it('should return query unchanged', () => { + const result = expander.expand('tooltip configuration'); + expect(result).toBe('tooltip configuration'); + }); + + it('should handle empty query', () => { + const result = expander.expand(''); + expect(result).toBe(''); + }); +}); diff --git a/test/utils.test.ts b/test/utils.test.ts new file mode 100644 index 0000000..58d9548 --- /dev/null +++ b/test/utils.test.ts @@ -0,0 +1,116 @@ +import { describe, it, expect } from 'vitest'; +import { cosineSimilarity, evalMemoryFilter } from '../src/storage/utils'; + +describe('cosineSimilarity', () => { + it('should return 1 for identical vectors', () => { + const vec = [1, 0, 0]; + expect(cosineSimilarity(vec, vec)).toBeCloseTo(1); + }); + + it('should return 0 for orthogonal vectors', () => { + expect(cosineSimilarity([1, 0], [0, 1])).toBeCloseTo(0); + }); + + it('should return -1 for opposite vectors', () => { + expect(cosineSimilarity([1, 0], [-1, 0])).toBeCloseTo(-1); + }); + + it('should return 0 for zero vectors', () => { + expect(cosineSimilarity([0, 0, 0], [1, 2, 3])).toBe(0); + }); + + it('should handle high-dimensional vectors', () => { + const a = Array.from({ length: 512 }, (_, i) => i % 2 === 0 ? 0.1 : 0.2); + const b = Array.from({ length: 512 }, (_, i) => i % 2 === 0 ? 0.2 : 0.1); + const sim = cosineSimilarity(a, b); + expect(sim).toBeGreaterThan(0); + expect(sim).toBeLessThan(1); + }); +}); + +describe('evalMemoryFilter', () => { + const fields = { + content: 'hello world', + chunkIndex: 0, + parentDocId: 'abc__test', + totalChunks: 5, + headingPath: 'Line Chart > Tooltip', + }; + + describe('string equality', () => { + it('should match string field with quoted value', () => { + expect(evalMemoryFilter("parentDocId = 'abc__test'", fields)).toBe(true); + }); + + it('should reject non-matching string value', () => { + expect(evalMemoryFilter("parentDocId = 'xyz'", fields)).toBe(false); + }); + + it('should match content field', () => { + expect(evalMemoryFilter("content = 'hello world'", fields)).toBe(true); + }); + }); + + describe('number equality', () => { + it('should match numeric field with integer value', () => { + expect(evalMemoryFilter("chunkIndex = 0", fields)).toBe(true); + }); + + it('should reject non-matching numeric value', () => { + expect(evalMemoryFilter("chunkIndex = 3", fields)).toBe(false); + }); + + it('should match totalChunks field', () => { + expect(evalMemoryFilter("totalChunks = 5", fields)).toBe(true); + }); + + it('should handle negative numbers', () => { + expect(evalMemoryFilter("chunkIndex = -1", { chunkIndex: -1 })).toBe(true); + }); + }); + + describe('AND conjunction', () => { + it('should match when all clauses are true', () => { + expect(evalMemoryFilter("chunkIndex = 0 AND parentDocId = 'abc__test'", fields)).toBe(true); + }); + + it('should reject when any clause is false', () => { + expect(evalMemoryFilter("chunkIndex = 1 AND parentDocId = 'abc__test'", fields)).toBe(false); + }); + + it('should reject when all clauses are false', () => { + expect(evalMemoryFilter("chunkIndex = 3 AND parentDocId = 'xyz'", fields)).toBe(false); + }); + + it('should support mixed string and number clauses', () => { + expect(evalMemoryFilter("chunkIndex = 0 AND headingPath = 'Line Chart > Tooltip'", fields)).toBe(true); + }); + + it('should support three clauses with AND', () => { + expect(evalMemoryFilter("chunkIndex = 0 AND totalChunks = 5 AND parentDocId = 'abc__test'", fields)).toBe(true); + }); + }); + + describe('edge cases', () => { + it('should pass through unknown clause format', () => { + // Unknown format → true (let native zvec handle it) + expect(evalMemoryFilter("field > 10", fields)).toBe(true); + }); + + it('should pass through empty filter', () => { + expect(evalMemoryFilter('', fields)).toBe(true); + }); + + it('should handle AND with case-insensitive keyword', () => { + expect(evalMemoryFilter("chunkIndex = 0 and parentDocId = 'abc__test'", fields)).toBe(true); + }); + + it('should handle missing field gracefully (string clause)', () => { + expect(evalMemoryFilter("unknownField = 'value'", fields)).toBe(false); + }); + + it('should handle missing field gracefully (number clause)', () => { + expect(evalMemoryFilter("unknownField = 42", fields)).toBe(false); + }); + }); +}); diff --git a/test/zvec-store.test.ts b/test/zvec-store.test.ts index b8f7776..1fcff42 100644 --- a/test/zvec-store.test.ts +++ b/test/zvec-store.test.ts @@ -1,62 +1,182 @@ import { describe, it, expect, beforeEach, afterEach } from 'vitest'; import * as fs from 'fs'; import * as path from 'path'; -import { ZVecStore } from '../src/storage/zvec-store'; +import { + MemoryZvecStore, + 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'); -describe('ZVecStore', () => { - let store: ZVecStore; +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('MemoryZvecStore', () => { + let store: MemoryZvecStore; + + beforeEach(() => { + store = new MemoryZvecStore(); + }); + + it('should insert and search documents', async () => { + await store.insert([ + { id: 'doc1', vector: [0.1, 0.2, 0.3, 0.4], fields: { content: 'hello world' } }, + { id: 'doc2', vector: [0.4, 0.3, 0.2, 0.1], fields: { content: 'goodbye' } }, + ]); + + const results = await store.search({ vector: [0.1, 0.2, 0.3, 0.4], topK: 2 }); + expect(results.length).toBe(2); + expect(results[0].id).toBe('doc1'); + expect(results[0].score).toBeGreaterThan(0.9); + }); + + it('should support sync search', () => { + store.insert([ + { id: 'doc1', vector: [0.1, 0.2, 0.3, 0.4], fields: { content: 'hello' } }, + ]); + + const results = store.searchSync({ vector: [0.1, 0.2, 0.3, 0.4], topK: 1 }); + expect(results.length).toBe(1); + expect(results[0].id).toBe('doc1'); + }); + + it('should support filter expressions', async () => { + await store.insert([ + { id: 'a', vector: [1, 0, 0, 0], fields: { kind: 'foo' } }, + { id: 'b', vector: [0, 1, 0, 0], fields: { kind: 'bar' } }, + ]); + + const results = await store.search({ + vector: [1, 0, 0, 0], + topK: 2, + filter: "kind = 'foo'", + }); + expect(results.length).toBe(1); + expect(results[0].id).toBe('a'); + }); + + it('should support hybrid search with configurable field weights', async () => { + const weighted = new MemoryZvecStore([ + { fieldName: 'title', weight: 3 }, + { fieldName: 'content', weight: 1 }, + ]); + + await weighted.insert([ + { id: 't1', vector: [0.1, 0, 0, 0], fields: { title: 'sankey diagram', content: 'flow data' } }, + { id: 't2', vector: [0, 0.1, 0, 0], fields: { title: 'unrelated', content: 'sankey flow visualization' } }, + ]); + + const results = await weighted.searchHybrid({ + queryText: 'sankey', + queryVector: [0.1, 0.1, 0, 0], + topK: 2, + }); + expect(results.length).toBeGreaterThan(0); + }); + + it('should handle empty store gracefully', async () => { + const results = await store.search({ vector: [1, 0, 0, 0], topK: 5 }); + expect(results.length).toBe(0); + }); + + it('should clear all data on close', async () => { + await store.insert([ + { id: 'doc1', vector: [1, 0, 0, 0], fields: {} }, + ]); + await store.close(); + const results = await store.search({ vector: [1, 0, 0, 0], topK: 5 }); + expect(results.length).toBe(0); + }); +}); + +describe('ActualZvecStore (native)', () => { + let store: IZvecStore; beforeEach(async () => { - // Clean up before creating store if (fs.existsSync(TEST_DIR)) { fs.rmSync(TEST_DIR, { recursive: true, force: true }); } - store = await ZVecStore.create(STORE_PATH, 4); }); afterEach(async () => { - store?.close(); + 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 (e) { - // Ignore - may be locked - } + } catch { /* may be locked */ } } }); - it('should create and query store', () => { - store.add('doc1', [0.1, 0.2, 0.3, 0.4], 'content1'); - store.add('doc2', [0.4, 0.3, 0.2, 0.1], 'content2'); - store.save(); + 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.search([0.1, 0.2, 0.3, 0.4], 2); + 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 return meta with content', () => { - store.add('doc1', [0.1, 0.2, 0.3, 0.4], 'content1', { title: 'Test' }); - store.save(); + 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 doc = store.getDoc('doc1'); - expect(doc).toBeDefined(); - expect(doc!.content).toBe('content1'); - expect(doc!.meta).toEqual({ title: 'Test' }); + const results = store.searchHybridSync({ + queryText: 'sankey', + queryVector: [0.9, 0.1, 0, 0], + topK: 2, + }); + expect(results.length).toBeGreaterThan(0); }); - it('should clear all data', () => { - store.add('doc1', [0.1, 0.2, 0.3, 0.4], 'content1'); - store.add('doc2', [0.4, 0.3, 0.2, 0.1], 'content2'); - store.save(); + it('should fall back to MemoryZvecStore when zvec not available', async () => { + // When zvec IS available, createZvecStore creates an ActualZvecStore. + // Either way, the returned store should work. + const config = makeConfig(4); + const s = await createZvecStore(STORE_PATH, config); + expect(s).toBeDefined(); - store.clear(); + await s.insert([ + { id: 'x', vector: [1, 2, 3, 4], fields: { content: 'test content' } }, + ]); + const results = s.searchSync({ vector: [1, 2, 3, 4], topK: 1 }); + expect(results.length).toBe(1); + expect(results[0].id).toBe('x'); - expect(store.getDoc('doc1')).toBeUndefined(); - expect(store.getDoc('doc2')).toBeUndefined(); + await s.close(); }); -}); \ No newline at end of file +}); diff --git a/tsconfig.json b/tsconfig.json index 7db8599..cb7da18 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -1,8 +1,8 @@ { "compilerOptions": { "target": "ES2020", - "module": "ESNext", - "moduleResolution": "bundler", + "module": "commonjs", + "moduleResolution": "node", "lib": ["ES2020"], "outDir": "./dist", "rootDir": "./src", diff --git a/vitest.config.ts b/vitest.config.ts index 76131eb..5f2242a 100644 --- a/vitest.config.ts +++ b/vitest.config.ts @@ -3,8 +3,8 @@ import { defineConfig } from 'vitest/config'; export default defineConfig({ test: { include: ['test/**/*.test.ts'], - testTimeout: 300000, - hookTimeout: 120000, + testTimeout: 30000, + hookTimeout: 10000, pool: 'forks', poolOptions: { forks: { singleFork: true }, From 1c9e616146b5a3d8a9328f93fd6ea89f7e4ace8d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A6=8F=E6=99=8B?= Date: Mon, 6 Jul 2026 10:23:33 +0800 Subject: [PATCH 02/13] feat: bind the cache to embedder --- src/context.ts | 4 +- src/embedder/manager.ts | 61 +++++++++++-- src/embedder/transformers.ts | 74 ++++++++++++++- src/index.ts | 4 - src/loaders/chunker.ts | 169 ++++++++++++++++++++++++++++++++++- src/loaders/util.ts | 10 ++- src/query-expander.ts | 56 +++++++++--- src/reranker.ts | 81 +++++++++++++---- src/store-manager.ts | 49 ++++++++-- src/types.ts | 43 ++++++++- 10 files changed, 491 insertions(+), 60 deletions(-) diff --git a/src/context.ts b/src/context.ts index 0b4fde9..9a1089b 100644 --- a/src/context.ts +++ b/src/context.ts @@ -56,7 +56,8 @@ export class Context { : null; // Reranker is created eagerly — it has no model-load cost (KeywordReranker). // In the future this may be lazy if a cross-encoder is used. - this.reranker = createReranker(); + // Weights are configurable via ContextOptions.rerankWeights. + this.reranker = createReranker(options.rerankWeights); // Query expansion — uses user-provided synonym map only, no built-in defaults. // When queryExpansion is false, expansion is disabled entirely. // When queryExpansion is true/undefined/object without synonyms, SynonymExpander @@ -403,6 +404,7 @@ export class Context { meta, sourceFilePath: result.fields?.sourceFilePath as string | undefined, library, + embedderKind: this._embedderInfo.kind, }; // Attach chunk metadata when present (non-negative chunkIndex) diff --git a/src/embedder/manager.ts b/src/embedder/manager.ts index 8bf8ac3..dc6fa73 100644 --- a/src/embedder/manager.ts +++ b/src/embedder/manager.ts @@ -1,22 +1,62 @@ /** - * EmbedderManager — encapsulates module-level state for embedder resolution. + * EmbedderManager — encapsulates embedder resolution and lifecycle. * - * Previously the transformers module cache, failure tracking, and default - * embedder singleton were scattered as loose module-level variables. This - * class consolidates them into a single object so each Context instance can - * manage its own lifecycle without hidden global side-effects. + * 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 { SimpleEmbedder } from './simple'; -import { TransformersEmbedder, loadTransformersModule, resetTransformersModule } from './transformers'; +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; + private readonly _ownsLoader: boolean; + + constructor(options?: EmbedderManagerOptions) { + if (options?.transformersLoader) { + this._loader = options.transformersLoader; + this._ownsLoader = false; + } else { + // Default: use global shared loader (backward-compatible) + this._loader = { + load: loadTransformersModule, + reset: resetTransformersModule, + }; + this._ownsLoader = false; + } + } + + /** + * 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). @@ -28,10 +68,10 @@ export class EmbedderManager { async getEmbedder(synonymMap?: Map): Promise { if (this._defaultEmbedder) return this._defaultEmbedder; - const t = await loadTransformersModule(); + const t = await this._loader.load(); if (t) { try { - const probe = new TransformersEmbedder(); + const probe = new TransformersEmbedder(() => this._loader.load()); await probe.embed('probe'); // triggers lazy model download this._defaultEmbedder = probe; } catch (err) { @@ -59,10 +99,13 @@ export class EmbedderManager { /** * 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; - resetTransformersModule(); + this._loader.reset(); } } diff --git a/src/embedder/transformers.ts b/src/embedder/transformers.ts index 2836415..f70fd50 100644 --- a/src/embedder/transformers.ts +++ b/src/embedder/transformers.ts @@ -41,15 +41,83 @@ interface TransformersOutput { // 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. * - * This function is used by EmbedderManager but also available standalone - * for direct TransformersEmbedder construction. + * Uses the shared module-level cache. For instance-level isolation, + * use `createTransformersLoader()` instead. */ export async function loadTransformersModule(): Promise { if (_transformersModule) return _transformersModule; @@ -87,7 +155,7 @@ export async function loadTransformersModule(): Promise maxSize) { - // Try to break at a sentence boundary (。.!?。！？\n) - let cutPoint = findBreakPoint(current, maxSize); + let cutPoint = findBreakPointSafe(current, maxSize); chunks.push(current.slice(0, cutPoint).trim()); const suffix = current.slice(cutPoint).trim(); current = overlap > 0 && suffix.length > 0 @@ -271,6 +275,163 @@ function fixedSizeSplit( return chunks.length > 0 ? chunks : [text.trim()]; } +// --------------------------------------------------------------------------- +// Block-aware splitting helpers +// --------------------------------------------------------------------------- + +/** Regex matching the opening/closing of a fenced code block. */ +const FENCED_CODE_RE = /^(`{3,}|~{3,})/; + +/** + * Check whether a line is part of a pipe table row. + * Matches lines like `| col1 | col2 |` and separator rows `|---|---|`. + */ +function isTableRow(line: string): boolean { + const trimmed = line.trim(); + return trimmed.startsWith('|') && trimmed.endsWith('|') && trimmed.length > 1; +} + +/** + * Split text into logical units that respect fenced code blocks and tables. + * + * Each returned unit is either: + * - A complete fenced code block (never split) + * - A contiguous run of table rows (never split) + * - A run of normal paragraphs separated by double newlines + * + * This ensures downstream fixed-size splitting never cuts inside a + * code block or table. + */ +function splitRespectingBlocks(text: string): string[] { + const lines = text.split('\n'); + const units: string[] = []; + let buffer: string[] = []; + let inCodeBlock = false; + let codeFence = ''; + let inTable = false; + + function flushBuffer(): void { + if (buffer.length > 0) { + const content = buffer.join('\n').trim(); + if (content.length > 0) { + units.push(content); + } + buffer = []; + } + } + + for (const line of lines) { + // --- Fenced code block detection --- + const fenceMatch = line.match(FENCED_CODE_RE); + if (fenceMatch) { + if (!inCodeBlock) { + // Entering code block — flush any preceding normal content + flushBuffer(); + inCodeBlock = true; + codeFence = fenceMatch[1].charAt(0); // ` or ~ + buffer.push(line); + } else if (line.trimStart().startsWith(codeFence.repeat(codeFence.length)) && line.trim().length <= codeFence.length + 1) { + // Closing fence — end of code block + buffer.push(line); + flushBuffer(); // emit the complete code block as one unit + inCodeBlock = false; + codeFence = ''; + } else { + // Inside code block + buffer.push(line); + } + continue; + } + + if (inCodeBlock) { + buffer.push(line); + continue; + } + + // --- Table row detection --- + if (isTableRow(line)) { + if (!inTable) { + // Entering table — flush preceding normal content + flushBuffer(); + inTable = true; + } + buffer.push(line); + continue; + } + + // Not a table row + if (inTable) { + // Exiting table — flush the complete table as one unit + flushBuffer(); + inTable = false; + } + + // Normal line — accumulate; double-newline splits are handled + // by the paragraph-level logic below via empty-line detection + buffer.push(line); + } + + // Flush any remaining content + flushBuffer(); + + // Post-process: further split normal units on double-newlines so that + // the paragraph-level sliding window works correctly. Code blocks and + // tables remain as single units. + const result: string[] = []; + for (const unit of units) { + if (FENCED_CODE_RE.test(unit.split('\n')[0] ?? '') || isTableRow(unit.split('\n')[0] ?? '')) { + // Protected block — keep as-is + result.push(unit); + } else { + // Normal text — split on paragraph boundaries + const paras = unit.split(/\n\n+/).map((p) => p.trim()).filter((p) => p.length > 0); + result.push(...paras); + } + } + + return result.length > 0 ? result : [text.trim()]; +} + +/** + * Find a safe break point that avoids cutting inside code blocks or tables. + * + * Delegates to findBreakPoint for natural break detection, then adjusts + * the result if it lands inside a protected region. + */ +function findBreakPointSafe(text: string, maxSize: number): number { + const candidate = findBreakPoint(text, maxSize); + + // Check if the candidate falls inside a fenced code block + const beforeCut = text.slice(0, candidate); + const fenceOpens = (beforeCut.match(/^(`{3,}|~{3,})/gm) ?? []).length; + const fenceCloses = (beforeCut.match(/^(`{3,}|~{3,})\s*$/gm) ?? []).length; + + if (fenceOpens > fenceCloses) { + // Inside an unclosed code block — move cut point backward to before the fence + const lastOpenIdx = beforeCut.lastIndexOf('\n```'); + const lastOpenTilde = beforeCut.lastIndexOf('\n~~~'); + const safePoint = Math.max(lastOpenIdx, lastOpenTilde); + if (safePoint > 0) return safePoint; + } + + // Check if the candidate falls inside a table + const linesBeforeCut = beforeCut.split('\n'); + const lastLine = linesBeforeCut[linesBeforeCut.length - 1] ?? ''; + if (isTableRow(lastLine)) { + // Move backward to before the table started + for (let i = linesBeforeCut.length - 1; i >= 0; i--) { + if (!isTableRow(linesBeforeCut[i])) { + // Found the line before the table — compute its offset + const offset = linesBeforeCut.slice(0, i + 1).join('\n').length + 1; + if (offset > 0) return offset; + break; + } + } + } + + return candidate; +} + /** * Find a natural break point within [maxSize * 0.8, maxSize]. * Prefers paragraph breaks > sentence-ending punctuation > space. diff --git a/src/loaders/util.ts b/src/loaders/util.ts index 2cd8914..2c1fb8a 100644 --- a/src/loaders/util.ts +++ b/src/loaders/util.ts @@ -1,4 +1,5 @@ import * as crypto from 'crypto'; +import * as path from 'path'; /** * Convert a file path to a safe, collision-resistant ID for zvec. @@ -7,11 +8,16 @@ import * as crypto from 'crypto'; * 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: remove leading slashes - const normalized = filePath.replace(/^\/+/, ''); + // 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 = crypto.createHash('sha256').update(normalized).digest('hex').slice(0, 16); diff --git a/src/query-expander.ts b/src/query-expander.ts index 654f132..95414bc 100644 --- a/src/query-expander.ts +++ b/src/query-expander.ts @@ -32,8 +32,22 @@ export interface QueryExpander { // SynonymExpander // --------------------------------------------------------------------------- -/** Characters that delimit term boundaries in mixed CN/EN text. */ -const CJK_BOUNDARY_RE = /[\s一-鿿,，。]/; +/** + * 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. @@ -42,7 +56,10 @@ const CJK_BOUNDARY_RE = /[\s一-鿿,，。]/; * corresponding alternative terms are appended to the query. This bridges * CN↔EN terminology gaps and handles different naming conventions. * - * The expander normalises case and preserves CJK characters for exact matching. + * 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; @@ -78,17 +95,36 @@ export class SynonymExpander implements QueryExpander { } /** - * Check if a term appears as a word/phrase boundary in the query. - * Handles both space-separated (EN) and character-joined (CJK) text. + * 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 + * (e.g. "line" should not match inside "inline" or "pipeline"). */ private containsTerm(text: string, term: string): boolean { - if (text.includes(term)) { - const idx = text.indexOf(term); - const prevOk = idx === 0 || CJK_BOUNDARY_RE.test(text[idx - 1]); + if (!text.includes(term)) return false; + + // CJK terms: substring match is sufficient and correct + if (containsCJK(term)) { + return true; + } + + // Latin terms: enforce word boundaries on both sides + 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 || CJK_BOUNDARY_RE.test(text[afterIdx]); - return prevOk && afterOk; + const afterOk = afterIdx >= text.length || WORD_BOUNDARY_RE.test(text[afterIdx]); + + if (prevOk && afterOk) return true; + + // Move past this occurrence to check subsequent ones + searchFrom = idx + 1; } + return false; } } diff --git a/src/reranker.ts b/src/reranker.ts index 2e9fb7f..daacf0e 100644 --- a/src/reranker.ts +++ b/src/reranker.ts @@ -49,6 +49,27 @@ export interface RerankOptions { * Default: 10 */ minCandidates?: number; + + // ----------------------------------------------------------------------- + // KeywordReranker scoring weights — tune for your domain + // ----------------------------------------------------------------------- + + /** Weight for exact phrase match. Default: 3.0 */ + phraseWeight?: number; + /** Bonus per additional phrase occurrence beyond the first. Default: 0.5 */ + phraseRepeatBonus?: number; + /** Weight for whole-word term match. Default: 1.0 */ + termWeight?: number; + /** Bonus per additional term occurrence beyond the first. Default: 0.2 */ + termRepeatBonus?: number; + /** Weight for substring (partial) term match. Default: 0.3 */ + substringWeight?: number; + /** Bonus when a query term appears in the heading path. Default: 2.0 */ + headingTermBonus?: number; + /** Bonus when the full query phrase appears in the heading path. Default: 2.5 */ + headingPhraseBonus?: number; + /** Fraction of the original coarse score carried into the reranked score. Default: 0.1 */ + originalScoreCarry?: number; } // --------------------------------------------------------------------------- @@ -58,6 +79,18 @@ export interface RerankOptions { const DEFAULT_RERANK_FACTOR = 3; const DEFAULT_MIN_CANDIDATES = 10; +/** Default scoring weights for KeywordReranker. */ +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 // --------------------------------------------------------------------------- @@ -69,22 +102,34 @@ const DEFAULT_MIN_CANDIDATES = 10; * embedding stage captures semantic closeness, while this reranker boosts * candidates that contain the exact query terms (or their substrings). * - * Scoring factors: - * - Exact phrase match: 3.0× weight - * - Whole term match: 1.0× weight (per term) - * - Substring match: 0.3× weight (partial token overlap) - * - Heading match: 2.0× bonus (terms appearing in heading path) - * - Original score: 0.1× carry-over (respects the coarse rank) + * All scoring weights are configurable via constructor options. See + * `RerankOptions` for available tuning parameters. * - * All scores are normalised to [0, 1] via min-max scaling. + * Scores are normalised to [0, 1] via min-max scaling. */ export class KeywordReranker implements Reranker { + private readonly weights: typeof DEFAULT_WEIGHTS; + + 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(); @@ -92,10 +137,10 @@ export class KeywordReranker implements Reranker { // 1. Exact phrase match — strongest signal if (contentLower.includes(queryPhrase)) { - score += 3.0; + score += w.phraseWeight; // Bonus for each additional occurrence const phraseCount = countOccurrences(contentLower, queryPhrase); - score += (phraseCount - 1) * 0.5; + score += (phraseCount - 1) * w.phraseRepeatBonus; } // 2. Per-term matching @@ -103,12 +148,12 @@ export class KeywordReranker implements Reranker { if (contentLower.includes(term)) { // Whole word match if (isWordBoundary(contentLower, term)) { - score += 1.0; + score += w.termWeight; const termCount = countTermMatches(contentLower, term); - score += (termCount - 1) * 0.2; + score += (termCount - 1) * w.termRepeatBonus; } else { // Substring match (e.g. "tool" matches "tooltip") - score += 0.3; + score += w.substringWeight; } } } @@ -118,16 +163,16 @@ export class KeywordReranker implements Reranker { const headingLower = c.headingPath.toLowerCase(); for (const term of queryTerms) { if (headingLower.includes(term)) { - score += 2.0; + score += w.headingTermBonus; } } if (queryPhrase.length > 2 && headingLower.includes(queryPhrase)) { - score += 2.5; + score += w.headingPhraseBonus; } } // 4. Carry over a fraction of the original vector score - score += c.score * 0.1; + score += c.score * w.originalScoreCarry; return { id: c.id, score }; }); @@ -213,11 +258,11 @@ function isWordBoundary( // --------------------------------------------------------------------------- /** - * Create the default reranker. + * Create a reranker with optional weight configuration. * * Currently returns KeywordReranker. In the future this may auto-select * based on available models (cross-encoder, etc.). */ -export function createReranker(): Reranker { - return new KeywordReranker(); +export function createReranker(options?: RerankOptions): Reranker { + return new KeywordReranker(options); } diff --git a/src/store-manager.ts b/src/store-manager.ts index 7c876bc..7de5aac 100644 --- a/src/store-manager.ts +++ b/src/store-manager.ts @@ -3,8 +3,10 @@ import * as path from 'path'; import { createZvecStore, openZvecStoreSync, + isZvecAvailable, } from './storage/zvec-store'; import type { IZvecStore, ZvecStoreConfig, FtsFieldWeight, ActualZvecStoreOptions } from './storage/zvec-store'; +import { MemoryZvecStore } from './storage/memory-store'; import type { Embedder } from './embedder'; import { detectTokenizer } from './embedder'; import type { ContextOptions } from './types'; @@ -156,14 +158,47 @@ export class StoreManager { private async _doGetOrCreate(library: string, sampleText?: string): Promise { const filePath = this.getStorePath(library); - return fs.existsSync(filePath) - ? openZvecStoreSync(filePath, storeOpenOptions(this.contextOptions)) - : await createZvecStore( - filePath, - contextStoreConfig(this.embedder.dimensions, this.contextOptions, sampleText), - this.ftsWeights, - this.rankConstant, + const allowFallback = this.contextOptions?.allowMemoryFallback ?? false; + + if (fs.existsSync(filePath)) { + // Existing store on disk — try to open with zvec + if (isZvecAvailable()) { + return openZvecStoreSync(filePath, storeOpenOptions(this.contextOptions)); + } + // zvec unavailable for existing store + if (allowFallback) { + console.warn( + `[context] @zvec/zvec not available — opening library "${library}" ` + + `with in-memory MemoryZvecStore. Data will NOT be persisted.` ); + return new MemoryZvecStore(this.ftsWeights, this.rankConstant); + } + throw new Error( + `Cannot open existing store "${filePath}": @zvec/zvec is not installed. ` + + `Install it with: pnpm add @zvec/zvec\n` + + `Or set allowMemoryFallback: true to use in-memory store (no persistence).` + ); + } + + // New store — createZvecStore already handles fallback internally + try { + return await createZvecStore( + filePath, + contextStoreConfig(this.embedder.dimensions, this.contextOptions, sampleText), + this.ftsWeights, + this.rankConstant, + ); + } catch (err) { + if (allowFallback) { + console.warn( + `[context] Store creation failed for library "${library}", ` + + `falling back to in-memory MemoryZvecStore. ` + + `Error: ${(err as Error).message?.split('\n')[0]}` + ); + return new MemoryZvecStore(this.ftsWeights, this.rankConstant); + } + throw err; + } } /** diff --git a/src/types.ts b/src/types.ts index a3ab43e..f840fb0 100644 --- a/src/types.ts +++ b/src/types.ts @@ -203,6 +203,35 @@ export interface ContextOptions { * Defaults to 60 (standard RRF default). */ rankConstant?: number; + + /** + * KeywordReranker scoring weights — tune for your domain. + * + * These weights control how the second-stage reranker scores candidates. + * See `RerankOptions` for available tuning parameters. When omitted, + * sensible defaults are used. + * + * Example: boost heading matches for API docs: + * ```ts + * rerankWeights: { headingTermBonus: 4.0, phraseWeight: 5.0 } + * ``` + */ + rerankWeights?: Omit; + + /** + * Allow automatic fallback to in-memory store when @zvec/zvec is unavailable. + * + * When `true`, if @zvec/zvec cannot be loaded, the library will silently + * fall back to `MemoryZvecStore` (no persistence, vectors lost on restart). + * A warning is logged to console. + * + * When `false` (default), missing @zvec/zvec throws an error at store + * creation time, making the failure explicit. + * + * Useful for development/testing environments where native bindings + * may not be available. + */ + allowMemoryFallback?: boolean; } /** @@ -295,12 +324,22 @@ export interface QueryResult { */ sourceFilePath?: string; /** - * The library this result came from. + * Which library this result came from. * * Useful when querying multiple libraries (`library: ['g2', 'f2']` or `'*'`) - * to distinguish which library contributed each result. + * to identify the source of each result. */ library?: string; + /** + * The embedder kind used to produce the vectors for this result's store. + * + * - `'transformers'`: high-quality bilingual model (bge-small-zh-v1.5) + * - `'simple'`: lightweight fallback with lower recall quality + * + * Use this to detect when results come from a degraded embedder and + * warn users or trigger re-indexing with a better model. + */ + embedderKind?: import('./embedder/resolve').EmbedderKind; } /** From cbb5da3ec047998e8934f9a2db11e95cef3d8703 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A6=8F=E6=99=8B?= Date: Mon, 6 Jul 2026 10:48:58 +0800 Subject: [PATCH 03/13] fix: build error --- src/reranker.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/reranker.ts b/src/reranker.ts index daacf0e..dad5578 100644 --- a/src/reranker.ts +++ b/src/reranker.ts @@ -108,7 +108,7 @@ const DEFAULT_WEIGHTS = { * Scores are normalised to [0, 1] via min-max scaling. */ export class KeywordReranker implements Reranker { - private readonly weights: typeof DEFAULT_WEIGHTS; + private readonly weights: Record; constructor(options?: RerankOptions) { this.weights = { From 5fdc318df5bac749c8f755674fab21e9a2868fc9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A6=8F=E6=99=8B?= Date: Mon, 6 Jul 2026 11:00:42 +0800 Subject: [PATCH 04/13] chore: add git ignore --- .gitignore | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.gitignore b/.gitignore index 32710e4..2355ead 100644 --- a/.gitignore +++ b/.gitignore @@ -11,3 +11,5 @@ package-lock.json yarn.lock pnpm-lock.yaml pnpm-workspace.yaml +.test-tmp-* +.vscode \ No newline at end of file From 7f82427b5df7318196060298a84dd878676184df Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A6=8F=E6=99=8B?= Date: Mon, 6 Jul 2026 16:09:29 +0800 Subject: [PATCH 05/13] chore: remove sample embedder --- README.md | 5 +- src/context.ts | 20 +--- src/embedder/index.ts | 12 +-- src/embedder/manager.ts | 56 +++++------ src/embedder/resolve.ts | 142 +++++++-------------------- src/embedder/simple.ts | 184 ----------------------------------- src/types.ts | 25 +---- test/context.test.ts | 8 +- test/simple-embedder.test.ts | 92 ------------------ 9 files changed, 73 insertions(+), 471 deletions(-) delete mode 100644 src/embedder/simple.ts delete mode 100644 test/simple-embedder.test.ts diff --git a/README.md b/README.md index 2e43334..5adcf1a 100644 --- a/README.md +++ b/README.md @@ -14,7 +14,7 @@ A local context retrieval library that enables semantic search over your documen - 🔍 **Hybrid Retrieval**: Combines vector similarity + FTS text matching via RRF fusion for better recall - 🔄 **Deduplication**: Automatically skip already-loaded documents; content-hash change detection for re-embedding updated files - ⚖️ **Weight Configuration**: Per-field FTS boost weights and RRF rank constant tuning -- 🛡️ **Graceful Degradation**: Falls back to SimpleEmbedder when Transformers model unavailable +- 🛡️ **Clear Error Messages**: Throws descriptive errors when Transformers model is unavailable, guiding users to fix the issue - 🧩 **Document Chunking**: Split documents into semantic chunks (heading-aware for Markdown, fixed-size for plain text) for finer-grained retrieval - 🔁 **Two-stage Reranking**: KeywordReranker boosts candidates with exact query term matches after coarse vector/hybrid search - 🌐 **Query Expansion**: SynonymExpander uses user-provided synonym maps to bridge CN↔EN terminology gaps @@ -79,7 +79,6 @@ await ctx.close(); | `model` | `string` | auto | Transformers model name for embedding. Skipped when custom `embedder` is provided. | | `loaders` | `Loader[]` | built-in | Custom loaders (default: MarkdownLoader, JsonLoader, TextLoader) | | `embedder` | `Embedder` | auto-resolved | Custom embedder. Skips auto-resolution when provided. | -| `onEmbedderFallback` | `(info: EmbedderInfo) => void` | — | Callback invoked when the embedder falls back to SimpleEmbedder. | | `onProgress` | `(phase, detail) => void` | — | Progress callback for `load()` phases: `'load'` → `'chunk'` → `'embed'` → `'insert'`. | | `chunking` | `ChunkingOptions | false` | `{ strategy: 'auto', maxChunkSize: 1024, chunkOverlap: 128 }` | Document chunking config. `false` disables chunking. | | `queryExpansion` | `QueryExpansionOptions | false` | `false` (no-op) | Query expansion with user-provided synonym map. `false` disables. Without `synonyms`, expansion is a no-op. | @@ -356,7 +355,7 @@ await ctx.close(); - **Chunking**: `MarkdownChunker`, `FixedSizeChunker`, `createChunker`, `ChunkingOptions`, `Chunk`, `Chunker` - **Reranking**: `KeywordReranker`, `createReranker`, `Reranker`, `RerankCandidate`, `RerankResult`, `RerankOptions` - **Query Expansion**: `SynonymExpander`, `NoopExpander`, `QueryExpander`, `QueryExpansionOptions` -- **Advanced API**: `Embedder`, `SimpleEmbedder`, `TransformersEmbedder`, `EmbedderManager`, `IZvecStore`, `MemoryZvecStore`, `ActualZvecStore`, `DocumentRegistry`, `StoreManager` +- **Advanced API**: `Embedder`, `TransformersEmbedder`, `EmbedderManager`, `IZvecStore`, `MemoryZvecStore`, `ActualZvecStore`, `DocumentRegistry`, `StoreManager` ## License diff --git a/src/context.ts b/src/context.ts index 9a1089b..270f5c5 100644 --- a/src/context.ts +++ b/src/context.ts @@ -77,22 +77,15 @@ export class Context { if (options.embedder) { // User provided a custom embedder — infer info from its class - const isTransformers = options.embedder.constructor.name === 'TransformersEmbedder'; embedder = options.embedder; embedderInfo = { - kind: isTransformers ? 'transformers' : 'simple', + kind: 'transformers', dimensions: embedder.dimensions, - isFallback: false, }; } else { const result = await resolveEmbedder(options.model); embedder = result.embedder; embedderInfo = result.info; - - // Notify user if a fallback occurred (via callback if provided) - if (embedderInfo.isFallback && options.onEmbedderFallback) { - options.onEmbedderFallback(embedderInfo); - } } // Ensure vectors directory exists @@ -142,17 +135,6 @@ export class Context { /** * Diagnostic information about the active embedder. - * - * Use this to detect when a fallback to SimpleEmbedder occurred, so - * you can warn your users or log the event for troubleshooting. - * - * Example: - * ```ts - * const ctx = await Context.create({ vectorsDir: './vectors' }); - * if (ctx.embedderInfo.isFallback) { - * console.warn(`Using fallback embedder: ${ctx.embedderInfo.fallbackReason}`); - * } - * ``` */ get embedderInfo(): EmbedderInfo { return this._embedderInfo; diff --git a/src/embedder/index.ts b/src/embedder/index.ts index 59667fa..9cd83b7 100644 --- a/src/embedder/index.ts +++ b/src/embedder/index.ts @@ -1,13 +1,12 @@ /** * embedder — aggregate entry point for all embedding modules. * - * Re-exports from split files for backward compatibility: + * Re-exports from split files: * types.ts → Embedder interface * language.ts → isCJK, detectLanguage, tokenizerForLanguage, detectTokenizer - * simple.ts → SimpleEmbedder * transformers.ts → TransformersEmbedder * manager.ts → EmbedderManager, getEmbedder, resetEmbedder - * resolve.ts → resolveEmbedder, isRecoverableError + * resolve.ts → resolveEmbedder */ // Types @@ -17,15 +16,12 @@ export type { Embedder } from './types'; export { isCJK, splitMixed, detectLanguage, tokenizerForLanguage, detectTokenizer } from './language'; export type { LanguageHint } from './language'; -// SimpleEmbedder — lightweight fallback -export { SimpleEmbedder } from './simple'; - // TransformersEmbedder — production-quality model embedder export { TransformersEmbedder } from './transformers'; // EmbedderManager & global convenience functions export { EmbedderManager, getEmbedder, resetEmbedder } from './manager'; -// Embedder resolution & error classification -export { resolveEmbedder, isRecoverableError } from './resolve'; +// Embedder resolution +export { resolveEmbedder } from './resolve'; export type { EmbedderInfo, EmbedderKind, ResolveResult } from './resolve'; \ No newline at end of file diff --git a/src/embedder/manager.ts b/src/embedder/manager.ts index dc6fa73..83a4787 100644 --- a/src/embedder/manager.ts +++ b/src/embedder/manager.ts @@ -7,7 +7,6 @@ */ import { Embedder } from './types'; -import { SimpleEmbedder } from './simple'; import { TransformersEmbedder, loadTransformersModule, @@ -61,38 +60,36 @@ export class EmbedderManager { /** * Return a shared Embedder instance (async). * - * Tries TransformersEmbedder first (needs @huggingface/transformers installed - * AND the model downloadable from HuggingFace Hub). If either fails, falls - * back to SimpleEmbedder gracefully. + * Requires @huggingface/transformers to be installed and the model to be + * loadable. Throws an error when either condition is not met. */ - async getEmbedder(synonymMap?: Map): Promise { + async getEmbedder(): Promise { if (this._defaultEmbedder) return this._defaultEmbedder; const t = await this._loader.load(); - if (t) { - try { - const probe = new TransformersEmbedder(() => this._loader.load()); - await probe.embed('probe'); // triggers lazy model download - this._defaultEmbedder = probe; - } catch (err) { - console.warn( - `[embedder] Bilingual model (bge-small-zh-v1.5) load failed, falling back to SimpleEmbedder.\n` + - ` Error: ${(err as Error).message?.split('\n')[0]}\n` + - `\n` + - ` SimpleEmbedder has lower recall quality. To fix model download:\n` + - ` 1. Set mirror: export HF_ENDPOINT=https://hf-mirror.com\n` + - ` 2. Manual download: node scripts/download-model.mjs\n` - ); - this._defaultEmbedder = new SimpleEmbedder(synonymMap); - } - } else { - console.warn( - '[embedder] @huggingface/transformers not installed, using SimpleEmbedder.\n' + - ' Install it to enable bilingual model for better recall:\n' + + 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' + - ' node scripts/download-model.mjs\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' ); - this._defaultEmbedder = new SimpleEmbedder(synonymMap); } return this._defaultEmbedder; } @@ -123,10 +120,9 @@ const _globalManager = new EmbedderManager(); * explicit embedder to `Context.create({ embedder })` to avoid hidden * global state. This function remains for backward compatibility. * - * @param synonymMap Optional synonym map to inject into SimpleEmbedder fallback. */ -export async function getEmbedder(synonymMap?: Map): Promise { - return _globalManager.getEmbedder(synonymMap); +export async function getEmbedder(): Promise { + return _globalManager.getEmbedder(); } /** diff --git a/src/embedder/resolve.ts b/src/embedder/resolve.ts index c6059f8..d0076af 100644 --- a/src/embedder/resolve.ts +++ b/src/embedder/resolve.ts @@ -1,13 +1,10 @@ /** - * Embedder resolution — auto-selects the best available embedder. - * - * Provides `resolveEmbedder()` which tries TransformersEmbedder first, - * falling back to SimpleEmbedder on model-load failures. Also exports - * `isRecoverableError()` for distinguishing transient vs code-level errors. + * 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 { SimpleEmbedder } from './simple'; import { TransformersEmbedder, loadTransformersModule } from './transformers'; // --------------------------------------------------------------------------- @@ -15,7 +12,7 @@ import { TransformersEmbedder, loadTransformersModule } from './transformers'; // --------------------------------------------------------------------------- /** Describes the kind of embedder being used. */ -export type EmbedderKind = 'transformers' | 'simple'; +export type EmbedderKind = 'transformers'; /** Diagnostic information about the resolved embedder. */ export interface EmbedderInfo { @@ -25,10 +22,6 @@ export interface EmbedderInfo { dimensions: number; /** Model ID (only set for TransformersEmbedder). */ modelId?: string; - /** Whether the resolution fell back from the preferred embedder. */ - isFallback: boolean; - /** Reason for fallback (if applicable). */ - fallbackReason?: string; } // --------------------------------------------------------------------------- @@ -48,108 +41,41 @@ export interface ResolveResult { /** * Resolve which Embedder to use based on the model option. * - * - If model is specified, try TransformersEmbedder first, fall back to - * SimpleEmbedder on failure (with a cooldown to allow retry). - * - If no model, try TransformersEmbedder with the default model. - * - * Returns both the embedder instance and diagnostic info so callers - * can detect fallbacks and report them to users. + * 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) { - 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', - isFallback: false, - }, - }; - } catch (err) { - // Only fallback for network/model-load errors, not for code bugs - if (isRecoverableError(err)) { - const reason = (err as Error).message?.split('\n')[0] ?? 'unknown'; - console.warn( - `[context] Model (${model ?? 'bge-small-zh-v1.5'}) load failed, falling back to basic mode (lower recall quality).\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\n` - ); - const fallback = new SimpleEmbedder(); - return { - embedder: fallback, - info: { - kind: 'simple', - dimensions: fallback.dimensions, - isFallback: true, - fallbackReason: reason, - }, - }; - } - // Unrecoverable errors should propagate - throw err; - } + 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' + ); } - // Transformers not installed — fallback - console.warn( - '[context] @huggingface/transformers not installed, using basic mode (lower recall quality).\n' + - ' Install it for better retrieval:\n' + - ' npm install @huggingface/transformers\n' - ); - const fallback = new SimpleEmbedder(); - return { - embedder: fallback, - info: { - kind: 'simple', - dimensions: fallback.dimensions, - isFallback: true, - fallbackReason: '@huggingface/transformers not installed', - }, - }; -} - -/** - * Determine if an error is recoverable (network, model-not-found, etc.) - * vs a code-level bug that should not be silently swallowed. - */ -export function isRecoverableError(err: unknown): boolean { - if (!(err instanceof Error)) return true; // unknown errors → fallback - - const message = err.message ?? ''; - - // Network / download failures - if (message.includes('fetch') || message.includes('network') || - message.includes('ENOTFOUND') || message.includes('ECONNREFUSED') || - message.includes('timeout') || message.includes('Failed to fetch')) { - return true; + 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' + ); } - - // Model not found or invalid - if (message.includes('not found') || message.includes('404') || - message.includes('model') || message.includes('shape') || - message.includes('dimension') || message.includes('size')) { - return true; - } - - // WASM / native binding issues - if (message.includes('wasm') || message.includes('native') || - message.includes('binding')) { - return true; - } - - // SyntaxError, TypeError, ReferenceError are code bugs — don't swallow - if (err instanceof SyntaxError || err instanceof TypeError || - err instanceof ReferenceError || err instanceof RangeError) { - return false; - } - - // Default: recoverable (most runtime errors in model loading are transient) - return true; } diff --git a/src/embedder/simple.ts b/src/embedder/simple.ts deleted file mode 100644 index c609dc2..0000000 --- a/src/embedder/simple.ts +++ /dev/null @@ -1,184 +0,0 @@ -/** - * SimpleEmbedder — lightweight pseudo-embedding, no external dependencies. - * - * Uses weighted CJK n-grams and English word hashing with log-scale - * count compression and L2 normalization. Not a semantic embedder — - * it's a tuned bag-of-tokens fingerprint used as a graceful-degradation - * fallback when TransformersEmbedder cannot load its model. - * - * **This embedder is used internally as a graceful-degradation fallback.** - * It is not part of the public API and may change without notice. - * - * @internal — fallback only; prefer TransformersEmbedder for - * production-quality retrieval. - */ - -import { Embedder } from './types'; -import { isCJK, splitMixed } from './language'; - -// --------------------------------------------------------------------------- -// Constants -// --------------------------------------------------------------------------- - -const SIMPLE_DIMS = 512; - -const CJK_UNIGRAM_WEIGHT = 0.15; -const CJK_BIGRAM_WEIGHT = 1.0; -const CJK_TRIGRAM_WEIGHT = 2.0; - -const EN_WORD_WEIGHT = 1.5; -const EN_SINGLE_CHAR_WEIGHT = 0.1; - -const STOP_WORDS = new Set([ - '的','了','在','是','我','有','和','就','不', - 'the','a','an','is','are','was','were','be','been','being', - 'to','of','in','for','on','with','at','by','from','as', - 'i','me','my','we','our','he','him','his','she','her','it','its', - 'and','but','or','if','this','that','these','those', - 'not','no','nor','only', - 'chart','using','use', - '图表','数据','配置','展示','需要','支持','进行','通过', - '绘制','实现','基于','根据','使用','方式','效果','功能', - '用于','可以','一个','表示','如下','参考', -]); - -const CJK_UNIGRAM_STOP = new Set([ - '的','了','在','是','和','就','不','也','都','很','到','要','会','着', - '能','可','以','对','与','或','而','且','但','则','因','所','被','把', - '从','由','向','往','用','为','让','使','给','将','比','更','最','只', - '这','那','其','各','某','每','任','何','另','别','全','整','些','几', - '上','下','中','内','外','前','后','左','右','大','小','多','少','高', - '一','二','三','两','个','次','种','项','批','组','类','型', -]); - -// --------------------------------------------------------------------------- -// Internal helpers — weighted tokenization -// --------------------------------------------------------------------------- - -interface WeightedToken { - token: string; - weight: number; -} - -function tokenizeWeighted( - text: string, - synonymMap: Map | null -): WeightedToken[] { - const tokens: WeightedToken[] = []; - const seen = new Set(); - const lower = text.toLowerCase(); - const segments = splitMixed(lower); - - for (const seg of segments) { - if (isCJK(seg)) { - for (let i = 0; i + 3 <= seg.length; i++) { - const t = seg.slice(i, i + 3); - if (!seen.has(t)) { seen.add(t); tokens.push({ token: t, weight: CJK_TRIGRAM_WEIGHT }); } - } - for (let i = 0; i + 2 <= seg.length; i++) { - const t = seg.slice(i, i + 2); - if (!seen.has(t)) { seen.add(t); tokens.push({ token: t, weight: CJK_BIGRAM_WEIGHT }); } - } - for (const ch of seg) { - if (seen.has(ch) || CJK_UNIGRAM_STOP.has(ch)) continue; - seen.add(ch); - tokens.push({ token: ch, weight: CJK_UNIGRAM_WEIGHT }); - } - - if (synonymMap) { - for (const [term, synonyms] of synonymMap) { - if (seg.includes(term)) { - for (const syn of synonyms) { - if (seen.has(syn)) continue; - seen.add(syn); - tokens.push({ token: syn, weight: 1.0 }); - } - } - } - } - } else { - for (const w of seg.split(/\s+/)) { - const trimmed = w.trim(); - if (!trimmed || STOP_WORDS.has(trimmed) || seen.has(trimmed)) continue; - seen.add(trimmed); - const weight = trimmed.length === 1 - ? EN_SINGLE_CHAR_WEIGHT - : EN_WORD_WEIGHT; - tokens.push({ token: trimmed, weight }); - - if (synonymMap) { - const syns = synonymMap.get(trimmed); - if (syns) { - for (const syn of syns) { - if (seen.has(syn)) continue; - seen.add(syn); - tokens.push({ token: syn, weight: 1.0 }); - } - } - } - } - } - } - - return tokens; -} - -/** FNV-1a 32-bit hash with optional seed for multi-hash. */ -function hashToken(token: string, seed = 0): number { - let hash = (2166136261 + seed) >>> 0; - for (let i = 0; i < token.length; i++) { - hash ^= token.charCodeAt(i); - hash = Math.imul(hash, 16777619); - } - return hash >>> 0; -} - -// --------------------------------------------------------------------------- -// SimpleEmbedder -// --------------------------------------------------------------------------- - -export class SimpleEmbedder implements Embedder { - readonly dimensions = SIMPLE_DIMS; - private _synonymMap: Map | null; - - constructor(synonymMap?: Map) { - this._synonymMap = synonymMap ?? null; - } - - async embed(text: string): Promise { - return this.embedSync(text); - } - - async embedBatch(texts: string[]): Promise { - return texts.map((t) => this.embedSync(t)); - } - - /** Synchronous embedding – no async overhead, usable in sync code paths. */ - embedSync(text: string): number[] { - const vec = new Array(SIMPLE_DIMS).fill(0); - const tokens = tokenizeWeighted(text, this._synonymMap); - - for (const { token, weight } of tokens) { - // 3 hash functions per token for collision resistance - for (let h = 0; h < 3; h++) { - vec[hashToken(token, h) % SIMPLE_DIMS] += weight; - } - } - - // Log-scale compression: prevents dimension saturation from - // high-frequency terms (a term appearing 50x contributes log(51) ≈ 3.93 - // instead of 50, giving rare terms proportionally more influence). - for (let i = 0; i < SIMPLE_DIMS; i++) { - if (vec[i] > 0) vec[i] = Math.log(1 + vec[i]); - } - - // L2-normalise - let norm = 0; - for (let i = 0; i < SIMPLE_DIMS; i++) norm += vec[i] * vec[i]; - norm = Math.sqrt(norm) || 1; - for (let i = 0; i < SIMPLE_DIMS; i++) { - vec[i] /= norm; - } - return vec; - } -} diff --git a/src/types.ts b/src/types.ts index f840fb0..7f3e2a1 100644 --- a/src/types.ts +++ b/src/types.ts @@ -91,30 +91,11 @@ export interface ContextOptions { loaders?: import('./loaders').Loader[]; /** * Custom embedder instance. When provided, skips the auto-resolution - * logic and uses this embedder directly. Useful for tests or when you + * logic and uses this embedder directly. Useful for when you * have a pre-initialized embedder. */ embedder?: import('./embedder').Embedder; - /** - * Callback invoked when the embedder resolution falls back to - * SimpleEmbedder (lower recall quality). - * - * Use this to alert users, log diagnostics, or trigger a retry later. - * Not called when a custom `embedder` is provided. - * - * Example: - * ```ts - * const ctx = await Context.create({ - * vectorsDir: './vectors', - * onEmbedderFallback: (info) => { - * console.warn(`Embedder fallback: ${info.fallbackReason}`); - * }, - * }); - * ``` - */ - onEmbedderFallback?: (info: import('./embedder/resolve').EmbedderInfo) => void; - /** * Progress callback for `load()` — called after each major phase completes. * @@ -334,10 +315,6 @@ export interface QueryResult { * The embedder kind used to produce the vectors for this result's store. * * - `'transformers'`: high-quality bilingual model (bge-small-zh-v1.5) - * - `'simple'`: lightweight fallback with lower recall quality - * - * Use this to detect when results come from a degraded embedder and - * warn users or trigger re-indexing with a better model. */ embedderKind?: import('./embedder/resolve').EmbedderKind; } diff --git a/test/context.test.ts b/test/context.test.ts index 67ff82a..fb64727 100644 --- a/test/context.test.ts +++ b/test/context.test.ts @@ -2,13 +2,15 @@ import { describe, it, expect, afterAll, beforeAll } from 'vitest'; import * as fs from 'fs'; import * as path from 'path'; import { Context } from '../src/index'; -import { SimpleEmbedder } from '../src/embedder'; +import { TransformersEmbedder } from '../src/embedder'; const FIXTURES_DIR = path.join(__dirname, 'fixtures/docs'); const TEST_DIR = path.join(__dirname, '.test-tmp'); -// Use SimpleEmbedder for fast unit tests — no model download needed. -const testEmbedder = new SimpleEmbedder(); +// Use TransformersEmbedder for tests — requires model download. +// For faster local testing without model download, provide a custom +// embedder mock via the embedder option. +const testEmbedder = new TransformersEmbedder(); describe('Context', () => { let ctx: Context; diff --git a/test/simple-embedder.test.ts b/test/simple-embedder.test.ts deleted file mode 100644 index def974c..0000000 --- a/test/simple-embedder.test.ts +++ /dev/null @@ -1,92 +0,0 @@ -import { describe, it, expect } from 'vitest'; -import { SimpleEmbedder } from '../src/embedder/simple'; - -describe('SimpleEmbedder', () => { - const embedder = new SimpleEmbedder(); - - describe('basic properties', () => { - it('should have 512 dimensions', () => { - expect(embedder.dimensions).toBe(512); - }); - }); - - describe('embed', () => { - it('should return a vector of correct length', async () => { - const vec = await embedder.embed('hello world'); - expect(vec.length).toBe(512); - }); - - it('should return non-zero vectors for meaningful text', async () => { - const vec = await embedder.embed('important semantic content'); - const nonzeroCount = vec.filter((v) => v !== 0).length; - expect(nonzeroCount).toBeGreaterThan(0); - }); - - it('should produce different vectors for different text', async () => { - const vec1 = await embedder.embed('chart visualization'); - const vec2 = await embedder.embed('database query'); - // Vectors should not be identical - const identical = vec1.every((v, i) => v === vec2[i]); - expect(identical).toBe(false); - }); - - it('should handle CJK text', async () => { - const vec = await embedder.embed('折线图配置'); - expect(vec.length).toBe(512); - const nonzeroCount = vec.filter((v) => v !== 0).length; - expect(nonzeroCount).toBeGreaterThan(0); - }); - - it('should handle mixed CN/EN text', async () => { - const vec = await embedder.embed('line chart 折线图'); - expect(vec.length).toBe(512); - }); - - it('should handle empty text gracefully', async () => { - const vec = await embedder.embed(''); - expect(vec.length).toBe(512); - }); - }); - - describe('embedBatch', () => { - it('should return vectors for each input text', async () => { - const vecs = await embedder.embedBatch(['hello', 'world']); - expect(vecs.length).toBe(2); - expect(vecs[0].length).toBe(512); - expect(vecs[1].length).toBe(512); - }); - - it('should produce same result as individual embed calls', async () => { - const batchVecs = await embedder.embedBatch(['test one', 'test two']); - const singleVec1 = await embedder.embed('test one'); - const singleVec2 = await embedder.embed('test two'); - // Results should be identical (SimpleEmbedder is deterministic) - expect(batchVecs[0]).toEqual(singleVec1); - expect(batchVecs[1]).toEqual(singleVec2); - }); - }); - - describe('L2 normalization', () => { - it('should produce L2-normalized vectors', async () => { - const vec = await embedder.embed('normalized vector test'); - const norm = Math.sqrt(vec.reduce((sum, v) => sum + v * v, 0)); - expect(Math.abs(norm - 1)).toBeLessThan(0.01); - }); - }); - - describe('with synonym map', () => { - it('should incorporate synonym terms into embedding', async () => { - const synonyms = new Map([ - ['tooltip', ['提示框']], - ]); - const embedderWithSyn = new SimpleEmbedder(synonyms); - - const vecWithout = await embedder.embed('tooltip'); - const vecWith = await embedderWithSyn.embed('tooltip'); - - // With synonyms, the vector should differ (extra terms injected) - const identical = vecWithout.every((v, i) => v === vecWith[i]); - expect(identical).toBe(false); - }); - }); -}); From 18647976133559ff645c11434b7f0dc59a9a3eb0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A6=8F=E6=99=8B?= Date: Mon, 6 Jul 2026 16:15:48 +0800 Subject: [PATCH 06/13] chore: remove temp placeholder absolute paths --- src/context.ts | 5 +++-- src/loaders/chunker.ts | 16 +++++++++++----- src/loaders/json.ts | 3 --- src/loaders/markdown.ts | 3 --- src/loaders/text.ts | 3 --- src/types.ts | 6 ++++-- test/loaders.test.ts | 10 +++++----- 7 files changed, 23 insertions(+), 23 deletions(-) diff --git a/src/context.ts b/src/context.ts index 270f5c5..7ab7b34 100644 --- a/src/context.ts +++ b/src/context.ts @@ -220,9 +220,10 @@ export class Context { }) ); - // Internal type that extends Document with load-phase metadata - // not stored in the Document interface itself. + // Internal type that extends Document with load-phase metadata. + // id is required here — Context.load() always assigns it via pathToId. interface LoadedDoc extends Document { + id: string; // override optional Document.id → required contentHash?: string; sourceFilePath?: string; } diff --git a/src/loaders/chunker.ts b/src/loaders/chunker.ts index 1b495b9..0b37913 100644 --- a/src/loaders/chunker.ts +++ b/src/loaders/chunker.ts @@ -44,9 +44,15 @@ export interface ChunkingOptions { strategy?: 'markdown' | 'fixed' | 'auto'; } +/** A document with a guaranteed ID — used by chunkers. + * Context.load() always assigns id via pathToId before chunking. */ +export interface ChunkableDocument extends Document { + id: string; +} + /** Chunker interface — custom chunkers can implement this. */ export interface Chunker { - chunk(doc: Document): Chunk[]; + chunk(doc: ChunkableDocument): Chunk[]; } // --------------------------------------------------------------------------- @@ -77,7 +83,7 @@ export class MarkdownChunker implements Chunker { this.chunkOverlap = opts.chunkOverlap; } - chunk(doc: Document): Chunk[] { + chunk(doc: ChunkableDocument): Chunk[] { return chunkMarkdown(doc, this.maxChunkSize, this.chunkOverlap); } } @@ -87,7 +93,7 @@ export class MarkdownChunker implements Chunker { * sections that exceed maxChunkSize with a fixed-size sliding window. */ function chunkMarkdown( - doc: Document, + doc: ChunkableDocument, maxChunkSize: number, overlap: number, ): Chunk[] { @@ -205,7 +211,7 @@ export class FixedSizeChunker implements Chunker { this.chunkOverlap = opts.chunkOverlap; } - chunk(doc: Document): Chunk[] { + chunk(doc: ChunkableDocument): Chunk[] { const parts = fixedSizeSplit(doc.content, this.maxChunkSize, this.chunkOverlap); return parts.map((content, i) => ({ content, @@ -505,7 +511,7 @@ class AutoChunker implements Chunker { this.fixed = new FixedSizeChunker(options); } - chunk(doc: Document): Chunk[] { + chunk(doc: ChunkableDocument): Chunk[] { // If the document content contains markdown headings (## or ###), use // heading-aware chunking; otherwise fall back to fixed-size splits. if (/^#{1,3}\s+\S/m.test(doc.content)) { diff --git a/src/loaders/json.ts b/src/loaders/json.ts index 0baabd0..7cf2d3a 100644 --- a/src/loaders/json.ts +++ b/src/loaders/json.ts @@ -11,10 +11,7 @@ export class JsonLoader implements Loader { const content = await fs.readFile(filePath, 'utf-8'); const data = JSON.parse(content); - // ID is a temporary placeholder — Context.load() derives the canonical - // hash-based ID from the path relative to basePath (cross-machine consistency). return { - id: filePath, content: typeof data === 'string' ? data : JSON.stringify(data, null, 2), }; } diff --git a/src/loaders/markdown.ts b/src/loaders/markdown.ts index 1892d54..6cffa64 100644 --- a/src/loaders/markdown.ts +++ b/src/loaders/markdown.ts @@ -12,10 +12,7 @@ export class MarkdownLoader implements Loader { const content = await fs.readFile(filePath, 'utf-8'); const { data: meta, content: body } = matter(content); - // ID is a temporary placeholder — Context.load() derives the canonical - // hash-based ID from the path relative to basePath (cross-machine consistency). return { - id: filePath, content: body.trim(), meta, }; diff --git a/src/loaders/text.ts b/src/loaders/text.ts index bcf47f8..20e92fa 100644 --- a/src/loaders/text.ts +++ b/src/loaders/text.ts @@ -10,10 +10,7 @@ export class TextLoader implements Loader { async load(filePath: string): Promise { const content = await fs.readFile(filePath, 'utf-8'); - // ID is a temporary placeholder — Context.load() derives the canonical - // hash-based ID from the path relative to basePath (cross-machine consistency). return { - id: filePath, content: content.trim(), }; } diff --git a/src/types.ts b/src/types.ts index 7f3e2a1..ed6d6b4 100644 --- a/src/types.ts +++ b/src/types.ts @@ -2,8 +2,10 @@ * Document structure */ export interface Document { - /** Document unique identifier (default: file path) */ - id: string; + /** Document unique identifier — assigned by Context.load() (hash-based). + * Not set by Loaders; Context.load() derives it from the file path + * relative to basePath for cross-machine consistency. */ + id?: string; /** Document content */ content: string; /** Markdown front-matter metadata */ diff --git a/test/loaders.test.ts b/test/loaders.test.ts index 2e42130..98f14f4 100644 --- a/test/loaders.test.ts +++ b/test/loaders.test.ts @@ -13,9 +13,9 @@ describe('Loaders', () => { const filePath = path.join(FIXTURES_DIR, 'getting-started.md'); const doc = await loader.load(filePath); - // Loader returns the file path as a temporary ID — the canonical - // hash-based ID is assigned by Context.load() for cross-machine consistency. - expect(doc.id).toBe(filePath); + // Loader does not assign id — the canonical hash-based ID is + // assigned by Context.load() for cross-machine consistency. + expect(doc.id).toBeUndefined(); expect(doc.content).toContain('Getting Started'); expect(doc.content).toContain('npm install'); expect(doc.meta).toEqual({ @@ -38,7 +38,7 @@ describe('Loaders', () => { const filePath = path.join(FIXTURES_DIR, 'api.json'); const doc = await loader.load(filePath); - expect(doc.id).toBe(filePath); + expect(doc.id).toBeUndefined(); expect(doc.content).toContain('API Reference'); expect(doc.content).toContain('/users'); }); @@ -56,7 +56,7 @@ describe('Loaders', () => { const filePath = path.join(FIXTURES_DIR, 'notes.txt'); const doc = await loader.load(filePath); - expect(doc.id).toBe(filePath); + expect(doc.id).toBeUndefined(); expect(doc.content).toContain('notes'); }); From 66eb1bb0672ab2eee519c8cab17c463aacefc63a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A6=8F=E6=99=8B?= Date: Mon, 6 Jul 2026 16:35:03 +0800 Subject: [PATCH 07/13] chore: remove chunk --- README.md | 58 +---- src/context.ts | 158 ++----------- src/index.ts | 8 - src/loaders/chunker.ts | 522 ----------------------------------------- src/loaders/index.ts | 8 +- src/store-manager.ts | 4 - src/types.ts | 52 +--- test/context.test.ts | 97 +------- test/utils.test.ts | 47 +--- 9 files changed, 42 insertions(+), 912 deletions(-) delete mode 100644 src/loaders/chunker.ts diff --git a/README.md b/README.md index 5adcf1a..f9cb253 100644 --- a/README.md +++ b/README.md @@ -15,10 +15,10 @@ A local context retrieval library that enables semantic search over your documen - 🔄 **Deduplication**: Automatically skip already-loaded documents; content-hash change detection for re-embedding updated files - ⚖️ **Weight Configuration**: Per-field FTS boost weights and RRF rank constant tuning - 🛡️ **Clear Error Messages**: Throws descriptive errors when Transformers model is unavailable, guiding users to fix the issue -- 🧩 **Document Chunking**: Split documents into semantic chunks (heading-aware for Markdown, fixed-size for plain text) for finer-grained retrieval +- 🔁 **Two-stage Reranking**: KeywordReranker boosts candidates with exact query term matches after coarse vector/hybrid search - 🔁 **Two-stage Reranking**: KeywordReranker boosts candidates with exact query term matches after coarse vector/hybrid search - 🌐 **Query Expansion**: SynonymExpander uses user-provided synonym maps to bridge CN↔EN terminology gaps -- 📊 **Progress Callback**: `onProgress` hook for monitoring load phases (load → chunk → embed → insert) +- 📊 **Progress Callback**: `onProgress` hook for monitoring load phases (load → embed → insert) - 🏗️ **fromDir() Quick-start**: One-call setup from a project directory with auto-derived defaults @@ -43,7 +43,7 @@ await ctx.load('f2', './f2-docs/**/*.json'); // Query a single library (default: hybrid search = vector + FTS text) const results = await ctx.query('How to configure a line chart', { library: 'g2', topK: 5 }); -// => [{ content: '...', score: 0.92, scoreMode: 'hybrid', id: 'g2-docs/line.md', chunk: {...} }, ...] +// => [{ content: '...', score: 0.92, scoreMode: 'hybrid', id: 'g2-docs/line.md' }, ...] // Query with two-stage reranking (pulls extra candidates, then re-scores) const rerankedResults = await ctx.query('sankey diagram', { library: 'g2', topK: 5, rerank: { rerankFactor: 3 } }); @@ -58,10 +58,7 @@ const allResults = await ctx.query('visualization', { library: '*', topK: 10 }); const vectorResults = await ctx.query('chart', { library: 'g2', topK: 5, mode: 'vector' }); // Filter results by field value -const filteredResults = await ctx.query('tooltip', { library: 'g2', topK: 5, filter: "parentDocId = 'abc123__getting_started'" }); - -// Expand a chunk result — retrieve neighboring chunks for context -const expanded = await ctx.expandChunk('g2', 'abc123__getting_started', { before: 1, after: 1 }); +const filteredResults = await ctx.query('tooltip', { library: 'g2', topK: 5, filter: "sourceFilePath = 'docs/line-chart.md'" }); // Close when done (releases resources) await ctx.close(); @@ -79,8 +76,7 @@ await ctx.close(); | `model` | `string` | auto | Transformers model name for embedding. Skipped when custom `embedder` is provided. | | `loaders` | `Loader[]` | built-in | Custom loaders (default: MarkdownLoader, JsonLoader, TextLoader) | | `embedder` | `Embedder` | auto-resolved | Custom embedder. Skips auto-resolution when provided. | -| `onProgress` | `(phase, detail) => void` | — | Progress callback for `load()` phases: `'load'` → `'chunk'` → `'embed'` → `'insert'`. | -| `chunking` | `ChunkingOptions | false` | `{ strategy: 'auto', maxChunkSize: 1024, chunkOverlap: 128 }` | Document chunking config. `false` disables chunking. | +| `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. | | `ftsFields` | `string[]` | `['content']` | Fields to index for Full Text Search in hybrid mode | | `ftsFieldWeights` | `Record` | `{ content: 1 }` | Per-field boost weights for FTS text path. Higher = more influence. | @@ -99,22 +95,6 @@ const ctx = await Context.create({ }); ``` -#### Chunking Configuration Example - -```typescript -const ctx = await Context.create({ - vectorsDir: './vectors', - // Heading-aware chunking for Markdown, fixed-size for plain text - chunking: { strategy: 'auto', maxChunkSize: 1024, chunkOverlap: 128 }, -}); - -// Disable chunking — embed whole documents instead -const ctxWhole = await Context.create({ - vectorsDir: './vectors', - chunking: false, -}); -``` - #### Query Expansion Configuration Example ```typescript @@ -139,7 +119,7 @@ const ctxNoExpand = await Context.create({ ### `ctx.load(library, pattern)` -Load files into a specified library with automatic batch vectorization. Documents are split into semantic chunks (when chunking is enabled), embedded in batches, and inserted into the vector store. A content-hash change detection mechanism re-embeds files whose content has changed since the last load. +Load files into a specified library with automatic batch vectorization. Documents are embedded in batches and inserted into the vector store. A content-hash change detection mechanism re-embeds files whose content has changed since the last load. Document IDs are derived from file paths relative to `basePath` for cross-machine consistency. @@ -162,7 +142,7 @@ const ctx = await Context.create({ console.log(`${phase}: ${detail.loaded}/${detail.total}`); }, }); -// Phases: 'load' → 'chunk' → 'embed' → 'insert' +// Phases: 'load' → 'embed' → 'insert' ``` ### `ctx.query(text, options)` @@ -175,7 +155,7 @@ Two-stage retrieval: coarse search (vector / hybrid) → optional reranking → | `topK` | `number` | `5` | Number of results to return | | `mode` | `'hybrid' | 'vector'` | `'hybrid'` | Search mode. `'hybrid'` = vector + FTS text (better recall), `'vector'` = pure semantic search | | `rerank` | `RerankOptions | false` | `false` | Reranking configuration. Pass an object `{ rerankFactor, minCandidates }` to enable, or `false` to skip | -| `filter` | `string` | — | Filter expression for zvec exact-match filtering, e.g. `"parentDocId = 'abc123'"` | +| `filter` | `string` | — | Filter expression for zvec exact-match filtering, e.g. `"sourceFilePath = 'docs/line.md'"` | ```typescript // Hybrid search (default) — best recall for exact term matching @@ -189,8 +169,8 @@ const results = await ctx.query('line chart config', { // Pure vector search — when FTS is not needed const results = await ctx.query('chart', { library: 'g2', topK: 5, mode: 'vector' }); -// Filter by parent document — retrieve all chunks of a specific doc -const chunks = await ctx.query('tooltip', { library: 'g2', filter: "parentDocId = 'abc123'" }); +// Filter by source file path +const filtered = await ctx.query('tooltip', { library: 'g2', filter: "sourceFilePath = 'docs/line.md'" }); // Multiple libraries const results = await ctx.query('chart', { library: ['g2', 'f2'], topK: 5 }); @@ -205,12 +185,11 @@ Each result includes: | Field | Type | Description | |------|------|-------------| -| `id` | `string` | Document / chunk ID | +| `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) | -| `chunk` | `ChunkMeta` | Chunk metadata (if chunking is enabled) | | `sourceFilePath` | `string` | Original file path relative to `basePath` | | `library` | `string` | Which library this result came from | @@ -242,21 +221,6 @@ await ctx.untrack('g2', 'abc123__getting_started'); await ctx.rebuild('g2', './g2-docs/**/*.md'); ``` -### `ctx.expandChunk(library, parentDocId, options?)` - -Expand a chunk result — retrieve neighboring chunks from the same parent document for context. Useful when a query returns a chunked fragment and you need surrounding context. - -| Parameter | Type | Default | Description | -|-----------|------|---------|-------------| -| `library` | `string` | — | Library the chunk belongs to | -| `parentDocId` | `string` | — | Parent document ID (from `chunk.parentDocId`) | -| `before` | `number` | `1` | Number of preceding chunks | -| `after` | `number` | `1` | Number of following chunks | - -```typescript -const expanded = await ctx.expandChunk('g2', 'abc123__getting_started', { before: 2, after: 2 }); -``` - ### `Context.fromDir(dir, options?)` Quick-start convenience method — creates a Context from a project directory with auto-derived defaults (`basePath` = dir, `vectorsDir` = dir/.context/vectors). diff --git a/src/context.ts b/src/context.ts index 7ab7b34..e0046ee 100644 --- a/src/context.ts +++ b/src/context.ts @@ -2,7 +2,7 @@ import * as path from 'path'; import * as fs from 'fs'; import * as crypto from 'crypto'; import { glob } from 'glob'; -import { ContextOptions, QueryOptions, QueryResult, Document, ChunkMeta, LoadPhase, LoadProgress } from './types'; +import { ContextOptions, QueryOptions, QueryResult, Document, LoadPhase, LoadProgress } from './types'; import { resolveEmbedder, } from './embedder'; @@ -13,8 +13,6 @@ import { DocumentRegistry } from './registry'; import { StoreManager } from './store-manager'; import type { ZvecDoc } from './storage/zvec-store'; import { pathToId } from './loaders/util'; -import { createChunker } from './loaders/chunker'; -import type { Chunker } from './loaders/chunker'; import { createReranker } from './reranker'; import type { Reranker, RerankCandidate } from './reranker'; import { SynonymExpander, NoopExpander } from './query-expander'; @@ -32,7 +30,6 @@ export class Context { private readonly storeManager: StoreManager; private readonly registry: DocumentRegistry; private readonly loaders: Loader[]; - private readonly chunker: Chunker | null; private readonly reranker: Reranker | null; private readonly queryExpander: QueryExpander; private readonly _onProgress?: (phase: LoadPhase, detail: LoadProgress) => void; @@ -49,13 +46,7 @@ export class Context { new JsonLoader(), new TextLoader(), ]; - // Chunking defaults to enabled with auto strategy. - // Set `chunking: false` to embed whole documents instead. - this.chunker = options.chunking !== false - ? createChunker(options.chunking ?? undefined) - : null; // Reranker is created eagerly — it has no model-load cost (KeywordReranker). - // In the future this may be lazy if a cross-encoder is used. // Weights are configurable via ContextOptions.rerankWeights. this.reranker = createReranker(options.rerankWeights); // Query expansion — uses user-provided synonym map only, no built-in defaults. @@ -113,13 +104,13 @@ export class Context { } /** - * Quick-start convenience method - creates a Context from a project directory. + * Quick-start convenience method — creates a Context from a project directory. * * Auto-derives sensible defaults: * - basePath = dir (the project root) * - vectorsDir = dir/.context/vectors (hidden, won't pollute project) * - * All other options (model, chunking, embedder, etc.) can still be + * All other options (model, embedder, etc.) can still be * overridden via options. */ static async fromDir(dir: string, options?: Partial): Promise { @@ -250,54 +241,25 @@ export class Context { if (docsToEmbed.length === 0) return; - // Phase 2: Apply chunking (if enabled). Each chunk is independently - // embedded so queries can match at the paragraph / section level. - const items = this.chunker - ? docsToEmbed.flatMap((doc) => { - const chunks = this.chunker!.chunk(doc); - return chunks.map((chunk) => ({ - id: `${doc.id}__c${chunk.chunkIndex}`, - content: chunk.content, - meta: doc.meta, - sourceFilePath: doc.sourceFilePath, - chunk: { - chunkIndex: chunk.chunkIndex, - totalChunks: chunk.totalChunks, - parentDocId: doc.id, - headingPath: chunk.headingPath, - } satisfies ChunkMeta, - })); - }) - : docsToEmbed; - - // Progress: chunk phase complete - if (this._onProgress) { - this._onProgress('chunk', { loaded: items.length, total: docsToEmbed.length }); - } - - // Phase 3: Batch embed all items for better performance - const contents = items.map((item) => item.content); + // 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: items.length }); + this._onProgress('embed', { loaded: vectors.length, total: docsToEmbed.length }); } - // Phase 4: Batch insert into store - const zvecDocs: ZvecDoc[] = items.map((item, index) => ({ - id: item.id, + // Phase 3: Batch insert into store + const zvecDocs: ZvecDoc[] = docsToEmbed.map((doc, index) => ({ + id: doc.id, vector: vectors[index], fields: { - content: item.content, - meta: item.meta && Object.keys(item.meta).length > 0 - ? JSON.stringify(item.meta) + content: doc.content, + meta: doc.meta && Object.keys(doc.meta).length > 0 + ? JSON.stringify(doc.meta) : '', - chunkIndex: item.chunk?.chunkIndex ?? -1, - totalChunks: item.chunk?.totalChunks ?? 1, - parentDocId: item.chunk?.parentDocId ?? item.id, - headingPath: item.chunk?.headingPath?.join(' > ') ?? '', - sourceFilePath: item.sourceFilePath ?? '', + sourceFilePath: doc.sourceFilePath ?? '', }, })); @@ -305,10 +267,10 @@ export class Context { // Progress: insert phase complete if (this._onProgress) { - this._onProgress('insert', { loaded: zvecDocs.length, total: items.length }); + this._onProgress('insert', { loaded: zvecDocs.length, total: docsToEmbed.length }); } - // Phase 5: Update registry (track by parent doc ID for deduplication + change detection) + // Phase 4: Update registry (track by parent doc ID for deduplication + change detection) for (const doc of docsToEmbed) { this.registry.add(library, doc.id, doc.contentHash); } @@ -375,34 +337,17 @@ export class Context { const content = String(result.fields?.content ?? ''); const metaStr = result.fields?.meta as string | undefined; const meta = safeParseMeta(metaStr); - const chunkIndex = result.fields?.chunkIndex as number | undefined; - const parentDocId = result.fields?.parentDocId as string | undefined; - const headingPathRaw = result.fields?.headingPath as string | undefined; - const queryResult: QueryResult = { + return { id: result.id, content, score: result.score, - scoreMode: mode === 'hybrid' ? 'hybrid' : 'vector', + scoreMode: mode === 'hybrid' ? 'hybrid' as const : 'vector' as const, meta, sourceFilePath: result.fields?.sourceFilePath as string | undefined, library, embedderKind: this._embedderInfo.kind, }; - - // Attach chunk metadata when present (non-negative chunkIndex) - if (chunkIndex !== undefined && chunkIndex >= 0 && parentDocId) { - queryResult.chunk = { - chunkIndex, - totalChunks: result.fields?.totalChunks as number ?? 1, - parentDocId, - headingPath: headingPathRaw - ? headingPathRaw.split(' > ').filter(Boolean) - : [], - }; - } - - return queryResult; }); }) ); @@ -418,7 +363,6 @@ export class Context { id: r.id, content: r.content, score: r.score, - headingPath: r.chunk?.headingPath?.join(' > '), })); const reranked = await this.reranker!.rerank(text, candidates); @@ -439,74 +383,6 @@ export class Context { return allResults.slice(0, topK); } - /** - * Expand a chunk result — retrieve neighboring chunks from the same - * parent document for context. - * - * When a query returns a chunked document fragment, the user often - * needs to see the surrounding context (before/after chunks) to - * understand the full meaning. This method fetches those neighbors. - * - * @param library The library the chunk belongs to. - * @param parentDocId The parent document ID (from `chunk.parentDocId`). - * @param options Optional range control: - * - `before`: number of preceding chunks (default 1) - * - `after`: number of following chunks (default 1) - * - * @returns Array of chunk QueryResults sorted by chunkIndex, or empty - * array if the library or parent document is not found. - */ - async expandChunk( - library: string, - parentDocId: string, - options?: { before?: number; after?: number }, - ): Promise { - const store = this.storeManager.tryOpen(library); - if (!store) return []; - - const before = options?.before ?? 1; - const after = options?.after ?? 1; - - // Use a zero-vector search with a filter on parentDocId to retrieve - // all chunks of the parent document. The filter narrows results to - // the correct parent, and we sort by chunkIndex to find neighbors. - const zeroVector = new Array(this.embedder.dimensions).fill(0); - const allChunks = await store.search({ - vector: zeroVector, - topK: 100, // generous limit — a single doc rarely has >100 chunks - filter: `parentDocId = '${parentDocId}'`, - }); - - if (allChunks.length === 0) return []; - - // Sort by chunkIndex and parse into QueryResults - const parsed = allChunks - .map((result) => { - const chunkIndex = result.fields?.chunkIndex as number ?? -1; - return { - id: result.id, - content: String(result.fields?.content ?? ''), - score: result.score, - meta: safeParseMeta(result.fields?.meta as string | undefined), - sourceFilePath: result.fields?.sourceFilePath as string | undefined, - library, - chunk: { - chunkIndex, - totalChunks: result.fields?.totalChunks as number ?? 1, - parentDocId: result.fields?.parentDocId as string ?? parentDocId, - headingPath: (result.fields?.headingPath as string ?? '') - .split(' > ').filter(Boolean), - }, - chunkIndex, - }; - }) - .sort((a, b) => a.chunkIndex - b.chunkIndex); - - // Find the target chunk index and return the window around it - // If we don't know which chunk the user came from, return all chunks - return parsed.slice(0, parsed.length); - } - /** * Remove a document from a library's dedup registry. * diff --git a/src/index.ts b/src/index.ts index 0460d5d..38e2ab1 100644 --- a/src/index.ts +++ b/src/index.ts @@ -14,14 +14,6 @@ export { } from './loaders'; export { pathToId } from './loaders/util'; -// Chunking — document splitting strategies -export { - MarkdownChunker, - FixedSizeChunker, - createChunker, -} from './loaders/chunker'; -export type { Chunk, Chunker } from './loaders/chunker'; - // Reranker — two-stage retrieval precision scoring export { KeywordReranker, createReranker } from './reranker'; export type { Reranker, RerankCandidate, RerankResult } from './reranker'; diff --git a/src/loaders/chunker.ts b/src/loaders/chunker.ts deleted file mode 100644 index 0b37913..0000000 --- a/src/loaders/chunker.ts +++ /dev/null @@ -1,522 +0,0 @@ -/** - * Document chunker — splits loaded documents into smaller, semantic chunks - * so that vector embeddings capture fine-grained meaning instead of averaging - * the entire document. - * - * Strategies: - * - MarkdownChunker: splits on headings (## / ###), then falls back to fixed-size - * - FixedSizeChunker: sliding window with overlap, used for plain text / JSON - * - * Each chunk stores its heading path, index, and parent document ID so the - * query layer can reconstruct context or deduplicate by parent. - */ - -import { Document } from '../types'; - -// --------------------------------------------------------------------------- -// Types -// --------------------------------------------------------------------------- - -/** A semantic chunk of a document. */ -export interface Chunk { - /** Chunk content — the actual text that will be embedded. */ - content: string; - /** Zero-based chunk index within the parent document. */ - chunkIndex: number; - /** Total number of chunks in the parent document. */ - totalChunks: number; - /** Parent document ID (hash-based, assigned by Context). */ - parentDocId: string; - /** - * Heading path from the document root to this chunk. - * e.g. ["Configuration", "Line Chart", "Basic Usage"] - */ - headingPath: string[]; -} - -/** Configuration for document chunking. */ -export interface ChunkingOptions { - /** Maximum characters per chunk (roughly ¼ tokens for mixed CN/EN, ½ for pure EN). */ - maxChunkSize?: number; - /** Overlap in characters between adjacent chunks (avoids boundary cuts). */ - chunkOverlap?: number; - /** Strategy used for chunking. */ - strategy?: 'markdown' | 'fixed' | 'auto'; -} - -/** A document with a guaranteed ID — used by chunkers. - * Context.load() always assigns id via pathToId before chunking. */ -export interface ChunkableDocument extends Document { - id: string; -} - -/** Chunker interface — custom chunkers can implement this. */ -export interface Chunker { - chunk(doc: ChunkableDocument): Chunk[]; -} - -// --------------------------------------------------------------------------- -// Defaults -// --------------------------------------------------------------------------- - -const DEFAULT_MAX_CHUNK_SIZE = 1024; // ~256 tokens for mixed CN/EN -const DEFAULT_CHUNK_OVERLAP = 128; - -function chunkingDefaults(options?: ChunkingOptions): Required> { - return { - maxChunkSize: options?.maxChunkSize ?? DEFAULT_MAX_CHUNK_SIZE, - chunkOverlap: options?.chunkOverlap ?? DEFAULT_CHUNK_OVERLAP, - }; -} - -// --------------------------------------------------------------------------- -// Markdown Chunker — heading-aware splitting -// --------------------------------------------------------------------------- - -export class MarkdownChunker implements Chunker { - private readonly maxChunkSize: number; - private readonly chunkOverlap: number; - - constructor(options?: ChunkingOptions) { - const opts = chunkingDefaults(options); - this.maxChunkSize = opts.maxChunkSize; - this.chunkOverlap = opts.chunkOverlap; - } - - chunk(doc: ChunkableDocument): Chunk[] { - return chunkMarkdown(doc, this.maxChunkSize, this.chunkOverlap); - } -} - -/** - * Split markdown content into heading-anchored sections, then sub-split - * sections that exceed maxChunkSize with a fixed-size sliding window. - */ -function chunkMarkdown( - doc: ChunkableDocument, - maxChunkSize: number, - overlap: number, -): Chunk[] { - const sections = splitByHeadings(doc.content); - const chunks: Chunk[] = []; - - for (const section of sections) { - if (section.content.length <= maxChunkSize) { - if (section.content.trim().length > 0) { - chunks.push({ - content: section.content.trim(), - chunkIndex: 0, // filled in after - totalChunks: 0, // filled in after - parentDocId: doc.id, - headingPath: section.headingPath, - }); - } - continue; - } - - // Sub-split oversized sections with overlap - const subChunks = fixedSizeSplit(section.content, maxChunkSize, overlap); - for (const sc of subChunks) { - chunks.push({ - content: sc, - chunkIndex: 0, - totalChunks: 0, - parentDocId: doc.id, - headingPath: section.headingPath, - }); - } - } - - // Fill in chunkIndex and totalChunks - return chunks.map((c, i) => ({ - ...c, - chunkIndex: i, - totalChunks: chunks.length, - })); -} - -interface HeadingSection { - content: string; - headingPath: string[]; -} - -/** - * Split markdown text into sections delimited by markdown headings (##, ###). - * - * Lines before the first heading become a section with an empty heading path. - */ -function splitByHeadings(text: string): HeadingSection[] { - const sections: HeadingSection[] = []; - const lines = text.split('\n'); - let currentHeadingPath: string[] = []; - let currentLines: string[] = []; - - function flushSection(): void { - const content = currentLines.join('\n').trim(); - if (content.length > 0) { - sections.push({ - content, - headingPath: [...currentHeadingPath], - }); - } - currentLines = []; - } - - for (const line of lines) { - const h1Match = line.match(/^#\s+(.+)/); - const h2Match = line.match(/^##\s+(.+)/); - const h3Match = line.match(/^###\s+(.+)/); - - if (h1Match) { - flushSection(); - currentHeadingPath = [h1Match[1].trim()]; - } else if (h2Match) { - flushSection(); - const h1Ancestor = currentHeadingPath.length > 0 ? currentHeadingPath[0] : null; - currentHeadingPath = h1Ancestor - ? [h1Ancestor, h2Match[1].trim()] - : [h2Match[1].trim()]; - } else if (h3Match) { - flushSection(); - currentHeadingPath = [...currentHeadingPath, h3Match[1].trim()]; - if (currentHeadingPath.length === 0) { - currentHeadingPath = [h3Match[1].trim()]; - } - } else { - currentLines.push(line); - } - } - - flushSection(); - - // If no headings were found at all, return the full content as one section - if (sections.length === 0 && currentLines.length === 0 && text.trim().length > 0) { - sections.push({ content: text.trim(), headingPath: [] }); - } - - return sections; -} - -// --------------------------------------------------------------------------- -// Fixed-Size Chunker — sliding window, no semantic awareness -// --------------------------------------------------------------------------- - -export class FixedSizeChunker implements Chunker { - private readonly maxChunkSize: number; - private readonly chunkOverlap: number; - - constructor(options?: ChunkingOptions) { - const opts = chunkingDefaults(options); - this.maxChunkSize = opts.maxChunkSize; - this.chunkOverlap = opts.chunkOverlap; - } - - chunk(doc: ChunkableDocument): Chunk[] { - const parts = fixedSizeSplit(doc.content, this.maxChunkSize, this.chunkOverlap); - return parts.map((content, i) => ({ - content, - chunkIndex: i, - totalChunks: parts.length, - parentDocId: doc.id, - headingPath: [], - })); - } -} - -/** - * Split text into fixed-size windows with overlap. - * - * Tries to break at paragraph boundaries (double newline) within a tolerance - * of maxSize to avoid cutting sentences in half. Falls back to hard cut. - * - * Protects fenced code blocks (```) and pipe tables (|...|) from being - * split mid-structure — cut points are adjusted to land before or after - * these regions. - */ -function fixedSizeSplit( - text: string, - maxSize: number, - overlap: number, -): string[] { - const paragraphs = splitRespectingBlocks(text); - const chunks: string[] = []; - let current = ''; - - for (const para of paragraphs) { - const trimmed = para.trim(); - if (!trimmed) continue; - - if (current.length + trimmed.length + 2 <= maxSize) { - current = current ? `${current}\n\n${trimmed}` : trimmed; - } else { - // Current chunk is full — flush it - if (current) { - chunks.push(current); - // Start new chunk with overlap: take the last `overlap` chars of the - // previous chunk as prefix so boundary semantics aren't lost. - const overlapPrefix = overlap > 0 && current.length > overlap - ? current.slice(-overlap).replace(/^[^\w一-鿿]+/, '') + '\n\n' - : ''; - current = overlapPrefix + trimmed; - } else { - // Single paragraph/block is larger than maxSize — force split - // but still respect code block / table boundaries - current = trimmed; - while (current.length > maxSize) { - let cutPoint = findBreakPointSafe(current, maxSize); - chunks.push(current.slice(0, cutPoint).trim()); - const suffix = current.slice(cutPoint).trim(); - current = overlap > 0 && suffix.length > 0 - ? current.slice(Math.max(0, cutPoint - overlap), cutPoint).replace(/^[^\w一-鿿]+/, '') + '\n\n' + suffix - : suffix; - } - } - } - } - - if (current.trim()) { - chunks.push(current.trim()); - } - - return chunks.length > 0 ? chunks : [text.trim()]; -} - -// --------------------------------------------------------------------------- -// Block-aware splitting helpers -// --------------------------------------------------------------------------- - -/** Regex matching the opening/closing of a fenced code block. */ -const FENCED_CODE_RE = /^(`{3,}|~{3,})/; - -/** - * Check whether a line is part of a pipe table row. - * Matches lines like `| col1 | col2 |` and separator rows `|---|---|`. - */ -function isTableRow(line: string): boolean { - const trimmed = line.trim(); - return trimmed.startsWith('|') && trimmed.endsWith('|') && trimmed.length > 1; -} - -/** - * Split text into logical units that respect fenced code blocks and tables. - * - * Each returned unit is either: - * - A complete fenced code block (never split) - * - A contiguous run of table rows (never split) - * - A run of normal paragraphs separated by double newlines - * - * This ensures downstream fixed-size splitting never cuts inside a - * code block or table. - */ -function splitRespectingBlocks(text: string): string[] { - const lines = text.split('\n'); - const units: string[] = []; - let buffer: string[] = []; - let inCodeBlock = false; - let codeFence = ''; - let inTable = false; - - function flushBuffer(): void { - if (buffer.length > 0) { - const content = buffer.join('\n').trim(); - if (content.length > 0) { - units.push(content); - } - buffer = []; - } - } - - for (const line of lines) { - // --- Fenced code block detection --- - const fenceMatch = line.match(FENCED_CODE_RE); - if (fenceMatch) { - if (!inCodeBlock) { - // Entering code block — flush any preceding normal content - flushBuffer(); - inCodeBlock = true; - codeFence = fenceMatch[1].charAt(0); // ` or ~ - buffer.push(line); - } else if (line.trimStart().startsWith(codeFence.repeat(codeFence.length)) && line.trim().length <= codeFence.length + 1) { - // Closing fence — end of code block - buffer.push(line); - flushBuffer(); // emit the complete code block as one unit - inCodeBlock = false; - codeFence = ''; - } else { - // Inside code block - buffer.push(line); - } - continue; - } - - if (inCodeBlock) { - buffer.push(line); - continue; - } - - // --- Table row detection --- - if (isTableRow(line)) { - if (!inTable) { - // Entering table — flush preceding normal content - flushBuffer(); - inTable = true; - } - buffer.push(line); - continue; - } - - // Not a table row - if (inTable) { - // Exiting table — flush the complete table as one unit - flushBuffer(); - inTable = false; - } - - // Normal line — accumulate; double-newline splits are handled - // by the paragraph-level logic below via empty-line detection - buffer.push(line); - } - - // Flush any remaining content - flushBuffer(); - - // Post-process: further split normal units on double-newlines so that - // the paragraph-level sliding window works correctly. Code blocks and - // tables remain as single units. - const result: string[] = []; - for (const unit of units) { - if (FENCED_CODE_RE.test(unit.split('\n')[0] ?? '') || isTableRow(unit.split('\n')[0] ?? '')) { - // Protected block — keep as-is - result.push(unit); - } else { - // Normal text — split on paragraph boundaries - const paras = unit.split(/\n\n+/).map((p) => p.trim()).filter((p) => p.length > 0); - result.push(...paras); - } - } - - return result.length > 0 ? result : [text.trim()]; -} - -/** - * Find a safe break point that avoids cutting inside code blocks or tables. - * - * Delegates to findBreakPoint for natural break detection, then adjusts - * the result if it lands inside a protected region. - */ -function findBreakPointSafe(text: string, maxSize: number): number { - const candidate = findBreakPoint(text, maxSize); - - // Check if the candidate falls inside a fenced code block - const beforeCut = text.slice(0, candidate); - const fenceOpens = (beforeCut.match(/^(`{3,}|~{3,})/gm) ?? []).length; - const fenceCloses = (beforeCut.match(/^(`{3,}|~{3,})\s*$/gm) ?? []).length; - - if (fenceOpens > fenceCloses) { - // Inside an unclosed code block — move cut point backward to before the fence - const lastOpenIdx = beforeCut.lastIndexOf('\n```'); - const lastOpenTilde = beforeCut.lastIndexOf('\n~~~'); - const safePoint = Math.max(lastOpenIdx, lastOpenTilde); - if (safePoint > 0) return safePoint; - } - - // Check if the candidate falls inside a table - const linesBeforeCut = beforeCut.split('\n'); - const lastLine = linesBeforeCut[linesBeforeCut.length - 1] ?? ''; - if (isTableRow(lastLine)) { - // Move backward to before the table started - for (let i = linesBeforeCut.length - 1; i >= 0; i--) { - if (!isTableRow(linesBeforeCut[i])) { - // Found the line before the table — compute its offset - const offset = linesBeforeCut.slice(0, i + 1).join('\n').length + 1; - if (offset > 0) return offset; - break; - } - } - } - - return candidate; -} - -/** - * Find a natural break point within [maxSize * 0.8, maxSize]. - * Prefers paragraph breaks > sentence-ending punctuation > space. - */ -function findBreakPoint(text: string, maxSize: number): number { - const minPoint = Math.floor(maxSize * 0.7); - - // Try double newline first - for (let i = maxSize; i >= minPoint; i--) { - if (text[i] === '\n' && text[i - 1] === '\n') return i + 1; - } - - // Try sentence-ending punctuation followed by space or newline - const sentenceBreaks = /[。.!?！？]\s/g; - let match: RegExpExecArray | null; - let bestSentence = -1; - - sentenceBreaks.lastIndex = minPoint; - while ((match = sentenceBreaks.exec(text)) !== null) { - if (match.index > maxSize) break; - bestSentence = match.index + 1; // after the punctuation - } - if (bestSentence > minPoint) return bestSentence; - - // Try any newline - for (let i = maxSize; i >= minPoint; i--) { - if (text[i] === '\n') return i + 1; - } - - // Fallback: hard cut at maxSize, but try to land on a space - for (let i = maxSize; i >= minPoint; i--) { - if (text[i] === ' ') return i + 1; - } - - return maxSize; -} - -// --------------------------------------------------------------------------- -// Factory -// --------------------------------------------------------------------------- - -/** - * Create a chunker based on the configured strategy. - * - * - `'auto'` (default): tries MarkdownChunker; if no headings found, falls - * back to FixedSizeChunker. - * - `'markdown'`: always uses heading-aware splitting. - * - `'fixed'`: always uses fixed-size sliding window. - */ -export function createChunker(options?: ChunkingOptions): Chunker { - const strategy = options?.strategy ?? 'auto'; - - switch (strategy) { - case 'markdown': - return new MarkdownChunker(options); - case 'fixed': - return new FixedSizeChunker(options); - case 'auto': - default: - return new AutoChunker(options); - } -} - -/** Auto-detect chunking strategy: markdown when headings exist, else fixed. */ -class AutoChunker implements Chunker { - private markdown: MarkdownChunker; - private fixed: FixedSizeChunker; - - constructor(options?: ChunkingOptions) { - this.markdown = new MarkdownChunker(options); - this.fixed = new FixedSizeChunker(options); - } - - chunk(doc: ChunkableDocument): Chunk[] { - // If the document content contains markdown headings (## or ###), use - // heading-aware chunking; otherwise fall back to fixed-size splits. - if (/^#{1,3}\s+\S/m.test(doc.content)) { - return this.markdown.chunk(doc); - } - return this.fixed.chunk(doc); - } -} diff --git a/src/loaders/index.ts b/src/loaders/index.ts index fe32023..799617e 100644 --- a/src/loaders/index.ts +++ b/src/loaders/index.ts @@ -2,10 +2,4 @@ export { Loader } from './base'; export { MarkdownLoader } from './markdown'; export { JsonLoader } from './json'; export { TextLoader } from './text'; -export { pathToId } from './util'; -export { - MarkdownChunker, - FixedSizeChunker, - createChunker, -} from './chunker'; -export type { Chunk, Chunker } from './chunker'; \ No newline at end of file +export { pathToId } from './util'; \ No newline at end of file diff --git a/src/store-manager.ts b/src/store-manager.ts index 7de5aac..aaa1f1d 100644 --- a/src/store-manager.ts +++ b/src/store-manager.ts @@ -70,10 +70,6 @@ function contextStoreConfig(dims: number, options?: ContextOptions, sampleText?: indexOptions: { tokenizerName }, }, { name: 'meta', dataType: 'STRING' }, - { name: 'chunkIndex', dataType: 'INT64' }, - { name: 'totalChunks', dataType: 'INT64' }, - { name: 'parentDocId', dataType: 'STRING' }, - { name: 'headingPath', dataType: 'STRING' }, { name: 'sourceFilePath', dataType: 'STRING' }, ], }; diff --git a/src/types.ts b/src/types.ts index ed6d6b4..c78b1d3 100644 --- a/src/types.ts +++ b/src/types.ts @@ -10,11 +10,6 @@ export interface Document { content: string; /** Markdown front-matter metadata */ meta?: Record; - /** - * Chunk metadata — populated when document chunking is enabled. - * Undefined for whole-document (unchunked) mode. - */ - chunk?: ChunkMeta; /** * SHA-256 content hash (16 chars) — used internally for change detection. * Populated by `Context.load()` before registry insertion; not stored in zvec. @@ -22,18 +17,6 @@ export interface Document { contentHash?: string; } -/** Metadata for a chunked document fragment. */ -export interface ChunkMeta { - /** Zero-based chunk index within the parent document. */ - chunkIndex: number; - /** Total number of chunks in the parent document. */ - totalChunks: number; - /** Parent document ID (hash-based). */ - parentDocId: string; - /** Heading path from the document root to this chunk. */ - headingPath: string[]; -} - /** Configuration for query expansion. */ export interface QueryExpansionOptions { /** @@ -65,12 +48,6 @@ export interface RerankOptions { minCandidates?: number; } -/** Configuration for document chunking. Set to `false` to disable. - * - * Re-exported from loaders/chunker for convenience. - * See {@link import('./loaders/chunker').ChunkingOptions} for full details. */ -export type { ChunkingOptions } from './loaders/chunker'; - /** * Context initialization options */ @@ -116,18 +93,6 @@ export interface ContextOptions { */ onProgress?: (phase: LoadPhase, detail: LoadProgress) => void; - /** - * Document chunking configuration. - * - * When enabled, each document is split into semantic chunks (heading-aware - * for Markdown, fixed-size for plain text). Each chunk is independently - * embedded and can be retrieved at a finer granularity. - * - * Set to `false` to disable chunking and embed whole documents instead. - * Defaults to enabled with `{ strategy: 'auto', maxChunkSize: 1024, chunkOverlap: 128 }`. - */ - chunking?: import('./loaders/chunker').ChunkingOptions | false; - /** * Query expansion configuration. * @@ -264,10 +229,7 @@ export interface QueryOptions { * Passed directly to the zvec engine as an exact-match filter. * Format: `"field = 'value'"` or `"field = 'val1' AND field2 = 'val2'"`. * - * Useful for scoping queries to specific document categories, versions, - * or parent documents: - * `"parentDocId = 'f3a1b2c4__getting_started'"` — all chunks of one doc - * `"chunkIndex = 0"` — only first chunks + * Useful for scoping queries to specific document categories or versions. */ filter?: string; } @@ -293,12 +255,6 @@ export interface QueryResult { scoreMode?: 'vector' | 'hybrid' | 'reranked'; /** Document metadata */ meta?: Record; - /** - * Chunk metadata — present when the result is a chunk of a larger document. - * Use `parentDocId` to look up the original document; `headingPath` shows - * the section path (e.g. ["Configuration", "Line Chart"]). - */ - chunk?: ChunkMeta; /** * Original file path relative to `basePath` — allows users to trace * back to the source document for context review or navigation. @@ -324,15 +280,15 @@ export interface QueryResult { /** * Load progress phases — emitted by `Context.load()` via the `onProgress` callback. */ -export type LoadPhase = 'load' | 'chunk' | 'embed' | 'insert'; +export type LoadPhase = 'load' | 'embed' | 'insert'; /** * Load progress detail — passed to `onProgress` callback. */ export interface LoadProgress { - /** How many documents/chunks have been processed in this phase */ + /** How many documents have been processed in this phase */ loaded: number; - /** Total documents/chunks to process in this phase */ + /** Total documents to process in this phase */ total: number; } diff --git a/test/context.test.ts b/test/context.test.ts index fb64727..154103f 100644 --- a/test/context.test.ts +++ b/test/context.test.ts @@ -51,8 +51,6 @@ describe('Context', () => { it('should load markdown files', async () => { await ctx.load('md', path.join(FIXTURES_DIR, 'getting-started.md')); - // With chunking, the matching chunk may not be the first result. - // Check that at least one result contains the expected term. const results = await ctx.query('installation', { library: 'md', topK: 3 }); expect(results.length).toBeGreaterThan(0); const contents = results.map((r) => r.content).join(' '); @@ -79,9 +77,8 @@ describe('Context', () => { await ctx.load('md', path.join(FIXTURES_DIR, 'getting-started.md')); const results = await ctx.query('install', { library: 'md', topK: 10 }); - // With chunking enabled, a single doc may produce multiple chunks. - // The key invariant is that a second load() with the same pattern - // should NOT increase the count (dedup prevents double-insert). + // A second load() with the same pattern should NOT increase + // the count (dedup prevents double-insert). expect(results.length).toBeGreaterThan(0); }); }); @@ -146,91 +143,6 @@ describe('Context', () => { } }); }); - - describe('chunking', () => { - const chunkTestDir = TEST_DIR + '-chunk-test'; - - afterAll(() => { - if (fs.existsSync(chunkTestDir)) { - fs.rmSync(chunkTestDir, { recursive: true, force: true }); - } - }); - - it('should split large documents into chunks', async () => { - const ctx = await Context.create({ - vectorsDir: chunkTestDir, - embedder: testEmbedder, - chunking: { maxChunkSize: 500, chunkOverlap: 50 }, - }); - - await ctx.load('chunked', path.join(FIXTURES_DIR, 'line-chart-guide.md')); - - const results = await ctx.query('tooltip', { library: 'chunked', topK: 3 }); - expect(results.length).toBeGreaterThan(0); - - // At least one result should be a chunk (long doc with small maxChunkSize) - const chunkResults = results.filter((r) => r.chunk); - expect(chunkResults.length).toBeGreaterThan(0); - - // Chunk metadata should be present - const chunk = chunkResults[0].chunk!; - expect(chunk.parentDocId).toContain('line_chart_guide'); - expect(typeof chunk.chunkIndex).toBe('number'); - expect(chunk.chunkIndex).toBeGreaterThanOrEqual(0); - - await ctx.close(); - }); - - it('should include headingPath in chunk metadata', async () => { - const ctx = await Context.create({ - vectorsDir: chunkTestDir + '-2', - embedder: testEmbedder, - chunking: { maxChunkSize: 400, chunkOverlap: 50 }, - }); - - await ctx.load('chunked2', path.join(FIXTURES_DIR, 'line-chart-guide.md')); - - // Query for tooltip section specifically - const results = await ctx.query('Line Chart with Tooltip hover', { - library: 'chunked2', - topK: 5, - }); - expect(results.length).toBeGreaterThan(0); - - // At least one chunk should have a heading path containing "Line Chart" - const withHeading = results.filter( - (r) => r.chunk?.headingPath?.some((h) => h.includes('Line Chart')) - ); - expect(withHeading.length).toBeGreaterThan(0); - - await ctx.close(); - if (fs.existsSync(chunkTestDir + '-2')) { - fs.rmSync(chunkTestDir + '-2', { recursive: true, force: true }); - } - }); - - it('should support disabling chunking', async () => { - const ctx = await Context.create({ - vectorsDir: chunkTestDir + '-nochunk', - embedder: testEmbedder, - chunking: false, - }); - - await ctx.load('nochunk', path.join(FIXTURES_DIR, 'line-chart-guide.md')); - - const results = await ctx.query('tooltip', { library: 'nochunk', topK: 3 }); - expect(results.length).toBeGreaterThan(0); - - // When chunking is disabled, no results should have chunk metadata - const chunkResults = results.filter((r) => r.chunk); - expect(chunkResults.length).toBe(0); - - await ctx.close(); - if (fs.existsSync(chunkTestDir + '-nochunk')) { - fs.rmSync(chunkTestDir + '-nochunk', { recursive: true, force: true }); - } - }); - }); }); describe('Context with reranking', () => { @@ -246,7 +158,6 @@ describe('Context with reranking', () => { const ctx = await Context.create({ vectorsDir: rerankTestDir, embedder: testEmbedder, - chunking: { maxChunkSize: 400, chunkOverlap: 50 }, }); await ctx.load('rerank', path.join(FIXTURES_DIR, 'line-chart-guide.md')); @@ -271,7 +182,6 @@ describe('Context with reranking', () => { const ctx = await Context.create({ vectorsDir: rerankTestDir + '-disabled', embedder: testEmbedder, - chunking: { maxChunkSize: 400, chunkOverlap: 50 }, }); await ctx.load('rerank2', path.join(FIXTURES_DIR, 'line-chart-guide.md')); @@ -308,7 +218,6 @@ describe('Context with query expansion', () => { const ctx = await Context.create({ vectorsDir: expandTestDir, embedder: testEmbedder, - chunking: { maxChunkSize: 500, chunkOverlap: 50 }, }); await ctx.load('expand', path.join(FIXTURES_DIR, 'line-chart-guide.md')); @@ -331,7 +240,6 @@ describe('Context with query expansion', () => { vectorsDir: expandTestDir + '-disabled', embedder: testEmbedder, queryExpansion: false, - chunking: false, }); await ctx.load('expand2', path.join(FIXTURES_DIR, 'getting-started.md')); @@ -358,7 +266,6 @@ describe('Context with query expansion', () => { 'animation': ['动效', 'animate', 'transition'], }, }, - chunking: { maxChunkSize: 500, chunkOverlap: 50 }, }); await ctx.load('expand3', path.join(FIXTURES_DIR, 'line-chart-guide.md')); diff --git a/test/utils.test.ts b/test/utils.test.ts index 58d9548..0f15aee 100644 --- a/test/utils.test.ts +++ b/test/utils.test.ts @@ -31,19 +31,16 @@ describe('cosineSimilarity', () => { describe('evalMemoryFilter', () => { const fields = { content: 'hello world', - chunkIndex: 0, - parentDocId: 'abc__test', - totalChunks: 5, - headingPath: 'Line Chart > Tooltip', + sourceFilePath: 'docs/getting-started.md', }; describe('string equality', () => { it('should match string field with quoted value', () => { - expect(evalMemoryFilter("parentDocId = 'abc__test'", fields)).toBe(true); + expect(evalMemoryFilter("sourceFilePath = 'docs/getting-started.md'", fields)).toBe(true); }); it('should reject non-matching string value', () => { - expect(evalMemoryFilter("parentDocId = 'xyz'", fields)).toBe(false); + expect(evalMemoryFilter("sourceFilePath = 'xyz'", fields)).toBe(false); }); it('should match content field', () => { @@ -51,43 +48,17 @@ describe('evalMemoryFilter', () => { }); }); - describe('number equality', () => { - it('should match numeric field with integer value', () => { - expect(evalMemoryFilter("chunkIndex = 0", fields)).toBe(true); - }); - - it('should reject non-matching numeric value', () => { - expect(evalMemoryFilter("chunkIndex = 3", fields)).toBe(false); - }); - - it('should match totalChunks field', () => { - expect(evalMemoryFilter("totalChunks = 5", fields)).toBe(true); - }); - - it('should handle negative numbers', () => { - expect(evalMemoryFilter("chunkIndex = -1", { chunkIndex: -1 })).toBe(true); - }); - }); - describe('AND conjunction', () => { it('should match when all clauses are true', () => { - expect(evalMemoryFilter("chunkIndex = 0 AND parentDocId = 'abc__test'", fields)).toBe(true); + expect(evalMemoryFilter("content = 'hello world' AND sourceFilePath = 'docs/getting-started.md'", fields)).toBe(true); }); it('should reject when any clause is false', () => { - expect(evalMemoryFilter("chunkIndex = 1 AND parentDocId = 'abc__test'", fields)).toBe(false); + expect(evalMemoryFilter("content = 'hello' AND sourceFilePath = 'docs/getting-started.md'", fields)).toBe(false); }); it('should reject when all clauses are false', () => { - expect(evalMemoryFilter("chunkIndex = 3 AND parentDocId = 'xyz'", fields)).toBe(false); - }); - - it('should support mixed string and number clauses', () => { - expect(evalMemoryFilter("chunkIndex = 0 AND headingPath = 'Line Chart > Tooltip'", fields)).toBe(true); - }); - - it('should support three clauses with AND', () => { - expect(evalMemoryFilter("chunkIndex = 0 AND totalChunks = 5 AND parentDocId = 'abc__test'", fields)).toBe(true); + expect(evalMemoryFilter("content = 'foo' AND sourceFilePath = 'xyz'", fields)).toBe(false); }); }); @@ -102,15 +73,11 @@ describe('evalMemoryFilter', () => { }); it('should handle AND with case-insensitive keyword', () => { - expect(evalMemoryFilter("chunkIndex = 0 and parentDocId = 'abc__test'", fields)).toBe(true); + expect(evalMemoryFilter("content = 'hello world' and sourceFilePath = 'docs/getting-started.md'", fields)).toBe(true); }); it('should handle missing field gracefully (string clause)', () => { expect(evalMemoryFilter("unknownField = 'value'", fields)).toBe(false); }); - - it('should handle missing field gracefully (number clause)', () => { - expect(evalMemoryFilter("unknownField = 42", fields)).toBe(false); - }); }); }); From 35b880d32c240b63e86f2aca3303019b0d19ad1f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A6=8F=E6=99=8B?= Date: Mon, 6 Jul 2026 17:18:13 +0800 Subject: [PATCH 08/13] chore: add utils --- README.md | 100 ++---- src/context.ts | 176 ++-------- src/embedder/index.ts | 10 +- src/index.ts | 13 +- src/storage/store.ts | 315 ++++++++++++++++++ src/store-manager.ts | 268 --------------- src/types.ts | 27 +- src/utils/common.ts | 91 +++++ src/{query-expander.ts => utils/expander.ts} | 14 +- src/utils/index.ts | 32 ++ src/{ => utils}/reranker.ts | 45 +-- .../language.ts => utils/tokenizer.ts} | 7 +- test/context.test.ts | 15 - test/language.test.ts | 2 +- test/reranker.test.ts | 4 +- test/synonym-expander.test.ts | 2 +- 16 files changed, 534 insertions(+), 587 deletions(-) create mode 100644 src/storage/store.ts delete mode 100644 src/store-manager.ts create mode 100644 src/utils/common.ts rename src/{query-expander.ts => utils/expander.ts} (86%) create mode 100644 src/utils/index.ts rename src/{ => utils}/reranker.ts (77%) rename src/{embedder/language.ts => utils/tokenizer.ts} (90%) diff --git a/README.md b/README.md index f9cb253..c5352fc 100644 --- a/README.md +++ b/README.md @@ -41,12 +41,9 @@ const ctx2 = await Context.fromDir('/path/to/project'); await ctx.load('g2', './g2-docs/**/*.md'); await ctx.load('f2', './f2-docs/**/*.json'); -// Query a single library (default: hybrid search = vector + FTS text) +// Query a single library (default: hybrid search + reranking) const results = await ctx.query('How to configure a line chart', { library: 'g2', topK: 5 }); -// => [{ content: '...', score: 0.92, scoreMode: 'hybrid', id: 'g2-docs/line.md' }, ...] - -// Query with two-stage reranking (pulls extra candidates, then re-scores) -const rerankedResults = await ctx.query('sankey diagram', { library: 'g2', topK: 5, rerank: { rerankFactor: 3 } }); +// => [{ content: '...', score: 0.92, scoreMode: 'reranked', id: 'g2-docs/line.md' }, ...] // Query multiple libraries (array form) const crossResults = await ctx.query('chart configuration', { library: ['g2', 'f2'], topK: 5 }); @@ -54,12 +51,6 @@ const crossResults = await ctx.query('chart configuration', { library: ['g2', 'f // Query all loaded libraries const allResults = await ctx.query('visualization', { library: '*', topK: 10 }); -// Pure vector search (skip FTS text path) -const vectorResults = await ctx.query('chart', { library: 'g2', topK: 5, mode: 'vector' }); - -// Filter results by field value -const filteredResults = await ctx.query('tooltip', { library: 'g2', topK: 5, filter: "sourceFilePath = 'docs/line-chart.md'" }); - // Close when done (releases resources) await ctx.close(); ``` @@ -73,14 +64,10 @@ await ctx.close(); |-----------|------|---------|-------------| | `vectorsDir` | `string` | — | **Required**. Directory to store vector files | | `basePath` | `string` | `process.cwd()` | Base path for resolving document IDs. Set for cross-machine consistent IDs. | -| `model` | `string` | auto | Transformers model name for embedding. Skipped when custom `embedder` is provided. | -| `loaders` | `Loader[]` | built-in | Custom loaders (default: MarkdownLoader, JsonLoader, TextLoader) | -| `embedder` | `Embedder` | auto-resolved | Custom embedder. Skips auto-resolution when provided. | | `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. | | `ftsFields` | `string[]` | `['content']` | Fields to index for Full Text Search in hybrid mode | | `ftsFieldWeights` | `Record` | `{ content: 1 }` | Per-field boost weights for FTS text path. Higher = more influence. | -| `tokenizer` | `'jieba' | 'standard' | 'auto'` | `'auto'` | FTS tokenizer. `jieba` for CN, `standard` for EN, `auto` picks safe default. | | `rankConstant` | `number` | `60` | RRF rank constant for hybrid search fusion. Lower = "winner-takes-all", higher = more even. | #### Weight Configuration Example @@ -147,31 +134,17 @@ const ctx = await Context.create({ ### `ctx.query(text, options)` -Two-stage retrieval: coarse search (vector / hybrid) → optional reranking → final topK results. +Two-stage retrieval: coarse search (vector / hybrid) → reranking → final topK results. | Parameter | Type | Default | Description | |-----------|------|---------|-------------| | `library` | `string | string[]` | — | Library name(s). Single: `'g2'`, Multiple: `['g2', 'f2']`, All: `'*'`. Comma-separated `'g2,f2'` also supported. | | `topK` | `number` | `5` | Number of results to return | -| `mode` | `'hybrid' | 'vector'` | `'hybrid'` | Search mode. `'hybrid'` = vector + FTS text (better recall), `'vector'` = pure semantic search | -| `rerank` | `RerankOptions | false` | `false` | Reranking configuration. Pass an object `{ rerankFactor, minCandidates }` to enable, or `false` to skip | -| `filter` | `string` | — | Filter expression for zvec exact-match filtering, e.g. `"sourceFilePath = 'docs/line.md'"` | ```typescript -// Hybrid search (default) — best recall for exact term matching +// Semantic search — hybrid (vector + FTS) + reranking by default const results = await ctx.query('sankey diagram', { library: 'g2', topK: 5 }); -// Hybrid search with reranking enabled — two-stage retrieval for precision -const results = await ctx.query('line chart config', { - library: 'g2', topK: 5, rerank: { rerankFactor: 3, minCandidates: 10 } -}); - -// Pure vector search — when FTS is not needed -const results = await ctx.query('chart', { library: 'g2', topK: 5, mode: 'vector' }); - -// Filter by source file path -const filtered = await ctx.query('tooltip', { library: 'g2', filter: "sourceFilePath = 'docs/line.md'" }); - // Multiple libraries const results = await ctx.query('chart', { library: ['g2', 'f2'], topK: 5 }); @@ -233,7 +206,7 @@ Quick-start convenience method — creates a Context from a project directory wi ```typescript const ctx = await Context.fromDir('/path/to/project'); // With custom overrides -const ctx = await Context.fromDir('/path/to/project', { model: 'custom-model' }); +const ctx = await Context.fromDir('/path/to/project', { ftsFieldWeights: { content: 2 } }); ``` ### `ctx.remove(library, id)` — **Deprecated** @@ -276,37 +249,35 @@ await ctx.close(); +--------+--------+ | (SynonymExpander)| | +--------+--------+ +--------v--------+ | - | Chunker | v - | (Markdown/Fixed)+--------+ +--------v--------+ - +--------+--------+ | | Embedder | - | | +--------+--------+ - +--------v--------+ | | - | EmbedBatch | | +--------v--------+ - +--------+--------+ | | Vectorize | - | | +--------+--------+ - +--------v--------+ | | - | .zvec | +--------+-----------+ - +-----------------+ | - v - +-----------+-----------+ - | | - +-------v-------+ +-------v-------+ - | FTS Text Path | | Vector Path | - |(ftsFieldWeights| | | - | tokenizer) | | | - +-------+-------+ +-------+-------+ - | | - +-----------+-----------+ - | - +-----------v-----------+ - | RRF Fusion | - | (rankConstant) | - +-----------+-----------+ - | - +-----------v-----------+ - | KeywordReranker | - | (optional, 2nd stage) | - +-----------+-----------+ + | EmbedBatch | v + +--------+--------+ +--------v--------+ + | | Embedder | + +--------v--------+ +--------+--------+ + | .zvec | | + +-----------------+ +--------v--------+ + | Vectorize | + +--------+--------+ + | + +--------+-----------+ + | + +-----------v-----------+ + | | + +-------v-------+ +-------v-------+ + | FTS Text Path | | Vector Path | + |(ftsFieldWeights| | | + +-------+-------+ +-------+-------+ + | | + +-----------+-----------+ + | + +-----------v-----------+ + | RRF Fusion | + | (rankConstant) | + +-----------+-----------+ + | + +-----------v-----------+ + | KeywordReranker | + | (optional, 2nd stage) | + +-----------+-----------+ | Query Result @@ -316,10 +287,9 @@ await ctx.close(); ### Module Structure - **Public API**: `Context`, `QueryOptions`, `QueryResult`, `Document`, `Loader`, `MarkdownLoader`, `JsonLoader`, `TextLoader`, `pathToId` -- **Chunking**: `MarkdownChunker`, `FixedSizeChunker`, `createChunker`, `ChunkingOptions`, `Chunk`, `Chunker` - **Reranking**: `KeywordReranker`, `createReranker`, `Reranker`, `RerankCandidate`, `RerankResult`, `RerankOptions` - **Query Expansion**: `SynonymExpander`, `NoopExpander`, `QueryExpander`, `QueryExpansionOptions` -- **Advanced API**: `Embedder`, `TransformersEmbedder`, `EmbedderManager`, `IZvecStore`, `MemoryZvecStore`, `ActualZvecStore`, `DocumentRegistry`, `StoreManager` +- **Advanced API**: `Embedder`, `TransformersEmbedder`, `EmbedderManager`, `IZvecStore`, `MemoryZvecStore`, `ActualZvecStore`, `DocumentRegistry`, `Store` ## License diff --git a/src/context.ts b/src/context.ts index e0046ee..f2ba5e6 100644 --- a/src/context.ts +++ b/src/context.ts @@ -1,22 +1,24 @@ import * as path from 'path'; import * as fs from 'fs'; -import * as crypto from 'crypto'; import { glob } from 'glob'; import { ContextOptions, QueryOptions, QueryResult, Document, LoadPhase, LoadProgress } from './types'; -import { - resolveEmbedder, -} from './embedder'; -import type { Embedder } from './embedder'; -import type { EmbedderInfo } from './embedder'; +import { resolveEmbedder } from './embedder'; +import type { Embedder, EmbedderInfo } from './embedder'; import { Loader, MarkdownLoader, JsonLoader, TextLoader } from './loaders'; import { DocumentRegistry } from './registry'; -import { StoreManager } from './store-manager'; +import { Store } from './storage/store'; import type { ZvecDoc } from './storage/zvec-store'; import { pathToId } from './loaders/util'; -import { createReranker } from './reranker'; -import type { Reranker, RerankCandidate } from './reranker'; -import { SynonymExpander, NoopExpander } from './query-expander'; -import type { QueryExpander } from './query-expander'; +import { + createReranker, + SynonymExpander, + NoopExpander, + safeParseMeta, + resolveLibraries, + computeContentHash, + selectSampleFiles, +} from './utils'; +import type { Reranker, RerankCandidate, QueryExpander } from './utils'; // --------------------------------------------------------------------------- // Context class @@ -27,7 +29,7 @@ export class Context { private readonly basePath: string; private readonly embedder: Embedder; private readonly _embedderInfo: EmbedderInfo; - private readonly storeManager: StoreManager; + private readonly store: Store; private readonly registry: DocumentRegistry; private readonly loaders: Loader[]; private readonly reranker: Reranker | null; @@ -39,20 +41,14 @@ export class Context { this.basePath = options.basePath ?? process.cwd(); this.embedder = embedder; this._embedderInfo = embedderInfo; - this.storeManager = new StoreManager(options.vectorsDir, embedder, options); + this.store = new Store(options.vectorsDir, embedder, options); this.registry = new DocumentRegistry(); - this.loaders = options.loaders ?? [ + this.loaders = [ new MarkdownLoader(), new JsonLoader(), new TextLoader(), ]; - // Reranker is created eagerly — it has no model-load cost (KeywordReranker). - // Weights are configurable via ContextOptions.rerankWeights. this.reranker = createReranker(options.rerankWeights); - // Query expansion — uses user-provided synonym map only, no built-in defaults. - // When queryExpansion is false, expansion is disabled entirely. - // When queryExpansion is true/undefined/object without synonyms, SynonymExpander - // is created with an empty map (effectively a no-op). this.queryExpander = options.queryExpansion === false ? new NoopExpander() : new SynonymExpander( @@ -63,21 +59,9 @@ export class Context { } static async create(options: ContextOptions): Promise { - let embedder: Embedder; - let embedderInfo: EmbedderInfo; - - if (options.embedder) { - // User provided a custom embedder — infer info from its class - embedder = options.embedder; - embedderInfo = { - kind: 'transformers', - dimensions: embedder.dimensions, - }; - } else { - const result = await resolveEmbedder(options.model); - embedder = result.embedder; - embedderInfo = result.info; - } + const result = await resolveEmbedder(); + const embedder = result.embedder; + const embedderInfo = result.info; // Ensure vectors directory exists if (!fs.existsSync(options.vectorsDir)) { @@ -110,8 +94,7 @@ export class Context { * - basePath = dir (the project root) * - vectorsDir = dir/.context/vectors (hidden, won't pollute project) * - * All other options (model, embedder, etc.) can still be - * overridden via options. + * All other options can still be overridden via options. */ static async fromDir(dir: string, options?: Partial): Promise { const absoluteDir = path.resolve(dir); @@ -175,7 +158,7 @@ export class Context { } } - const store = await this.storeManager.getOrCreate(library, sampleText); + const openedStore = await this.store.create(library, sampleText); // Load registry from disk if not already loaded for this library if (!this.registry.hasLibrary(library)) { @@ -263,7 +246,7 @@ export class Context { }, })); - await store.insert(zvecDocs); + await this.store.addDoc(library, zvecDocs); // Progress: insert phase complete if (this._onProgress) { @@ -321,17 +304,15 @@ export class Context { // eliminates the serial wait time when querying multiple libraries. const perLibraryResults = await Promise.all( libraries.map(async (library) => { - const store = this.storeManager.tryOpen(library); - if (!store) return [] as QueryResult[]; - - const searchResults = mode === 'hybrid' - ? await store.searchHybrid({ - queryText: expandedText, - queryVector: vector, - topK: searchTopK, - filter: options.filter, - }) - : await store.search({ vector, topK: searchTopK, filter: options.filter }); + const searchResults = await this.store.queryDoc(library, { + mode, + queryText: expandedText, + queryVector: vector, + topK: searchTopK, + filter: options.filter, + }); + + if (searchResults.length === 0) return [] as QueryResult[]; return searchResults.map((result) => { const content = String(result.fields?.content ?? ''); @@ -353,7 +334,7 @@ export class Context { ); // Flatten and sort by coarse score - let allResults = perLibraryResults.flat(); + const allResults = perLibraryResults.flat(); allResults.sort((a, b) => b.score - a.score); // Stage 2: Rerank the candidate pool for precision (when enabled). @@ -403,7 +384,7 @@ export class Context { this.registry.remove(library, id); this.registry.saveToDisk(this.vectorsDir, library); - await this.storeManager.close(library); + await this.store.close(library); } /** @@ -420,8 +401,8 @@ export class Context { */ async rebuild(library: string, pattern: string | string[]): Promise { // Close and delete the existing store - await this.storeManager.close(library); - await this.storeManager.deleteStore(library); + await this.store.close(library); + await this.store.deleteStore(library); // Clear the registry so all docs will be re-loaded this.registry.removeLibrary(library); @@ -447,91 +428,6 @@ export class Context { * exit or before re-creating a new instance). */ async close(): Promise { - await this.storeManager.closeAll(); - } -} - -// --------------------------------------------------------------------------- -// Internal helpers -// --------------------------------------------------------------------------- - -/** - * Safely parse meta JSON string. Returns undefined on invalid JSON. - */ -function safeParseMeta(metaStr: string | undefined): Record | undefined { - if (!metaStr) return undefined; - try { - return JSON.parse(metaStr); - } catch { - return undefined; - } -} - - -/** - * Resolve library names from the query option. - * - * - '*' queries all loaded libraries. - * - Array of names queries multiple specific libraries. - * - Comma-separated string is supported for backward compatibility. - * - Single string is the normal case. - */ -function resolveLibraries( - librarySpec: string | string[], - registry: DocumentRegistry -): string[] { - // Array form: direct - if (Array.isArray(librarySpec)) { - return librarySpec.filter(Boolean); - } - - // Wildcard: all libraries - if (librarySpec === '*') { - return registry.getLibraryNames(); - } - - // Comma-separated: backward compatibility - if (librarySpec.includes(',')) { - return librarySpec.split(',').map((s) => s.trim()).filter(Boolean); - } - - // Single library - return [librarySpec]; -} - -/** - * Compute a short content hash for change detection. - * - * Uses SHA-256 truncated to 16 hex chars (64-bit) — compact enough - * for registry storage, collision-resistant enough for dedup. - */ -function computeContentHash(content: string): string { - return crypto.createHash('sha256').update(content).digest('hex').slice(0, 16); -} - -/** - * 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. - */ -function selectSampleFiles(files: string[], maxCount: number): string[] { - if (files.length <= maxCount) return files; - - const result: string[] = []; - // Always include first and last - result.push(files[0]); - // Add evenly-spaced samples from the middle - 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]); - } - } - // Always include last (if not already included) - if (result[result.length - 1] !== files[files.length - 1] && result.length < maxCount) { - result.push(files[files.length - 1]); + await this.store.closeAll(); } - return result; } diff --git a/src/embedder/index.ts b/src/embedder/index.ts index 9cd83b7..8fbce55 100644 --- a/src/embedder/index.ts +++ b/src/embedder/index.ts @@ -3,18 +3,20 @@ * * Re-exports from split files: * types.ts → Embedder interface - * language.ts → isCJK, detectLanguage, tokenizerForLanguage, detectTokenizer * 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 -export { isCJK, splitMixed, detectLanguage, tokenizerForLanguage, detectTokenizer } from './language'; -export type { LanguageHint } from './language'; +// Language detection & CJK utilities — re-exported from utils for backward compatibility +export { isCJK, splitMixed, detectLanguage, tokenizerForLanguage, detectTokenizer } from '../utils/tokenizer'; +export type { LanguageHint } from '../utils/tokenizer'; // TransformersEmbedder — production-quality model embedder export { TransformersEmbedder } from './transformers'; diff --git a/src/index.ts b/src/index.ts index 38e2ab1..3e910b8 100644 --- a/src/index.ts +++ b/src/index.ts @@ -15,12 +15,12 @@ export { export { pathToId } from './loaders/util'; // Reranker — two-stage retrieval precision scoring -export { KeywordReranker, createReranker } from './reranker'; -export type { Reranker, RerankCandidate, RerankResult } from './reranker'; +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 './query-expander'; -export type { QueryExpander } from './query-expander'; +export { SynonymExpander, NoopExpander } from './utils/expander'; +export type { QueryExpander } from './utils/expander'; // --------------------------------------------------------------------------- // Advanced API — for extending or customizing internals @@ -67,6 +67,7 @@ export type { // DocumentRegistry — dedup tracking export { DocumentRegistry } from './registry'; -// StoreManager — zvec store lifecycle management -export { StoreManager } from './store-manager'; +// Store — zvec store lifecycle management +export { Store } from './storage/store'; +export type { StoreQueryParams } from './storage/store'; diff --git a/src/storage/store.ts b/src/storage/store.ts new file mode 100644 index 0000000..0f917a9 --- /dev/null +++ b/src/storage/store.ts @@ -0,0 +1,315 @@ +import * as fs from 'fs'; +import * as path from 'path'; +import { + createZvecStore, + openZvecStoreSync, + isZvecAvailable, +} from './zvec-store'; +import type { IZvecStore, ZvecStoreConfig, FtsFieldWeight, ActualZvecStoreOptions } from './zvec-store'; +import type { ZvecDoc, ZvecQueryResult, ZvecSearchParams, ZvecHybridParams } from './types'; +import { MemoryZvecStore } from './memory-store'; +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 resolveFtsFields(options?: ContextOptions): string[] { + return options?.ftsFields ?? DEFAULT_FTS_FIELDS; +} + +function resolveFtsWeights(options?: ContextOptions): FtsFieldWeight[] { + const ftsFields = resolveFtsFields(options); + const weightMap = options?.ftsFieldWeights; + + if (weightMap) { + return ftsFields.map((fieldName) => ({ + fieldName, + weight: weightMap[fieldName] ?? 1.0, + })); + } + + return ftsFields.map((fieldName) => ({ fieldName, weight: 1.0 })); +} + +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' }, + ], + }; +} + +function storeOpenOptions(options?: ContextOptions): ActualZvecStoreOptions { + return { + vectorField: DEFAULT_VECTOR_FIELD, + ftsFields: resolveFtsFields(options), + rankConstant: options?.rankConstant ?? DEFAULT_RANK_CONSTANT, + }; +} + +// --------------------------------------------------------------------------- +// 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; +} + +// --------------------------------------------------------------------------- +// Store — manages zvec store lifecycle per library +// --------------------------------------------------------------------------- + +/** + * Store manages creation, caching, and querying of zvec store instances. + * + * 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) + */ +export class Store { + private readonly vectorsDir: string; + private readonly embedder: Embedder; + private readonly ftsWeights: FtsFieldWeight[]; + private readonly rankConstant: number; + 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(); + + constructor(vectorsDir: string, embedder: Embedder, options?: ContextOptions) { + this.vectorsDir = vectorsDir; + this.embedder = embedder; + this.ftsWeights = resolveFtsWeights(options); + this.rankConstant = options?.rankConstant ?? DEFAULT_RANK_CONSTANT; + this.contextOptions = 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); + if (cached) return cached; + + const pendingPromise = this.pending.get(library); + if (pendingPromise) return pendingPromise; + + 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); + } + } + + /** + * Batch-insert documents into a library's store. + * + * The store must have been created via `create()` first. + * + * @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); + } + + /** + * 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 []; + + if (params.mode === 'hybrid') { + return store.searchHybrid({ + queryText: params.queryText ?? '', + queryVector: params.queryVector, + topK: params.topK, + filter: params.filter, + }); + } + + return store.search({ + vector: params.queryVector, + topK: params.topK, + filter: params.filter, + }); + } + + // ── Lifecycle helpers (used by Context internally) ───────────────────── + + /** + * Close and remove a single library's store from cache. + * + * Used by `Context.untrack()` and `Context.rebuild()`. + */ + async close(library: string): Promise { + const store = this.stores.get(library); + if (store) { + await store.close(); + this.stores.delete(library); + } + } + + /** + * Delete a library's store file from disk (after closing it). + * + * Used by `Context.rebuild()` to physically remove vector data before + * re-creating the store with fresh embeddings. + */ + async deleteStore(library: string): Promise { + await this.close(library); + const filePath = this._getStorePath(library); + if (fs.existsSync(filePath)) { + fs.unlinkSync(filePath); + } + } + + /** + * 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(); + } + this.stores.clear(); + } + + // ── Private helpers ──────────────────────────────────────────────────── + + private _getStorePath(library: string): string { + return path.join(this.vectorsDir, `${library}.zvec`); + } + + private async _doCreate(library: string, sampleText?: string): Promise { + const filePath = this._getStorePath(library); + const allowFallback = this.contextOptions?.allowMemoryFallback ?? false; + + if (fs.existsSync(filePath)) { + if (isZvecAvailable()) { + return openZvecStoreSync(filePath, storeOpenOptions(this.contextOptions)); + } + if (allowFallback) { + console.warn( + `[context] @zvec/zvec not available — opening library "${library}" ` + + `with in-memory MemoryZvecStore. Data will NOT be persisted.` + ); + return new MemoryZvecStore(this.ftsWeights, this.rankConstant); + } + throw new Error( + `Cannot open existing store "${filePath}": @zvec/zvec is not installed. ` + + `Install it with: pnpm add @zvec/zvec\n` + + `Or set allowMemoryFallback: true to use in-memory store (no persistence).` + ); + } + + try { + return await createZvecStore( + filePath, + contextStoreConfig(this.embedder.dimensions, sampleText), + this.ftsWeights, + this.rankConstant, + ); + } catch (err) { + if (allowFallback) { + console.warn( + `[context] Store creation failed for library "${library}", ` + + `falling back to in-memory MemoryZvecStore. ` + + `Error: ${(err as Error).message?.split('\n')[0]}` + ); + return new MemoryZvecStore(this.ftsWeights, this.rankConstant); + } + throw err; + } + } + + /** + * 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) && isZvecAvailable()) { + const store = openZvecStoreSync(filePath, storeOpenOptions(this.contextOptions)); + this.stores.set(library, store); + return store; + } + return undefined; + } +} diff --git a/src/store-manager.ts b/src/store-manager.ts deleted file mode 100644 index aaa1f1d..0000000 --- a/src/store-manager.ts +++ /dev/null @@ -1,268 +0,0 @@ -import * as fs from 'fs'; -import * as path from 'path'; -import { - createZvecStore, - openZvecStoreSync, - isZvecAvailable, -} from './storage/zvec-store'; -import type { IZvecStore, ZvecStoreConfig, FtsFieldWeight, ActualZvecStoreOptions } from './storage/zvec-store'; -import { MemoryZvecStore } from './storage/memory-store'; -import type { Embedder } from './embedder'; -import { detectTokenizer } from './embedder'; -import type { ContextOptions } from './types'; - -// --------------------------------------------------------------------------- -// Default schema constants -// --------------------------------------------------------------------------- - -const DEFAULT_VECTOR_FIELD = 'embedding'; -const DEFAULT_FTS_FIELDS = ['content']; -const DEFAULT_RANK_CONSTANT = 60; - -function resolveFtsFields(options?: ContextOptions): string[] { - return options?.ftsFields ?? DEFAULT_FTS_FIELDS; -} - -function resolveFtsWeights(options?: ContextOptions): FtsFieldWeight[] { - const ftsFields = resolveFtsFields(options); - const weightMap = options?.ftsFieldWeights; - - if (weightMap) { - return ftsFields.map((fieldName) => ({ - fieldName, - weight: weightMap[fieldName] ?? 1.0, - })); - } - - // Default: all FTS fields have weight 1.0 - return ftsFields.map((fieldName) => ({ fieldName, weight: 1.0 })); -} - -function resolveTokenizer(options?: ContextOptions, sampleText?: string): string { - const tokenizer = options?.tokenizer ?? 'auto'; - if (tokenizer === 'auto') { - // When a document sample is available, auto-detect the best tokenizer - // based on character distribution (CJK → jieba, Latin → standard). - // Falls back to 'jieba' as the safe default for mixed-language content - // when no sample has been loaded yet. - if (sampleText && sampleText.trim().length > 0) { - return detectTokenizer(sampleText); - } - return 'jieba'; - } - return tokenizer; -} - -function contextStoreConfig(dims: number, options?: ContextOptions, sampleText?: string): ZvecStoreConfig { - const ftsFields = resolveFtsFields(options); - const tokenizerName = resolveTokenizer(options, 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' }, - ], - }; -} - -function storeOpenOptions(options?: ContextOptions): ActualZvecStoreOptions { - return { - vectorField: DEFAULT_VECTOR_FIELD, - ftsFields: resolveFtsFields(options), - rankConstant: options?.rankConstant ?? DEFAULT_RANK_CONSTANT, - }; -} - -// --------------------------------------------------------------------------- -// StoreManager — manages zvec store lifecycle per library -// --------------------------------------------------------------------------- - -/** - * StoreManager handles creation, caching, and lazy-loading of IZvecStore - * instances. Each library gets its own `.zvec` file on disk. - */ -export class StoreManager { - private readonly vectorsDir: string; - private readonly embedder: Embedder; - private readonly ftsWeights: FtsFieldWeight[]; - private readonly rankConstant: number; - 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(); - - constructor(vectorsDir: string, embedder: Embedder, options?: ContextOptions) { - this.vectorsDir = vectorsDir; - this.embedder = embedder; - this.ftsWeights = resolveFtsWeights(options); - this.rankConstant = options?.rankConstant ?? DEFAULT_RANK_CONSTANT; - this.contextOptions = options; - } - - /** Get the FTS field weights for MemoryZvecStore fallback. */ - getFtsWeights(): FtsFieldWeight[] { - return this.ftsWeights; - } - - /** Get the RRF rank constant for hybrid search. */ - getRankConstant(): number { - return this.rankConstant; - } - - private getStorePath(library: string): string { - return path.join(this.vectorsDir, `${library}.zvec`); - } - - /** - * Get or create a zvec store for a library. - * - * Uses a Promise-based lock so that concurrent calls for the same library - * share a single creation attempt instead of racing to produce duplicate - * store instances. - * - * @param library Library name. - * @param sampleText Optional document sample for auto-detecting FTS tokenizer - * when `tokenizer` is `'auto'`. Only used for new stores. - */ - async getOrCreate(library: string, sampleText?: string): Promise { - const cached = this.stores.get(library); - if (cached) return cached; - - const pending = this.pending.get(library); - if (pending) return pending; - - const promise = this._doGetOrCreate(library, sampleText); - this.pending.set(library, promise); - try { - const store = await promise; - this.stores.set(library, store); - return store; - } finally { - this.pending.delete(library); - } - } - - private async _doGetOrCreate(library: string, sampleText?: string): Promise { - const filePath = this.getStorePath(library); - const allowFallback = this.contextOptions?.allowMemoryFallback ?? false; - - if (fs.existsSync(filePath)) { - // Existing store on disk — try to open with zvec - if (isZvecAvailable()) { - return openZvecStoreSync(filePath, storeOpenOptions(this.contextOptions)); - } - // zvec unavailable for existing store - if (allowFallback) { - console.warn( - `[context] @zvec/zvec not available — opening library "${library}" ` + - `with in-memory MemoryZvecStore. Data will NOT be persisted.` - ); - return new MemoryZvecStore(this.ftsWeights, this.rankConstant); - } - throw new Error( - `Cannot open existing store "${filePath}": @zvec/zvec is not installed. ` + - `Install it with: pnpm add @zvec/zvec\n` + - `Or set allowMemoryFallback: true to use in-memory store (no persistence).` - ); - } - - // New store — createZvecStore already handles fallback internally - try { - return await createZvecStore( - filePath, - contextStoreConfig(this.embedder.dimensions, this.contextOptions, sampleText), - this.ftsWeights, - this.rankConstant, - ); - } catch (err) { - if (allowFallback) { - console.warn( - `[context] Store creation failed for library "${library}", ` + - `falling back to in-memory MemoryZvecStore. ` + - `Error: ${(err as Error).message?.split('\n')[0]}` - ); - return new MemoryZvecStore(this.ftsWeights, this.rankConstant); - } - throw err; - } - } - - /** - * Get an existing store (already opened or cached). - * Returns undefined if the store has not been opened yet. - */ - getCached(library: string): IZvecStore | undefined { - return this.stores.get(library); - } - - /** - * Try to lazily open an existing store on disk. - * Returns the store if found, undefined otherwise. - */ - tryOpen(library: string): IZvecStore | undefined { - if (this.stores.has(library)) { - return this.stores.get(library)!; - } - - const filePath = this.getStorePath(library); - if (fs.existsSync(filePath)) { - const store = openZvecStoreSync(filePath, storeOpenOptions(this.contextOptions)); - this.stores.set(library, store); - return store; - } - - return undefined; - } - - /** - * Check whether a store file exists on disk for a library. - */ - existsOnDisk(library: string): boolean { - return fs.existsSync(this.getStorePath(library)); - } - - /** - * Close and remove a store from cache. - */ - async close(library: string): Promise { - const store = this.stores.get(library); - if (store) { - await store.close(); - this.stores.delete(library); - } - } - - /** - * Delete the store file from disk (after closing it). - * - * Used by `Context.rebuild()` to physically remove vector data before - * re-creating the store with fresh embeddings. - */ - async deleteStore(library: string): Promise { - await this.close(library); - const filePath = this.getStorePath(library); - if (fs.existsSync(filePath)) { - fs.unlinkSync(filePath); - } - } - - /** - * Close all stores. - */ - async closeAll(): Promise { - for (const [, store] of this.stores) { - await store.close(); - } - this.stores.clear(); - } -} diff --git a/src/types.ts b/src/types.ts index c78b1d3..e846c68 100644 --- a/src/types.ts +++ b/src/types.ts @@ -54,8 +54,6 @@ export interface RerankOptions { export interface ContextOptions { /** Directory to store vector files */ vectorsDir: string; - /** Transformers model name */ - model?: string; /** * Base path for resolving document IDs. * @@ -66,14 +64,6 @@ export interface ContextOptions { * Defaults to `process.cwd()`. */ basePath?: string; - /** Custom loaders to register (default: MarkdownLoader, JsonLoader, TextLoader) */ - loaders?: import('./loaders').Loader[]; - /** - * Custom embedder instance. When provided, skips the auto-resolution - * logic and uses this embedder directly. Useful for when you - * have a pre-initialized embedder. - */ - embedder?: import('./embedder').Embedder; /** * Progress callback for `load()` — called after each major phase completes. @@ -130,17 +120,6 @@ export interface ContextOptions { */ ftsFieldWeights?: Record; - /** - * FTS tokenizer for full-text search indexing. - * - * - `'jieba'` (default): Chinese word segmentation, best for CN/mixed content - * - `'standard'`: English stemmer + stop words, best for pure English content - * - `'auto'`: same as `'jieba'` (safe default for bilingual datasets) - * - * Only affects newly created stores. Existing stores keep their schema. - */ - tokenizer?: 'jieba' | 'standard' | 'auto'; - /** * RRF rank constant for hybrid search fusion. * @@ -164,7 +143,7 @@ export interface ContextOptions { * rerankWeights: { headingTermBonus: 4.0, phraseWeight: 5.0 } * ``` */ - rerankWeights?: Omit; + rerankWeights?: Omit; /** * Allow automatic fallback to in-memory store when @zvec/zvec is unavailable. @@ -213,12 +192,12 @@ export interface QueryOptions { /** * Reranking configuration. * - * When enabled, the query pipeline uses two-stage retrieval: + * When enabled (default), the query pipeline uses two-stage retrieval: * 1. Coarse search (vector / hybrid) returns topK × rerankFactor candidates * 2. Reranker scores each candidate against the query for precision * 3. Final sort by reranked score → topK results * - * Set to `false` to skip reranking (default). Pass an object to configure + * Set to `false` to skip reranking. Pass an object to configure * the rerank factor and minimum candidate pool size. */ rerank?: RerankOptions | false; diff --git a/src/utils/common.ts b/src/utils/common.ts new file mode 100644 index 0000000..748c781 --- /dev/null +++ b/src/utils/common.ts @@ -0,0 +1,91 @@ +import * as crypto from 'crypto'; +import type { DocumentRegistry } from '../registry'; + +// --------------------------------------------------------------------------- +// JSON / meta helpers +// --------------------------------------------------------------------------- + +/** + * Safely parse a meta JSON string. Returns undefined on invalid JSON. + */ +export function safeParseMeta(metaStr: string | undefined): Record | undefined { + if (!metaStr) return undefined; + try { + return JSON.parse(metaStr); + } catch { + return undefined; + } +} + +// --------------------------------------------------------------------------- +// Library name resolution +// --------------------------------------------------------------------------- + +/** + * Resolve library names from the query option. + * + * - '*' queries all loaded libraries. + * - Array of names queries multiple specific libraries. + * - Comma-separated string is supported for backward compatibility. + * - Single string is the normal case. + */ +export function resolveLibraries( + librarySpec: string | string[], + registry: DocumentRegistry, +): string[] { + if (Array.isArray(librarySpec)) { + return librarySpec.filter(Boolean); + } + + if (librarySpec === '*') { + return registry.getLibraryNames(); + } + + if (librarySpec.includes(',')) { + return librarySpec.split(',').map((s) => s.trim()).filter(Boolean); + } + + return [librarySpec]; +} + +// --------------------------------------------------------------------------- +// Content hashing +// --------------------------------------------------------------------------- + +/** + * Compute a short content hash for change detection. + * + * Uses SHA-256 truncated to 16 hex chars (64-bit) — compact enough + * for registry storage, collision-resistant enough for dedup. + */ +export function computeContentHash(content: string): string { + return crypto.createHash('sha256').update(content).digest('hex').slice(0, 16); +} + +// --------------------------------------------------------------------------- +// 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; +} diff --git a/src/query-expander.ts b/src/utils/expander.ts similarity index 86% rename from src/query-expander.ts rename to src/utils/expander.ts index 95414bc..608fbb9 100644 --- a/src/query-expander.ts +++ b/src/utils/expander.ts @@ -8,9 +8,6 @@ * * Synonyms are entirely user-provided — no built-in defaults. Pass your * own map to `SynonymExpander` to define domain-specific term bridges. - * - * Example (with a custom synonym map): - * { "tooltip": ["提示框", "hover"] } → "tooltip config 提示框 hover 配置" */ // --------------------------------------------------------------------------- @@ -72,14 +69,11 @@ export class SynonymExpander implements QueryExpander { const queryLower = query.toLowerCase().trim(); const additions = new Set(); - // Find matching synonym entries in the query for (const [term, syns] of Object.entries(this.synonyms)) { const termLower = term.toLowerCase(); - // Exact match (whole term exists in query) if (this.containsTerm(queryLower, termLower)) { for (const syn of syns) { - // Don't add synonyms that are already in the query if (!this.containsTerm(queryLower, syn.toLowerCase())) { additions.add(syn); } @@ -89,8 +83,6 @@ export class SynonymExpander implements QueryExpander { if (additions.size === 0) return query; - // Append additions — space-separated keeps compatibility with both - // embedding models and FTS tokenizers. return `${query} ${[...additions].join(' ')}`; } @@ -98,18 +90,15 @@ export class SynonymExpander implements QueryExpander { * 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 - * (e.g. "line" should not match inside "inline" or "pipeline"). + * - Latin terms use strict word-boundary detection to avoid false positives. */ private containsTerm(text: string, term: string): boolean { if (!text.includes(term)) return false; - // CJK terms: substring match is sufficient and correct if (containsCJK(term)) { return true; } - // Latin terms: enforce word boundaries on both sides let searchFrom = 0; while (searchFrom <= text.length - term.length) { const idx = text.indexOf(term, searchFrom); @@ -121,7 +110,6 @@ export class SynonymExpander implements QueryExpander { if (prevOk && afterOk) return true; - // Move past this occurrence to check subsequent ones searchFrom = idx + 1; } diff --git a/src/utils/index.ts b/src/utils/index.ts new file mode 100644 index 0000000..704dc7f --- /dev/null +++ b/src/utils/index.ts @@ -0,0 +1,32 @@ +/** + * 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, library resolution, hashing, sampling +export { + safeParseMeta, + resolveLibraries, + computeContentHash, + selectSampleFiles, +} from './common'; + +// Tokenizer selection — language detection & FTS tokenizer auto-configuration +export { + isCJK, + splitMixed, + 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'; diff --git a/src/reranker.ts b/src/utils/reranker.ts similarity index 77% rename from src/reranker.ts rename to src/utils/reranker.ts index dad5578..a510cad 100644 --- a/src/reranker.ts +++ b/src/utils/reranker.ts @@ -5,9 +5,6 @@ * 1. Coarse search (vector / hybrid) → topK × rerankFactor candidates * 2. Reranker scores each candidate against the query * 3. Final sort by reranked score → topK results - * - * This two-stage approach combines fast vector pre-filtering with a more - * expensive but more precise scoring model on the shortlist. */ // --------------------------------------------------------------------------- @@ -38,37 +35,15 @@ export interface Reranker { /** Configuration for reranking. */ export interface RerankOptions { - /** - * How many extra candidates to pull from the coarse search stage. - * e.g. with topK=5 and rerankFactor=3, 15 candidates are reranked. - * Default: 3 - */ rerankFactor?: number; - /** - * Minimum number of candidates to rerank (floor). - * Default: 10 - */ minCandidates?: number; - - // ----------------------------------------------------------------------- - // KeywordReranker scoring weights — tune for your domain - // ----------------------------------------------------------------------- - - /** Weight for exact phrase match. Default: 3.0 */ phraseWeight?: number; - /** Bonus per additional phrase occurrence beyond the first. Default: 0.5 */ phraseRepeatBonus?: number; - /** Weight for whole-word term match. Default: 1.0 */ termWeight?: number; - /** Bonus per additional term occurrence beyond the first. Default: 0.2 */ termRepeatBonus?: number; - /** Weight for substring (partial) term match. Default: 0.3 */ substringWeight?: number; - /** Bonus when a query term appears in the heading path. Default: 2.0 */ headingTermBonus?: number; - /** Bonus when the full query phrase appears in the heading path. Default: 2.5 */ headingPhraseBonus?: number; - /** Fraction of the original coarse score carried into the reranked score. Default: 0.1 */ originalScoreCarry?: number; } @@ -79,7 +54,6 @@ export interface RerankOptions { const DEFAULT_RERANK_FACTOR = 3; const DEFAULT_MIN_CANDIDATES = 10; -/** Default scoring weights for KeywordReranker. */ const DEFAULT_WEIGHTS = { phraseWeight: 3.0, phraseRepeatBonus: 0.5, @@ -98,13 +72,7 @@ const DEFAULT_WEIGHTS = { /** * Reranker that scores candidates by keyword / phrase overlap with the query. * - * This provides a complementary signal to vector similarity: the coarse - * embedding stage captures semantic closeness, while this reranker boosts - * candidates that contain the exact query terms (or their substrings). - * - * All scoring weights are configurable via constructor options. See - * `RerankOptions` for available tuning parameters. - * + * All scoring weights are configurable via constructor options. * Scores are normalised to [0, 1] via min-max scaling. */ export class KeywordReranker implements Reranker { @@ -138,7 +106,6 @@ export class KeywordReranker implements Reranker { // 1. Exact phrase match — strongest signal if (contentLower.includes(queryPhrase)) { score += w.phraseWeight; - // Bonus for each additional occurrence const phraseCount = countOccurrences(contentLower, queryPhrase); score += (phraseCount - 1) * w.phraseRepeatBonus; } @@ -146,19 +113,17 @@ export class KeywordReranker implements Reranker { // 2. Per-term matching for (const term of queryTerms) { if (contentLower.includes(term)) { - // Whole word match if (isWordBoundary(contentLower, term)) { score += w.termWeight; const termCount = countTermMatches(contentLower, term); score += (termCount - 1) * w.termRepeatBonus; } else { - // Substring match (e.g. "tool" matches "tooltip") score += w.substringWeight; } } } - // 3. Heading path bonus — terms in section headings are more relevant + // 3. Heading path bonus if (c.headingPath) { const headingLower = c.headingPath.toLowerCase(); for (const term of queryTerms) { @@ -196,16 +161,13 @@ export class KeywordReranker implements Reranker { /** Split a query into meaningful tokens. */ function tokenizeQuery(query: string): string[] { - // Split on whitespace and CJK-friendly boundaries const raw = query .split(/[\s,，。.!！？?、：:；;]+/) .filter(Boolean); - // For mixed CN/EN queries, also extract CJK bigrams as tokens const tokens: string[] = []; for (const r of raw) { tokens.push(r); - // For CJK segments, add bigrams as additional tokens if (/[一-鿿]{3,}/.test(r)) { for (let i = 0; i + 2 <= r.length; i++) { tokens.push(r.slice(i, i + 2)); @@ -259,9 +221,6 @@ function isWordBoundary( /** * Create a reranker with optional weight configuration. - * - * Currently returns KeywordReranker. In the future this may auto-select - * based on available models (cross-encoder, etc.). */ export function createReranker(options?: RerankOptions): Reranker { return new KeywordReranker(options); diff --git a/src/embedder/language.ts b/src/utils/tokenizer.ts similarity index 90% rename from src/embedder/language.ts rename to src/utils/tokenizer.ts index d4cc67b..8114108 100644 --- a/src/embedder/language.ts +++ b/src/utils/tokenizer.ts @@ -1,8 +1,8 @@ /** * Language detection and tokenizer selection utilities. * - * Used by StoreManager to auto-configure FTS tokenization based on - * the character distribution of loaded documents. + * Pure functions that determine the best FTS tokenizer based on + * the character distribution of sample text content. */ // --------------------------------------------------------------------------- @@ -51,7 +51,6 @@ export function splitMixed(text: string): string[] { /** * Detect the dominant language category of a text sample. * - * Used by StoreManager to pick the right FTS tokenizer: * - 'cjk' → jieba (Chinese word segmentation) * - 'latin' → standard (English stemmer + stop words) * - 'mixed' → jieba (the more general choice for mixed scripts) @@ -83,8 +82,6 @@ export function detectLanguage(text: string): LanguageHint { /** * Pick the best FTS tokenizer name for a given language hint. - * - * This is the function StoreManager calls to auto-configure tokenization. */ export function tokenizerForLanguage(hint: LanguageHint): string { switch (hint) { diff --git a/test/context.test.ts b/test/context.test.ts index 154103f..2a1e15a 100644 --- a/test/context.test.ts +++ b/test/context.test.ts @@ -2,16 +2,10 @@ import { describe, it, expect, afterAll, beforeAll } from 'vitest'; import * as fs from 'fs'; import * as path from 'path'; import { Context } from '../src/index'; -import { TransformersEmbedder } from '../src/embedder'; const FIXTURES_DIR = path.join(__dirname, 'fixtures/docs'); const TEST_DIR = path.join(__dirname, '.test-tmp'); -// Use TransformersEmbedder for tests — requires model download. -// For faster local testing without model download, provide a custom -// embedder mock via the embedder option. -const testEmbedder = new TransformersEmbedder(); - describe('Context', () => { let ctx: Context; @@ -23,7 +17,6 @@ describe('Context', () => { ctx = await Context.create({ vectorsDir: TEST_DIR, - embedder: testEmbedder, }); }); @@ -134,7 +127,6 @@ describe('Context', () => { const closeTestDir = TEST_DIR + '-close-test'; const ctx2 = await Context.create({ vectorsDir: closeTestDir, - embedder: testEmbedder, }); await ctx2.load('close-test', path.join(FIXTURES_DIR, 'getting-started.md')); await ctx2.close(); @@ -157,7 +149,6 @@ describe('Context with reranking', () => { it('should rerank results with keyword scoring', async () => { const ctx = await Context.create({ vectorsDir: rerankTestDir, - embedder: testEmbedder, }); await ctx.load('rerank', path.join(FIXTURES_DIR, 'line-chart-guide.md')); @@ -181,7 +172,6 @@ describe('Context with reranking', () => { it('should support disabling reranking', async () => { const ctx = await Context.create({ vectorsDir: rerankTestDir + '-disabled', - embedder: testEmbedder, }); await ctx.load('rerank2', path.join(FIXTURES_DIR, 'line-chart-guide.md')); @@ -217,7 +207,6 @@ describe('Context with query expansion', () => { it('should expand CN query to match EN content', async () => { const ctx = await Context.create({ vectorsDir: expandTestDir, - embedder: testEmbedder, }); await ctx.load('expand', path.join(FIXTURES_DIR, 'line-chart-guide.md')); @@ -238,7 +227,6 @@ describe('Context with query expansion', () => { it('should support disabling query expansion', async () => { const ctx = await Context.create({ vectorsDir: expandTestDir + '-disabled', - embedder: testEmbedder, queryExpansion: false, }); @@ -260,7 +248,6 @@ describe('Context with query expansion', () => { it('should expand EN query to match CN concepts', async () => { const ctx = await Context.create({ vectorsDir: expandTestDir + '-en', - embedder: testEmbedder, queryExpansion: { synonyms: { 'animation': ['动效', 'animate', 'transition'], @@ -299,7 +286,6 @@ describe('Context with weight configuration', () => { it('should create Context with custom ftsFieldWeights', async () => { const ctx = await Context.create({ vectorsDir: weightTestDir, - embedder: testEmbedder, ftsFields: ['content'], ftsFieldWeights: { content: 2 }, rankConstant: 30, @@ -317,7 +303,6 @@ describe('Context with weight configuration', () => { // When zvec is not available, MemoryZvecStore is used with FtsFieldWeights const ctx = await Context.create({ vectorsDir: weightTestDir + '-mem', - embedder: testEmbedder, ftsFields: ['content'], ftsFieldWeights: { content: 3 }, }); diff --git a/test/language.test.ts b/test/language.test.ts index da8b5dc..e249309 100644 --- a/test/language.test.ts +++ b/test/language.test.ts @@ -1,5 +1,5 @@ import { describe, it, expect } from 'vitest'; -import { isCJK, detectLanguage, tokenizerForLanguage, detectTokenizer } from '../src/embedder/language'; +import { isCJK, detectLanguage, tokenizerForLanguage, detectTokenizer } from '../src/utils/tokenizer'; describe('isCJK', () => { it('should detect Chinese characters', () => { diff --git a/test/reranker.test.ts b/test/reranker.test.ts index e997c14..977bbab 100644 --- a/test/reranker.test.ts +++ b/test/reranker.test.ts @@ -1,6 +1,6 @@ import { describe, it, expect } from 'vitest'; -import { KeywordReranker, createReranker } from '../src/reranker'; -import type { RerankCandidate } from '../src/reranker'; +import { KeywordReranker, createReranker } from '../src/utils/reranker'; +import type { RerankCandidate } from '../src/utils/reranker'; describe('KeywordReranker', () => { const reranker = new KeywordReranker(); diff --git a/test/synonym-expander.test.ts b/test/synonym-expander.test.ts index 4c8eec0..a5a7b80 100644 --- a/test/synonym-expander.test.ts +++ b/test/synonym-expander.test.ts @@ -1,5 +1,5 @@ import { describe, it, expect } from 'vitest'; -import { SynonymExpander, NoopExpander } from '../src/query-expander'; +import { SynonymExpander, NoopExpander } from '../src/utils/expander'; describe('SynonymExpander', () => { describe('with user-provided synonyms', () => { From 375d536c6a8c054a29ca7db8c9effdb503b2112c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A6=8F=E6=99=8B?= Date: Mon, 6 Jul 2026 18:14:45 +0800 Subject: [PATCH 09/13] chore: context opt --- README.md | 88 ++----------------- src/context.ts | 205 ++++++++++++------------------------------- src/storage/store.ts | 5 -- src/types.ts | 21 +---- src/utils/common.ts | 32 ------- src/utils/index.ts | 3 +- test/context.test.ts | 12 --- 7 files changed, 68 insertions(+), 298 deletions(-) diff --git a/README.md b/README.md index c5352fc..2211ab7 100644 --- a/README.md +++ b/README.md @@ -8,18 +8,10 @@ A local context retrieval library that enables semantic search over your documen ## Features -- 📄 **Multi-format Support**: Supports Markdown, JSON, Text, and other file formats -- 📚 **Multi-library Support**: Manage documents by library -- ⚡ **Auto-indexing**: Automatic vectorization on load with batch embedding for performance -- 🔍 **Hybrid Retrieval**: Combines vector similarity + FTS text matching via RRF fusion for better recall -- 🔄 **Deduplication**: Automatically skip already-loaded documents; content-hash change detection for re-embedding updated files -- ⚖️ **Weight Configuration**: Per-field FTS boost weights and RRF rank constant tuning -- 🛡️ **Clear Error Messages**: Throws descriptive errors when Transformers model is unavailable, guiding users to fix the issue -- 🔁 **Two-stage Reranking**: KeywordReranker boosts candidates with exact query term matches after coarse vector/hybrid search -- 🔁 **Two-stage Reranking**: KeywordReranker boosts candidates with exact query term matches after coarse vector/hybrid search -- 🌐 **Query Expansion**: SynonymExpander uses user-provided synonym maps to bridge CN↔EN terminology gaps -- 📊 **Progress Callback**: `onProgress` hook for monitoring load phases (load → embed → insert) -- 🏗️ **fromDir() Quick-start**: One-call setup from a project directory with auto-derived defaults +- 📄 **Multi-format Support**: Markdown, JSON, Text 文档自动加载与向量化 +- 🔍 **Hybrid Retrieval**: 向量语义 + FTS 全文检索双路召回，RRF 融合排序 +- 🔁 **Two-stage Reranking**: KeywordReranker 精排，关键词命中优先 +- 🌐 **Query Expansion**: 用户自定义同义词表，CN↔EN 跨语言召回增强 ## Quick Start @@ -34,23 +26,14 @@ import { Context } from '@antv/context'; // Standard creation — specify vectorsDir const ctx = await Context.create({ vectorsDir: './vectors' }); -// Quick-start from a project directory (auto-derives basePath & vectorsDir) -const ctx2 = await Context.fromDir('/path/to/project'); - // Load documents into a specific library with automatic vectorization await ctx.load('g2', './g2-docs/**/*.md'); await ctx.load('f2', './f2-docs/**/*.json'); -// Query a single library (default: hybrid search + reranking) +// Query a library (default: hybrid search + reranking) const results = await ctx.query('How to configure a line chart', { library: 'g2', topK: 5 }); // => [{ content: '...', score: 0.92, scoreMode: 'reranked', id: 'g2-docs/line.md' }, ...] -// Query multiple libraries (array form) -const crossResults = await ctx.query('chart configuration', { library: ['g2', 'f2'], topK: 5 }); - -// Query all loaded libraries -const allResults = await ctx.query('visualization', { library: '*', topK: 10 }); - // Close when done (releases resources) await ctx.close(); ``` @@ -138,18 +121,12 @@ Two-stage retrieval: coarse search (vector / hybrid) → reranking → final top | Parameter | Type | Default | Description | |-----------|------|---------|-------------| -| `library` | `string | string[]` | — | Library name(s). Single: `'g2'`, Multiple: `['g2', 'f2']`, All: `'*'`. Comma-separated `'g2,f2'` also supported. | +| `library` | `string` | — | Library name to query. | | `topK` | `number` | `5` | Number of results to return | ```typescript // Semantic search — hybrid (vector + FTS) + reranking by default const results = await ctx.query('sankey diagram', { library: 'g2', topK: 5 }); - -// Multiple libraries -const results = await ctx.query('chart', { library: ['g2', 'f2'], topK: 5 }); - -// All libraries -const results = await ctx.query('chart', { library: '*', topK: 5 }); ``` #### Query Result Fields @@ -164,60 +141,7 @@ Each result includes: | `scoreMode` | `'vector' | 'hybrid' | 'reranked'` | How the score was computed | | `meta` | `Record` | Front-matter metadata (if present) | | `sourceFilePath` | `string` | Original file path relative to `basePath` | -| `library` | `string` | Which library this result came from | - -### `ctx.untrack(library, id)` - -Remove a document from a library's dedup registry. **Important**: zvec does not support single-document deletion, so vector data remains in the store. `untrack()` only removes the dedup entry — the actual vectors remain until you call `rebuild()`. - -| Parameter | Type | Description | -|-----------|------|---------| -| `library` | `string` | Library name | -| `id` | `string` | Document ID to untrack from dedup tracking | - -```typescript -await ctx.untrack('g2', 'abc123__getting_started'); -``` - -### `ctx.rebuild(library, pattern)` - -Rebuild a library's vector store from scratch. Deletes the existing `.zvec` store file, clears the dedup registry, and re-embeds all matching documents. Use this after `untrack()` to actually remove vectors. - -| Parameter | Type | Description | -|-----------|------|---------| -| `library` | `string` | Library name to rebuild | -| `pattern` | `string | string[]` | Glob pattern(s) for re-loading documents | - -```typescript -// Rebuild after untracking documents -await ctx.untrack('g2', 'abc123__getting_started'); -await ctx.rebuild('g2', './g2-docs/**/*.md'); -``` - -### `Context.fromDir(dir, options?)` - -Quick-start convenience method — creates a Context from a project directory with auto-derived defaults (`basePath` = dir, `vectorsDir` = dir/.context/vectors). - -| Parameter | Type | Description | -|-----------|------|---------| -| `dir` | `string` | Project directory path | -| `options` | `Partial` | Optional overrides for auto-derived defaults | -```typescript -const ctx = await Context.fromDir('/path/to/project'); -// With custom overrides -const ctx = await Context.fromDir('/path/to/project', { ftsFieldWeights: { content: 2 } }); -``` - -### `ctx.remove(library, id)` — **Deprecated** - -> Use `untrack()` instead. This alias only removes the dedup tracking entry — vector data remains in the store. -> To physically remove data, call `untrack()` then `rebuild()`. - -```typescript -// Deprecated — use untrack() + rebuild() instead -await ctx.remove('g2', 'abc123__getting_started'); -``` ### `ctx.close()` diff --git a/src/context.ts b/src/context.ts index f2ba5e6..6df2b5c 100644 --- a/src/context.ts +++ b/src/context.ts @@ -1,7 +1,14 @@ import * as path from 'path'; import * as fs from 'fs'; import { glob } from 'glob'; -import { ContextOptions, QueryOptions, QueryResult, Document, LoadPhase, LoadProgress } from './types'; +import { + ContextOptions, + QueryOptions, + QueryResult, + Document, + LoadPhase, + LoadProgress, +} from './types'; import { resolveEmbedder } from './embedder'; import type { Embedder, EmbedderInfo } from './embedder'; import { Loader, MarkdownLoader, JsonLoader, TextLoader } from './loaders'; @@ -14,7 +21,6 @@ import { SynonymExpander, NoopExpander, safeParseMeta, - resolveLibraries, computeContentHash, selectSampleFiles, } from './utils'; @@ -43,18 +49,16 @@ export class Context { this._embedderInfo = embedderInfo; this.store = new Store(options.vectorsDir, embedder, options); this.registry = new DocumentRegistry(); - this.loaders = [ - new MarkdownLoader(), - new JsonLoader(), - new TextLoader(), - ]; + 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.queryExpander = + options.queryExpansion === false + ? new NoopExpander() + : new SynonymExpander( + options.queryExpansion && typeof options.queryExpansion === 'object' + ? options.queryExpansion.synonyms + : undefined, + ); this._onProgress = options.onProgress; } @@ -76,7 +80,8 @@ export class Context { // files and loading them into the registry prevents duplicate // re-embedding of unchanged documents. if (fs.existsSync(options.vectorsDir)) { - const indexFiles = fs.readdirSync(options.vectorsDir) + const indexFiles = fs + .readdirSync(options.vectorsDir) .filter((f) => f.endsWith('.index.json')); for (const indexFile of indexFiles) { const library = indexFile.replace('.index.json', ''); @@ -87,26 +92,6 @@ export class Context { return ctx; } - /** - * Quick-start convenience method — creates a Context from a project directory. - * - * Auto-derives sensible defaults: - * - basePath = dir (the project root) - * - vectorsDir = dir/.context/vectors (hidden, won't pollute project) - * - * All other options can still be overridden via options. - */ - static async fromDir(dir: string, options?: Partial): Promise { - const absoluteDir = path.resolve(dir); - const vectorsDir = options?.vectorsDir ?? path.join(absoluteDir, '.context', 'vectors'); - - return Context.create({ - ...options, - basePath: options?.basePath ?? absoluteDir, - vectorsDir, - }); - } - /** * Diagnostic information about the active embedder. */ @@ -145,7 +130,7 @@ export class Context { try { const sampleFiles = selectSampleFiles(files, 5); const samples = await Promise.allSettled( - sampleFiles.map((f) => fs.promises.readFile(f, 'utf-8')) + sampleFiles.map((f) => fs.promises.readFile(f, 'utf-8')), ); const validSamples = samples .filter((r) => r.status === 'fulfilled') @@ -191,7 +176,7 @@ export class Context { if (this.registry.has(library, docId, contentHash)) return null; return { ...doc, id: docId, contentHash, sourceFilePath: relativePath }; - }) + }), ); // Internal type that extends Document with load-phase metadata. @@ -213,7 +198,7 @@ export class Context { } if (failCount > 0) { console.warn( - `[context] ${failCount}/${files.length} file(s) failed to load in library "${library}" and were skipped.` + `[context] ${failCount}/${files.length} file(s) failed to load in library "${library}" and were skipped.`, ); } @@ -239,9 +224,7 @@ export class Context { vector: vectors[index], fields: { content: doc.content, - meta: doc.meta && Object.keys(doc.meta).length > 0 - ? JSON.stringify(doc.meta) - : '', + meta: doc.meta && Object.keys(doc.meta).length > 0 ? JSON.stringify(doc.meta) : '', sourceFilePath: doc.sourceFilePath ?? '', }, })); @@ -267,17 +250,12 @@ export class Context { * matching via RRF fusion for better recall. Use `mode: 'vector'` for * pure semantic search when FTS is not needed. * - * Supports querying a single library, multiple libraries (array), - * or all loaded libraries ('*' wildcard). - * * @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 libraries = resolveLibraries(options.library, this.registry); - - if (libraries.length === 0) return []; + const library = options.library; // Expand the query with synonyms / cross-language bridging terms. // The expanded text is used for both embedding and FTS, so a single @@ -291,51 +269,41 @@ export class Context { // 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 searchTopK = rerankEnabled - ? Math.max(topK * rerankFactor, minCandidates) - : topK; - - // Query all libraries concurrently and merge results. - // Each library's store is independent — running searches in parallel - // eliminates the serial wait time when querying multiple libraries. - const perLibraryResults = await Promise.all( - libraries.map(async (library) => { - const searchResults = await this.store.queryDoc(library, { - mode, - queryText: expandedText, - queryVector: vector, - topK: searchTopK, - filter: options.filter, - }); - - if (searchResults.length === 0) return [] as QueryResult[]; - - return searchResults.map((result) => { - const content = String(result.fields?.content ?? ''); - const metaStr = result.fields?.meta as string | undefined; - const meta = safeParseMeta(metaStr); - - 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, - library, - embedderKind: this._embedderInfo.kind, - }; - }); - }) - ); + 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 searchTopK = rerankEnabled ? Math.max(topK * rerankFactor, minCandidates) : topK; + + const searchResults = await this.store.queryDoc(library, { + mode, + queryText: expandedText, + queryVector: vector, + topK: searchTopK, + filter: options.filter, + }); - // Flatten and sort by coarse score - const allResults = perLibraryResults.flat(); - allResults.sort((a, b) => b.score - a.score); + if (searchResults.length === 0) return []; + + const allResults: QueryResult[] = searchResults.map((result) => { + const content = String(result.fields?.content ?? ''); + const metaStr = result.fields?.meta as string | undefined; + const meta = safeParseMeta(metaStr); + + 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, + }; + }); // Stage 2: Rerank the candidate pool for precision (when enabled). // Pull extra candidates, re-score each against the query, then keep topK. @@ -364,63 +332,6 @@ export class Context { return allResults.slice(0, topK); } - /** - * Remove a document from a library's dedup registry. - * - * **Important**: this only removes the document from the deduplication - * tracking — the underlying vectors remain in the store because zvec - * does not support single-document deletion. To actually remove the - * vector data, call `rebuild(library)` after untracking documents. - * - * Typical workflow for updating a document: - * 1. `ctx.untrack(library, docId)` — remove from dedup tracking - * 2. `ctx.rebuild(library)` — delete store + re-embed remaining docs - * - * @param library Library name. - * @param id Document ID to untrack. - */ - async untrack(library: string, id: string): Promise { - if (!this.registry.has(library, id)) return; - - this.registry.remove(library, id); - this.registry.saveToDisk(this.vectorsDir, library); - await this.store.close(library); - } - - /** - * Rebuild a library's vector store from scratch. - * - * Deletes the existing `.zvec` store file, clears the dedup registry, - * and re-embeds all documents that match the given glob pattern(s). - * - * Use this after `untrack()` to physically remove unwanted vectors, or - * when the store schema needs to change (e.g. new FTS fields). - * - * @param library Library name to rebuild. - * @param pattern Glob pattern(s) for re-loading documents. - */ - async rebuild(library: string, pattern: string | string[]): Promise { - // Close and delete the existing store - await this.store.close(library); - await this.store.deleteStore(library); - - // Clear the registry so all docs will be re-loaded - this.registry.removeLibrary(library); - this.registry.saveToDisk(this.vectorsDir, library); - - // Re-load all matching documents - await this.load(library, pattern); - } - - /** - * @deprecated Use `untrack()` instead. This alias only removes the - * dedup tracking entry — vector data remains in the store. - * To physically remove data, call `untrack()` then `rebuild()`. - */ - async remove(library: string, id: string): Promise { - return this.untrack(library, id); - } - /** * Close all stores and release resources. * diff --git a/src/storage/store.ts b/src/storage/store.ts index 0f917a9..b7bafb9 100644 --- a/src/storage/store.ts +++ b/src/storage/store.ts @@ -214,8 +214,6 @@ export class Store { /** * Close and remove a single library's store from cache. - * - * Used by `Context.untrack()` and `Context.rebuild()`. */ async close(library: string): Promise { const store = this.stores.get(library); @@ -227,9 +225,6 @@ export class Store { /** * Delete a library's store file from disk (after closing it). - * - * Used by `Context.rebuild()` to physically remove vector data before - * re-creating the store with fresh embeddings. */ async deleteStore(library: string): Promise { await this.close(library); diff --git a/src/types.ts b/src/types.ts index e846c68..b21c5f4 100644 --- a/src/types.ts +++ b/src/types.ts @@ -166,16 +166,9 @@ export interface ContextOptions { */ export interface QueryOptions { /** - * Library name(s) to query. - * - * - Single library: `'g2'` - * - Multiple libraries: `['g2', 'f2']` - * - All libraries: `'*'` - * - * Comma-separated strings (`'g2,f2'`) are also supported for backward - * compatibility but the array form is preferred. + * Library name to query. */ - library: string | string[]; + library: string; /** Number of results to return */ topK?: number; /** @@ -234,20 +227,12 @@ export interface QueryResult { scoreMode?: 'vector' | 'hybrid' | 'reranked'; /** Document metadata */ meta?: Record; - /** - * Original file path relative to `basePath` — allows users to trace + /** Original file path relative to `basePath` — allows users to trace * back to the source document for context review or navigation. * * Populated during `load()` and stored as a zvec field. */ sourceFilePath?: string; - /** - * Which library this result came from. - * - * Useful when querying multiple libraries (`library: ['g2', 'f2']` or `'*'`) - * to identify the source of each result. - */ - library?: string; /** * The embedder kind used to produce the vectors for this result's store. * diff --git a/src/utils/common.ts b/src/utils/common.ts index 748c781..5ee2e8b 100644 --- a/src/utils/common.ts +++ b/src/utils/common.ts @@ -1,5 +1,4 @@ import * as crypto from 'crypto'; -import type { DocumentRegistry } from '../registry'; // --------------------------------------------------------------------------- // JSON / meta helpers @@ -17,37 +16,6 @@ export function safeParseMeta(metaStr: string | undefined): Record s.trim()).filter(Boolean); - } - - return [librarySpec]; -} - // --------------------------------------------------------------------------- // Content hashing // --------------------------------------------------------------------------- diff --git a/src/utils/index.ts b/src/utils/index.ts index 704dc7f..475edc2 100644 --- a/src/utils/index.ts +++ b/src/utils/index.ts @@ -5,10 +5,9 @@ * a specific domain module (embedder, loaders, storage). */ -// Common helpers — JSON parsing, library resolution, hashing, sampling +// Common helpers — JSON parsing, hashing, sampling export { safeParseMeta, - resolveLibraries, computeContentHash, selectSampleFiles, } from './common'; diff --git a/test/context.test.ts b/test/context.test.ts index 2a1e15a..33ecc6a 100644 --- a/test/context.test.ts +++ b/test/context.test.ts @@ -98,11 +98,6 @@ describe('Context', () => { expect(results.length).toBe(0); }); - it('should support array library parameter', async () => { - const results = await ctx.query('install', { library: ['md', 'json'], topK: 5 }); - expect(results.length).toBeGreaterThan(0); - }); - it('should default to hybrid search mode', async () => { // Default mode is 'hybrid' — combines vector + text path const results = await ctx.query('install', { library: 'md', topK: 1 }); @@ -115,13 +110,6 @@ describe('Context', () => { }); }); - describe('cross-library query', () => { - it('should query multiple comma-separated libraries', async () => { - const results = await ctx.query('install', { library: 'md,json', topK: 5 }); - expect(results.length).toBeGreaterThan(0); - }); - }); - describe('close', () => { it('should close all stores without error', async () => { const closeTestDir = TEST_DIR + '-close-test'; From 062f1d815bb3d86a6e3c92d74ed3c2614e6a442b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A6=8F=E6=99=8B?= Date: Mon, 6 Jul 2026 19:13:47 +0800 Subject: [PATCH 10/13] chore: remove memory store --- README.md | 2 +- src/index.ts | 2 - src/storage/actual-store.ts | 24 +--- src/storage/memory-store.ts | 265 ------------------------------------ src/storage/store.ts | 80 +++-------- src/storage/types.ts | 6 - src/storage/utils.ts | 51 ------- src/storage/zvec-store.ts | 8 +- src/types.ts | 16 --- test/context.test.ts | 24 ---- test/utils.test.ts | 58 +------- test/zvec-store.test.ts | 96 ------------- 12 files changed, 29 insertions(+), 603 deletions(-) delete mode 100644 src/storage/memory-store.ts diff --git a/README.md b/README.md index 2211ab7..b33a989 100644 --- a/README.md +++ b/README.md @@ -213,7 +213,7 @@ await ctx.close(); - **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`, `MemoryZvecStore`, `ActualZvecStore`, `DocumentRegistry`, `Store` +- **Advanced API**: `Embedder`, `TransformersEmbedder`, `EmbedderManager`, `IZvecStore`, `ActualZvecStore`, `DocumentRegistry`, `Store` ## License diff --git a/src/index.ts b/src/index.ts index 3e910b8..f0f3873 100644 --- a/src/index.ts +++ b/src/index.ts @@ -44,7 +44,6 @@ export type { EmbedderInfo, EmbedderKind } from './embedder'; // Zvec store — custom vector storage backends export { IZvecStore, - MemoryZvecStore, ActualZvecStore, createZvecStore, openZvecStore, @@ -58,7 +57,6 @@ export type { ZvecQueryResult, ZvecSearchParams, ZvecHybridParams, - FtsFieldWeight, ZvecFieldSchema, ZvecStoreConfig, ActualZvecStoreOptions, diff --git a/src/storage/actual-store.ts b/src/storage/actual-store.ts index 018f774..3424a67 100644 --- a/src/storage/actual-store.ts +++ b/src/storage/actual-store.ts @@ -337,40 +337,31 @@ export class ActualZvecStore implements IZvecStore { /** * Create a zvec store with the given config. * - * When @zvec/zvec is available, creates a persisted ActualZvecStore. - * Otherwise falls back to an in-memory MemoryZvecStore. + * Requires @zvec/zvec to be installed. Throws an error when it is not available. */ export async function createZvecStore( path: string, config: ZvecStoreConfig, - memoryFtsWeights?: import('./types').FtsFieldWeight[], - rankConstant?: number, ): Promise { - const z = loadZvecSync(); - if (z) { - return ActualZvecStore.create(path, config); - } - const { MemoryZvecStore } = await import('./memory-store'); - return new MemoryZvecStore(memoryFtsWeights, rankConstant); + 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 { - // Both ActualZvecStore.open() and MemoryZvecStore are internally - // synchronous — this async wrapper is for API consistency only. return openZvecStoreSync(path, options); } /** * Synchronously open a zvec store. * - * This avoids the async Promise wrapper so it can be used inside synchronous - * code paths (e.g. `retrieve()`). + * Requires @zvec/zvec to be installed. Throws an error when it is not available. */ export function openZvecStoreSync( path: string, @@ -381,9 +372,8 @@ export function openZvecStoreSync( return ActualZvecStore.openSync(path, options); } throw new Error( - 'Cannot open zvec store: @zvec/zvec is not installed and MemoryZvecStore ' + - 'has no persistence. Install @zvec/zvec or use createZvecStore() to create ' + - 'a new MemoryZvecStore.' + '@zvec/zvec is not installed. Install it with:\n' + + ' pnpm add @zvec/zvec' ); } diff --git a/src/storage/memory-store.ts b/src/storage/memory-store.ts deleted file mode 100644 index fa9e54b..0000000 --- a/src/storage/memory-store.ts +++ /dev/null @@ -1,265 +0,0 @@ -/** - * MemoryZvecStore — pure JS fallback (cosine similarity + linear scan + text). - */ - -import type { - ZvecDoc, - ZvecQueryResult, - ZvecSearchParams, - ZvecHybridParams, - IZvecStore, - FtsFieldWeight, -} from './types'; -import { cosineSimilarity, evalMemoryFilter } from './utils'; - -const DEFAULT_FTS_WEIGHTS: FtsFieldWeight[] = [ - { fieldName: 'content', weight: 1.0 } -]; - -const DEFAULT_RANK_CONSTANT = 60; - -export class MemoryZvecStore implements IZvecStore { - private docs: ZvecDoc[] = []; - private _ftsWeights: FtsFieldWeight[]; - private _rankConstant: number; - - constructor(ftsFieldWeights?: FtsFieldWeight[], rankConstant?: number) { - this._ftsWeights = ftsFieldWeights && ftsFieldWeights.length > 0 - ? ftsFieldWeights - : DEFAULT_FTS_WEIGHTS; - this._rankConstant = rankConstant ?? DEFAULT_RANK_CONSTANT; - } - - async insert(docs: ZvecDoc[]): Promise { - this.docs.push(...docs); - } - - async search(params: ZvecSearchParams): Promise { - return doSyncSearch(this.docs, params); - } - - async searchHybrid(params: ZvecHybridParams): Promise { - return doSyncHybridSearch(this.docs, params, this._ftsWeights, this._rankConstant); - } - - searchSync(params: ZvecSearchParams): ZvecQueryResult[] { - return doSyncSearch(this.docs, params); - } - - searchHybridSync(params: ZvecHybridParams): ZvecQueryResult[] { - return doSyncHybridSearch(this.docs, params, this._ftsWeights, this._rankConstant); - } - - async close(): Promise { - this.docs = []; - } -} - -function doSyncSearch( - docs: ZvecDoc[], - params: ZvecSearchParams -): ZvecQueryResult[] { - const { vector, topK, filter } = params; - const scored: ZvecQueryResult[] = []; - for (const doc of docs) { - if (filter && !evalMemoryFilter(filter, doc.fields)) continue; - scored.push({ - id: doc.id, - score: cosineSimilarity(vector, doc.vector), - fields: doc.fields - }); - } - scored.sort((a, b) => b.score - a.score); - return scored.slice(0, topK); -} - -function doSyncHybridSearch( - docs: ZvecDoc[], - params: ZvecHybridParams, - ftsWeights: FtsFieldWeight[], - rankConstant: number, -): ZvecQueryResult[] { - const { queryText, queryVector, topK, filter } = params; - const rrScores = new Map(); - - // Filtered docs — apply once - const candidates = filter - ? docs.filter((d) => evalMemoryFilter(filter, d.fields)) - : docs; - - // 1. Text path: field-boosted scoring with word-boundary matching + TF weighting. - // - // Scoring rules (per query term per field): - // - Exact word-boundary match (EN: surrounded by space/punctuation, - // CJK: standalone character group): 1.0 × field_weight - // - Substring match (no word boundary): 0.3 × field_weight - // - TF multiplier: log(1 + count) so multiple occurrences boost - // the score but don't saturate it. - // - CJK bigrams from the query also score (0.5 × field_weight per match). - // - // This produces better FTS results than raw substring matching: - // "tooltip" won't falsely match "tooltipconfigxx" - // "折线" will match "折线图" via bigram scoring - const terms = queryText.toLowerCase().split(/\s+/).filter(Boolean); - const cjkBigrams = extractCJKBigrams(queryText.toLowerCase()); - - const textRanked = candidates - .map((doc) => { - let score = 0; - for (const fw of ftsWeights) { - const fieldVal = String(doc.fields[fw.fieldName] || '').toLowerCase(); - for (const term of terms) { - const count = countTermMatches(fieldVal, term); - if (count > 0) { - // Word-boundary match count vs total occurrences - const boundaryCount = countBoundaryMatches(fieldVal, term); - const subCount = count - boundaryCount; - // TF-weighted: log(1 + count) prevents saturation - const boundaryScore = boundaryCount > 0 - ? Math.log(1 + boundaryCount) * fw.weight - : 0; - const subScore = subCount > 0 - ? Math.log(1 + subCount) * 0.3 * fw.weight - : 0; - score += boundaryScore + subScore; - } - } - // CJK bigram scoring — bridges partial character matches - for (const bigram of cjkBigrams) { - const count = countTermMatches(fieldVal, bigram); - if (count > 0) { - score += Math.log(1 + count) * 0.5 * fw.weight; - } - } - } - return { id: doc.id, score }; - }) - .filter((r) => r.score > 0) - .sort((a, b) => b.score - a.score); - - // Standard RRF: score = 1 / (rankConstant + rank), rank starts at 1 - for (let i = 0; i < textRanked.length; i++) { - rrScores.set(textRanked[i].id, 1 / (rankConstant + (i + 1))); - } - - // 2. Vector path: cosine similarity - const vecRanked = candidates - .map((doc) => ({ - id: doc.id, - score: cosineSimilarity(queryVector, doc.vector), - })) - .sort((a, b) => b.score - a.score); - - for (let i = 0; i < vecRanked.length; i++) { - const existing = rrScores.get(vecRanked[i].id) ?? 0; - rrScores.set(vecRanked[i].id, existing + 1 / (rankConstant + (i + 1))); - } - - // 3. Merge by RRF score (no division — scores are already correctly scaled) - const docMap = new Map(candidates.map((d) => [d.id, d])); - return [...rrScores.entries()] - .sort((a, b) => b[1] - a[1]) - .slice(0, topK) - .map(([id, score]) => { - const doc = docMap.get(id)!; - return { id, score, fields: doc.fields }; - }); -} - -// --------------------------------------------------------------------------- -// FTS helper functions — word-boundary matching + CJK bigram extraction -// --------------------------------------------------------------------------- - -/** - * Count total occurrences of a term in text (substring match). - */ -function countTermMatches(text: string, term: string): number { - let count = 0; - let pos = 0; - while ((pos = text.indexOf(term, pos)) !== -1) { - count++; - pos += term.length; - } - return count; -} - -/** - * Count occurrences of a term that land on a word boundary. - * - * A word boundary means the character before and after the match is - * a space, punctuation, or string boundary. CJK characters are treated - * as individual "words" — a CJK term naturally forms a boundary within - * a CJK string. - */ -function countBoundaryMatches(text: string, term: string): number { - let count = 0; - let pos = 0; - while ((pos = text.indexOf(term, pos)) !== -1) { - if (isWordBoundary(text, pos, term.length)) { - count++; - } - pos += term.length; - } - return count; -} - -/** - * Check if a substring occurrence at the given position is at a word boundary. - */ -function isWordBoundary(text: string, pos: number, len: number): boolean { - const before = pos === 0 || /[\s\n.,;:!?，。！？、：；"'(（【《\-_]/.test(text[pos - 1]); - const afterIdx = pos + len; - const after = afterIdx >= text.length || /[\s\n.,;:!?，。！？、：；"')）】》\-_]/.test(text[afterIdx]); - return before && after; -} - -/** - * Extract CJK bigrams (2-char sequences) from text for partial matching. - * - * For example, "折线图配置" produces bigrams: ["折线", "线图", "图配", "配置"] - * These bigrams help the FTS path find documents that contain partial - * character sequences of the query, improving recall for CJK searches. - */ -function extractCJKBigrams(text: string): string[] { - const bigrams: string[] = []; - // Extract continuous CJK segments, then generate bigrams within each - const segments = splitCJKSegments(text); - for (const seg of segments) { - for (let i = 0; i + 2 <= seg.length; i++) { - bigrams.push(seg.slice(i, i + 2)); - } - } - return bigrams; -} - -/** - * Split text into continuous CJK character segments. - * Non-CJK characters break the segment. - */ -function splitCJKSegments(text: string): string[] { - const segments: string[] = []; - let current = ''; - for (const ch of text) { - if (isCJKChar(ch)) { - current += ch; - } else { - if (current.length >= 2) segments.push(current); - current = ''; - } - } - if (current.length >= 2) segments.push(current); - return segments; -} - -/** - * Check if a character is CJK. - */ -function isCJKChar(ch: string): boolean { - const cp = ch.codePointAt(0)!; - return ( - (cp >= 0x4e00 && cp <= 0x9fff) || - (cp >= 0x3400 && cp <= 0x4dbf) || - (cp >= 0x3040 && cp <= 0x30ff) || - (cp >= 0xac00 && cp <= 0xd7af) - ); -} diff --git a/src/storage/store.ts b/src/storage/store.ts index b7bafb9..e25a927 100644 --- a/src/storage/store.ts +++ b/src/storage/store.ts @@ -3,11 +3,9 @@ import * as path from 'path'; import { createZvecStore, openZvecStoreSync, - isZvecAvailable, } from './zvec-store'; -import type { IZvecStore, ZvecStoreConfig, FtsFieldWeight, ActualZvecStoreOptions } from './zvec-store'; -import type { ZvecDoc, ZvecQueryResult, ZvecSearchParams, ZvecHybridParams } from './types'; -import { MemoryZvecStore } from './memory-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'; @@ -20,24 +18,6 @@ const DEFAULT_VECTOR_FIELD = 'embedding'; const DEFAULT_FTS_FIELDS = ['content']; const DEFAULT_RANK_CONSTANT = 60; -function resolveFtsFields(options?: ContextOptions): string[] { - return options?.ftsFields ?? DEFAULT_FTS_FIELDS; -} - -function resolveFtsWeights(options?: ContextOptions): FtsFieldWeight[] { - const ftsFields = resolveFtsFields(options); - const weightMap = options?.ftsFieldWeights; - - if (weightMap) { - return ftsFields.map((fieldName) => ({ - fieldName, - weight: weightMap[fieldName] ?? 1.0, - })); - } - - return ftsFields.map((fieldName) => ({ fieldName, weight: 1.0 })); -} - function resolveTokenizer(sampleText?: string): string { // Always auto-detect based on sample text content. // Falls back to 'jieba' as safe default for mixed-language content @@ -73,7 +53,7 @@ function contextStoreConfig(dims: number, sampleText?: string): ZvecStoreConfig function storeOpenOptions(options?: ContextOptions): ActualZvecStoreOptions { return { vectorField: DEFAULT_VECTOR_FIELD, - ftsFields: resolveFtsFields(options), + ftsFields: options?.ftsFields ?? DEFAULT_FTS_FIELDS, rankConstant: options?.rankConstant ?? DEFAULT_RANK_CONSTANT, }; } @@ -112,7 +92,6 @@ export interface StoreQueryParams { export class Store { private readonly vectorsDir: string; private readonly embedder: Embedder; - private readonly ftsWeights: FtsFieldWeight[]; private readonly rankConstant: number; private readonly contextOptions?: ContextOptions; private readonly stores: Map = new Map(); @@ -122,7 +101,6 @@ export class Store { constructor(vectorsDir: string, embedder: Embedder, options?: ContextOptions) { this.vectorsDir = vectorsDir; this.embedder = embedder; - this.ftsWeights = resolveFtsWeights(options); this.rankConstant = options?.rankConstant ?? DEFAULT_RANK_CONSTANT; this.contextOptions = options; } @@ -254,44 +232,15 @@ export class Store { private async _doCreate(library: string, sampleText?: string): Promise { const filePath = this._getStorePath(library); - const allowFallback = this.contextOptions?.allowMemoryFallback ?? false; if (fs.existsSync(filePath)) { - if (isZvecAvailable()) { - return openZvecStoreSync(filePath, storeOpenOptions(this.contextOptions)); - } - if (allowFallback) { - console.warn( - `[context] @zvec/zvec not available — opening library "${library}" ` + - `with in-memory MemoryZvecStore. Data will NOT be persisted.` - ); - return new MemoryZvecStore(this.ftsWeights, this.rankConstant); - } - throw new Error( - `Cannot open existing store "${filePath}": @zvec/zvec is not installed. ` + - `Install it with: pnpm add @zvec/zvec\n` + - `Or set allowMemoryFallback: true to use in-memory store (no persistence).` - ); + return openZvecStoreSync(filePath, storeOpenOptions(this.contextOptions)); } - try { - return await createZvecStore( - filePath, - contextStoreConfig(this.embedder.dimensions, sampleText), - this.ftsWeights, - this.rankConstant, - ); - } catch (err) { - if (allowFallback) { - console.warn( - `[context] Store creation failed for library "${library}", ` + - `falling back to in-memory MemoryZvecStore. ` + - `Error: ${(err as Error).message?.split('\n')[0]}` - ); - return new MemoryZvecStore(this.ftsWeights, this.rankConstant); - } - throw err; - } + return await createZvecStore( + filePath, + contextStoreConfig(this.embedder.dimensions, sampleText), + ); } /** @@ -300,10 +249,15 @@ export class Store { */ private _tryOpenFromDisk(library: string): IZvecStore | undefined { const filePath = this._getStorePath(library); - if (fs.existsSync(filePath) && isZvecAvailable()) { - const store = openZvecStoreSync(filePath, storeOpenOptions(this.contextOptions)); - this.stores.set(library, store); - return store; + 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; + } } return undefined; } diff --git a/src/storage/types.ts b/src/storage/types.ts index ef24060..f2faecb 100644 --- a/src/storage/types.ts +++ b/src/storage/types.ts @@ -43,12 +43,6 @@ export interface IZvecStore { close(): Promise; } -/** Weighted field for MemoryZvecStore FTS text scoring in hybrid search. */ -export interface FtsFieldWeight { - fieldName: string; - weight: number; -} - /** Field schema entry for building a zvec collection. */ export interface ZvecFieldSchema { name: string; diff --git a/src/storage/utils.ts b/src/storage/utils.ts index e28504d..c7aa273 100644 --- a/src/storage/utils.ts +++ b/src/storage/utils.ts @@ -15,54 +15,3 @@ export function cosineSimilarity(a: number[], b: number[]): number { const denom = Math.sqrt(na) * Math.sqrt(nb); return denom === 0 ? 0 : dot / denom; } - -/** - * Evaluate a filter expression against document fields. - * - * Supports: - * - String equality: `field = 'value'` - * - Number equality: `field = 0` - * - AND conjunction: `field = 'val1' AND field2 = 'val2'` - * - * Unknown / malformed expressions return true (pass-through) so that - * MemoryZvecStore remains usable even when filters are only meant for - * the native zvec engine. - */ -export function evalMemoryFilter( - filter: string, - fields: Record -): boolean { - // Split on AND (case-insensitive, with surrounding spaces) - const clauses = filter.split(/\s+AND\s+/i); - - return clauses.every((clause) => evalSingleClause(clause.trim(), fields)); -} - -/** - * Evaluate a single filter clause. - * - * Forms: - * `fieldName = 'stringVal'` → string equality - * `fieldName = numericVal` → number equality (coerces field to Number) - */ -function evalSingleClause( - clause: string, - fields: Record -): boolean { - // String equality: field = 'value' - const stringMatch = clause.match(/^(\w+)\s*=\s*'([^']*)'$/); - if (stringMatch) { - const [, fieldName, expected] = stringMatch; - return String(fields[fieldName] ?? '') === expected; - } - - // Number equality: field = 0 (or any integer/float) - const numberMatch = clause.match(/^(\w+)\s*=\s*(-?\d+(?:\.\d+)?)$/); - if (numberMatch) { - const [, fieldName, expected] = numberMatch; - return Number(fields[fieldName]) === Number(expected); - } - - // Unknown clause format — pass through (let native zvec handle it) - return true; -} diff --git a/src/storage/zvec-store.ts b/src/storage/zvec-store.ts index 153f2fe..cc8b07a 100644 --- a/src/storage/zvec-store.ts +++ b/src/storage/zvec-store.ts @@ -2,10 +2,9 @@ * zvec-store — aggregate entry point for all storage modules. * * Re-exports from the split files for backward compatibility: - * types.ts → public types & interfaces - * memory-store.ts → MemoryZvecStore + * types.ts → public types & interfaces * actual-store.ts → ActualZvecStore, factories, schema builder - * utils.ts → cosineSimilarity, evalMemoryFilter + * utils.ts → cosineSimilarity */ export type { @@ -14,14 +13,11 @@ export type { ZvecSearchParams, ZvecHybridParams, IZvecStore, - FtsFieldWeight, ZvecFieldSchema, ZvecStoreConfig, ActualZvecStoreOptions, } from './types'; -export { MemoryZvecStore } from './memory-store'; - export { ActualZvecStore, createZvecStore, diff --git a/src/types.ts b/src/types.ts index b21c5f4..97e37f8 100644 --- a/src/types.ts +++ b/src/types.ts @@ -144,21 +144,6 @@ export interface ContextOptions { * ``` */ rerankWeights?: Omit; - - /** - * Allow automatic fallback to in-memory store when @zvec/zvec is unavailable. - * - * When `true`, if @zvec/zvec cannot be loaded, the library will silently - * fall back to `MemoryZvecStore` (no persistence, vectors lost on restart). - * A warning is logged to console. - * - * When `false` (default), missing @zvec/zvec throws an error at store - * creation time, making the failure explicit. - * - * Useful for development/testing environments where native bindings - * may not be available. - */ - allowMemoryFallback?: boolean; } /** @@ -262,7 +247,6 @@ export type { ZvecQueryResult, ZvecSearchParams, ZvecHybridParams, - FtsFieldWeight, ZvecFieldSchema, ZvecStoreConfig, } from './storage/zvec-store'; diff --git a/test/context.test.ts b/test/context.test.ts index 33ecc6a..ce60855 100644 --- a/test/context.test.ts +++ b/test/context.test.ts @@ -286,28 +286,4 @@ describe('Context with weight configuration', () => { await ctx.close(); }); - - it('should use hybrid search with MemoryZvecStore when zvec unavailable', async () => { - // When zvec is not available, MemoryZvecStore is used with FtsFieldWeights - const ctx = await Context.create({ - vectorsDir: weightTestDir + '-mem', - ftsFields: ['content'], - ftsFieldWeights: { content: 3 }, - }); - - await ctx.load('mem-weighted', path.join(FIXTURES_DIR, 'getting-started.md')); - - // Hybrid mode (default) should use both vector + text path - const hybridResults = await ctx.query('Getting Started', { library: 'mem-weighted', topK: 1 }); - expect(hybridResults.length).toBeGreaterThan(0); - - // Vector-only mode should work too - const vectorResults = await ctx.query('Getting Started', { library: 'mem-weighted', topK: 1, mode: 'vector' }); - expect(vectorResults.length).toBeGreaterThan(0); - - await ctx.close(); - if (fs.existsSync(weightTestDir + '-mem')) { - fs.rmSync(weightTestDir + '-mem', { recursive: true, force: true }); - } - }); }); \ No newline at end of file diff --git a/test/utils.test.ts b/test/utils.test.ts index 0f15aee..c7ff7af 100644 --- a/test/utils.test.ts +++ b/test/utils.test.ts @@ -1,5 +1,5 @@ import { describe, it, expect } from 'vitest'; -import { cosineSimilarity, evalMemoryFilter } from '../src/storage/utils'; +import { cosineSimilarity } from '../src/storage/utils'; describe('cosineSimilarity', () => { it('should return 1 for identical vectors', () => { @@ -26,58 +26,4 @@ describe('cosineSimilarity', () => { expect(sim).toBeGreaterThan(0); expect(sim).toBeLessThan(1); }); -}); - -describe('evalMemoryFilter', () => { - const fields = { - content: 'hello world', - sourceFilePath: 'docs/getting-started.md', - }; - - describe('string equality', () => { - it('should match string field with quoted value', () => { - expect(evalMemoryFilter("sourceFilePath = 'docs/getting-started.md'", fields)).toBe(true); - }); - - it('should reject non-matching string value', () => { - expect(evalMemoryFilter("sourceFilePath = 'xyz'", fields)).toBe(false); - }); - - it('should match content field', () => { - expect(evalMemoryFilter("content = 'hello world'", fields)).toBe(true); - }); - }); - - describe('AND conjunction', () => { - it('should match when all clauses are true', () => { - expect(evalMemoryFilter("content = 'hello world' AND sourceFilePath = 'docs/getting-started.md'", fields)).toBe(true); - }); - - it('should reject when any clause is false', () => { - expect(evalMemoryFilter("content = 'hello' AND sourceFilePath = 'docs/getting-started.md'", fields)).toBe(false); - }); - - it('should reject when all clauses are false', () => { - expect(evalMemoryFilter("content = 'foo' AND sourceFilePath = 'xyz'", fields)).toBe(false); - }); - }); - - describe('edge cases', () => { - it('should pass through unknown clause format', () => { - // Unknown format → true (let native zvec handle it) - expect(evalMemoryFilter("field > 10", fields)).toBe(true); - }); - - it('should pass through empty filter', () => { - expect(evalMemoryFilter('', fields)).toBe(true); - }); - - it('should handle AND with case-insensitive keyword', () => { - expect(evalMemoryFilter("content = 'hello world' and sourceFilePath = 'docs/getting-started.md'", fields)).toBe(true); - }); - - it('should handle missing field gracefully (string clause)', () => { - expect(evalMemoryFilter("unknownField = 'value'", fields)).toBe(false); - }); - }); -}); +}); \ No newline at end of file diff --git a/test/zvec-store.test.ts b/test/zvec-store.test.ts index 1fcff42..1370cd1 100644 --- a/test/zvec-store.test.ts +++ b/test/zvec-store.test.ts @@ -2,7 +2,6 @@ import { describe, it, expect, beforeEach, afterEach } from 'vitest'; import * as fs from 'fs'; import * as path from 'path'; import { - MemoryZvecStore, createZvecStore, isZvecAvailable, } from '../src/storage/zvec-store'; @@ -23,84 +22,6 @@ function makeConfig(dims: number = 4): ZvecStoreConfig { }; } -describe('MemoryZvecStore', () => { - let store: MemoryZvecStore; - - beforeEach(() => { - store = new MemoryZvecStore(); - }); - - it('should insert and search documents', async () => { - await store.insert([ - { id: 'doc1', vector: [0.1, 0.2, 0.3, 0.4], fields: { content: 'hello world' } }, - { id: 'doc2', vector: [0.4, 0.3, 0.2, 0.1], fields: { content: 'goodbye' } }, - ]); - - const results = await store.search({ vector: [0.1, 0.2, 0.3, 0.4], topK: 2 }); - expect(results.length).toBe(2); - expect(results[0].id).toBe('doc1'); - expect(results[0].score).toBeGreaterThan(0.9); - }); - - it('should support sync search', () => { - store.insert([ - { id: 'doc1', vector: [0.1, 0.2, 0.3, 0.4], fields: { content: 'hello' } }, - ]); - - const results = store.searchSync({ vector: [0.1, 0.2, 0.3, 0.4], topK: 1 }); - expect(results.length).toBe(1); - expect(results[0].id).toBe('doc1'); - }); - - it('should support filter expressions', async () => { - await store.insert([ - { id: 'a', vector: [1, 0, 0, 0], fields: { kind: 'foo' } }, - { id: 'b', vector: [0, 1, 0, 0], fields: { kind: 'bar' } }, - ]); - - const results = await store.search({ - vector: [1, 0, 0, 0], - topK: 2, - filter: "kind = 'foo'", - }); - expect(results.length).toBe(1); - expect(results[0].id).toBe('a'); - }); - - it('should support hybrid search with configurable field weights', async () => { - const weighted = new MemoryZvecStore([ - { fieldName: 'title', weight: 3 }, - { fieldName: 'content', weight: 1 }, - ]); - - await weighted.insert([ - { id: 't1', vector: [0.1, 0, 0, 0], fields: { title: 'sankey diagram', content: 'flow data' } }, - { id: 't2', vector: [0, 0.1, 0, 0], fields: { title: 'unrelated', content: 'sankey flow visualization' } }, - ]); - - const results = await weighted.searchHybrid({ - queryText: 'sankey', - queryVector: [0.1, 0.1, 0, 0], - topK: 2, - }); - expect(results.length).toBeGreaterThan(0); - }); - - it('should handle empty store gracefully', async () => { - const results = await store.search({ vector: [1, 0, 0, 0], topK: 5 }); - expect(results.length).toBe(0); - }); - - it('should clear all data on close', async () => { - await store.insert([ - { id: 'doc1', vector: [1, 0, 0, 0], fields: {} }, - ]); - await store.close(); - const results = await store.search({ vector: [1, 0, 0, 0], topK: 5 }); - expect(results.length).toBe(0); - }); -}); - describe('ActualZvecStore (native)', () => { let store: IZvecStore; @@ -162,21 +83,4 @@ describe('ActualZvecStore (native)', () => { }); expect(results.length).toBeGreaterThan(0); }); - - it('should fall back to MemoryZvecStore when zvec not available', async () => { - // When zvec IS available, createZvecStore creates an ActualZvecStore. - // Either way, the returned store should work. - const config = makeConfig(4); - const s = await createZvecStore(STORE_PATH, config); - expect(s).toBeDefined(); - - await s.insert([ - { id: 'x', vector: [1, 2, 3, 4], fields: { content: 'test content' } }, - ]); - const results = s.searchSync({ vector: [1, 2, 3, 4], topK: 1 }); - expect(results.length).toBe(1); - expect(results[0].id).toBe('x'); - - await s.close(); - }); }); From 593101c63f8636108b31e2902f8a02ef31d34707 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A6=8F=E6=99=8B?= Date: Mon, 6 Jul 2026 21:49:58 +0800 Subject: [PATCH 11/13] chore: remove registry --- src/context.ts | 100 +++++++++---------------- src/index.ts | 3 - src/registry.ts | 144 ------------------------------------ src/storage/actual-store.ts | 26 ++++++- src/storage/store.ts | 29 +++++++- src/storage/types.ts | 7 ++ src/types.ts | 2 +- src/utils/common.ts | 2 +- 8 files changed, 93 insertions(+), 220 deletions(-) delete mode 100644 src/registry.ts diff --git a/src/context.ts b/src/context.ts index 6df2b5c..cb4e5a8 100644 --- a/src/context.ts +++ b/src/context.ts @@ -12,7 +12,6 @@ import { import { resolveEmbedder } from './embedder'; import type { Embedder, EmbedderInfo } from './embedder'; import { Loader, MarkdownLoader, JsonLoader, TextLoader } from './loaders'; -import { DocumentRegistry } from './registry'; import { Store } from './storage/store'; import type { ZvecDoc } from './storage/zvec-store'; import { pathToId } from './loaders/util'; @@ -31,24 +30,20 @@ import type { Reranker, RerankCandidate, QueryExpander } from './utils'; // --------------------------------------------------------------------------- export class Context { - private readonly vectorsDir: string; private readonly basePath: string; private readonly embedder: Embedder; private readonly _embedderInfo: EmbedderInfo; private readonly store: Store; - private readonly registry: DocumentRegistry; 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.vectorsDir = options.vectorsDir; this.basePath = options.basePath ?? process.cwd(); this.embedder = embedder; this._embedderInfo = embedderInfo; this.store = new Store(options.vectorsDir, embedder, options); - this.registry = new DocumentRegistry(); this.loaders = [new MarkdownLoader(), new JsonLoader(), new TextLoader()]; this.reranker = createReranker(options.rerankWeights); this.queryExpander = @@ -74,21 +69,6 @@ export class Context { const ctx = new Context(options, embedder, embedderInfo); - // Auto-recover registry from existing index files on disk. - // When a process restarts, the in-memory registry is empty, but the - // zvec store and `.index.json` files persist. Scanning for these - // files and loading them into the registry prevents duplicate - // re-embedding of unchanged documents. - if (fs.existsSync(options.vectorsDir)) { - const indexFiles = fs - .readdirSync(options.vectorsDir) - .filter((f) => f.endsWith('.index.json')); - for (const indexFile of indexFiles) { - const library = indexFile.replace('.index.json', ''); - ctx.registry.loadFromDisk(options.vectorsDir, library); - } - } - return ctx; } @@ -106,12 +86,9 @@ export class Context { /** * Load documents into a library with automatic vectorization. * - * Documents that have already been loaded (same id) are skipped to - * prevent duplicate vectors in the store. Uses batch embedding for - * better performance when loading many files at once. - * - * Document IDs are derived from the file path relative to `basePath`, - * ensuring the same document gets the same ID across different machines. + * 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. @@ -121,10 +98,6 @@ export class Context { const files = await glob(patterns, { absolute: true }); // Sample multiple files for FTS tokenizer auto-detection. - // Only used when `tokenizer` is 'auto' and the store does not yet exist. - // Sampling up to 5 files from different positions in the file list - // ensures we don't pick the wrong tokenizer when the first file is - // an English README but most content is Chinese. let sampleText: string | undefined; if (files.length > 0) { try { @@ -143,55 +116,35 @@ export class Context { } } - const openedStore = await this.store.create(library, sampleText); + await this.store.create(library, sampleText); - // Load registry from disk if not already loaded for this library - if (!this.registry.hasLibrary(library)) { - this.registry.loadFromDisk(this.vectorsDir, library); + // Internal type that extends Document with load-phase metadata. + interface LoadedDoc extends Document { + id: string; + contentHash: string; + sourceFilePath: string; } - // Phase 1: Load all documents concurrently and filter out duplicates. - // Each file is an independent I/O operation — running them in parallel - // eliminates the serial disk-read bottleneck. We use allSettled so - // one broken file does not kill the entire batch. - // - // Change detection: files whose content hash differs from the stored - // hash are re-embedded (content was updated since last load). + // Phase 1: Load all files concurrently and collect candidates. const loadSettled = await Promise.allSettled( files.map(async (filePath) => { const loader = this.getLoader(filePath); if (!loader) return null; const doc = await loader.load(filePath); - - // Derive ID from relative path for cross-machine consistency const relativePath = path.relative(this.basePath, filePath); const docId = pathToId(relativePath); - - // Compute content hash for change detection const contentHash = computeContentHash(doc.content); - // Deduplication: skip documents whose content hasn't changed. - // If the hash differs, the file was updated — re-embed it. - if (this.registry.has(library, docId, contentHash)) return null; - return { ...doc, id: docId, contentHash, sourceFilePath: relativePath }; }), ); - // Internal type that extends Document with load-phase metadata. - // id is required here — Context.load() always assigns it via pathToId. - interface LoadedDoc extends Document { - id: string; // override optional Document.id → required - contentHash?: string; - sourceFilePath?: string; - } - - const docsToEmbed: LoadedDoc[] = []; + const candidates: LoadedDoc[] = []; let failCount = 0; for (const r of loadSettled) { if (r.status === 'fulfilled' && r.value !== null) { - docsToEmbed.push(r.value); + candidates.push(r.value); } else if (r.status === 'rejected') { failCount++; } @@ -202,6 +155,26 @@ export class Context { ); } + 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 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) + } + // Progress: load phase complete if (this._onProgress) { this._onProgress('load', { loaded: docsToEmbed.length, total: files.length }); @@ -218,7 +191,7 @@ export class Context { this._onProgress('embed', { loaded: vectors.length, total: docsToEmbed.length }); } - // Phase 3: Batch insert into store + // 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], @@ -226,6 +199,7 @@ export class Context { content: doc.content, meta: doc.meta && Object.keys(doc.meta).length > 0 ? JSON.stringify(doc.meta) : '', sourceFilePath: doc.sourceFilePath ?? '', + contentHash: doc.contentHash, }, })); @@ -235,12 +209,6 @@ export class Context { if (this._onProgress) { this._onProgress('insert', { loaded: zvecDocs.length, total: docsToEmbed.length }); } - - // Phase 4: Update registry (track by parent doc ID for deduplication + change detection) - for (const doc of docsToEmbed) { - this.registry.add(library, doc.id, doc.contentHash); - } - this.registry.saveToDisk(this.vectorsDir, library); } /** diff --git a/src/index.ts b/src/index.ts index f0f3873..0567fd0 100644 --- a/src/index.ts +++ b/src/index.ts @@ -62,9 +62,6 @@ export type { ActualZvecStoreOptions, } from './storage/zvec-store'; -// DocumentRegistry — dedup tracking -export { DocumentRegistry } from './registry'; - // Store — zvec store lifecycle management export { Store } from './storage/store'; export type { StoreQueryParams } from './storage/store'; diff --git a/src/registry.ts b/src/registry.ts deleted file mode 100644 index 861cd2a..0000000 --- a/src/registry.ts +++ /dev/null @@ -1,144 +0,0 @@ -import * as fs from 'fs'; -import * as path from 'path'; - -/** - * DocumentRegistry — manages document metadata for each library. - * - * Tracks loaded document IDs along with their content hashes for - * change-detection. When a document's content hash changes (e.g. the - * file was edited), the dedup check will allow it to be re-embedded. - */ - -/** Registry entry: document ID + content hash for change detection. */ -interface RegistryEntry { - id: string; - /** SHA-256 hash of the document content (first 16 chars). Empty = legacy entry without hash. */ - contentHash: string; -} - -export class DocumentRegistry { - private readonly entries: Map> = new Map(); - - /** - * Check whether a document ID has already been loaded into a library, - * optionally verifying that the content hash matches. - * - * - If `contentHash` is provided and the existing entry has a hash, - * returns `true` only when both the ID and hash match (unchanged file). - * - If `contentHash` is provided but differs from the stored hash, - * returns `false` (file has been updated — needs re-embedding). - * - If `contentHash` is not provided, returns `true` when the ID exists - * (legacy behaviour — no change detection). - */ - has(library: string, id: string, contentHash?: string): boolean { - const lib = this.entries.get(library); - if (!lib) return false; - - const entry = lib.get(id); - if (!entry) return false; - - // If hash is provided, do change detection - if (contentHash !== undefined && entry.contentHash) { - return entry.contentHash === contentHash; - } - - // No hash provided or stored — just check existence - return true; - } - - /** - * Register a document as loaded for a library, with its content hash. - */ - add(library: string, id: string, contentHash?: string): void { - let lib = this.entries.get(library); - if (!lib) { - lib = new Map(); - this.entries.set(library, lib); - } - lib.set(id, { id, contentHash: contentHash ?? '' }); - } - - /** - * Get all loaded document IDs for a library. - */ - getIds(library: string): Set { - const lib = this.entries.get(library); - return lib ? new Set(lib.keys()) : new Set(); - } - - /** - * Get all registry entries for a library (including content hashes). - */ - getEntries(library: string): Map { - return this.entries.get(library) ?? new Map(); - } - - /** - * Check whether a library has any loaded documents. - */ - hasLibrary(library: string): boolean { - return (this.entries.get(library)?.size ?? 0) > 0; - } - - /** - * Remove a specific document from a library's registry. - */ - remove(library: string, id: string): void { - this.entries.get(library)?.delete(id); - } - - /** - * Get all library names that have loaded documents. - */ - getLibraryNames(): string[] { - return [...this.entries.keys()]; - } - - /** - * Remove all documents for a library. - */ - removeLibrary(library: string): void { - this.entries.delete(library); - } - - /** - * Persist registry state to disk (for process restart recovery). - * - * Format: `{ id, contentHash }` entries for change detection across - * restarts. - */ - saveToDisk(vectorsDir: string, library: string): void { - const indexPath = path.join(vectorsDir, `${library}.index.json`); - const lib = this.entries.get(library); - if (!lib) { - // No entries — write empty array - fs.writeFileSync(indexPath, JSON.stringify([])); - return; - } - const data = [...lib.values()]; - fs.writeFileSync(indexPath, JSON.stringify(data)); - } - - /** - * Load registry state from disk. - */ - loadFromDisk(vectorsDir: string, library: string): void { - const indexPath = path.join(vectorsDir, `${library}.index.json`); - if (fs.existsSync(indexPath)) { - const raw: RegistryEntry[] | string[] = JSON.parse(fs.readFileSync(indexPath, 'utf-8')); - const lib = new Map(); - - for (const item of raw) { - if (typeof item === 'string') { - // Legacy format: plain ID strings (no content hash) - lib.set(item, { id: item, contentHash: '' }); - } else { - // New format: { id, contentHash } entries - lib.set(item.id, item); - } - } - - this.entries.set(library, lib); - } - } -} diff --git a/src/storage/actual-store.ts b/src/storage/actual-store.ts index 3424a67..e9dfaac 100644 --- a/src/storage/actual-store.ts +++ b/src/storage/actual-store.ts @@ -33,6 +33,8 @@ interface ZvecSDK { /** 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; @@ -212,7 +214,29 @@ export class ActualZvecStore implements IZvecStore { vectors: { [this._vectorField]: d.vector }, fields: d.fields })); - this._collection.insertSync(records); + // 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 { diff --git a/src/storage/store.ts b/src/storage/store.ts index e25a927..f12f065 100644 --- a/src/storage/store.ts +++ b/src/storage/store.ts @@ -46,6 +46,7 @@ function contextStoreConfig(dims: number, sampleText?: string): ZvecStoreConfig }, { name: 'meta', dataType: 'STRING' }, { name: 'sourceFilePath', dataType: 'STRING' }, + { name: 'contentHash', dataType: 'STRING' }, ], }; } @@ -92,7 +93,6 @@ export interface StoreQueryParams { export class Store { private readonly vectorsDir: string; private readonly embedder: Embedder; - private readonly rankConstant: number; private readonly contextOptions?: ContextOptions; private readonly stores: Map = new Map(); /** In-flight creation promises — prevents duplicate stores from concurrent calls. */ @@ -101,7 +101,6 @@ export class Store { constructor(vectorsDir: string, embedder: Embedder, options?: ContextOptions) { this.vectorsDir = vectorsDir; this.embedder = embedder; - this.rankConstant = options?.rankConstant ?? DEFAULT_RANK_CONSTANT; this.contextOptions = options; } @@ -140,9 +139,10 @@ export class Store { } /** - * Batch-insert documents into a library's store. + * Batch-insert documents into a library's store (upsert semantics). * - * The store must have been created via `create()` first. + * 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). @@ -155,6 +155,27 @@ export class Store { await store.insert(docs); } + /** + * 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); + } + /** * Query a library's store for similar documents. * diff --git a/src/storage/types.ts b/src/storage/types.ts index f2faecb..0387fc8 100644 --- a/src/storage/types.ts +++ b/src/storage/types.ts @@ -32,6 +32,13 @@ export interface ZvecHybridParams { 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). */ diff --git a/src/types.ts b/src/types.ts index 97e37f8..12873fb 100644 --- a/src/types.ts +++ b/src/types.ts @@ -12,7 +12,7 @@ export interface Document { meta?: Record; /** * SHA-256 content hash (16 chars) — used internally for change detection. - * Populated by `Context.load()` before registry insertion; not stored in zvec. + * Populated by `Context.load()` and stored in zvec as the single source of truth. */ contentHash?: string; } diff --git a/src/utils/common.ts b/src/utils/common.ts index 5ee2e8b..59becc5 100644 --- a/src/utils/common.ts +++ b/src/utils/common.ts @@ -24,7 +24,7 @@ export function safeParseMeta(metaStr: string | undefined): Record Date: Mon, 6 Jul 2026 21:50:27 +0800 Subject: [PATCH 12/13] chore: update readme --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index b33a989..e461b0d 100644 --- a/README.md +++ b/README.md @@ -213,7 +213,7 @@ await ctx.close(); - **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`, `DocumentRegistry`, `Store` +- **Advanced API**: `Embedder`, `TransformersEmbedder`, `EmbedderManager`, `IZvecStore`, `ActualZvecStore`, `Store` ## License From 94d674969d17dfe2326ab9dad09f40767fb64fb4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A6=8F=E6=99=8B?= Date: Tue, 7 Jul 2026 14:22:48 +0800 Subject: [PATCH 13/13] chore: update hook timeout --- package.json | 2 +- vitest.config.ts | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/package.json b/package.json index b264da5..cb85fa1 100644 --- a/package.json +++ b/package.json @@ -6,7 +6,7 @@ "types": "dist/index.d.ts", "scripts": { "build": "tsc", - "test": "vitest run --coverage", + "test": "HF_ENDPOINT=https://hf-mirror.com vitest run --coverage", "lint": "eslint src/ --ext .ts", "lint:fix": "eslint src/ --ext .ts --fix", "format": "prettier --write 'src/**/*.ts' 'test/**/*.ts'" diff --git a/vitest.config.ts b/vitest.config.ts index 5f2242a..75d5680 100644 --- a/vitest.config.ts +++ b/vitest.config.ts @@ -4,7 +4,7 @@ export default defineConfig({ test: { include: ['test/**/*.test.ts'], testTimeout: 30000, - hookTimeout: 10000, + hookTimeout: 30000, pool: 'forks', poolOptions: { forks: { singleFork: true },