From e68c753e39b8214d1c62e10b37bdb60983d0fb8b Mon Sep 17 00:00:00 2001 From: Chris Sherwood Date: Thu, 14 May 2026 15:01:24 -0700 Subject: [PATCH] =?UTF-8?q?feat(KB):=20surface=20embedding-disk=20estimate?= =?UTF-8?q?=20in=20curated=20tier-change=20modal=20(RFC=20#883=20=C2=A71)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When a user picks a tier in TierSelectionModal, show how much additional disk space the AI Assistant will need if the new ZIMs are indexed, plus a policy-aware footer explaining whether they'll auto-index (Always) or wait for opt-in (Manual). Estimates consume #891's KbRatioRegistry via a new POST /api/rag/estimate-batch endpoint. Backend - New POST /api/rag/estimate-batch route + RagController.estimateBatch - VineJS schema accepting array of {filename, sizeBytes}, capped at 500 - KbRatioRegistry.estimateBatch aggregates via the existing prefix-match lookup, returns {totalChunks, totalBytes, hasUnknown} - New BYTES_PER_CHUNK_ON_DISK constant (~8 KB: 3 KB vector + ~3 KB chunk text + ~2 KB payload/index overhead). Tunable; will be replaced by Phase 4 self-calibration once we have real measurements. - Controller normalizes incoming filenames via path.basename so callers that send full paths or URLs still match registry prefixes correctly. Frontend - api.estimateEmbeddingBatch() client method - TierSelectionModal: when localSelectedSlug is set, resolve the tier's resources (incl. inherited tiers), POST to /estimate-batch, and render a new info block with the +~X GB figure + ingest-policy copy. Also fetches rag.defaultIngestPolicy so the same block surfaces whether indexing will fire automatically or wait for the user. - resourceFilename() helper extracts the basename from the resource URL so the registry lookup hits the right prefix regardless of mirror. Tests - 4 new cases in tests/unit/kb_ratio_lookup.spec.ts covering the estimateBatch aggregator: standard sum, unknown-flagging, video-only ZIM (0 chunks but known, hasUnknown stays false), empty input. Stacks on feat/kb-ratio-registry (#891) — consumes the registry table seeded by that PR. Once #891 merges to rc, this PR auto-rebases. Out of scope for this PR (deferred to follow-ups): - Per-batch opt-in checkbox (RFC §1's '☑ Also index these for AI') needs a per-batch policy override path and is a separate PR - Guardrail modal at 50 GB / 10% free / 6 hr thresholds (RFC §7) is also separate; this PR is informational, not gating - Time-to-embed estimate awaits a chunks-per-second metric per host --- admin/app/controllers/rag_controller.ts | 17 +++- admin/app/models/kb_ratio_registry.ts | 18 +++- admin/app/utils/kb_ratio_lookup.ts | 53 +++++++++++ admin/app/validators/rag.ts | 14 +++ .../inertia/components/TierSelectionModal.tsx | 89 +++++++++++++++++-- admin/inertia/lib/api.ts | 11 +++ admin/start/routes.ts | 1 + admin/tests/unit/kb_ratio_lookup.spec.ts | 51 ++++++++++- 8 files changed, 246 insertions(+), 8 deletions(-) diff --git a/admin/app/controllers/rag_controller.ts b/admin/app/controllers/rag_controller.ts index 205fb69..3d01ca4 100644 --- a/admin/app/controllers/rag_controller.ts +++ b/admin/app/controllers/rag_controller.ts @@ -1,11 +1,13 @@ import { RagService } from '#services/rag_service' import { EmbedFileJob } from '#jobs/embed_file_job' +import KbRatioRegistry from '#models/kb_ratio_registry' import { inject } from '@adonisjs/core' import type { HttpContext } from '@adonisjs/core/http' import app from '@adonisjs/core/services/app' import { randomBytes } from 'node:crypto' import { sanitizeFilename } from '../utils/fs.js' -import { deleteFileSchema, getJobStatusSchema } from '#validators/rag' +import { basename } from 'node:path' +import { deleteFileSchema, estimateBatchSchema, getJobStatusSchema } from '#validators/rag' import logger from '@adonisjs/core/services/logger' @inject() @@ -122,4 +124,17 @@ export default class RagController { const result = await this.ragService.checkQdrantHealth() return response.status(200).json(result) } + + public async estimateBatch({ request, response }: HttpContext) { + const { files } = await request.validateUsing(estimateBatchSchema) + // The registry matches on basename prefixes; if a caller passes a full path + // (e.g. /app/storage/zim/wikipedia_en_simple_…), strip directories first so + // patterns like `wikipedia_en_simple_` still match. + const normalized = files.map((f) => ({ + filename: basename(f.filename), + sizeBytes: f.sizeBytes, + })) + const result = await KbRatioRegistry.estimateBatch(normalized) + return response.status(200).json(result) + } } diff --git a/admin/app/models/kb_ratio_registry.ts b/admin/app/models/kb_ratio_registry.ts index 97cd1f2..13731b8 100644 --- a/admin/app/models/kb_ratio_registry.ts +++ b/admin/app/models/kb_ratio_registry.ts @@ -1,6 +1,12 @@ import { DateTime } from 'luxon' import { BaseModel, column, SnakeCaseNamingStrategy } from '@adonisjs/lucid/orm' -import { findChunksPerMb, estimateChunkCount } from '../utils/kb_ratio_lookup.js' +import { + findChunksPerMb, + estimateChunkCount, + estimateBatch, + type BatchEstimate, + type BatchEstimateInput, +} from '../utils/kb_ratio_lookup.js' /** * Self-calibrating registry of `{filename-prefix → chunks_per_mb}` ratios used @@ -48,4 +54,14 @@ export default class KbRatioRegistry extends BaseModel { const rows = await this.all() return estimateChunkCount(filename, fileSizeBytes, rows) } + + /** + * Aggregate an embedding-disk-cost estimate across a batch of files. Used by + * the curated-tier-change UI to show "you're about to add ~X GB of + * embeddings on top of the ZIM downloads" before the user commits. + */ + static async estimateBatch(files: BatchEstimateInput[]): Promise { + const rows = await this.all() + return estimateBatch(files, rows) + } } diff --git a/admin/app/utils/kb_ratio_lookup.ts b/admin/app/utils/kb_ratio_lookup.ts index 19b22b6..1abd5c6 100644 --- a/admin/app/utils/kb_ratio_lookup.ts +++ b/admin/app/utils/kb_ratio_lookup.ts @@ -3,6 +3,59 @@ export interface RatioRow { chunks_per_mb: number } +/** + * Bytes of on-disk storage one embedded chunk consumes inside Qdrant. + * + * Rough composition for our pipeline: + * - vector: 768 dims × float32 = 3,072 B + * - chunk text payload: ~3,000 B (target 1,500 tokens × 2 chars/token) + * - source/metadata payload + Qdrant indexes: ~2,000 B + * + * Used for surfacing pre-ingest disk-cost estimates; the actual figure + * varies with collection params and will be replaced by self-calibration + * (RFC #883 Phase 4) once we have real measurements. + */ +export const BYTES_PER_CHUNK_ON_DISK = 8_000 + +export interface BatchEstimateInput { + filename: string + sizeBytes: number +} + +export interface BatchEstimate { + totalChunks: number + totalBytes: number + hasUnknown: boolean +} + +/** + * Aggregate an embedding-disk-cost estimate across a batch of files (curated + * tier add, multi-upload, sync preview, etc). `hasUnknown` is true when at + * least one file did not match any registry row — the totals only include + * matched files, so callers should annotate "estimate excludes unknown files" + * when surfacing the figure. + */ +export function estimateBatch( + files: BatchEstimateInput[], + rows: RatioRow[] +): BatchEstimate { + let totalChunks = 0 + let hasUnknown = false + for (const f of files) { + const chunks = estimateChunkCount(f.filename, f.sizeBytes, rows) + if (chunks === null) { + hasUnknown = true + } else { + totalChunks += chunks + } + } + return { + totalChunks, + totalBytes: totalChunks * BYTES_PER_CHUNK_ON_DISK, + hasUnknown, + } +} + /** * Pick the chunks_per_mb estimate for a filename by longest-prefix match. * diff --git a/admin/app/validators/rag.ts b/admin/app/validators/rag.ts index a9124b4..8326f44 100644 --- a/admin/app/validators/rag.ts +++ b/admin/app/validators/rag.ts @@ -11,3 +11,17 @@ export const deleteFileSchema = vine.compile( source: vine.string(), }) ) + +export const estimateBatchSchema = vine.compile( + vine.object({ + files: vine + .array( + vine.object({ + filename: vine.string().minLength(1).maxLength(255), + sizeBytes: vine.number().min(0), + }) + ) + .minLength(1) + .maxLength(500), + }) +) diff --git a/admin/inertia/components/TierSelectionModal.tsx b/admin/inertia/components/TierSelectionModal.tsx index 0a67169..a55a244 100644 --- a/admin/inertia/components/TierSelectionModal.tsx +++ b/admin/inertia/components/TierSelectionModal.tsx @@ -1,13 +1,25 @@ -import { Fragment, useState, useEffect } from 'react' +import { Fragment, useState, useEffect, useMemo } from 'react' import { Dialog, Transition } from '@headlessui/react' import { IconX, IconCheck, IconInfoCircle } from '@tabler/icons-react' +import { useQuery } from '@tanstack/react-query' import type { CategoryWithStatus, SpecTier, SpecResource } from '../../types/collections' import { resolveTierResources } from '~/lib/collections' import { formatBytes } from '~/lib/util' +import api from '~/lib/api' import classNames from 'classnames' import DynamicIcon, { DynamicIconName } from './DynamicIcon' import StyledButton from './StyledButton' +/** + * Filename for the embed-estimate registry lookup. Strips the URL path so + * patterns like `wikipedia_en_simple_` continue to match upstream filenames + * regardless of mirror domain. + */ +function resourceFilename(resource: SpecResource): string { + const last = resource.url.split('/').pop() + return last && last.length > 0 ? last : resource.id +} + interface TierSelectionModalProps { isOpen: boolean onClose: () => void @@ -33,13 +45,47 @@ const TierSelectionModal: React.FC = ({ } }, [isOpen, category, selectedTierSlug]) - if (!category) return null - - // Get all resources for a tier (including inherited resources) + // Get all resources for a tier (including inherited resources). Defined as a + // hook-safe closure (always callable, returns [] when no category) so the + // memo below can depend on `category` without breaking hook order. const getAllResourcesForTier = (tier: SpecTier): SpecResource[] => { + if (!category) return [] return resolveTierResources(tier, category.tiers) } + // Pre-compute the selected tier's resources outside the JSX so hooks below + // don't re-run on every render. Empty array when no selection. + const selectedTierResources = useMemo(() => { + if (!category || !localSelectedSlug) return [] + const tier = category.tiers.find((t) => t.slug === localSelectedSlug) + return tier ? resolveTierResources(tier, category.tiers) : [] + }, [category, localSelectedSlug]) + + const embedEstimateRequest = useMemo( + () => + selectedTierResources.map((r) => ({ + filename: resourceFilename(r), + sizeBytes: Math.round(r.size_mb * 1024 * 1024), + })), + [selectedTierResources] + ) + + const { data: embedEstimate } = useQuery({ + queryKey: ['embedEstimateBatch', embedEstimateRequest], + queryFn: () => api.estimateEmbeddingBatch(embedEstimateRequest), + enabled: embedEstimateRequest.length > 0, + staleTime: 5 * 60_000, + }) + + const { data: ingestPolicySetting } = useQuery({ + queryKey: ['ingestPolicy'], + queryFn: () => api.getSetting('rag.defaultIngestPolicy'), + }) + const ingestPolicy: 'Always' | 'Manual' = + ingestPolicySetting?.value === 'Manual' ? 'Manual' : 'Always' + + if (!category) return null + const getTierTotalSize = (tier: SpecTier): number => { return getAllResourcesForTier(tier).reduce((acc, r) => acc + r.size_mb * 1024 * 1024, 0) } @@ -203,8 +249,41 @@ const TierSelectionModal: React.FC = ({ })} + {/* Embedding-cost preview — visible whenever a tier is + selected. The estimate uses #891's ratio registry to + project how much extra disk space the AI Assistant will + need for these files on top of the raw downloads. */} + {localSelectedSlug && embedEstimate && embedEstimate.totalBytes > 0 && ( +
+
+ +
+

+ +~{formatBytes(embedEstimate.totalBytes, 1)} + {' '}of additional storage if these are indexed for the AI Assistant + {embedEstimate.hasUnknown && ( + (estimate excludes some files we have no prior data for) + )} + . +

+

+ {ingestPolicy === 'Always' ? ( + <> + Your Auto-index setting is Always, so these files will be indexed automatically once downloaded. You can change this in the Knowledge Base settings. + + ) : ( + <> + Your Auto-index setting is Manual, so these files will sit unindexed until you opt in from the Knowledge Base settings. + + )} +

+
+
+
+ )} + {/* Info note */} -
+

You can change your selection at any time. Click Submit to confirm your choice. diff --git a/admin/inertia/lib/api.ts b/admin/inertia/lib/api.ts index a859cf5..0b3359f 100644 --- a/admin/inertia/lib/api.ts +++ b/admin/inertia/lib/api.ts @@ -835,6 +835,17 @@ class API { })() } + async estimateEmbeddingBatch(files: { filename: string; sizeBytes: number }[]) { + return catchInternal(async () => { + const response = await this.client.post<{ + totalChunks: number + totalBytes: number + hasUnknown: boolean + }>('/rag/estimate-batch', { files }) + return response.data + })() + } + // Wikipedia selector methods async getWikipediaState(): Promise { diff --git a/admin/start/routes.ts b/admin/start/routes.ts index 4d214f3..2038cb7 100644 --- a/admin/start/routes.ts +++ b/admin/start/routes.ts @@ -149,6 +149,7 @@ router router.post('/sync', [RagController, 'scanAndSync']) router.post('/re-embed-all', [RagController, 'reembedAll']) router.post('/reset-and-rebuild', [RagController, 'resetAndRebuild']) + router.post('/estimate-batch', [RagController, 'estimateBatch']) router.get('/health', [RagController, 'health']) }) .prefix('/api/rag') diff --git a/admin/tests/unit/kb_ratio_lookup.spec.ts b/admin/tests/unit/kb_ratio_lookup.spec.ts index 08c3350..330e2a3 100644 --- a/admin/tests/unit/kb_ratio_lookup.spec.ts +++ b/admin/tests/unit/kb_ratio_lookup.spec.ts @@ -1,7 +1,12 @@ import * as assert from 'node:assert/strict' import { test } from 'node:test' -import { estimateChunkCount, findChunksPerMb } from '../../app/utils/kb_ratio_lookup.js' +import { + BYTES_PER_CHUNK_ON_DISK, + estimateBatch, + estimateChunkCount, + findChunksPerMb, +} from '../../app/utils/kb_ratio_lookup.js' const SEEDED_ROWS = [ { pattern: 'devdocs_', chunks_per_mb: 1100 }, @@ -60,3 +65,47 @@ test('estimateChunkCount returns null when no match and no fallback', () => { null ) }) + +test('estimateBatch sums chunks and bytes for matched files', () => { + const files = [ + // 100 MB devdocs -> 110,000 chunks + { filename: 'devdocs_en_python_2026-02.zim', sizeBytes: 100 * 1024 * 1024 }, + // 500 MB wikipedia_en_simple -> 135,000 chunks + { filename: 'wikipedia_en_simple_all_nopic_2026-02.zim', sizeBytes: 500 * 1024 * 1024 }, + ] + const out = estimateBatch(files, SEEDED_ROWS) + assert.equal(out.totalChunks, 110_000 + 135_000) + assert.equal(out.totalBytes, (110_000 + 135_000) * BYTES_PER_CHUNK_ON_DISK) + assert.equal(out.hasUnknown, false) +}) + +test('estimateBatch sets hasUnknown when a file has no registry match', () => { + // Drop the empty-string fallback so the unknown file truly has no match + const rows = SEEDED_ROWS.filter((r) => r.pattern !== '') + const files = [ + { filename: 'devdocs_en_python_2026-02.zim', sizeBytes: 100 * 1024 * 1024 }, + { filename: 'something_unknown_2026-02.zim', sizeBytes: 50 * 1024 * 1024 }, + ] + const out = estimateBatch(files, rows) + assert.equal(out.totalChunks, 110_000) // only the matched file + assert.equal(out.hasUnknown, true) +}) + +test('estimateBatch handles video-only ZIMs (0 chunks/MB) without flagging hasUnknown', () => { + // A 5 GB video ZIM matches the registry with 0 chunks/MB; that is a + // *known* value, not an unknown — totals should be 0 and hasUnknown false. + const files = [ + { filename: 'lrnselfreliance_en_all_2025-12.zim', sizeBytes: 5 * 1024 * 1024 * 1024 }, + ] + const out = estimateBatch(files, SEEDED_ROWS) + assert.equal(out.totalChunks, 0) + assert.equal(out.totalBytes, 0) + assert.equal(out.hasUnknown, false) +}) + +test('estimateBatch on empty input returns zeros', () => { + const out = estimateBatch([], SEEDED_ROWS) + assert.equal(out.totalChunks, 0) + assert.equal(out.totalBytes, 0) + assert.equal(out.hasUnknown, false) +})