refactor: update reranker

Author: hustcc

README.md

@@ -208,14 +208,6 @@ await ctx.close();
 +------------------------------------------------------------------------+
 ```
 
-### Module Structure
-
-- **Public API**: `Context`, `QueryOptions`, `QueryResult`, `Document`, `Loader`, `MarkdownLoader`, `JsonLoader`, `TextLoader`, `pathToId`
-- **Reranking**: `KeywordReranker`, `createReranker`, `Reranker`, `RerankCandidate`, `RerankResult`, `RerankOptions`
-- **Query Expansion**: `SynonymExpander`, `NoopExpander`, `QueryExpander`, `QueryExpansionOptions`
-- **Advanced API**: `Embedder`, `TransformersEmbedder`, `EmbedderManager`, `IZvecStore`, `ActualZvecStore`, `Store`
-
-
 ## License
 
 MIT

src/context.ts

@@ -5,10 +5,7 @@ import {
   ContextOptions,
   QueryOptions,
   QueryResult,
-  Document,
   LoadedDoc,
-  LoadPhase,
-  LoadProgress,
 } from './types';
 import { Embedder } from './embedder';
 
@@ -18,24 +15,18 @@ import { pathToId } from './utils/doc';
 import { Store } from './storage/store';
 import type { ZvecDoc } from './storage/zvec-store';
 import {
-  createReranker,
   safeJsonParse,
   computeContentHash,
   loadSampleText,
 } from './utils';
 import { expand } from './expander';
-import type { Reranker, RerankCandidate } from './utils';
-
-// ---------------------------------------------------------------------------
-// Context class
-// ---------------------------------------------------------------------------
+import { createReranker, applyRerank } from './reranker';
 
 export class Context {
   private readonly options: ContextOptions;
   private readonly embedder: Embedder;
   readonly embedderInfo: EmbedderInfo;
   private readonly store: Store;
-  private readonly reranker: Reranker | null;
 
   private constructor(options: ContextOptions, embedder: Embedder, embedderInfo: EmbedderInfo) {
     this.options = {
@@ -46,7 +37,6 @@ export class Context {
     this.embedder = embedder;
     this.embedderInfo = embedderInfo;
     this.store = new Store(this.options.vectorsDir!, embedder, this.options);
-    this.reranker = createReranker(this.options.rerankWeights);
   }
 
   static async create(options: ContextOptions): Promise<Context> {
@@ -163,23 +153,9 @@ export class Context {
       };
     });
 
-    if (rerankEnabled && queryResults.length > topK) {
-      const docs: RerankCandidate[] = queryResults.map((r) => ({
-        id: r.id,
-        content: r.content,
-        score: r.score,
-      }));
-
-      const reranked = await this.reranker!.rerank(text, docs);
-
-      const scoreMap = new Map(reranked.map((r) => [r.id, r.score]));
-      for (const result of queryResults) {
-        const newScore = scoreMap.get(result.id);
-        if (newScore !== undefined) {
-          result.score = newScore;
-          result.scoreMode = 'reranked';
-        }
-      }
+    if (rerankEnabled) {
+      const reranker = createReranker(this.options.rerankWeights);
+      await applyRerank(reranker, text, queryResults, topK);
     }
 
     queryResults.sort((a, b) => b.score - a.score);

src/index.ts

@@ -1,16 +1,2 @@
-// ---------------------------------------------------------------------------
-// Public API
-// ---------------------------------------------------------------------------
-
 export { Context } from './context';
-export * from './types';
-
-// Reranker
-export { KeywordReranker, createReranker } from './utils/reranker';
-export type { RerankCandidate, RerankResult } from './utils/reranker';
-
-// Query expansion
-export { expand } from './expander';
-
-// Embedder
-export { Embedder } from './embedder';
+export * from './types';

src/reranker/index.ts

@@ -0,0 +1,13 @@
+import type { Reranker, RerankCandidate, RerankResult, RerankOptions } from './types';
+import { KeywordReranker } from './keyword-reranker';
+import { applyRerank } from './rerank';
+
+export type { Reranker, RerankCandidate, RerankResult, RerankOptions };
+export { KeywordReranker, applyRerank };
+
+/**
+ * Create a reranker with optional weight configuration.
+ */
+export function createReranker(options?: RerankOptions): Reranker {
+  return new KeywordReranker(options);
+}

src/utils/reranker.ts src/reranker/keyword-reranker.tssrc/utils/reranker.ts renamed to src/reranker/keyword-reranker.ts

@@ -1,51 +1,4 @@
-/**
- * Reranker — second-stage precision scoring for search results.
- *
- * The pipeline:
- *   1. Coarse search (vector / hybrid) → topK × rerankFactor candidates
- *   2. Reranker scores each candidate against the query
- *   3. Final sort by reranked score → topK results
- */
-
-// ---------------------------------------------------------------------------
-// Types
-// ---------------------------------------------------------------------------
-
-/** A candidate document for reranking. */
-export interface RerankCandidate {
-  id: string;
-  content: string;
-  /** Original score from the coarse search stage. */
-  score: number;
-  /** Heading path as a string (e.g. "Line Chart > Tooltip"). */
-  headingPath?: string;
-}
-
-/** A reranked result. */
-export interface RerankResult {
-  id: string;
-  /** Final score after reranking (higher is better). */
-  score: number;
-}
-
-/** Reranker type */
-export type Reranker = {
-  rerank(query: string, candidates: RerankCandidate[]): Promise<RerankResult[]>;
-};
-
-/** Configuration for reranking. */
-export interface RerankOptions {
-  rerankFactor?: number;
-  minCandidates?: number;
-  phraseWeight?: number;
-  phraseRepeatBonus?: number;
-  termWeight?: number;
-  termRepeatBonus?: number;
-  substringWeight?: number;
-  headingTermBonus?: number;
-  headingPhraseBonus?: number;
-  originalScoreCarry?: number;
-}
+import type { Reranker, RerankCandidate, RerankResult, RerankOptions } from './types';
 
 // ---------------------------------------------------------------------------
 // Defaults
@@ -72,7 +25,7 @@ const DEFAULT_WEIGHTS = {
  * All scoring weights are configurable via constructor options.
  * Scores are normalised to [0, 1] via min-max scaling.
  */
-export class KeywordReranker {
+export class KeywordReranker implements Reranker {
   private readonly weights: Record<string, number>;
 
   constructor(options?: RerankOptions) {
@@ -210,15 +163,4 @@ function isWordBoundary(
   const afterIdx = idx + term.length;
   const after = afterIdx >= text.length || /[\s\n.,;:!?，。！？、：；"')）】》\-_]/.test(text[afterIdx]);
   return before && after;
-}
-
-// ---------------------------------------------------------------------------
-// Factory
-// ---------------------------------------------------------------------------
-
-/**
- * Create a reranker with optional weight configuration.
- */
-export function createReranker(options?: RerankOptions): Reranker {
-  return new KeywordReranker(options);
-}
+}

src/reranker/rerank.ts

@@ -0,0 +1,39 @@
+import type { Reranker, RerankCandidate } from './types';
+import type { QueryResult } from '../types';
+
+/**
+ * Apply reranking to query results.
+ *
+ * @param reranker - The reranker instance
+ * @param query - The original query text
+ * @param results - Query results to rerank
+ * @param topK - Original topK (only rerank if results > topK)
+ * @returns The results with scores updated if reranking was applied
+ */
+export async function applyRerank(
+  reranker: Reranker,
+  query: string,
+  results: QueryResult[],
+  topK: number,
+): Promise<QueryResult[]> {
+  if (results.length <= topK) return results;
+
+  const candidates: RerankCandidate[] = results.map((r) => ({
+    id: r.id,
+    content: r.content,
+    score: r.score,
+  }));
+
+  const reranked = await reranker.rerank(query, candidates);
+
+  const scoreMap = new Map(reranked.map((r) => [r.id, r.score]));
+  for (const result of results) {
+    const newScore = scoreMap.get(result.id);
+    if (newScore !== undefined) {
+      result.score = newScore;
+      result.scoreMode = 'reranked';
+    }
+  }
+
+  return results;
+}

src/reranker/types.ts

@@ -0,0 +1,43 @@
+/**
+ * Reranker — second-stage precision scoring for search results.
+ */
+
+// ---------------------------------------------------------------------------
+// Types
+// ---------------------------------------------------------------------------
+
+/** A candidate document for reranking. */
+export interface RerankCandidate {
+  id: string;
+  content: string;
+  /** Original score from the coarse search stage. */
+  score: number;
+  /** Heading path as a string (e.g. "Line Chart > Tooltip"). */
+  headingPath?: string;
+}
+
+/** A reranked result. */
+export interface RerankResult {
+  id: string;
+  /** Final score after reranking (higher is better). */
+  score: number;
+}
+
+/** Reranker type */
+export type Reranker = {
+  rerank(query: string, candidates: RerankCandidate[]): Promise<RerankResult[]>;
+};
+
+/** Configuration for reranking. */
+export interface RerankOptions {
+  rerankFactor?: number;
+  minCandidates?: number;
+  phraseWeight?: number;
+  phraseRepeatBonus?: number;
+  termWeight?: number;
+  termRepeatBonus?: number;
+  substringWeight?: number;
+  headingTermBonus?: number;
+  headingPhraseBonus?: number;
+  originalScoreCarry?: number;
+}

src/utils/index.ts

@@ -4,6 +4,4 @@ export { safeJsonParse } from './common';
 export { loadSampleText } from './sample';
 export { isCJK, detectLanguage, tokenizerForLanguage, detectTokenizer } from './tokenizer';
 export type { LanguageHint } from './tokenizer';
-export { containsCJK } from './str';
-export { KeywordReranker, createReranker } from './reranker';
-export type { Reranker, RerankCandidate, RerankResult } from './reranker';
+export { containsCJK } from './str';

test/reranker.test.ts

@@ -1,6 +1,6 @@
 import { describe, it, expect } from 'vitest';
-import { KeywordReranker, createReranker } from '../src/utils/reranker';
-import type { RerankCandidate } from '../src/utils/reranker';
+import { KeywordReranker, createReranker } from '../src/reranker';
+import type { RerankCandidate } from '../src/reranker';
 
 describe('KeywordReranker', () => {
   const reranker = new KeywordReranker();