@statewalker/indexer-sqlite 0.0.0-stage → 0.1.1

package/src/ngram-provider.ts

@@ -0,0 +1,153 @@
+import { NgramStore } from "@statewalker/db-sqlite-index/ngram";
+import type { DocumentPath, PathSelector } from "@statewalker/indexer-api";
+import { toAsyncIterable } from "@statewalker/indexer-core";
+import {
+  NGRAM_TYPE,
+  type NgramBlock,
+  type NgramConfig,
+  type NgramIndex,
+  type NgramMatch,
+  type NgramQuery,
+  type NgramResult,
+  type NgramSearchSummary,
+} from "@statewalker/indexer-ngram";
+import { newEnumeration } from "./enumeration.js";
+import type { SqliteModalityProvider } from "./types.js";
+
+/**
+ * Blocks per `NgramStore.upsert`: each upsert is one transaction, and a
+ * Durable Object retries a write that runs past its ~15 s deadline.
+ */
+const UPSERT_BATCH = 500;
+
+/** 1-based line of each offset, from the positions of `\n` in `text`. */
+function withLines(text: string, matches: { start: number; end: number }[]): NgramMatch[] {
+  const breaks: number[] = [];
+  for (let i = text.indexOf("\n"); i !== -1; i = text.indexOf("\n", i + 1)) breaks.push(i);
+  let line = 1;
+  let b = 0;
+  return [...matches]
+    .sort((x, y) => x.start - y.start)
+    .map(({ start, end }) => {
+      while (b < breaks.length && (breaks[b] as number) < start) {
+        b++;
+        line++;
+      }
+      return { start, end, line };
+    });
+}
+
+/**
+ * Wraps the id-based `NgramStore` (own transactions; never call it inside
+ * `inTransaction`). Tables live under `{prefix}ng_`. Size and enumeration join
+ * the store's `idTable`. Score = match count; ties keep the store's
+ * ascending-id order.
+ */
+export function newSqliteNgramProvider(): SqliteModalityProvider<NgramConfig, NgramIndex> {
+  return {
+    type: NGRAM_TYPE,
+
+    async drop({ db, prefix }) {
+      await new NgramStore({ db, prefix: `${prefix}ng_` }).drop(); // drops either store's tables
+    },
+
+    async open({ db, prefix, blockIds }, config) {
+      const storeKind = config.store ?? "fts5";
+      const store = new NgramStore({ db, prefix: `${prefix}ng_`, store: storeKind });
+      await store.init();
+      const list = newEnumeration(db, blockIds, `SELECT id FROM ${store.idTable}`);
+      let closed = false;
+      const ensureOpen = () => {
+        if (closed) throw new Error("NgramIndex is closed");
+      };
+
+      const index: NgramIndex = {
+        async getIndexInfo() {
+          return { store: storeKind, ...(config.metadata ? { metadata: config.metadata } : {}) };
+        },
+
+        async *search(params: NgramQuery): AsyncGenerator<NgramResult, NgramSearchSummary> {
+          ensureOpen();
+          const allowed = params.paths?.length
+            ? new Set(
+                await blockIds.select(
+                  db,
+                  params.paths.map((path) => ({ path })),
+                ),
+              )
+            : undefined;
+          const result = await store.search(params.pattern, {
+            regex: params.regex ?? false,
+            caseSensitive: params.caseSensitive ?? true,
+            limit: params.topK ?? 100,
+            ...(allowed ? { filter: (id: number) => allowed.has(id) } : {}),
+          });
+          const hits = result.hits.filter((h) => h.matches.length > 0);
+          const ids = hits.map((h) => h.id);
+          const [texts, refs] = await Promise.all([store.getTexts(ids), blockIds.refs(db, ids)]);
+          // Array sort is stable: equal counts keep the store's ascending-id order.
+          const ranked = hits.sort((x, y) => y.matches.length - x.matches.length);
+          for (const h of ranked) {
+            const ref = refs.get(h.id);
+            if (!ref) continue;
+            const count = h.matches.length;
+            yield {
+              ...ref,
+              score: count,
+              count,
+              matches: withLines(texts.get(h.id) ?? "", h.matches),
+            };
+          }
+          return { truncated: result.truncated };
+        },
+
+        async addDocument(blocks: NgramBlock[]) {
+          ensureOpen();
+          for (let i = 0; i < blocks.length; i += UPSERT_BATCH) {
+            const batch = blocks.slice(i, i + UPSERT_BATCH);
+            const ids = await blockIds.ensure(db, batch);
+            await store.upsert(
+              ids.map((id, j) => ({ id, text: (batch[j] as NgramBlock).content })),
+            );
+          }
+        },
+
+        async addDocuments(batches) {
+          for await (const batch of batches) await this.addDocument(batch);
+        },
+
+        async deleteDocuments(selectors) {
+          ensureOpen();
+          const sels: PathSelector[] = [];
+          for await (const s of toAsyncIterable(selectors)) sels.push(s);
+          const ids = await blockIds.select(db, sels);
+          if (ids.length === 0) return;
+          await store.delete(ids);
+        },
+
+        getSize: (p?: DocumentPath) => list.size(p),
+        getDocumentPaths: (p?: DocumentPath) => list.paths(p),
+        getDocumentBlocksRefs: (p?: DocumentPath) => list.refs(p),
+
+        async *getDocumentsBlocks(p?: DocumentPath): AsyncGenerator<NgramBlock> {
+          const entries = await list.entries(p);
+          const texts = await store.getTexts(entries.map((e) => e.id));
+          for (const e of entries) {
+            const content = texts.get(e.id);
+            if (content !== undefined) yield { path: e.path, blockId: e.blockId, content };
+          }
+        },
+
+        async close() {
+          closed = true;
+        },
+        async flush() {},
+        async deleteIndex() {
+          closed = true;
+          await store.drop();
+        },
+      };
+      return index;
+    },
+  };
+}

package/src/path-sql.ts

@@ -0,0 +1,17 @@
+/**
+ * SQL condition matching `column` under any of `prefixes`, on path-component
+ * boundaries (`/foo` matches `/foo` and `/foo/x`, never `/foobar`) and
+ * case-sensitively. Uses `substr` rather than `LIKE`: SQLite's `LIKE` ignores
+ * ASCII case, and `PRAGMA case_sensitive_like` is not authorised in Durable
+ * Objects. `column` must be table-qualified (`b.path`): `json_each` has a `path`
+ * column of its own. No prefixes → `1` (no constraint).
+ */
+export function pathPrefixSql(
+  column: string,
+  prefixes: readonly string[],
+): { sql: string; params: string[] } {
+  if (prefixes.length === 0) return { sql: "1", params: [] };
+  // One JSON parameter, however many prefixes: Durable Objects cap bound variables at 100.
+  const sql = `EXISTS (SELECT 1 FROM json_each(?) p WHERE CASE WHEN substr(p.value, -1) = '/' THEN substr(${column}, 1, length(p.value)) = p.value ELSE ${column} = p.value OR substr(${column}, 1, length(p.value) + 1) = p.value || '/' END)`;
+  return { sql, params: [JSON.stringify(prefixes)] };
+}

package/src/types.ts

@@ -0,0 +1,26 @@
+import type { Db } from "@statewalker/db-api";
+import type { ScoredHit, SearchIndex } from "@statewalker/indexer-api";
+import type { BlockIds } from "./block-ids.js";
+
+export interface SqliteModalityContext {
+  /** The database. Inside `inTransaction(db, (tx) => …)` use `tx` and `ctx.blockIds` methods with `tx`. */
+  db: Db;
+  /** Unique table-name prefix for this sub-index, e.g. `idx_docs__body_`. */
+  prefix: string;
+  /** The Index's shared block-id table. */
+  blockIds: BlockIds;
+}
+
+/**
+ * A modality backed by tables in one SQLite `Db`. Unlike `ModalityProvider`,
+ * opening is async (tables are created, existing data is adopted) and the
+ * provider can drop its tables without an open instance.
+ */
+export interface SqliteModalityProvider<
+  C = unknown,
+  I extends SearchIndex<unknown, unknown, ScoredHit> = SearchIndex<unknown, unknown, ScoredHit>,
+> {
+  readonly type: string;
+  open(ctx: SqliteModalityContext, config: C): Promise<I>;
+  drop(ctx: SqliteModalityContext): Promise<void>;
+}

package/src/vector-provider.ts

@@ -0,0 +1,132 @@
+import { dropVectorIndex, VectorIndex as VectorEngine } from "@statewalker/db-sqlite-index/vector";
+import type { DocumentPath, PathSelector } from "@statewalker/indexer-api";
+import { toAsyncIterable, validateDimensionality } from "@statewalker/indexer-core";
+import {
+  VECTOR_TYPE,
+  type VectorBlock,
+  type VectorConfig,
+  type VectorIndex,
+  type VectorQuery,
+  type VectorResult,
+} from "@statewalker/indexer-vector";
+import { newEnumeration } from "./enumeration.js";
+import type { SqliteModalityProvider } from "./types.js";
+
+/**
+ * Wraps the id-based `VectorIndex` engine. The engine owns its tables under
+ * `{prefix}vec_` and its own transactions (never call it inside `inTransaction`);
+ * this provider adds only the `{path, blockId}` mapping. Size and enumeration
+ * join the engine's `idTable`, so they list exactly what the engine holds.
+ */
+export function newSqliteVectorProvider(): SqliteModalityProvider<VectorConfig, VectorIndex> {
+  return {
+    type: VECTOR_TYPE,
+
+    async drop({ db, prefix }) {
+      await dropVectorIndex(db, `${prefix}vec_`);
+    },
+
+    async open({ db, prefix, blockIds }, config) {
+      const engine = new VectorEngine({
+        db,
+        prefix: `${prefix}vec_`,
+        dims: config.dimensionality,
+        metric: "cosine",
+      });
+      await engine.init();
+      const list = newEnumeration(db, blockIds, `SELECT id FROM ${engine.idTable}`);
+      let closed = false;
+      const ensureOpen = () => {
+        if (closed) throw new Error("VectorIndex is closed");
+      };
+      const info = {
+        dimensionality: config.dimensionality,
+        model: config.model,
+        metadata: config.metadata,
+      };
+
+      const index: VectorIndex = {
+        async getIndexInfo() {
+          return { ...info };
+        },
+
+        async *search(params: VectorQuery): AsyncGenerator<VectorResult> {
+          ensureOpen();
+          const k = params.topK ?? 100;
+          const allowed = params.paths?.length
+            ? new Set(
+                await blockIds.select(
+                  db,
+                  params.paths.map((path) => ({ path })),
+                ),
+              )
+            : undefined;
+          const filter = allowed ? (id: number) => allowed.has(id) : undefined;
+          // Several embeddings merge by max score, like every other backend.
+          const best = new Map<number, number>();
+          for (const q of params.embeddings ?? []) {
+            for (const hit of await engine.search(q, { k, ...(filter ? { filter } : {}) })) {
+              const score = 1 - hit.distance;
+              if (score > (best.get(hit.id) ?? Number.NEGATIVE_INFINITY)) best.set(hit.id, score);
+            }
+          }
+          const top = [...best.entries()].sort((a, b) => b[1] - a[1]).slice(0, k);
+          const refs = await blockIds.refs(
+            db,
+            top.map(([id]) => id),
+          );
+          for (const [id, score] of top) {
+            const ref = refs.get(id);
+            if (ref) yield { ...ref, score };
+          }
+        },
+
+        async addDocument(blocks: VectorBlock[]) {
+          ensureOpen();
+          if (blocks.length === 0) return;
+          for (const b of blocks) validateDimensionality(info, b.embedding);
+          const ids = await blockIds.ensure(db, blocks);
+          await engine.upsert(
+            ids.map((id, i) => ({ id, vector: (blocks[i] as VectorBlock).embedding })),
+          );
+        },
+
+        async addDocuments(batches) {
+          for await (const batch of batches) await this.addDocument(batch);
+        },
+
+        async deleteDocuments(selectors) {
+          ensureOpen();
+          const sels: PathSelector[] = [];
+          for await (const s of toAsyncIterable(selectors)) sels.push(s);
+          const ids = await blockIds.select(db, sels);
+          if (ids.length === 0) return;
+          await engine.delete(ids);
+        },
+
+        getSize: (p?: DocumentPath) => list.size(p),
+        getDocumentPaths: (p?: DocumentPath) => list.paths(p),
+        getDocumentBlocksRefs: (p?: DocumentPath) => list.refs(p),
+
+        async *getDocumentsBlocks(p?: DocumentPath): AsyncGenerator<VectorBlock> {
+          const entries = await list.entries(p);
+          const vectors = await engine.getVectors(entries.map((e) => e.id));
+          for (const e of entries) {
+            const embedding = vectors.get(e.id);
+            if (embedding) yield { path: e.path, blockId: e.blockId, embedding };
+          }
+        },
+
+        async close() {
+          closed = true;
+        },
+        async flush() {},
+        async deleteIndex() {
+          closed = true;
+          await dropVectorIndex(db, `${prefix}vec_`);
+        },
+      };
+      return index;
+    },
+  };
+}