diff --git a/USAGE.md b/USAGE.md index e361dca8..c519fa80 100644 --- a/USAGE.md +++ b/USAGE.md @@ -324,6 +324,7 @@ npx tsx hybrid_fast_seed.ts --filesdir [options] - `--overwrite` - Drop existing tables and rebuild from scratch - `--rebuild-concepts` - Rebuild concept index without re-processing documents - `--auto-reseed` - Automatically fix and re-process incomplete catalog records +- `--populate-summaries[=targets]` - Fill in missing summaries only (see below) **Examples:** ```bash @@ -340,6 +341,47 @@ npx tsx hybrid_fast_seed.ts --filesdir ~/Documents --rebuild-concepts npx tsx hybrid_fast_seed.ts --filesdir ~/Documents --auto-reseed ``` +#### Rebuilding summaries only + +`--populate-summaries` fills in missing summaries in an existing database. It +needs no `--filesdir`: documents are never re-loaded, re-chunked, or re-extracted. + +| Table | What is regenerated | Source | +|-------|--------------------|--------| +| `catalog.summary` | Document overview (and the embedding vector, which encodes it) | Existing chunks, in reading order | +| `concepts.summary` | One-sentence definition | Concept name | +| `categories.summary` | One-sentence description | Category name | + +A summary counts as missing when it is empty, or when it is a seeding fallback +(`Document overview (N pages)` for documents, the generated description for +categories). + +```bash +# See what is missing - no LLM calls, no writes +npx tsx hybrid_fast_seed.ts --populate-summaries --dry-run + +# Fill in everything that is missing +RUST_LOG=error npx tsx hybrid_fast_seed.ts --populate-summaries + +# Restrict to one or more tables +npx tsx hybrid_fast_seed.ts --populate-summaries=concepts,categories + +# Try it on a handful of rows first +npx tsx hybrid_fast_seed.ts --populate-summaries=concepts --summary-max-items 50 + +# Regenerate every summary, not just the missing ones +npx tsx hybrid_fast_seed.ts --populate-summaries --force-summaries +``` + +**Additional options:** `--summary-batch-size N` (names per LLM request, default 30), +`--summary-flush-size N` (summaries buffered before writing, default 250), +`--summary-max-items N` (cap rows per table), `--summary-model ID` (defaults to the +model configured in `src/concepts/summary_generator.ts`). + +Results are written back in batches with a merge-insert keyed on `id`, so all +other columns and the table schema are left untouched. An interrupted run keeps +everything it already wrote ā€” re-run the same command to continue. + **šŸ“ Logging:** Each run creates a timestamped log in `logs/seed-YYYY-MM-DDTHH-MM-SS.log`. --- diff --git a/docs/development.md b/docs/development.md index de196e00..b4cdba66 100644 --- a/docs/development.md +++ b/docs/development.md @@ -123,6 +123,24 @@ npm run test:integration # Integration tests only | `--auto-reseed` | Re-process documents with incomplete metadata | | `--max-docs N` | Process at most N new documents (for batching) | | `--with-wordnet` | Enable WordNet enrichment (disabled by default) | +| `--populate-summaries[=targets]` | Fill in missing summaries only, no `--filesdir` needed | +| `--force-summaries` | With `--populate-summaries`: regenerate every summary | +| `--dry-run` | With `--populate-summaries`: report only, no LLM calls or writes | + +**Rebuild summaries without re-seeding:** + +```bash +# Report what is missing across catalog, concepts and categories +npx tsx hybrid_fast_seed.ts --populate-summaries --dry-run + +# Fill the gaps (targets default to all three tables) +RUST_LOG=error npx tsx hybrid_fast_seed.ts --populate-summaries=concepts +``` + +Document overviews are regenerated from existing chunks (and the catalog vector +re-embedded, since it encodes the summary); concept and category summaries are +generated from their names. Writes go through a merge-insert keyed on `id`, so +every other column survives and the run is resumable. **Seed specific documents:** diff --git a/hybrid_fast_seed.ts b/hybrid_fast_seed.ts index 1963ce2f..a3c30dad 100644 --- a/hybrid_fast_seed.ts +++ b/hybrid_fast_seed.ts @@ -28,9 +28,11 @@ import { findDocumentFilesRecursively, getDatabaseSize, truncateFilePath, + backfillSummaries, + parseSummaryTargets, type DataCompletenessCheck } from './src/infrastructure/seeding/index.js'; -import { generateCategorySummaries } from './src/concepts/summary_generator.js'; +import { generateCategorySummaries, generateDocumentOverview } from './src/concepts/summary_generator.js'; import { parseFilenameMetadata } from './src/infrastructure/utils/filename-metadata-parser.js'; import { SeedingCheckpoint } from './src/infrastructure/checkpoint/seeding-checkpoint.js'; import { StageCache, type CachedDocumentData } from './src/infrastructure/checkpoint/stage-cache.js'; @@ -166,8 +168,8 @@ console.log = (...args: any[]) => { }; const argv: minimist.ParsedArgs = minimist(process.argv.slice(2), { - boolean: ["overwrite", "rebuild-concepts", "auto-reseed", "clean-checkpoint", "resume", "with-wordnet", "clear-cache", "cache-only", "no-cache"], - string: ["dbpath", "filesdir", "max-docs", "parallel", "cache-dir"] + boolean: ["overwrite", "rebuild-concepts", "auto-reseed", "clean-checkpoint", "resume", "with-wordnet", "clear-cache", "cache-only", "no-cache", "force-summaries", "dry-run"], + string: ["dbpath", "filesdir", "max-docs", "parallel", "cache-dir", "populate-summaries", "summary-batch-size", "summary-flush-size", "summary-max-items", "summary-model"] }); const databaseDir = path.resolve(argv["dbpath"] || process.env.CONCEPT_RAG_DB_PATH || path.join(process.env.HOME || process.env.USERPROFILE || "~", ".concept_rag")); @@ -187,6 +189,16 @@ const cacheOnly = argv["cache-only"]; const useCache = !argv["no-cache"]; // Default: true (use cache), --no-cache disables const customCacheDir = argv["cache-dir"]; +// Summary backfill flags (standalone mode - operates on an existing database) +const summaryMode = argv["populate-summaries"] !== undefined; +const summaryTargetsArg = argv["populate-summaries"]; +const forceSummaries = argv["force-summaries"]; +const dryRun = argv["dry-run"]; +const summaryBatchSize = argv["summary-batch-size"] ? parseInt(argv["summary-batch-size"], 10) : undefined; +const summaryFlushSize = argv["summary-flush-size"] ? parseInt(argv["summary-flush-size"], 10) : undefined; +const summaryMaxItems = argv["summary-max-items"] ? parseInt(argv["summary-max-items"], 10) : undefined; +const summaryModel = argv["summary-model"] || undefined; + function showUsageAndExit() { console.error("Please provide a directory with files (--filesdir) to process"); console.error("Usage: npx tsx hybrid_fast_seed.ts --filesdir [--dbpath ] [options]"); @@ -206,6 +218,17 @@ function showUsageAndExit() { console.error(" --max-docs N: Process at most N NEW documents (skips already processed, enables batching)"); console.error(" --parallel N: Process N documents concurrently for concept extraction (default: 10, max: 25)"); console.error(""); + console.error("Summary backfill (no --filesdir needed, runs against --dbpath and exits):"); + console.error(" --populate-summaries[=TARGETS]: Fill in missing summaries in an existing database."); + console.error(" TARGETS is a comma-separated list of catalog,concepts,categories"); + console.error(" (default: all three)"); + console.error(" --force-summaries: Regenerate every summary, not just the missing ones"); + console.error(" --dry-run: Report what is missing without calling the LLM or writing"); + console.error(" --summary-batch-size N: Concept/category names per LLM request (default: 30)"); + console.error(" --summary-flush-size N: Summaries buffered before writing back (default: 250)"); + console.error(" --summary-max-items N: Process at most N rows per table (trial runs)"); + console.error(" --summary-model ID: Model override (default: the configured summary model)"); + console.error(""); console.error("Cache options:"); console.error(" --no-cache: Disable stage cache (don't use cached LLM results)"); console.error(" --clear-cache: Clear stage cache before processing"); @@ -350,33 +373,10 @@ async function verifyApiKey(): Promise { } // LLM API call for summarization +// Shared with the --populate-summaries backfill so both produce identical overviews. +// No rate limiting here: overviews are generated across parallel workers during seeding. async function callOpenRouterChat(text: string): Promise { - const response = await fetch('https://openrouter.ai/api/v1/chat/completions', { - method: 'POST', - headers: { - 'Authorization': `Bearer ${openrouterApiKey}`, - 'Content-Type': 'application/json', - 'HTTP-Referer': 'https://github.com/adiom-data/lance-mcp', - 'X-Title': 'Lance MCP Server' - }, - body: JSON.stringify({ - model: 'x-ai/grok-4-fast', // Grok-4-fast: blazing fast for simple summaries - messages: [{ - role: 'user', - content: `Write a high-level one sentence content overview based on the text below. WRITE THE CONTENT OVERVIEW ONLY:\n\n${text.slice(0, 8000)}` - }], - max_tokens: 100, - temperature: 0.3 - }) - }); - - if (!response.ok) { - const errorData = await response.text(); - throw new Error(`LLM API error: ${response.status} - ${errorData}`); - } - - const data = await response.json(); - return data.choices[0].message.content.trim(); + return generateDocumentOverview(text, { apiKey: openrouterApiKey }); } async function generateContentOverview(rawDocs: Document[]): Promise { @@ -2055,7 +2055,95 @@ async function rebuildConceptIndexFromExistingData( await createCategoriesTable(db, catalogDocs); } +/** + * Standalone mode for --populate-summaries. + * + * Fills in missing summaries in an already-seeded database: document overviews + * are regenerated from existing chunks, concept and category summaries from + * their names. No documents are loaded, chunked or re-extracted. + */ +async function populateMissingSummaries() { + let targets; + try { + targets = parseSummaryTargets(summaryTargetsArg); + } catch (error: any) { + console.error(`āŒ ${error.message}`); + process.exit(1); + } + + console.log("\nšŸ“ POPULATE SUMMARIES"); + console.log("=".repeat(70)); + console.log(`šŸ“‚ Database: ${databaseDir}`); + console.log(`šŸŽÆ Targets: ${targets.join(', ')}`); + if (forceSummaries) console.log("āš ļø Force mode: regenerating ALL summaries"); + if (dryRun) console.log("šŸ” Dry run: no LLM calls, no writes"); + if (summaryMaxItems) console.log(`šŸ”¢ Max items per table: ${summaryMaxItems}`); + console.log(""); + + if (!fs.existsSync(databaseDir)) { + console.error(`āŒ Database not found: ${databaseDir}`); + process.exit(1); + } + + if (!dryRun) { + if (!openrouterApiKey) { + console.error("Please set OPENROUTER_API_KEY environment variable"); + process.exit(1); + } + await verifyApiKey(); + } + + const db = await lancedb.connect(databaseDir); + + const report = await backfillSummaries(db, { + targets, + apiKey: openrouterApiKey, + model: summaryModel, + batchSize: summaryBatchSize, + flushSize: summaryFlushSize, + maxItems: summaryMaxItems, + force: forceSummaries, + dryRun, + onProgress: ({ table, completed, total }) => { + const width = 30; + const percentage = total > 0 ? Math.round((completed / total) * 100) : 100; + const filled = Math.round((percentage / 100) * width); + const bar = 'ā–ˆ'.repeat(filled) + 'ā–‘'.repeat(width - filled); + process.stdout.write(`\r [${bar}] ${table}: ${completed}/${total} `); + if (completed >= total) process.stdout.write('\r' + ' '.repeat(80) + '\r'); + } + }); + + console.log("\n" + "=".repeat(70)); + for (const result of report.results) { + if (result.skippedReason) { + console.log(`ā­ļø ${result.table}: skipped (${result.skippedReason})`); + continue; + } + const suffix = report.dryRun + ? '(dry run - nothing written)' + : `ā†’ ${result.written} written${result.failed > 0 ? `, ${result.failed} failed` : ''}`; + console.log(` ${result.table}: ${result.missing}/${result.total} missing ${suffix}`); + } + + if (report.dryRun) { + console.log("\nšŸ’” Re-run without --dry-run to generate the missing summaries"); + } else { + console.log("\nāœ… Summary backfill complete!"); + } +} + async function hybridFastSeed() { + // Summary backfill runs against an existing database and exits + if (summaryMode) { + await populateMissingSummaries(); + return; + } + + if (dryRun) { + console.warn("āš ļø --dry-run only applies to --populate-summaries - continuing with a normal seeding run"); + } + const sourceDirs = await validateArgs(); // Preflight check: verify API key before any database operations diff --git a/scripts/populate_summaries.ts b/scripts/populate_summaries.ts deleted file mode 100644 index a7d151ba..00000000 --- a/scripts/populate_summaries.ts +++ /dev/null @@ -1,476 +0,0 @@ -#!/usr/bin/env npx tsx -/** - * Populate summary fields for concepts and categories - * - * Uses grok-4-fast to generate one-sentence summaries in batches - * - * RESUME SUPPORT: The script automatically resumes from where it left off - * by finding items that already have summaries and skipping them. - * Progress is saved after each batch for crash recovery. - * - * Usage: - * # Suppress LanceDB warnings for clean progress bar: - * RUST_LOG=error npx tsx scripts/populate_summaries.ts [db-path] - * - * # Example: - * RUST_LOG=error npx tsx scripts/populate_summaries.ts ~/.concept_rag - * - * Options: - * --concepts-only Only populate concept summaries - * --categories-only Only populate category summaries - * --batch-size=N Number of items per LLM request (default: 20) - * --dry-run Show what would be done without making changes - * --force Regenerate ALL summaries (ignore existing) - */ - -import { connect, Table } from '@lancedb/lancedb'; -import * as path from 'path'; - -const OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"; -const MODEL = "x-ai/grok-4-fast"; -const DEFAULT_BATCH_SIZE = 20; -const MIN_REQUEST_INTERVAL = 1000; // 1 second between requests -const SAVE_EVERY_N_BATCHES = 5; // Save progress every N batches - -interface SummaryResult { - name: string; - summary: string; -} - -let lastRequestTime = 0; - -async function rateLimitDelay(): Promise { - const now = Date.now(); - const timeSinceLastRequest = now - lastRequestTime; - if (timeSinceLastRequest < MIN_REQUEST_INTERVAL) { - const delayNeeded = MIN_REQUEST_INTERVAL - timeSinceLastRequest; - await new Promise(resolve => setTimeout(resolve, delayNeeded)); - } - lastRequestTime = Date.now(); -} - -/** - * Generate ASCII progress bar - */ -function progressBar(current: number, total: number, width: number = 30): string { - const percentage = Math.round((current / total) * 100); - const filled = Math.round((percentage / 100) * width); - const empty = width - filled; - return 'ā–ˆ'.repeat(filled) + 'ā–‘'.repeat(empty); -} - -async function generateSummaries( - items: string[], - type: 'concept' | 'category', - apiKey: string -): Promise { - await rateLimitDelay(); - - const prompt = type === 'concept' - ? `Generate a one-sentence summary for each of the following concepts. The summary should be a clear, concise definition that would help someone understand what the concept means. - -Format your response as JSON array: -[{"name": "concept name", "summary": "one sentence summary"}] - -Concepts to summarize: -${items.map((item, i) => `${i + 1}. ${item}`).join('\n')} - -Return ONLY the JSON array, no other text.` - : `Generate a one-sentence summary for each of the following categories/domains. The summary should describe what topics and subjects fall under this category. - -Format your response as JSON array: -[{"name": "category name", "summary": "one sentence summary"}] - -Categories to summarize: -${items.map((item, i) => `${i + 1}. ${item}`).join('\n')} - -Return ONLY the JSON array, no other text.`; - - try { - const response = await fetch(`${OPENROUTER_BASE_URL}/chat/completions`, { - method: 'POST', - headers: { - 'Authorization': `Bearer ${apiKey}`, - 'Content-Type': 'application/json', - 'HTTP-Referer': 'https://github.com/concept-rag', - 'X-Title': 'Concept-RAG Summary Generation' - }, - body: JSON.stringify({ - model: MODEL, - messages: [ - { role: 'user', content: prompt } - ], - temperature: 0.3, - max_tokens: 4000 - }) - }); - - if (!response.ok) { - const error = await response.text(); - throw new Error(`API error: ${response.status} - ${error}`); - } - - const data = await response.json(); - const content = data.choices[0]?.message?.content || ''; - - // Parse JSON from response - const jsonMatch = content.match(/\[[\s\S]*\]/); - if (!jsonMatch) { - console.error(' āš ļø Failed to parse JSON from response:', content.substring(0, 200)); - return []; - } - - const results: SummaryResult[] = JSON.parse(jsonMatch[0]); - return results; - } catch (error) { - console.error(' āŒ Error generating summaries:', error); - return []; - } -} - -/** - * Convert Arrow Vectors and other LanceDB types to native arrays - */ -function toArray(val: any): string[] { - if (!val) return ['']; - if (Array.isArray(val)) return val.length > 0 ? val : ['']; - if (typeof val === 'object' && 'toArray' in val) { - const arr = Array.from(val.toArray()) as string[]; - return arr.length > 0 ? arr : ['']; - } - return ['']; -} - -function toNumberArray(val: any): number[] { - if (!val) return [0]; - if (Array.isArray(val)) return val.length > 0 ? val : [0]; - if (typeof val === 'object' && 'toArray' in val) { - const arr = Array.from(val.toArray()) as number[]; - return arr.length > 0 ? arr : [0]; - } - return [0]; -} - -function toVectorArray(val: any): number[] { - if (Array.isArray(val)) return val; - if (typeof val === 'object' && 'toArray' in val) { - return Array.from(val.toArray()) as number[]; - } - return []; -} - -/** - * Save concepts table with updated summaries (incremental save) - */ -async function saveConceptsIncremental( - db: any, - allConcepts: any[], - updates: Map -): Promise { - const migratedConcepts = allConcepts.map(row => { - const updatedSummary = updates.get(row.id); - return { - id: row.id, - name: row.name || row.concept || '', // Support both 'name' (new) and 'concept' (legacy) - summary: updatedSummary ?? row.summary ?? '', - catalog_ids: toNumberArray(row.catalog_ids), - chunk_ids: toNumberArray(row.chunk_ids), - related_concept_ids: toNumberArray(row.related_concept_ids), - synonyms: toArray(row.synonyms), - broader_terms: toArray(row.broader_terms), - narrower_terms: toArray(row.narrower_terms), - weight: row.weight || 0, - vector: toVectorArray(row.vector) - }; - }); - - await db.dropTable('concepts'); - await db.createTable('concepts', migratedConcepts, { mode: 'overwrite' }); -} - -/** - * Save categories table with updated summaries (incremental save) - */ -async function saveCategoriesIncremental( - db: any, - allCategories: any[], - updates: Map -): Promise { - const migratedCategories = allCategories.map(row => { - const updatedSummary = updates.get(row.id); - return { - id: row.id, - category: row.category || '', - description: row.description || '', - summary: updatedSummary ?? row.summary ?? '', - parent_category_id: row.parent_category_id ?? null, - aliases: toArray(row.aliases), - related_categories: toNumberArray(row.related_categories), - document_count: row.document_count || 0, - chunk_count: row.chunk_count || 0, - concept_count: row.concept_count || 0, - vector: toVectorArray(row.vector) - }; - }); - - await db.dropTable('categories'); - await db.createTable('categories', migratedCategories, { mode: 'overwrite' }); -} - -/** - * Preflight check: Verify OpenRouter API key is valid before starting any operations. - * Makes a minimal chat completion request to detect 401/403 errors early. - */ -async function verifyApiKey(apiKey: string): Promise { - console.log('šŸ”‘ Verifying OpenRouter API key...'); - - try { - // Use chat completions endpoint with minimal request (requires auth, unlike /models) - const response = await fetch('https://openrouter.ai/api/v1/chat/completions', { - method: 'POST', - headers: { - 'Authorization': `Bearer ${apiKey}`, - 'Content-Type': 'application/json', - 'HTTP-Referer': 'https://github.com/adiom-data/lance-mcp', - 'X-Title': 'Concept-RAG API Verification' - }, - body: JSON.stringify({ - model: 'openai/gpt-4o-mini', // Cheapest model for validation - messages: [{ role: 'user', content: 'hi' }], - max_tokens: 1 // Minimal tokens to minimize cost - }) - }); - - if (!response.ok) { - const errorData = await response.text(); - let errorMessage: string; - - try { - const parsed = JSON.parse(errorData); - errorMessage = parsed.error?.message || errorData; - } catch { - errorMessage = errorData; - } - - if (response.status === 401 || response.status === 403) { - console.error(''); - console.error('āŒ API KEY VALIDATION FAILED'); - console.error('ā”'.repeat(50)); - console.error(` Status: ${response.status} ${response.statusText}`); - console.error(` Error: ${errorMessage}`); - console.error(''); - console.error(' Your OpenRouter API key is invalid or expired.'); - console.error(' Please check your OPENROUTER_API_KEY in .envrc'); - console.error('ā”'.repeat(50)); - process.exit(1); - } - - console.warn(`āš ļø API check returned ${response.status}: ${errorMessage}`); - console.warn(' Proceeding anyway - API may still work for completions.'); - return; - } - - console.log('āœ… API key verified'); - } catch (error) { - console.warn(`āš ļø Could not verify API key: ${(error as Error).message}`); - console.warn(' Proceeding anyway - check your network connection if errors occur.'); - } -} - -async function populateSummaries( - dbPath: string, - options: { - conceptsOnly?: boolean; - categoriesOnly?: boolean; - batchSize?: number; - dryRun?: boolean; - force?: boolean; - } = {} -) { - const apiKey = process.env.OPENROUTER_API_KEY; - if (!apiKey) { - console.error('āŒ OPENROUTER_API_KEY environment variable not set'); - process.exit(1); - } - - // Preflight check: verify API key before any database operations - await verifyApiKey(apiKey); - - const batchSize = options.batchSize || DEFAULT_BATCH_SIZE; - - console.log('\nšŸ“ POPULATE SUMMARIES'); - console.log('='.repeat(70)); - console.log(`Database: ${dbPath}`); - console.log(`Model: ${MODEL}`); - console.log(`Batch size: ${batchSize}`); - if (options.dryRun) console.log('šŸ” DRY RUN MODE - no changes will be made'); - if (options.force) console.log('āš ļø FORCE MODE - regenerating ALL summaries'); - console.log('šŸ’¾ Progress saved after each batch (resume-safe)'); - if (!process.env.RUST_LOG) { - console.log('šŸ’” Tip: Run with RUST_LOG=error to suppress LanceDB warnings'); - } - console.log(''); - - const db = await connect(dbPath); - - // ========== CONCEPTS ========== - if (!options.categoriesOnly) { - console.log('šŸ§  Processing concepts...'); - const conceptsTable = await db.openTable('concepts'); - const allConcepts = await conceptsTable.query().limit(100000).toArray(); - - // Find concepts without summaries (or all if --force) - const needsSummary = options.force - ? allConcepts - : allConcepts.filter(c => !c.summary || c.summary === ''); - - const alreadyDone = allConcepts.length - needsSummary.length; - - if (alreadyDone > 0 && !options.force) { - console.log(` āœ“ Already have summaries: ${alreadyDone}/${allConcepts.length}`); - console.log(` ā†’ Resuming from item ${alreadyDone + 1}...`); - } - console.log(` šŸ“‹ Need to process: ${needsSummary.length} concepts`); - - if (needsSummary.length > 0 && !options.dryRun) { - const updates = new Map(); - const totalBatches = Math.ceil(needsSummary.length / batchSize); - - for (let i = 0; i < needsSummary.length; i += batchSize) { - const batch = needsSummary.slice(i, i + batchSize); - const names = batch.map(c => c.name); - const batchNum = Math.floor(i / batchSize) + 1; - const processed = Math.min(i + batchSize, needsSummary.length); - - // Show progress bar - const bar = progressBar(processed, needsSummary.length); - process.stdout.write(`\r šŸ§  [${bar}] ${processed}/${needsSummary.length} (batch ${batchNum}/${totalBatches}) `); - - const results = await generateSummaries(names, 'concept', apiKey); - - // Match results back to concepts - for (const result of results) { - const concept = batch.find(c => - c.name.toLowerCase() === result.name.toLowerCase() - ); - if (concept && result.summary) { - updates.set(concept.id, result.summary); - } - } - - // Save progress after each batch - if (updates.size > 0) { - await saveConceptsIncremental(db, allConcepts, updates); - // Refresh allConcepts with updated data - const refreshedTable = await db.openTable('concepts'); - const refreshedData = await refreshedTable.query().limit(100000).toArray(); - allConcepts.length = 0; - allConcepts.push(...refreshedData); - } - } - - // Clear progress bar and show completion - process.stdout.write('\r' + ' '.repeat(80) + '\r'); - console.log(` āœ… Generated ${updates.size} concept summaries`); - } else if (needsSummary.length === 0) { - console.log(` āœ… All ${allConcepts.length} concepts already have summaries`); - } - } - - // ========== CATEGORIES ========== - if (!options.conceptsOnly) { - console.log('\nšŸ“ Processing categories...'); - const categoriesTable = await db.openTable('categories'); - const allCategories = await categoriesTable.query().limit(10000).toArray(); - - // Find categories without summaries (or all if --force) - const needsSummary = options.force - ? allCategories - : allCategories.filter(c => !c.summary || c.summary === ''); - - const alreadyDone = allCategories.length - needsSummary.length; - - if (alreadyDone > 0 && !options.force) { - console.log(` āœ“ Already have summaries: ${alreadyDone}/${allCategories.length}`); - console.log(` ā†’ Resuming from item ${alreadyDone + 1}...`); - } - console.log(` šŸ“‹ Need to process: ${needsSummary.length} categories`); - - if (needsSummary.length > 0 && !options.dryRun) { - const updates = new Map(); - const totalBatches = Math.ceil(needsSummary.length / batchSize); - let lastSavedBatch = 0; - - for (let i = 0; i < needsSummary.length; i += batchSize) { - const batch = needsSummary.slice(i, i + batchSize); - const names = batch.map(c => c.category); - const batchNum = Math.floor(i / batchSize) + 1; - const processed = Math.min(i + batchSize, needsSummary.length); - - // Show progress bar - const bar = progressBar(processed, needsSummary.length); - process.stdout.write(`\r šŸ“‚ [${bar}] ${processed}/${needsSummary.length} (batch ${batchNum}/${totalBatches}) `); - - const results = await generateSummaries(names, 'category', apiKey); - - // Match results back to categories - for (const result of results) { - const category = batch.find(c => - c.category.toLowerCase() === result.name.toLowerCase() - ); - if (category && result.summary) { - updates.set(category.id, result.summary); - } - } - - // Save progress every N batches (or on last batch) - const isLastBatch = i + batchSize >= needsSummary.length; - const shouldSave = (batchNum - lastSavedBatch >= SAVE_EVERY_N_BATCHES) || isLastBatch; - - if (shouldSave && updates.size > 0) { - // Clear line and show save message - process.stdout.write(`\r šŸ’¾ Saving ${updates.size} summaries...` + ' '.repeat(50)); - - await saveCategoriesIncremental(db, allCategories, updates); - lastSavedBatch = batchNum; - - // Refresh allCategories with updated data - const refreshedTable = await db.openTable('categories'); - const refreshedData = await refreshedTable.query().limit(10000).toArray(); - allCategories.length = 0; - allCategories.push(...refreshedData); - - // Restore progress bar - process.stdout.write(`\r šŸ“‚ [${bar}] ${processed}/${needsSummary.length} (batch ${batchNum}/${totalBatches}) `); - } - } - - // Clear progress bar and show completion - process.stdout.write('\r' + ' '.repeat(80) + '\r'); - console.log(` āœ… Generated ${updates.size} category summaries`); - } else if (needsSummary.length === 0) { - console.log(` āœ… All ${allCategories.length} categories already have summaries`); - } - } - - console.log('\n' + '='.repeat(70)); - console.log('āœ… Summary population complete!'); -} - -// Parse arguments -const args = process.argv.slice(2); -const dbPath = args.find(a => !a.startsWith('--')) || path.join(process.env.HOME || '', '.concept_rag'); -const options = { - conceptsOnly: args.includes('--concepts-only'), - categoriesOnly: args.includes('--categories-only'), - batchSize: parseInt(args.find(a => a.startsWith('--batch-size='))?.split('=')[1] || '') || DEFAULT_BATCH_SIZE, - dryRun: args.includes('--dry-run'), - force: args.includes('--force') -}; - -populateSummaries(dbPath, options) - .then(() => process.exit(0)) - .catch(err => { - console.error('Error:', err); - process.exit(1); - }); diff --git a/src/concepts/summary_generator.ts b/src/concepts/summary_generator.ts index 8ef82268..4b20caba 100644 --- a/src/concepts/summary_generator.ts +++ b/src/concepts/summary_generator.ts @@ -1,12 +1,15 @@ /** - * Summary generation for concepts and categories using LLM + * Summary generation for documents, concepts and categories using LLM */ const OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"; -const DEFAULT_MODEL = "x-ai/grok-4-fast"; +const DEFAULT_MODEL = "x-ai/grok-4.6"; const DEFAULT_BATCH_SIZE = 30; const MIN_REQUEST_INTERVAL = 1000; // 1 second between requests +/** Maximum characters of document text sent to the LLM for an overview */ +export const OVERVIEW_INPUT_LIMIT = 8000; + export interface SummaryResult { name: string; summary: string; @@ -46,6 +49,67 @@ async function rateLimitDelay(): Promise { lastRequestTime = Date.now(); } +/** + * Generate a one-sentence content overview for a single document. + * + * This is the document-level summary stored in `catalog.summary` (ahead of the + * "Key Concepts"/"Categories" enrichment lines). Used by both the seeding + * pipeline and the `--populate-summaries` backfill so the wording stays + * identical across both paths. + * + * @param text - Document text (only the first {@link OVERVIEW_INPUT_LIMIT} chars are sent) + * @param options.apiKey - OpenRouter API key (defaults to OPENROUTER_API_KEY) + * @param options.model - Model override (defaults to this module's DEFAULT_MODEL) + * @param options.rateLimit - Space requests at least 1s apart (default: false). + * Seeding runs overviews across parallel workers, so it + * opts out; the backfill opts in. + * @throws If the API key is missing or the API returns a non-OK response + */ +export async function generateDocumentOverview( + text: string, + options: { + apiKey?: string; + model?: string; + rateLimit?: boolean; + } = {} +): Promise { + const apiKey = options.apiKey || process.env.OPENROUTER_API_KEY; + if (!apiKey) { + throw new Error('OPENROUTER_API_KEY not set'); + } + + if (options.rateLimit) { + await rateLimitDelay(); + } + + const response = await fetch(`${OPENROUTER_BASE_URL}/chat/completions`, { + method: 'POST', + headers: { + 'Authorization': `Bearer ${apiKey}`, + 'Content-Type': 'application/json', + 'HTTP-Referer': 'https://github.com/adiom-data/lance-mcp', + 'X-Title': 'Lance MCP Server' + }, + body: JSON.stringify({ + model: options.model || DEFAULT_MODEL, + messages: [{ + role: 'user', + content: `Write a high-level one sentence content overview based on the text below. WRITE THE CONTENT OVERVIEW ONLY:\n\n${text.slice(0, OVERVIEW_INPUT_LIMIT)}` + }], + max_tokens: 100, + temperature: 0.3 + }) + }); + + if (!response.ok) { + const errorData = await response.text(); + throw new Error(`LLM API error: ${response.status} - ${errorData}`); + } + + const data = await response.json(); + return data.choices[0].message.content.trim(); +} + /** * Generate summaries for a batch of items */ diff --git a/src/infrastructure/seeding/__tests__/summary-backfill.test.ts b/src/infrastructure/seeding/__tests__/summary-backfill.test.ts new file mode 100644 index 00000000..4c7208d9 --- /dev/null +++ b/src/infrastructure/seeding/__tests__/summary-backfill.test.ts @@ -0,0 +1,397 @@ +/** + * Unit Tests for Summary Backfill (--populate-summaries) + */ + +import { describe, it, expect, beforeEach, afterEach } from 'vitest'; +import * as fs from 'fs'; +import * as path from 'path'; +import * as os from 'os'; +import * as lancedb from '@lancedb/lancedb'; +import { + parseSummaryTargets, + splitCatalogSummary, + buildCatalogSummary, + assembleDocumentText, + isMissingCatalogSummary, + isMissingConceptSummary, + isMissingCategorySummary, + backfillSummaries, + SUMMARY_TARGETS +} from '../index.js'; + +describe('parseSummaryTargets', () => { + it('defaults to every table for a bare flag', () => { + expect(parseSummaryTargets(undefined)).toEqual([...SUMMARY_TARGETS]); + expect(parseSummaryTargets('')).toEqual([...SUMMARY_TARGETS]); + expect(parseSummaryTargets(true)).toEqual([...SUMMARY_TARGETS]); + expect(parseSummaryTargets('all')).toEqual([...SUMMARY_TARGETS]); + }); + + it('parses a comma-separated subset in canonical order', () => { + expect(parseSummaryTargets('concepts,catalog')).toEqual(['catalog', 'concepts']); + expect(parseSummaryTargets(' CATEGORIES ')).toEqual(['categories']); + }); + + it('drops duplicates', () => { + expect(parseSummaryTargets('concepts,concepts')).toEqual(['concepts']); + }); + + it('throws on an unknown target', () => { + expect(() => parseSummaryTargets('chunks')).toThrow(/Unknown summary target/); + }); +}); + +describe('splitCatalogSummary', () => { + it('separates the overview from the enrichment lines', () => { + const summary = 'A book about testing.\n\nKey Concepts: a, b\nCategories: x'; + expect(splitCatalogSummary(summary)).toEqual({ + overview: 'A book about testing.', + enrichment: 'Key Concepts: a, b\nCategories: x' + }); + }); + + it('treats a bare summary as all overview', () => { + expect(splitCatalogSummary('Just an overview.').overview).toBe('Just an overview.'); + expect(splitCatalogSummary('').overview).toBe(''); + }); +}); + +describe('buildCatalogSummary', () => { + it('matches the shape written during seeding', () => { + expect(buildCatalogSummary('An overview.', ['a', 'b'], ['x'])).toBe( + 'An overview.\n\nKey Concepts: a, b\nCategories: x' + ); + }); + + it('drops placeholder entries', () => { + expect(buildCatalogSummary('An overview.', ['', 'a'], [''])).toBe( + 'An overview.\n\nKey Concepts: a\nCategories:' + ); + }); + + it('round-trips through splitCatalogSummary', () => { + const summary = buildCatalogSummary('An overview.', ['a'], ['x']); + expect(splitCatalogSummary(summary).overview).toBe('An overview.'); + }); +}); + +describe('isMissingCatalogSummary', () => { + it('accepts a real overview', () => { + expect(isMissingCatalogSummary({ summary: 'A thorough guide to distributed systems.' })).toBe(false); + }); + + it('flags empty and too-short summaries', () => { + expect(isMissingCatalogSummary({ summary: '' })).toBe(true); + expect(isMissingCatalogSummary({})).toBe(true); + expect(isMissingCatalogSummary({ summary: 'short' })).toBe(true); + }); + + it('flags the seeder fallback and failure markers', () => { + expect(isMissingCatalogSummary({ summary: 'Document overview (42 pages)' })).toBe(true); + expect(isMissingCatalogSummary({ summary: 'LLM summarization failed for this document' })).toBe(true); + }); + + it('flags a row whose overview is missing but enrichment survives', () => { + expect(isMissingCatalogSummary({ summary: '\n\nKey Concepts: a\nCategories: x' })).toBe(true); + }); + + it('flags everything under force', () => { + expect(isMissingCatalogSummary({ summary: 'A thorough guide to distributed systems.' }, true)).toBe(true); + }); +}); + +describe('isMissingConceptSummary', () => { + it('flags empty summaries only', () => { + expect(isMissingConceptSummary({ summary: '' })).toBe(true); + expect(isMissingConceptSummary({ summary: ' ' })).toBe(true); + expect(isMissingConceptSummary({})).toBe(true); + expect(isMissingConceptSummary({ summary: 'A definition.' })).toBe(false); + }); + + it('flags everything under force', () => { + expect(isMissingConceptSummary({ summary: 'A definition.' }, true)).toBe(true); + }); +}); + +describe('isMissingCategorySummary', () => { + it('flags empty summaries', () => { + expect(isMissingCategorySummary({ summary: '', description: 'd' })).toBe(true); + }); + + it('flags a summary that is just the generated description', () => { + const description = 'Concepts and practices related to cryptography'; + expect(isMissingCategorySummary({ summary: description, description })).toBe(true); + }); + + it('accepts a distinct summary', () => { + expect( + isMissingCategorySummary({ + summary: 'Cryptography covers the design of ciphers and protocols.', + description: 'Concepts and practices related to cryptography' + }) + ).toBe(false); + }); +}); + +describe('assembleDocumentText', () => { + it('orders chunks by page number', () => { + const text = assembleDocumentText([ + { text: 'third', page_number: 3 }, + { text: 'first', page_number: 1 }, + { text: 'second', page_number: 2 } + ]); + + expect(text).toBe('first\n\nsecond\n\nthird'); + }); + + it('keeps storage order for chunks sharing a page number', () => { + // EPUBs store every chunk as page 1; the scan order is the reading order + const text = assembleDocumentText([ + { text: 'front matter', page_number: 1 }, + { text: 'chapter one', page_number: 1 }, + { text: 'chapter two', page_number: 1 } + ]); + + expect(text).toBe('front matter\n\nchapter one\n\nchapter two'); + }); + + it('skips empty chunks and truncates to the limit', () => { + const text = assembleDocumentText( + [{ text: '', page_number: 1 }, { text: 'abcdefghij', page_number: 2 }], + 4 + ); + + expect(text).toBe('abcd'); + }); +}); + +describe('backfillSummaries', () => { + let tempDir: string; + let db: lancedb.Connection; + + // Must match SimpleEmbeddingService's dimension: the backfill re-embeds catalog summaries + const vector = () => Array.from({ length: 384 }, (_, i) => (i % 10) / 10); + + beforeEach(async () => { + tempDir = await fs.promises.mkdtemp(path.join(os.tmpdir(), 'summary-backfill-test-')); + db = await lancedb.connect(path.join(tempDir, 'db')); + + await db.createTable('catalog', [ + { + id: 1, + hash: 'aaa', + source: '/docs/complete.pdf', + summary: 'An existing overview of a complete document.\n\nKey Concepts: kept\nCategories: kept', + concept_names: ['kept'], + category_names: ['kept'], + vector: vector() + }, + { + id: 2, + hash: 'bbb', + source: '/docs/fallback.pdf', + summary: 'Document overview (12 pages)', + concept_names: ['alpha', 'beta'], + category_names: ['engineering'], + vector: vector() + } + ]); + + await db.createTable('chunks', [ + { id: 10, hash: 'bbb', page_number: 2, text: 'second page text', vector: vector() }, + { id: 11, hash: 'bbb', page_number: 1, text: 'first page text', vector: vector() } + ]); + + await db.createTable('concepts', [ + { + id: 100, + name: 'alpha', + summary: '', + catalog_ids: [1], + catalog_titles: ['/docs/fallback.pdf'], + chunk_ids: [10], + adjacent_ids: [101], + related_ids: [101], + synonyms: ['a'], + broader_terms: [''], + narrower_terms: [''], + weight: 1, + vector: vector() + }, + { + id: 101, + name: 'beta', + summary: 'Beta already has a summary.', + catalog_ids: [1], + catalog_titles: ['/docs/fallback.pdf'], + chunk_ids: [11], + adjacent_ids: [100], + related_ids: [100], + synonyms: ['b'], + broader_terms: [''], + narrower_terms: [''], + weight: 1, + vector: vector() + } + ]); + + await db.createTable('categories', [ + { + id: 200, + category: 'engineering', + description: 'Concepts and practices related to engineering', + summary: 'Concepts and practices related to engineering', + parent_category_id: 0, + aliases: [''], + related_categories: [0], + document_count: 1, + chunk_count: 2, + concept_count: 2, + vector: vector() + } + ]); + }); + + afterEach(async () => { + await fs.promises.rm(tempDir, { recursive: true, force: true }); + }); + + const stubs = () => ({ + log: () => {}, + overviewGenerator: async (text: string) => `Overview of: ${text.slice(0, 16)}`, + batchGenerator: async (names: string[]) => + new Map(names.map(n => [n.toLowerCase(), `Summary of ${n}.`])) + }); + + it('reports missing summaries without writing under dry run', async () => { + const report = await backfillSummaries(db, { ...stubs(), dryRun: true }); + + expect(report.dryRun).toBe(true); + expect(report.results.map(r => [r.table, r.missing, r.written])).toEqual([ + ['catalog', 1, 0], + ['concepts', 1, 0], + ['categories', 1, 0] + ]); + + const concepts = await (await db.openTable('concepts')).query().limit(10).toArray(); + expect(concepts.find((c: any) => c.name === 'alpha')?.summary).toBe(''); + }); + + it('fills missing summaries across all three tables', async () => { + const report = await backfillSummaries(db, stubs()); + + expect(report.results.map(r => [r.table, r.written, r.failed])).toEqual([ + ['catalog', 1, 0], + ['concepts', 1, 0], + ['categories', 1, 0] + ]); + + const catalog = await (await db.openTable('catalog')).query().limit(10).toArray(); + const repaired: any = catalog.find((r: any) => r.id === 2); + // Overview regenerated from chunks, in page order, with enrichment rebuilt + expect(repaired.summary).toBe( + 'Overview of: first page text\n\nKey Concepts: alpha, beta\nCategories: engineering' + ); + + const untouched: any = catalog.find((r: any) => r.id === 1); + expect(untouched.summary).toContain('An existing overview'); + + const concepts = await (await db.openTable('concepts')).query().limit(10).toArray(); + expect(concepts.find((c: any) => c.name === 'alpha')?.summary).toBe('Summary of alpha.'); + expect(concepts.find((c: any) => c.name === 'beta')?.summary).toBe('Beta already has a summary.'); + + const categories = await (await db.openTable('categories')).query().limit(10).toArray(); + expect(categories[0].summary).toBe('Summary of engineering.'); + }); + + it('preserves every other column and the table schema', async () => { + const before = await (await db.openTable('concepts')).schema(); + + await backfillSummaries(db, { ...stubs(), targets: ['concepts'] }); + + // Re-open: LanceDB table handles are pinned to the version they were opened at + const conceptsTable = await db.openTable('concepts'); + const after = await conceptsTable.schema(); + expect(after.fields.map((f: any) => `${f.name}:${f.type}`)).toEqual( + before.fields.map((f: any) => `${f.name}:${f.type}`) + ); + + const rows = await conceptsTable.query().limit(10).toArray(); + expect(rows).toHaveLength(2); + + const alpha: any = rows.find((r: any) => r.name === 'alpha'); + expect(alpha.summary).toBe('Summary of alpha.'); + expect(Array.from(alpha.catalog_titles.toArray())).toEqual(['/docs/fallback.pdf']); + expect(Array.from(alpha.adjacent_ids.toArray())).toEqual([101]); + expect(Array.from(alpha.related_ids.toArray())).toEqual([101]); + expect(Array.from(alpha.chunk_ids.toArray())).toEqual([10]); + expect(Array.from(alpha.synonyms.toArray())).toEqual(['a']); + expect(Array.from(alpha.vector.toArray())).toHaveLength(384); + expect(alpha.weight).toBe(1); + }); + + it('refreshes the catalog vector, which embeds the summary text', async () => { + const before: any = (await (await db.openTable('catalog')).query().limit(10).toArray()).find( + (r: any) => r.id === 2 + ); + const beforeVector = Array.from(before.vector.toArray()); + + await backfillSummaries(db, { ...stubs(), targets: ['catalog'] }); + + const after: any = (await (await db.openTable('catalog')).query().limit(10).toArray()).find( + (r: any) => r.id === 2 + ); + expect(Array.from(after.vector.toArray())).not.toEqual(beforeVector); + // Untouched rows keep their original vector + const untouched: any = (await (await db.openTable('catalog')).query().limit(10).toArray()).find( + (r: any) => r.id === 1 + ); + expect(Array.from(untouched.vector.toArray())).toEqual(beforeVector); + }); + + it('regenerates everything under force', async () => { + const report = await backfillSummaries(db, { ...stubs(), force: true, targets: ['concepts'] }); + + expect(report.results[0].missing).toBe(2); + expect(report.results[0].written).toBe(2); + + const rows = await (await db.openTable('concepts')).query().limit(10).toArray(); + expect(rows.find((c: any) => c.name === 'beta')?.summary).toBe('Summary of beta.'); + }); + + it('caps work with maxItems', async () => { + const report = await backfillSummaries(db, { + ...stubs(), + force: true, + targets: ['concepts'], + maxItems: 1 + }); + + expect(report.results[0].written).toBe(1); + }); + + it('counts a document with no chunks as failed instead of throwing', async () => { + await (await db.openTable('chunks')).delete('hash = "bbb"'); + + const report = await backfillSummaries(db, { ...stubs(), targets: ['catalog'] }); + + expect(report.results[0]).toMatchObject({ missing: 1, written: 0, failed: 1 }); + }); + + it('skips tables that do not exist', async () => { + const emptyDb = await lancedb.connect(path.join(tempDir, 'empty')); + const report = await backfillSummaries(emptyDb, stubs()); + + expect(report.results.every(r => r.skippedReason === 'table not found')).toBe(true); + }); + + it('records a failure when the generator returns nothing for a name', async () => { + const report = await backfillSummaries(db, { + ...stubs(), + targets: ['concepts'], + batchGenerator: async () => new Map() + }); + + expect(report.results[0]).toMatchObject({ missing: 1, generated: 0, written: 0, failed: 1 }); + }); +}); diff --git a/src/infrastructure/seeding/index.ts b/src/infrastructure/seeding/index.ts index 3aa860a8..0d59e990 100644 --- a/src/infrastructure/seeding/index.ts +++ b/src/infrastructure/seeding/index.ts @@ -24,3 +24,20 @@ export { truncateFilePath, formatHashDisplay } from './string-utils.js'; + +export { + backfillSummaries, + parseSummaryTargets, + splitCatalogSummary, + buildCatalogSummary, + assembleDocumentText, + isMissingCatalogSummary, + isMissingConceptSummary, + isMissingCategorySummary, + SUMMARY_TARGETS, + type SummaryTarget, + type SummaryBackfillOptions, + type SummaryBackfillReport, + type TableBackfillResult, + type SummaryProgress +} from './summary-backfill.js'; diff --git a/src/infrastructure/seeding/summary-backfill.ts b/src/infrastructure/seeding/summary-backfill.ts new file mode 100644 index 00000000..7829c012 --- /dev/null +++ b/src/infrastructure/seeding/summary-backfill.ts @@ -0,0 +1,586 @@ +/** + * Summary Backfill + * + * Fills in missing summaries in an already-seeded database without re-running + * document loading, chunking or concept extraction. Backs the + * `--populate-summaries` flag of the seeding script. + * + * Three tables carry summaries, and each is repaired from data already in the + * database: + * - `catalog.summary` - document overview, regenerated from existing chunks, + * re-enriched with the row's concept/category names, and + * re-embedded (the catalog vector embeds this text). + * - `concepts.summary` - one-sentence definition generated from the concept name. + * - `categories.summary` - one-sentence description generated from the category name. + * + * Rows are written back with a merge-insert keyed on `id`, so every column the + * table already has (including `catalog_titles`, `adjacent_ids`, `related_ids`) + * survives untouched and the Arrow schema is never rewritten. + */ + +import * as lancedb from '@lancedb/lancedb'; +import { + generateSummaries, + generateDocumentOverview +} from '../../concepts/summary_generator.js'; +import { SimpleEmbeddingService } from '../embeddings/simple-embedding-service.js'; + +/** Tables that carry a summary field. */ +export const SUMMARY_TARGETS = ['catalog', 'concepts', 'categories'] as const; + +export type SummaryTarget = (typeof SUMMARY_TARGETS)[number]; + +/** Separator between the document overview and its enrichment lines in catalog.summary */ +const CONCEPTS_MARKER = '\n\nKey Concepts:'; + +/** Fallback text written by the seeder when overview generation fails */ +const OVERVIEW_FALLBACK_PREFIX = 'Document overview ('; + +/** An overview shorter than this is treated as junk rather than a summary */ +const MIN_OVERVIEW_LENGTH = 10; + +/** Characters of document text assembled from chunks before summarising */ +const DOCUMENT_TEXT_LIMIT = 10000; + +const DEFAULT_BATCH_SIZE = 30; +const DEFAULT_FLUSH_SIZE = 250; + +/** + * Generates summaries for a batch of concept/category names. + * Returns a map of lowercased name to summary. Injectable for testing. + */ +export type BatchSummaryGenerator = ( + names: string[], + type: 'concept' | 'category' +) => Promise>; + +/** + * Generates a one-sentence overview for a single document's text. + * Injectable for testing. + */ +export type OverviewGenerator = (text: string) => Promise; + +export interface SummaryBackfillOptions { + /** Tables to process (default: all three) */ + targets?: SummaryTarget[]; + /** OpenRouter API key (defaults to OPENROUTER_API_KEY) */ + apiKey?: string; + /** Model override for summary generation */ + model?: string; + /** Items per LLM request for concepts/categories */ + batchSize?: number; + /** Summaries buffered before writing back to the table */ + flushSize?: number; + /** Regenerate every summary instead of only the missing ones */ + force?: boolean; + /** Report what would change without calling the LLM or writing */ + dryRun?: boolean; + /** Process at most N items per table (useful for trial runs) */ + maxItems?: number; + /** Milestone logger (default: console.log) */ + log?: (message: string) => void; + /** Progress callback, called after each batch */ + onProgress?: (progress: SummaryProgress) => void; + /** Override batch summary generation (tests) */ + batchGenerator?: BatchSummaryGenerator; + /** Override document overview generation (tests) */ + overviewGenerator?: OverviewGenerator; +} + +export interface SummaryProgress { + table: SummaryTarget; + completed: number; + total: number; +} + +export interface TableBackfillResult { + table: SummaryTarget; + /** Rows in the table */ + total: number; + /** Rows found to be missing a summary */ + missing: number; + /** Summaries the LLM returned */ + generated: number; + /** Rows written back to the table */ + written: number; + /** Rows that were missing a summary but got no usable result */ + failed: number; + /** Set when the table was not processed at all */ + skippedReason?: string; +} + +export interface SummaryBackfillReport { + dryRun: boolean; + results: TableBackfillResult[]; +} + +const embeddingService = new SimpleEmbeddingService(); + +/** + * Parse the `--populate-summaries` value into a list of tables. + * + * Accepts an empty string or "all" for every table, or a comma-separated + * subset such as "concepts,catalog". + * + * @throws If a token is not a known table + */ +export function parseSummaryTargets(value?: string | boolean): SummaryTarget[] { + if (value === undefined || value === null || value === true || value === '') { + return [...SUMMARY_TARGETS]; + } + + const tokens = String(value) + .split(',') + .map(t => t.trim().toLowerCase()) + .filter(t => t.length > 0); + + if (tokens.length === 0 || tokens.includes('all')) { + return [...SUMMARY_TARGETS]; + } + + const unknown = tokens.filter(t => !SUMMARY_TARGETS.includes(t as SummaryTarget)); + if (unknown.length > 0) { + throw new Error( + `Unknown summary target(s): ${unknown.join(', ')}. Valid targets: ${SUMMARY_TARGETS.join(', ')}, all` + ); + } + + // Preserve canonical order and drop duplicates + return SUMMARY_TARGETS.filter(t => tokens.includes(t)); +} + +/** + * Split a catalog summary into its LLM overview and the enrichment lines + * ("Key Concepts: ..." / "Categories: ...") appended during seeding. + */ +export function splitCatalogSummary(summary: string): { overview: string; enrichment: string } { + const text = summary || ''; + const markerIndex = text.indexOf(CONCEPTS_MARKER); + + if (markerIndex === -1) { + return { overview: text.trim(), enrichment: '' }; + } + + return { + overview: text.slice(0, markerIndex).trim(), + enrichment: text.slice(markerIndex).trim() + }; +} + +/** + * Rebuild a catalog summary in the exact shape the seeder writes: + * overview, blank line, concept names, category names. + */ +export function buildCatalogSummary( + overview: string, + conceptNames: string[], + categoryNames: string[] +): string { + const concepts = conceptNames.filter(n => n && n.trim().length > 0); + const categories = categoryNames.filter(n => n && n.trim().length > 0); + + return `${overview}\n\nKey Concepts: ${concepts.join(', ')}\nCategories: ${categories.join(', ')}`.trim(); +} + +/** + * Is this catalog row missing a usable document overview? + * + * True for an empty summary, the seeder's "Document overview (N pages)" + * fallback, a failed-summarisation marker, or an overview too short to be one. + */ +export function isMissingCatalogSummary(row: { summary?: string }, force = false): boolean { + if (force) return true; + + const { overview } = splitCatalogSummary(row.summary ?? ''); + + return ( + overview.length < MIN_OVERVIEW_LENGTH || + overview.startsWith(OVERVIEW_FALLBACK_PREFIX) || + overview.includes('LLM summarization failed') + ); +} + +/** Is this concept row missing a summary? */ +export function isMissingConceptSummary(row: { summary?: string }, force = false): boolean { + if (force) return true; + return (row.summary ?? '').trim().length === 0; +} + +/** + * Is this category row missing a summary? + * + * Also true when the summary is just the generated description, which is what + * seeding falls back to when the LLM call is skipped or fails. + */ +export function isMissingCategorySummary( + row: { summary?: string; description?: string }, + force = false +): boolean { + if (force) return true; + + const summary = (row.summary ?? '').trim(); + if (summary.length === 0) return true; + + const description = (row.description ?? '').trim(); + return description.length > 0 && summary === description; +} + +/** + * Assemble document text from its chunks, in reading order, for summarisation. + * + * Chunks are ordered by page number with a *stable* sort, so chunks sharing a + * page keep the order they came back from the table in - which is insertion + * order, i.e. the order they were split from the document. Chunk ids are + * hash-based and carry no sequence, so they must not be used as a tiebreak: + * EPUBs store every chunk as page 1, and sorting those by id shuffles the book. + */ +export function assembleDocumentText( + chunks: Array<{ text?: string; page_number?: number }>, + maxChars: number = DOCUMENT_TEXT_LIMIT +): string { + const ordered = [...chunks].sort((a, b) => (a.page_number ?? 0) - (b.page_number ?? 0)); + + return ordered + .map(c => c.text ?? '') + .filter(t => t.length > 0) + .join('\n\n') + .slice(0, maxChars); +} + +/** Convert an Arrow-backed value to a plain JS value. */ +function toPlain(value: any): any { + if (value && typeof value === 'object' && typeof value.toArray === 'function') { + return Array.from(value.toArray()); + } + return value; +} + +/** Read a row's array field as a plain string array. */ +function toStringArray(value: any): string[] { + const plain = toPlain(value); + if (!Array.isArray(plain)) return []; + return plain.map(v => String(v ?? '')).filter(v => v.length > 0); +} + +/** + * Write updated rows back via merge-insert on `id`. + * + * Rows are reconstructed field-by-field from the live schema, so unknown or + * newly added columns pass through unchanged. + */ +async function writeUpdatedRows( + table: lancedb.Table, + fieldNames: string[], + updates: Array<{ row: any; values: Record }> +): Promise { + if (updates.length === 0) return 0; + + const data = updates.map(({ row, values }) => { + const record: Record = {}; + for (const field of fieldNames) { + record[field] = field in values ? values[field] : toPlain(row[field]); + } + return record; + }); + + await table.mergeInsert('id').whenMatchedUpdateAll().execute(data); + return data.length; +} + +/** + * Populate missing summaries across the catalog, concepts and categories tables. + * + * Every LLM result is written back in batches, so an interrupted run keeps the + * summaries it already produced and a re-run picks up where it stopped. + */ +export async function backfillSummaries( + db: lancedb.Connection, + options: SummaryBackfillOptions = {} +): Promise { + const targets = options.targets ?? [...SUMMARY_TARGETS]; + const batchSize = options.batchSize ?? DEFAULT_BATCH_SIZE; + const flushSize = options.flushSize ?? DEFAULT_FLUSH_SIZE; + const force = options.force ?? false; + const dryRun = options.dryRun ?? false; + const log = options.log ?? ((message: string) => console.log(message)); + + const generateBatch: BatchSummaryGenerator = + options.batchGenerator ?? + ((names, type) => + generateSummaries(names, type, { + apiKey: options.apiKey, + model: options.model, + batchSize, + onProgress: () => {} // progress is reported per batch by this module + })); + + const generateOverview: OverviewGenerator = + options.overviewGenerator ?? + (text => + generateDocumentOverview(text, { + apiKey: options.apiKey, + model: options.model, + rateLimit: true + })); + + const tableNames = await db.tableNames(); + const results: TableBackfillResult[] = []; + + for (const target of targets) { + if (!tableNames.includes(target)) { + results.push({ + table: target, + total: 0, + missing: 0, + generated: 0, + written: 0, + failed: 0, + skippedReason: 'table not found' + }); + continue; + } + + if (target === 'catalog') { + results.push( + await backfillCatalog(db, { + force, + dryRun, + flushSize, + maxItems: options.maxItems, + log, + onProgress: options.onProgress, + generateOverview + }) + ); + } else { + results.push( + await backfillNamedTable(db, target, { + force, + dryRun, + batchSize, + flushSize, + maxItems: options.maxItems, + log, + onProgress: options.onProgress, + generateBatch + }) + ); + } + } + + return { dryRun, results }; +} + +/** + * Regenerate document overviews for catalog rows, from their existing chunks. + */ +async function backfillCatalog( + db: lancedb.Connection, + opts: { + force: boolean; + dryRun: boolean; + flushSize: number; + maxItems?: number; + log: (message: string) => void; + onProgress?: (progress: SummaryProgress) => void; + generateOverview: OverviewGenerator; + } +): Promise { + const table = await db.openTable('catalog'); + const schema = await table.schema(); + const fieldNames = schema.fields.map((f: any) => f.name); + + const rows = await table.query().limit(1000000).toArray(); + const missingRows = rows.filter((r: any) => isMissingCatalogSummary(r, opts.force)); + const targetRows = opts.maxItems ? missingRows.slice(0, opts.maxItems) : missingRows; + + const result: TableBackfillResult = { + table: 'catalog', + total: rows.length, + missing: missingRows.length, + generated: 0, + written: 0, + failed: 0 + }; + + opts.log(`šŸ“š catalog: ${missingRows.length}/${rows.length} document(s) need a summary`); + + if (targetRows.length === 0 || opts.dryRun) { + return result; + } + + const chunksAvailable = (await db.tableNames()).includes('chunks'); + if (!chunksAvailable) { + result.skippedReason = 'chunks table not found (document text unavailable)'; + opts.log(` āš ļø ${result.skippedReason} - skipping catalog summaries`); + return result; + } + const chunksTable = await db.openTable('chunks'); + + let pending: Array<{ row: any; values: Record }> = []; + + const flush = async () => { + result.written += await writeUpdatedRows(table, fieldNames, pending); + pending = []; + }; + + for (let i = 0; i < targetRows.length; i++) { + const row: any = targetRows[i]; + const hash = String(row.hash ?? '').replace(/"/g, ''); + + try { + const chunks = hash + ? await chunksTable.query().where(`hash = "${hash}"`).limit(100000).toArray() + : []; + + if (chunks.length === 0) { + result.failed++; + opts.log(` āš ļø No chunks found for ${row.source ?? row.id} - cannot regenerate summary`); + continue; + } + + const text = assembleDocumentText( + chunks.map((c: any) => ({ + text: c.text, + page_number: typeof c.page_number === 'number' ? c.page_number : Number(c.page_number ?? 0) + })) + ); + + if (text.trim().length === 0) { + result.failed++; + opts.log(` āš ļø Chunks for ${row.source ?? row.id} contain no text`); + continue; + } + + const overview = (await opts.generateOverview(text)).trim(); + if (overview.length === 0) { + result.failed++; + continue; + } + + const summary = buildCatalogSummary( + overview, + toStringArray(row.concept_names), + toStringArray(row.category_names) + ); + + result.generated++; + pending.push({ + row, + values: { + summary, + // The catalog vector embeds the summary text, so it must be refreshed too + vector: embeddingService.generateEmbedding(summary) + } + }); + + if (pending.length >= opts.flushSize) { + await flush(); + } + } catch (error: any) { + result.failed++; + opts.log(` āŒ Summary failed for ${row.source ?? row.id}: ${error.message}`); + } + + opts.onProgress?.({ table: 'catalog', completed: i + 1, total: targetRows.length }); + } + + await flush(); + opts.log(` āœ… catalog: wrote ${result.written} summary/summaries`); + + return result; +} + +/** + * Generate summaries for the concepts or categories table from their names. + */ +async function backfillNamedTable( + db: lancedb.Connection, + target: 'concepts' | 'categories', + opts: { + force: boolean; + dryRun: boolean; + batchSize: number; + flushSize: number; + maxItems?: number; + log: (message: string) => void; + onProgress?: (progress: SummaryProgress) => void; + generateBatch: BatchSummaryGenerator; + } +): Promise { + const isConcepts = target === 'concepts'; + const nameField = isConcepts ? 'name' : 'category'; + const type = isConcepts ? 'concept' : 'category'; + const icon = isConcepts ? 'šŸ§ ' : 'šŸ“‚'; + + const table = await db.openTable(target); + const schema = await table.schema(); + const fieldNames = schema.fields.map((f: any) => f.name); + + const rows = await table.query().limit(1000000).toArray(); + const missingRows = rows.filter((r: any) => + isConcepts ? isMissingConceptSummary(r, opts.force) : isMissingCategorySummary(r, opts.force) + ); + const targetRows = (opts.maxItems ? missingRows.slice(0, opts.maxItems) : missingRows).filter( + (r: any) => String(r[nameField] ?? '').trim().length > 0 + ); + + const result: TableBackfillResult = { + table: target, + total: rows.length, + missing: missingRows.length, + generated: 0, + written: 0, + failed: 0 + }; + + opts.log(`${icon} ${target}: ${missingRows.length}/${rows.length} row(s) need a summary`); + + if (targetRows.length === 0 || opts.dryRun) { + return result; + } + + let pending: Array<{ row: any; values: Record }> = []; + + const flush = async () => { + result.written += await writeUpdatedRows(table, fieldNames, pending); + pending = []; + }; + + for (let i = 0; i < targetRows.length; i += opts.batchSize) { + const batch = targetRows.slice(i, i + opts.batchSize); + const names = batch.map((r: any) => String(r[nameField])); + + let summaries = new Map(); + try { + summaries = await opts.generateBatch(names, type); + } catch (error: any) { + opts.log(` āŒ Batch failed (${names.length} ${type}s): ${error.message}`); + } + + for (const row of batch) { + const summary = summaries.get(String(row[nameField]).toLowerCase().trim()); + if (summary && summary.trim().length > 0) { + result.generated++; + pending.push({ row, values: { summary: summary.trim() } }); + } else { + result.failed++; + } + } + + if (pending.length >= opts.flushSize) { + await flush(); + } + + opts.onProgress?.({ + table: target, + completed: Math.min(i + opts.batchSize, targetRows.length), + total: targetRows.length + }); + } + + await flush(); + opts.log(` āœ… ${target}: wrote ${result.written} summary/summaries`); + + return result; +}