feat(KB): surface embedding-disk estimate in curated tier-change modal (RFC #883 §1)
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
This commit is contained in:
Chris Sherwood
Jake Turner
parent
8eb8809154
commit
e68c753e39
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<BatchEstimate> {
|
||||
const rows = await this.all()
|
||||
return estimateBatch(files, rows)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
*
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
})
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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<TierSelectionModalProps> = ({
|
|||
}
|
||||
}, [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<SpecResource[]>(() => {
|
||||
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<TierSelectionModalProps> = ({
|
|||
})}
|
||||
</div>
|
||||
|
||||
{/* 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 && (
|
||||
<div className="mt-4 bg-surface-secondary border border-border-subtle rounded p-3 text-sm">
|
||||
<div className="flex items-start gap-2">
|
||||
<DynamicIcon icon="IconBrain" className="w-5 h-5 text-desert-green flex-shrink-0 mt-0.5" />
|
||||
<div className="flex-1">
|
||||
<p className="text-text-primary">
|
||||
<span className="font-medium">+~{formatBytes(embedEstimate.totalBytes, 1)}</span>
|
||||
{' '}of additional storage if these are indexed for the AI Assistant
|
||||
{embedEstimate.hasUnknown && (
|
||||
<span className="text-text-muted"> (estimate excludes some files we have no prior data for)</span>
|
||||
)}
|
||||
.
|
||||
</p>
|
||||
<p className="text-text-muted text-xs mt-1">
|
||||
{ingestPolicy === 'Always' ? (
|
||||
<>
|
||||
Your <strong>Auto-index</strong> setting is <strong>Always</strong>, so these files will be indexed automatically once downloaded. You can change this in the Knowledge Base settings.
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
Your <strong>Auto-index</strong> setting is <strong>Manual</strong>, so these files will sit unindexed until you opt in from the Knowledge Base settings.
|
||||
</>
|
||||
)}
|
||||
</p>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Info note */}
|
||||
<div className="mt-6 flex items-start gap-2 text-sm text-text-muted bg-blue-50 p-3 rounded">
|
||||
<div className="mt-4 flex items-start gap-2 text-sm text-text-muted bg-blue-50 p-3 rounded">
|
||||
<IconInfoCircle size={18} className="text-blue-500 flex-shrink-0 mt-0.5" />
|
||||
<p>
|
||||
You can change your selection at any time. Click Submit to confirm your choice.
|
||||
|
|
|
|||
|
|
@ -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<WikipediaState | undefined> {
|
||||
|
|
|
|||
|
|
@ -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')
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
})
|
||||
|
|
|
|||
Loading…
Reference in New Issue