@sorane/search 0.2.0

package/src/embeddings.ts

@@ -0,0 +1,79 @@
+import { resolve } from "node:path";
+import { readFileSync, existsSync } from "node:fs";
+
+export interface EmbeddingProvider {
+  readonly dimensions: number;
+  readonly modelId?: string;
+  readonly modelSha256?: string;
+  readonly quant?: string;
+  embed(text: string): Promise<number[]>;
+  embedBatch(texts: string[]): Promise<number[][]>;
+}
+
+export const DOC_PREFIX = "検索文書: ";
+export const QUERY_PREFIX = "検索クエリ: ";
+
+export interface RuriOptions {
+  readonly modelRoot?: string;
+  readonly modelId?: string;
+  readonly dtype?: string;
+}
+
+export class RuriEmbeddings implements EmbeddingProvider {
+  readonly dimensions = 256;
+  readonly modelId: string;
+  readonly modelDir: string;
+  private readonly dtype: string;
+  private readonly modelRoot: string;
+  private extractor: ((text: string, opts: object) => Promise<{ data: Float32Array | number[] }>) | null = null;
+
+  constructor(options: RuriOptions = {}) {
+    this.modelRoot = resolve(options.modelRoot ?? "vendor/models");
+    this.modelId = options.modelId ?? "ruri-v3-30m";
+    this.modelDir = resolve(this.modelRoot, this.modelId);
+    this.dtype = options.dtype ?? "q8";
+  }
+
+  get modelSha256(): string {
+    const p = resolve(this.modelDir, "version.txt");
+    return existsSync(p) ? readFileSync(p, "utf8").trim() : "";
+  }
+
+  get quant(): string {
+    return this.dtype;
+  }
+
+  private async ensure(): Promise<void> {
+    if (this.extractor) return;
+    if (!existsSync(this.modelDir)) {
+      throw new Error(
+        `model not found: ${this.modelDir}\n  run: npm run fetch-model`,
+      );
+    }
+    const { env, pipeline } = await import("@huggingface/transformers");
+    env.allowRemoteModels = false;
+    env.localModelPath = this.modelRoot;
+    env.useBrowserCache = false;
+    this.extractor = (await pipeline("feature-extraction", this.modelId, {
+      dtype: this.dtype as "q8",
+    })) as unknown as typeof this.extractor;
+  }
+
+  async embed(text: string): Promise<number[]> {
+    await this.ensure();
+    const out = await this.extractor!(text, { pooling: "mean", normalize: true });
+    const data = out.data;
+    if (data.length !== this.dimensions) {
+      throw new Error(`dimension mismatch: ${data.length} != ${this.dimensions}`);
+    }
+    return Array.from(data);
+  }
+
+  async embedBatch(texts: string[]): Promise<number[][]> {
+    const out: number[][] = [];
+    for (let i = 0; i < texts.length; i++) {
+      out.push(await this.embed(texts[i]!));
+    }
+    return out;
+  }
+}

package/src/emit-search-assets.ts

@@ -0,0 +1,85 @@
+import { mkdirSync } from "node:fs";
+import { join, resolve } from "node:path";
+import { deriveWebIndex, type WebSearchMode } from "./derive-web-index.ts";
+import { copySearchScript, vendorModel, vendorRuntime } from "./vendor-web.ts";
+
+export interface EmitSearchAssetsOptions {
+  readonly cwd: string;
+  readonly outDir: string;
+  readonly indexPath: string;
+  readonly mode?: WebSearchMode;
+  readonly modelRoot: string;
+  readonly modelId: string;
+  readonly bundleModel?: boolean;
+  readonly assetBaseUrl?: string;
+  readonly sourceToUrl: (source: string) => string;
+  readonly contentDir?: string;
+  readonly machineReadable?: boolean;
+  readonly repoRoot?: string;
+  readonly onProgress?: (message: string) => void;
+}
+
+export interface EmitSearchAssetsResult {
+  readonly written: boolean;
+  readonly chunks: number;
+  readonly bytes: number;
+  readonly mode?: WebSearchMode;
+  readonly model: boolean;
+  readonly runtime: boolean;
+}
+
+export async function emitSearchAssets(
+  opts: EmitSearchAssetsOptions,
+): Promise<EmitSearchAssetsResult> {
+  const log = opts.onProgress ?? (() => {});
+  const mode = opts.mode ?? "fts";
+  const assetsDir = join(opts.outDir, "assets");
+  mkdirSync(assetsDir, { recursive: true });
+
+  const webIdx = await deriveWebIndex(
+    opts.indexPath,
+    join(assetsDir, "search-index.json"),
+    opts.sourceToUrl,
+    mode,
+    {
+      contentDir: opts.contentDir,
+      machineReadable: opts.machineReadable,
+    },
+  );
+
+  if (!webIdx.written) {
+    log(`search-index.json: skipped (no index at ${opts.indexPath})`);
+    return { written: false, chunks: 0, bytes: 0, model: false, runtime: false };
+  }
+
+  const okScript = copySearchScript(opts.outDir, opts.repoRoot);
+  let okModel = false;
+  let okRuntime = false;
+
+  if (mode === "hybrid") {
+    const modelRoot = resolve(opts.cwd, opts.modelRoot);
+    okModel = opts.bundleModel !== false && vendorModel(modelRoot, opts.modelId, opts.outDir);
+    okRuntime = vendorRuntime(opts.outDir, opts.repoRoot);
+    log(
+      `search-index.json: ${webIdx.chunks} chunks, ${(webIdx.bytes / 1024 / 1024).toFixed(1)} MB` +
+        ` [hybrid] (model=${okModel ? "ok" : "missing"}, runtime=${okRuntime ? "ok" : "missing"}, script=${okScript ? "ok" : "missing"})`,
+    );
+    if (opts.assetBaseUrl) {
+      log(`search assets base URL: ${opts.assetBaseUrl}`);
+    }
+  } else {
+    log(
+      `search-index.json: ${webIdx.chunks} chunks, ${(webIdx.bytes / 1024).toFixed(1)} KB` +
+        ` [fts] (script=${okScript ? "ok" : "missing"})`,
+    );
+  }
+
+  return {
+    written: true,
+    chunks: webIdx.chunks,
+    bytes: webIdx.bytes,
+    mode: webIdx.mode,
+    model: okModel,
+    runtime: okRuntime,
+  };
+}

package/src/heading-slug.ts

@@ -0,0 +1,24 @@
+/** 見出しテキスト → アンカー id 用 slug（ページ内重複は連番）。 */
+
+export function slugifyHeading(text: string): string {
+  const base = text
+    .trim()
+    .toLowerCase()
+    .replace(/[*_`~]/g, "")
+    .replace(/\s+/g, "-")
+    .replace(/[^\w぀-ヿ㐀-鿿豈-﫿ｦ-ﾟ-]/g, "")
+    .replace(/-+/g, "-")
+    .replace(/^-+|-+$/g, "");
+  return base.length > 0 ? base : "section";
+}
+
+export class SlugLedger {
+  private readonly used = new Map<string, number>();
+
+  next(text: string): string {
+    const base = slugifyHeading(text);
+    const count = this.used.get(base) ?? 0;
+    this.used.set(base, count + 1);
+    return count === 0 ? base : `${base}-${count + 1}`;
+  }
+}

package/src/incremental.ts

@@ -0,0 +1,40 @@
+import { createHash } from "node:crypto";
+
+export function hashContent(content: string): string {
+  return createHash("sha256").update(content, "utf8").digest("hex");
+}
+
+export interface IncrementalPlan {
+  readonly added: string[];
+  readonly changed: string[];
+  readonly removed: string[];
+  readonly unchanged: string[];
+}
+
+export function planIncremental(
+  disk: Map<string, string>,
+  indexed: Map<string, string>,
+): IncrementalPlan {
+  const added: string[] = [];
+  const changed: string[] = [];
+  const unchanged: string[] = [];
+  const removed: string[] = [];
+
+  for (const [source, hash] of disk) {
+    const prev = indexed.get(source);
+    if (prev === undefined) added.push(source);
+    else if (prev === hash) unchanged.push(source);
+    else changed.push(source);
+  }
+  for (const source of indexed.keys()) {
+    if (!disk.has(source)) removed.push(source);
+  }
+
+  const sort = (xs: string[]) => xs.sort();
+  return {
+    added: sort(added),
+    changed: sort(changed),
+    removed: sort(removed),
+    unchanged: sort(unchanged),
+  };
+}

package/src/index.ts

@@ -0,0 +1,60 @@
+export { chunkDocument, MIN_BODY, MAX_BODY, type Chunk } from "./chunker.ts";
+export { hashContent, planIncremental, type IncrementalPlan } from "./incremental.ts";
+export { slugifyHeading, SlugLedger } from "./heading-slug.ts";
+export {
+  RuriEmbeddings,
+  DOC_PREFIX,
+  QUERY_PREFIX,
+  type EmbeddingProvider,
+  type RuriOptions,
+} from "./embeddings.ts";
+export {
+  IndexStore,
+  SCHEMA_VERSION,
+  type ChunkRow,
+  type MetaFilter,
+  type FtsHit,
+  type VecHit,
+  type Counts,
+  type IndexMeta,
+} from "./store.ts";
+export {
+  buildFtsQuery,
+  makeSnippet,
+  rrfFuse,
+  RRF_K,
+  checkModelMismatch,
+  searchFts,
+  searchHybrid,
+  search,
+  type SearchOptions,
+  type SearchResult,
+} from "./search.ts";
+export { walkMarkdown } from "./walk.ts";
+export { buildSearchIndex, type BuildIndexOptions, type BuildIndexResult } from "./build-index.ts";
+export {
+  buildWebIndex,
+  buildFtsWebIndex,
+  toSnippet,
+  defaultSourceUrl,
+  WEB_INDEX_SCHEMA_VERSION,
+  FTS_WEB_INDEX_SCHEMA_VERSION,
+  INT8_SCALE,
+  SNIPPET_LEN,
+  type WebChunk,
+  type WebIndex,
+  type FtsWebIndex,
+  type FtsWebChunk,
+} from "./web-export.ts";
+export { deriveWebIndex, type DeriveResult, type WebSearchMode } from "./derive-web-index.ts";
+export {
+  vendorModel,
+  vendorRuntime,
+  copySearchScript,
+  readSearchScript,
+} from "./vendor-web.ts";
+export {
+  emitSearchAssets,
+  type EmitSearchAssetsOptions,
+  type EmitSearchAssetsResult,
+} from "./emit-search-assets.ts";

package/src/search.ts

@@ -0,0 +1,131 @@
+import type { EmbeddingProvider } from "./embeddings.ts";
+import { QUERY_PREFIX } from "./embeddings.ts";
+import type { ChunkRow, FtsHit, IndexStore, MetaFilter, VecHit } from "./store.ts";
+
+export const RRF_K = 60;
+
+export function buildFtsQuery(query: string): string {
+  const segs = query
+    .split(/[぀-ゟ]+|[\s、。・，．:：;；!！?？()（）「」『』【】\[\]]+/)
+    .map((s) => s.replace(/"/g, "").trim())
+    .filter((s) => s.length >= 2);
+  if (segs.length === 0) return query.replace(/"/g, " ").trim();
+  return segs.map((s) => `"${s}"`).join(" OR ");
+}
+
+export function rrfFuse(rankings: number[][], k: number = RRF_K): Map<number, number> {
+  const scores = new Map<number, number>();
+  for (const ranking of rankings) {
+    for (let rank = 0; rank < ranking.length; rank++) {
+      const id = ranking[rank]!;
+      scores.set(id, (scores.get(id) ?? 0) + 1 / (k + rank + 1));
+    }
+  }
+  return scores;
+}
+
+export function makeSnippet(text: string, query: string, max: number = 160): string {
+  const flat = text.replace(/\s+/g, " ").trim();
+  if (flat.length <= max) return flat;
+  const q = query.trim();
+  let start = 0;
+  if (q) {
+    const idx = flat.indexOf(q);
+    if (idx >= 0) start = Math.max(0, idx - Math.floor(max / 4));
+  }
+  const end = Math.min(flat.length, start + max);
+  const body = flat.slice(start, end);
+  return (start > 0 ? "…" : "") + body + (end < flat.length ? "…" : "");
+}
+
+export function checkModelMismatch(
+  meta: Record<string, string>,
+  modelId: string,
+  dim: number,
+): string | null {
+  const issues: string[] = [];
+  if (meta.dim && Number(meta.dim) !== dim) {
+    issues.push(`dim index=${meta.dim} runtime=${dim}`);
+  }
+  if (meta.model_id && meta.model_id !== modelId) {
+    issues.push(`model index=${meta.model_id} runtime=${modelId}`);
+  }
+  return issues.length ? issues.join(", ") : null;
+}
+
+export interface SearchOptions {
+  readonly k?: number;
+  readonly filter?: MetaFilter;
+  readonly ftsOnly?: boolean;
+}
+
+export interface SearchResult extends ChunkRow {
+  readonly score: number;
+  readonly snippet: string;
+}
+
+export function searchFts(
+  store: IndexStore,
+  query: string,
+  opts: SearchOptions = {},
+): SearchResult[] {
+  const k = opts.k ?? 10;
+  const filter = opts.filter ?? {};
+  let hits: FtsHit[] = [];
+  try {
+    hits = store.ftsSearch(buildFtsQuery(query), k, filter);
+  } catch {
+    hits = [];
+  }
+  return hits.map((row, rank) => ({
+    ...row,
+    score: 1 / (rank + 1),
+    snippet: makeSnippet(row.text, query),
+  }));
+}
+
+export async function searchHybrid(
+  store: IndexStore,
+  embeddings: EmbeddingProvider,
+  query: string,
+  opts: SearchOptions = {},
+): Promise<SearchResult[]> {
+  const k = opts.k ?? 10;
+  const filter = opts.filter ?? {};
+  const pool = Math.max(k * 5, 50);
+
+  const queryVec = await embeddings.embed(QUERY_PREFIX + query);
+  const vecHits: VecHit[] = store.vecKnn(queryVec, pool, filter);
+  let ftsHits: FtsHit[] = [];
+  try {
+    ftsHits = store.ftsSearch(buildFtsQuery(query), pool, filter);
+  } catch {
+    ftsHits = [];
+  }
+
+  const byId = new Map<number, ChunkRow>();
+  for (const h of vecHits) byId.set(h.id, h);
+  for (const h of ftsHits) byId.set(h.id, h);
+
+  const fused = rrfFuse([vecHits.map((h) => h.id), ftsHits.map((h) => h.id)]);
+  const ranked = [...fused.entries()].sort((a, b) => b[1] - a[1]).slice(0, k);
+
+  return ranked.flatMap(([id, score]) => {
+    const row = byId.get(id);
+    if (!row) return [];
+    return [{ ...row, score, snippet: makeSnippet(row.text, query) }];
+  });
+}
+
+export async function search(
+  store: IndexStore,
+  embeddings: EmbeddingProvider | null,
+  query: string,
+  opts: SearchOptions = {},
+): Promise<SearchResult[]> {
+  const useHybrid = !opts.ftsOnly && embeddings !== null && store.hasVectors();
+  if (useHybrid && embeddings) {
+    return searchHybrid(store, embeddings, query, opts);
+  }
+  return searchFts(store, query, opts);
+}