diff --git a/admin/app/controllers/rag_controller.ts b/admin/app/controllers/rag_controller.ts index c836393..e8565e0 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() @@ -102,4 +104,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 aa1369b..99531a0 100644 --- a/admin/inertia/lib/api.ts +++ b/admin/inertia/lib/api.ts @@ -811,6 +811,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 fb0d47d..79e7234 100644 --- a/admin/start/routes.ts +++ b/admin/start/routes.ts @@ -147,6 +147,7 @@ router router.delete('/failed-jobs', [RagController, 'cleanupFailedJobs']) router.get('/job-status', [RagController, 'getJobStatus']) router.post('/sync', [RagController, 'scanAndSync']) + 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) +})