membot 0.7.0 → 0.10.0

package/src/ingest/ingest.ts

@@ -1,10 +1,12 @@
+import { DEFAULTS } from "../constants.ts";
 import type { AppContext } from "../context.ts";
 import { upsertBlob } from "../db/blobs.ts";
 import { insertChunksForVersion, rebuildFts } from "../db/chunks.ts";
 import { type FetcherKind, getCurrent, insertVersion, millisIso, type SourceType } from "../db/files.ts";
 import { asHelpful, HelpfulError } from "../errors.ts";
-import { logger } from "../output/logger.ts";
+import { pieFor } from "../output/progress.ts";
 import { chunkDeterministic } from "./chunker.ts";
+import { AsyncMutex, pMap } from "./concurrency.ts";
 import { convert } from "./converter/index.ts";
 import { describe } from "./describer.ts";
 import { embed } from "./embedder.ts";
@@ -33,6 +35,7 @@ export interface IngestEntryResult {
 	error?: string;
 	mime_type: string | null;
 	size_bytes: number;
+	chunk_count: number | null;
 	fetcher: FetcherKind;
 	source_sha256: string;
 }
@@ -51,17 +54,32 @@ export interface IngestResult {
  * without re-resolving anything. `onEntryStart` fires before the pipeline
  * touches an entry; `onEntryComplete` fires after the result (ok / unchanged
  * / failed) is known. Both are optional.
+ *
+ * The optional `workerId` arg threads the slot index through so the UI can
+ * show one status line per in-flight worker; callers that don't want that
+ * detail simply ignore it.
  */
 export interface IngestCallbacks {
-	onEntryStart?: (label: string) => void;
-	onEntryComplete?: (entry: IngestEntryResult) => void;
+	onEntryStart?: (label: string, workerId?: number) => void;
+	onEntryComplete?: (entry: IngestEntryResult, workerId?: number) => void;
 	/**
 	 * Fires for sub-step progress within a single entry (e.g. "embedding
 	 * 32/168"). The callback runs many times per entry and is intended for
 	 * driving an interactive spinner — non-interactive callers should ignore
 	 * it to avoid log spam.
 	 */
-	onEntryProgress?: (label: string, sublabel: string) => void;
+	onEntryProgress?: (label: string, sublabel: string, workerId?: number) => void;
+	/**
+	 * Fires once after the worker pool size has been determined, before the
+	 * first entry begins. Lets the progress reporter size its per-worker
+	 * status section.
+	 */
+	onWorkerCount?: (n: number) => void;
+	/**
+	 * Fires after each successful persist with the number of new chunks
+	 * written, so the progress reporter can track a running total.
+	 */
+	onChunks?: (n: number) => void;
 }
 
 /**
@@ -92,7 +110,26 @@ export async function ingest(input: IngestInput, ctx: AppContext): Promise<Inges
 	const total = countResolvedEntries(resolved);
 	ctx.progress.start(total, "ingest");
 	const callbacks: IngestCallbacks = {
-		onEntryStart: (label) => ctx.progress.tick(label),
+		// Tick on completion so the bar reflects done-and-persisted entries,
+		// not concurrently-in-flight ones. setLabel shows the in-flight file
+		// without advancing the count; sub-step suffix flows via update; per-
+		// worker status lines + chunk total light up if the reporter supports
+		// them (multi-line UI in TTY, no-op otherwise). The pie glyph fills
+		// in as the per-file pipeline marches read → … → persist.
+		onWorkerCount: (n) => ctx.progress.setWorkers(n),
+		onEntryStart: (label, workerId) => {
+			if (workerId !== undefined) ctx.progress.workerSet(workerId, `${pieFor(undefined)} ${label}`);
+			ctx.progress.setLabel(label);
+		},
+		onEntryComplete: (entry, workerId) => {
+			if (workerId !== undefined) ctx.progress.workerSet(workerId, "");
+			ctx.progress.tick(entry.logical_path);
+		},
+		onEntryProgress: (label, sublabel, workerId) => {
+			if (workerId !== undefined) ctx.progress.workerSet(workerId, `${pieFor(sublabel)} ${label} — ${sublabel}`);
+			ctx.progress.update(sublabel);
+		},
+		onChunks: (n) => ctx.progress.addChunks(n),
 	};
 	const result = await ingestResolved(resolved, input, ctx, callbacks);
 	const okCount = result.ok;
@@ -144,11 +181,12 @@ async function ingestInline(
 		status: "ok",
 		mime_type: "text/markdown",
 		size_bytes: bytes.byteLength,
+		chunk_count: null,
 		fetcher: "inline",
 		source_sha256: sha,
 	};
 	try {
-		const versionId = await persistVersion(
+		const persisted = await persistVersion(
 			ctx,
 			{
 				logicalPath,
@@ -168,7 +206,8 @@ async function ingestInline(
 			},
 			(sublabel) => callbacks?.onEntryProgress?.(logicalPath, sublabel),
 		);
-		result.version_id = versionId;
+		result.version_id = persisted.versionId;
+		result.chunk_count = persisted.chunkCount;
 	} catch (err) {
 		result.status = "failed";
 		result.error = errorMessage(err);
@@ -195,6 +234,7 @@ async function ingestUrl(
 		status: "ok",
 		mime_type: null,
 		size_bytes: 0,
+		chunk_count: null,
 		fetcher: "downloader",
 		source_sha256: "",
 	};
@@ -225,7 +265,7 @@ async function ingestUrl(
 			}
 		}
 
-		const versionId = await pipelineForBytes(
+		const persisted = await pipelineForBytes(
 			ctx,
 			{
 				logicalPath,
@@ -244,7 +284,8 @@ async function ingestUrl(
 			},
 			(sublabel) => callbacks?.onEntryProgress?.(url, sublabel),
 		);
-		result.version_id = versionId;
+		result.version_id = persisted.versionId;
+		result.chunk_count = persisted.chunkCount;
 	} catch (err) {
 		result.status = "failed";
 		result.error = errorMessage(err);
@@ -277,11 +318,28 @@ async function ingestLocalFiles(
 		});
 	}
 
-	const results: IngestEntryResult[] = [];
 	const isMulti = resolved.entries.length > 1;
-
-	for (const entry of resolved.entries) {
-		callbacks?.onEntryStart?.(entry.relPathFromBase);
+	// Cap worker count by the actual file count so tiny batches don't pay
+	// the cost of spawning N threads (each loads ~130MB of model weights);
+	// also clamp by config and the global MAX_WORKERS ceiling.
+	const configured = Math.min(DEFAULTS.MAX_WORKERS, Math.max(1, ctx.config.ingest.worker_concurrency));
+	const workerCount = Math.max(1, Math.min(configured, resolved.entries.length));
+	callbacks?.onWorkerCount?.(workerCount);
+	const persistMutex = new AsyncMutex();
+	let anyOk = false;
+
+	// Each pMap worker pulls a file from the shared queue and runs the
+	// entire pipeline end-to-end (read → unchanged check → convert →
+	// describe → chunk → embed → persist). The persist phase is gated by a
+	// single mutex because all workers share one DuckDB connection and
+	// DuckDB rejects nested BEGINs. The embed step itself fans out across
+	// the per-command embedder subprocess pool that `add` / `refresh`
+	// register via `withEmbedderPool()` — so the WASM call truly
+	// parallelizes across cores instead of serializing on the main JS
+	// event loop. When that pool isn't registered (single-shot SDK call,
+	// `embedding.workers = 1`), embed() runs inline against the in-process
+	// extractor with no IPC overhead.
+	const outcomes = await pMap(resolved.entries, workerCount, async (entry, _index, workerId) => {
 		const logicalPath = pickLogicalPath(input.logical_path, entry, isMulti);
 		const result: IngestEntryResult = {
 			source_path: entry.absPath,
@@ -290,10 +348,14 @@ async function ingestLocalFiles(
 			status: "ok",
 			mime_type: null,
 			size_bytes: 0,
+			chunk_count: null,
 			fetcher: "local",
 			source_sha256: "",
 		};
+		callbacks?.onEntryStart?.(entry.relPathFromBase, workerId);
+		const onPhase = (sublabel: string) => callbacks?.onEntryProgress?.(entry.relPathFromBase, sublabel, workerId);
 		try {
+			onPhase("reading");
 			const local = await readLocalFile(entry.absPath);
 			result.mime_type = local.mimeType;
 			result.size_bytes = local.sizeBytes;
@@ -304,48 +366,189 @@ async function ingestLocalFiles(
 				if (cur && cur.source_sha256 === local.sha256) {
 					result.status = "unchanged";
 					result.version_id = cur.version_id;
-					results.push(result);
-					callbacks?.onEntryComplete?.(result);
-					continue;
+					callbacks?.onEntryComplete?.(result, workerId);
+					return result;
 				}
 			}
 
-			const versionId = await pipelineForBytes(
-				ctx,
-				{
+			onPhase("converting");
+			const conversion = await convert(
+				local.bytes,
+				local.mimeType,
+				entry.absPath,
+				ctx.config.llm,
+				ctx.config.converters,
+			);
+			const markdown = conversion.markdown;
+
+			onPhase("describing");
+			const description = await describe(logicalPath, local.mimeType, markdown, ctx.config.llm);
+
+			onPhase("chunking");
+			const chunks = chunkDeterministic(markdown, ctx.config.chunker);
+			const searchTexts = chunks.map((c) => buildSearchText(logicalPath, description, c.content));
+
+			let embeddings: number[][];
+			try {
+				embeddings = await embed(searchTexts, ctx.config.embedding_model, {
+					onProgress: (done, total) => onPhase(`embedding ${done}/${total}`),
+				});
+			} catch (err) {
+				throw asHelpful(
+					err,
+					`while embedding chunks for ${logicalPath}`,
+					"Run `bun run prebuild` to apply the transformers WASM patch, or set a different config.embedding_model.",
+				);
+			}
+
+			const versionId = await persistMutex.lock(async () => {
+				onPhase("persisting");
+				return persistOne(ctx, {
 					logicalPath,
-					bytes: local.bytes,
-					mime: local.mimeType,
-					source: entry.absPath,
 					sourceType: "local",
 					sourcePath: entry.absPath,
 					sourceMtimeMs: local.mtimeMs,
 					sourceSha: local.sha256,
+					blobSha: local.sha256,
+					mime: local.mimeType,
+					bytes: local.bytes,
+					markdown,
+					description,
+					chunks,
+					searchTexts,
+					embeddings,
 					fetcher: "local",
 					downloader: null,
 					downloaderArgs: null,
 					refreshSec,
 					changeNote: input.change_note ?? null,
-				},
-				(sublabel) => callbacks?.onEntryProgress?.(entry.relPathFromBase, sublabel),
-			);
+				});
+			});
 			result.version_id = versionId;
+			result.chunk_count = chunks.length;
+			anyOk = true;
+			callbacks?.onChunks?.(chunks.length);
 		} catch (err) {
 			result.status = "failed";
 			result.error = errorMessage(err);
-		} finally {
-			// Release the DB lock between files in a directory/glob walk so
-			// concurrent processes can wedge in mid-batch. The next entry's
-			// first DB call reopens (cheap — same-process reopen).
-			await ctx.db.release();
 		}
-		results.push(result);
-		callbacks?.onEntryComplete?.(result);
+		callbacks?.onEntryComplete?.(result, workerId);
+		return result;
+	});
+
+	const results: IngestEntryResult[] = outcomes.map((o) => {
+		if (o.ok) return o.value;
+		// pMap caught a worker rejection — shouldn't happen since the worker
+		// catches its own errors, but surface defensively.
+		return {
+			source_path: "",
+			logical_path: "",
+			version_id: null,
+			status: "failed",
+			error: errorMessage(o.error),
+			mime_type: null,
+			size_bytes: 0,
+			chunk_count: null,
+			fetcher: "local",
+			source_sha256: "",
+		};
+	});
+
+	// Single FTS rebuild for the whole batch — replaces N per-entry rebuilds
+	// in the prior implementation. Skip when nothing was newly persisted.
+	if (anyOk) {
+		await rebuildFts(ctx.db);
 	}
 
 	return summarize(results);
 }
 
+/**
+ * Per-file persist payload. All inputs are precomputed by the worker; this
+ * helper just executes the transactional DB writes.
+ */
+interface PersistOneParams {
+	logicalPath: string;
+	sourceType: SourceType;
+	sourcePath: string | null;
+	sourceMtimeMs: number | null;
+	sourceSha: string;
+	blobSha: string | null;
+	mime: string;
+	bytes: Uint8Array | null;
+	markdown: string;
+	description: string;
+	chunks: { index: number; content: string }[];
+	searchTexts: string[];
+	embeddings: number[][];
+	fetcher: FetcherKind;
+	downloader: string | null;
+	downloaderArgs: Record<string, unknown> | null;
+	refreshSec: number | null;
+	changeNote: string | null;
+}
+
+/**
+ * Write blob + new (logical_path, version_id) row + its chunks under a
+ * single DuckDB transaction. ROLLBACK on failure keeps the row+chunks pair
+ * atomic; one COMMIT replaces ~N+2 autocommitted round-trips.
+ */
+async function persistOne(ctx: AppContext, p: PersistOneParams): Promise<string> {
+	const versionId = millisIso(Date.now());
+	const contentSha = sha256Hex(new TextEncoder().encode(p.markdown));
+	await ctx.db.exec("BEGIN TRANSACTION");
+	try {
+		if (p.bytes) {
+			await upsertBlob(ctx.db, {
+				sha256: p.sourceSha,
+				mime_type: p.mime,
+				size_bytes: p.bytes.byteLength,
+				bytes: p.bytes,
+			});
+		}
+		await insertVersion(ctx.db, {
+			logical_path: p.logicalPath,
+			version_id: versionId,
+			source_type: p.sourceType,
+			source_path: p.sourcePath,
+			source_mtime_ms: p.sourceMtimeMs,
+			source_sha256: p.sourceSha,
+			blob_sha256: p.blobSha,
+			content_sha256: contentSha,
+			content: p.markdown,
+			description: p.description,
+			mime_type: p.mime,
+			size_bytes: p.bytes?.byteLength ?? new TextEncoder().encode(p.markdown).byteLength,
+			fetcher: p.fetcher,
+			downloader: p.downloader,
+			downloader_args: p.downloaderArgs,
+			refresh_frequency_sec: p.refreshSec,
+			refreshed_at: new Date().toISOString(),
+			last_refresh_status: "ok",
+			change_note: p.changeNote,
+		});
+		await insertChunksForVersion(
+			ctx.db,
+			p.logicalPath,
+			versionId,
+			p.chunks.map((c, i) => ({
+				chunk_index: c.index,
+				chunk_content: c.content,
+				search_text: p.searchTexts[i] ?? buildSearchText(p.logicalPath, p.description, c.content),
+				embedding: p.embeddings[i] ?? new Array(p.embeddings[0]?.length ?? 0).fill(0),
+			})),
+		);
+		await ctx.db.exec("COMMIT");
+	} catch (err) {
+		await ctx.db.exec("ROLLBACK").catch(() => {
+			// Best effort — if ROLLBACK itself fails (already aborted, lock
+			// dropped, etc.) we still want the original error to surface.
+		});
+		throw err;
+	}
+	return versionId;
+}
+
 interface PipelineParams {
 	logicalPath: string;
 	bytes: Uint8Array;
@@ -373,7 +576,7 @@ async function pipelineForBytes(
 	ctx: AppContext,
 	p: PipelineParams,
 	onPhase?: (sublabel: string) => void,
-): Promise<string> {
+): Promise<{ versionId: string; chunkCount: number }> {
 	onPhase?.("storing blob");
 	await upsertBlob(ctx.db, {
 		sha256: p.sourceSha,
@@ -438,7 +641,7 @@ async function persistVersion(
 	ctx: AppContext,
 	p: PersistParams,
 	onPhase?: (sublabel: string) => void,
-): Promise<string> {
+): Promise<{ versionId: string; chunkCount: number }> {
 	onPhase?.("describing");
 	const description = await describe(p.logicalPath, p.mime, p.markdown, ctx.config.llm);
 	onPhase?.("chunking");
@@ -460,42 +663,49 @@ async function persistVersion(
 	onPhase?.("persisting");
 	const versionId = millisIso(Date.now());
 	const contentSha = p.contentSha ?? sha256Hex(new TextEncoder().encode(p.markdown));
-	await insertVersion(ctx.db, {
-		logical_path: p.logicalPath,
-		version_id: versionId,
-		source_type: p.sourceType,
-		source_path: p.sourcePath,
-		source_mtime_ms: p.sourceMtimeMs,
-		source_sha256: p.sourceSha,
-		blob_sha256: p.blobSha,
-		content_sha256: contentSha,
-		content: p.markdown,
-		description,
-		mime_type: p.mime,
-		size_bytes: p.bytes?.byteLength ?? new TextEncoder().encode(p.markdown).byteLength,
-		fetcher: p.fetcher,
-		downloader: p.downloader,
-		downloader_args: p.downloaderArgs,
-		refresh_frequency_sec: p.refreshSec,
-		refreshed_at: new Date().toISOString(),
-		last_refresh_status: "ok",
-		change_note: p.changeNote,
-	});
+	await ctx.db.exec("BEGIN TRANSACTION");
+	try {
+		await insertVersion(ctx.db, {
+			logical_path: p.logicalPath,
+			version_id: versionId,
+			source_type: p.sourceType,
+			source_path: p.sourcePath,
+			source_mtime_ms: p.sourceMtimeMs,
+			source_sha256: p.sourceSha,
+			blob_sha256: p.blobSha,
+			content_sha256: contentSha,
+			content: p.markdown,
+			description,
+			mime_type: p.mime,
+			size_bytes: p.bytes?.byteLength ?? new TextEncoder().encode(p.markdown).byteLength,
+			fetcher: p.fetcher,
+			downloader: p.downloader,
+			downloader_args: p.downloaderArgs,
+			refresh_frequency_sec: p.refreshSec,
+			refreshed_at: new Date().toISOString(),
+			last_refresh_status: "ok",
+			change_note: p.changeNote,
+		});
 
-	await insertChunksForVersion(
-		ctx.db,
-		p.logicalPath,
-		versionId,
-		chunks.map((c, i) => ({
-			chunk_index: c.index,
-			chunk_content: c.content,
-			search_text: searchTexts[i] ?? buildSearchText(p.logicalPath, description, c.content),
-			embedding: embeddings[i] ?? new Array(embeddings[0]?.length ?? 0).fill(0),
-		})),
-	);
+		await insertChunksForVersion(
+			ctx.db,
+			p.logicalPath,
+			versionId,
+			chunks.map((c, i) => ({
+				chunk_index: c.index,
+				chunk_content: c.content,
+				search_text: searchTexts[i] ?? buildSearchText(p.logicalPath, description, c.content),
+				embedding: embeddings[i] ?? new Array(embeddings[0]?.length ?? 0).fill(0),
+			})),
+		);
+		await ctx.db.exec("COMMIT");
+	} catch (err) {
+		await ctx.db.exec("ROLLBACK").catch(() => {});
+		throw err;
+	}
 	onPhase?.("indexing");
 	await rebuildFts(ctx.db);
-	return versionId;
+	return { versionId, chunkCount: chunks.length };
 }
 
 /**