From a36ba48aac293ea69c57a97db857bbc7f0134c28 Mon Sep 17 00:00:00 2001 From: Mike Clay Date: Sun, 16 Aug 2026 08:05:56 +0100 Subject: [PATCH] feat(seeding): add --populate-summaries to rebuild summaries only Summaries could only be repaired by re-seeding, which reloads, re-chunks and re-extracts concepts for every document. The stage cache keys the document overview and the concept extraction together, so there was no way to invalidate one without paying for the other. Adds a standalone mode to the seeding script that fills in missing summaries against an existing database, with no --filesdir: - catalog.summary regenerated from existing chunks in reading order, re-enriched with the row's concept/category names, and re-embedded (the catalog vector encodes it) - concepts.summary generated from the concept name - categories.summary generated from the category name A summary counts as missing when it is empty or a seeding fallback ("Document overview (N pages)", or the generated description for a category). --force-summaries regenerates everything, --dry-run reports without calling the LLM or writing. Writes go through a merge-insert keyed on id, with rows rebuilt from the live schema, so every other column and the Arrow schema are untouched and an interrupted run keeps what it already wrote. Chunks are ordered for summarisation by a stable sort on page number only. Chunk ids are hash-based and carry no sequence: EPUBs store every chunk as page 1, so using ids as a tiebreak shuffles the document. generateDocumentOverview moves into summary_generator.ts and the seeder delegates to it, so both paths share one prompt and cannot drift. Deletes scripts/populate_summaries.ts, which this replaces. It rewrote the concepts table from a hard-coded field list that predated the current schema, dropping catalog_titles and adjacent_ids and renaming related_ids - all three are read at query time by the concept repository. Verified against a copy of the live database: row counts and Arrow schemas identical across all three tables, summaries updated with no drift in any other column. Co-Authored-By: Claude Opus 5 (1M context) --- USAGE.md | 42 ++ docs/development.md | 18 + hybrid_fast_seed.ts | 146 ++++- scripts/populate_summaries.ts | 476 -------------- src/concepts/summary_generator.ts | 68 +- .../__tests__/summary-backfill.test.ts | 397 ++++++++++++ src/infrastructure/seeding/index.ts | 17 + .../seeding/summary-backfill.ts | 586 ++++++++++++++++++ 8 files changed, 1243 insertions(+), 507 deletions(-) delete mode 100644 scripts/populate_summaries.ts create mode 100644 src/infrastructure/seeding/__tests__/summary-backfill.test.ts create mode 100644 src/infrastructure/seeding/summary-backfill.ts 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; +}