feat: refactor embedder pkg to use db

Author: caiopizzol

.env.example

@@ -1,5 +1,5 @@
-# Database connection
-DATABASE_URL=postgres://postgres:postgres@localhost:5432/corpus
+# Database connection (use docker-compose up to start local postgres)
+DATABASE_URL=postgres://postgres:postgres@localhost:5432/docx_corpus
 
 # Local storage path
 STORAGE_PATH=./corpus
@@ -20,6 +20,12 @@ EXTRACT_OUTPUT_PREFIX=extracted   # Output directory prefix
 EXTRACT_BATCH_SIZE=100            # Documents per batch
 EXTRACT_WORKERS=4                 # Parallel worker processes
 
+# Embedder settings
+EMBED_INPUT_PREFIX=extracted      # Input directory prefix (extracted text)
+EMBED_MODEL=minilm                # Model: minilm, bge-m3, or voyage-lite
+EMBED_BATCH_SIZE=100              # Documents per batch
+VOYAGE_API_KEY=                   # Required for voyage-lite model
+
 # Cloudflare R2 (optional - for cloud storage)
 CLOUDFLARE_ACCOUNT_ID=
 R2_ACCESS_KEY_ID=

README.md

@@ -69,22 +69,35 @@ bun install
 
 ```
 packages/
-  shared/         # Shared utilities (progress bars, formatting)
+  shared/         # Shared utilities (DB client, storage, formatting)
   scraper/        # Core scraper logic (downloads WARC, validates .docx)
   extractor/      # Text extraction using Docling (Python)
+  embedder/       # Document embeddings using sentence-transformers (Python)
 apps/
   cli/            # Unified CLI - corpus <command>
   cdx-filter/     # AWS Lambda - filters CDX indexes for .docx URLs
   web/            # Landing page - docxcorp.us
+db/
+  schema.sql      # PostgreSQL schema (with pgvector)
+  migrations/     # Database migrations
 ```
 
-| Package/App    | Purpose                           | Runtime              |
-| -------------- | --------------------------------- | -------------------- |
-| **cli**        | Unified CLI entry point           | Bun                  |
-| **scraper**    | Download and validate .docx files | Bun                  |
-| **extractor**  | Extract text from .docx files     | Bun + Python         |
-| **cdx-filter** | Filter Common Crawl CDX indexes   | AWS Lambda (Node.js) |
-| **web**        | Landing page                      | Static HTML          |
+**Apps** (entry points)
+
+| App            | Purpose                         | Uses                     |
+| -------------- | ------------------------------- | ------------------------ |
+| **cli**        | `corpus` command                | scraper, extractor, embedder |
+| **cdx-filter** | Filter CDX indexes (Lambda)     | -                        |
+| **web**        | Landing page                    | -                        |
+
+**Packages** (libraries)
+
+| Package        | Purpose                           | Runtime      |
+| -------------- | --------------------------------- | ------------ |
+| **shared**     | DB client, storage, formatting    | Bun          |
+| **scraper**    | Download and validate .docx files | Bun          |
+| **extractor**  | Extract text (Docling)            | Bun + Python |
+| **embedder**   | Generate embeddings               | Bun + Python |
 
 ## Usage
 
@@ -134,6 +147,20 @@ bun run corpus extract --batch 50 --workers 8
 bun run corpus extract --verbose
 ```
 
+### 4. Generate embeddings
+
+```bash
+# Embed all extracted documents (default: minilm, 384 dims)
+bun run corpus embed
+
+# Use a different model
+bun run corpus embed --model bge-m3      # 1024 dims
+bun run corpus embed --model voyage-lite  # requires VOYAGE_API_KEY
+
+# Embed with batch limit
+bun run corpus embed --batch 100 --verbose
+```
+
 ### Docker
 
 Run the CLI in a container:
@@ -173,35 +200,62 @@ export R2_SECRET_ACCESS_KEY=xxx
 bun run corpus scrape --crawl CC-MAIN-2025-51 --batch 1000
 ```
 
+## Local Development
+
+Start PostgreSQL with pgvector locally:
+
+```bash
+docker compose up -d
+
+# Verify
+docker exec docx-corpus-postgres-1 psql -U postgres -d docx_corpus -c "\dt"
+```
+
+Run commands against local database:
+
+```bash
+DATABASE_URL=postgres://postgres:postgres@localhost:5432/docx_corpus \
+CLOUDFLARE_ACCOUNT_ID='' \
+bun run corpus status
+```
+
 ## Configuration
 
 All configuration via environment variables (`.env`):
 
 ```bash
-# Cloudflare R2 (required for both Lambda and scraper)
+# Database (required)
+DATABASE_URL=postgres://user:pass@host:5432/dbname
+
+# Cloudflare R2 (required for cloud storage)
 CLOUDFLARE_ACCOUNT_ID=
 R2_ACCESS_KEY_ID=
 R2_SECRET_ACCESS_KEY=
 R2_BUCKET_NAME=docx-corpus
 
-# Scraping
+# Local storage (used when R2 not configured)
 STORAGE_PATH=./corpus
-CRAWL_ID=CC-MAIN-2025-51
 
-# Performance tuning
-CONCURRENCY=50              # Parallel downloads
-RATE_LIMIT_RPS=50           # Requests per second (initial)
-MAX_RPS=100                 # Max requests per second
-MIN_RPS=10                  # Min requests per second
-TIMEOUT_MS=45000            # Request timeout in ms
-MAX_RETRIES=10              # Max retry attempts
-MAX_BACKOFF_MS=60000        # Max backoff delay (ms)
+# Scraping
+CRAWL_ID=CC-MAIN-2025-51
+CONCURRENCY=50
+RATE_LIMIT_RPS=50
+MAX_RPS=100
+MIN_RPS=10
+TIMEOUT_MS=45000
+MAX_RETRIES=10
 
 # Extractor
-EXTRACT_INPUT_PREFIX=documents  # Input directory prefix
-EXTRACT_OUTPUT_PREFIX=extracted # Output directory prefix
-EXTRACT_BATCH_SIZE=100          # Documents per batch
-EXTRACT_WORKERS=4               # Parallel workers
+EXTRACT_INPUT_PREFIX=documents
+EXTRACT_OUTPUT_PREFIX=extracted
+EXTRACT_BATCH_SIZE=100
+EXTRACT_WORKERS=4
+
+# Embedder
+EMBED_INPUT_PREFIX=extracted
+EMBED_MODEL=minilm           # minilm | bge-m3 | voyage-lite
+EMBED_BATCH_SIZE=100
+VOYAGE_API_KEY=              # Required for voyage-lite model
 ```
 
 ### Rate Limiting

apps/cli/commands/embed.ts

@@ -6,12 +6,11 @@ import {
   type EmbedConfig,
   type EmbeddingModel,
 } from "@docx-corpus/embedder";
-import { createLocalStorage, createR2Storage } from "@docx-corpus/shared";
+import { createDb, createLocalStorage, createR2Storage } from "@docx-corpus/shared";
 
 interface ParsedFlags {
   model?: EmbeddingModel;
   batchSize?: number;
-  workers?: number;
   verbose: boolean;
 }
 
@@ -35,11 +34,6 @@ function parseFlags(args: string[]): ParsedFlags {
         flags.batchSize = parseInt(next || "", 10);
         i++;
         break;
-      case "--workers":
-      case "-w":
-        flags.workers = parseInt(next || "", 10);
-        i++;
-        break;
       case "--verbose":
       case "-v":
         flags.verbose = true;
@@ -56,32 +50,31 @@ corpus embed - Generate embeddings for extracted documents
 Usage
   corpus embed [options]
 
+Reads extracted text from storage and writes embeddings to the database.
 Storage is auto-selected based on environment:
-  - With R2 credentials: reads from r2://extracted/, writes to r2://embeddings/
-  - Without R2 credentials: reads from ./corpus/extracted/, writes to ./corpus/embeddings/
+  - With R2 credentials: reads from r2://extracted/
+  - Without R2 credentials: reads from ./corpus/extracted/
 
-Already-embedded files are automatically skipped (tracked in index.jsonl).
+Embedding progress is tracked in the database (embedded_at column).
 
 Options
   --model, -m <name>      Embedding model (default: minilm)
                             minilm      - all-MiniLM-L6-v2 (fast, 384 dims)
                             bge-m3      - BAAI/bge-m3 (better quality, 1024 dims)
                             voyage-lite - Voyage 3.5 lite (best, requires API key)
   --batch, -b <n>         Limit to n documents (default: all)
-  --workers, -w <n>       Number of parallel workers (default: 4)
   --verbose, -v           Show detailed progress
   --help, -h              Show this help
 
 Environment Variables
+  DATABASE_URL            PostgreSQL connection string (required)
   STORAGE_PATH            Local storage path (default: ./corpus)
   CLOUDFLARE_ACCOUNT_ID   Cloudflare account ID (enables R2)
   R2_ACCESS_KEY_ID        R2 access key
   R2_SECRET_ACCESS_KEY    R2 secret key
   R2_BUCKET_NAME          R2 bucket (default: docx-corpus)
   EMBED_INPUT_PREFIX      Input prefix (default: extracted)
-  EMBED_OUTPUT_PREFIX     Output prefix (default: embeddings)
   EMBED_MODEL             Default model (default: minilm)
-  EMBED_WORKERS           Worker count (default: 4)
   VOYAGE_API_KEY          Voyage AI API key (required for voyage-lite)
 
 Examples
@@ -99,6 +92,13 @@ export async function runEmbed(args: string[]) {
 
   const flags = parseFlags(args);
   const envConfig = loadEmbedderConfig();
+
+  // Validate database URL
+  if (!envConfig.database.url) {
+    console.error("Error: DATABASE_URL environment variable is required");
+    process.exit(1);
+  }
+
   const useCloud = hasCloudflareCredentials(envConfig);
   const model = flags.model ?? envConfig.embed.model;
 
@@ -108,6 +108,9 @@ export async function runEmbed(args: string[]) {
     process.exit(1);
   }
 
+  // Create database client
+  const db = await createDb(envConfig.database.url);
+
   // Create storage based on credentials
   const storage = useCloud
     ? createR2Storage({
@@ -119,12 +122,11 @@ export async function runEmbed(args: string[]) {
     : createLocalStorage(envConfig.storage.localPath);
 
   const config: EmbedConfig = {
+    db,
     storage,
     inputPrefix: envConfig.embed.inputPrefix,
-    outputPrefix: envConfig.embed.outputPrefix,
     model,
-    batchSize: flags.batchSize ?? Infinity,
-    workers: flags.workers ?? envConfig.embed.workers,
+    batchSize: flags.batchSize ?? 1000000,
   };
 
   console.log("Document Embedder");
@@ -133,16 +135,17 @@ export async function runEmbed(args: string[]) {
     `Storage: ${useCloud ? `R2 (${envConfig.cloudflare.r2BucketName})` : `local (${envConfig.storage.localPath})`}`
   );
   console.log(`Input:   ${config.inputPrefix}/`);
-  console.log(`Output:  ${config.outputPrefix}/`);
+  console.log(`Output:  database (embedding column)`);
   console.log(`Model:   ${config.model}`);
-  console.log(`Workers: ${config.workers}`);
-  console.log(`Batch:   ${config.batchSize === Infinity ? "all" : config.batchSize}`);
+  console.log(`Batch:   ${config.batchSize >= 1000000 ? "all" : config.batchSize}`);
   console.log("");
 
   try {
     await processEmbeddings(config, flags.verbose);
   } catch (err) {
     console.error("Fatal error:", err);
     process.exit(1);
+  } finally {
+    await db.close();
   }
 }

apps/cli/commands/extract.ts

@@ -4,7 +4,7 @@ import {
   hasCloudflareCredentials,
   type ExtractConfig,
 } from "@docx-corpus/extractor";
-import { createLocalStorage, createR2Storage } from "@docx-corpus/shared";
+import { createDb, createLocalStorage, createR2Storage } from "@docx-corpus/shared";
 
 interface ParsedFlags {
   batchSize?: number;
@@ -52,7 +52,7 @@ Storage is auto-selected based on environment:
   - With R2 credentials: reads from r2://documents/, writes to r2://extracted/
   - Without R2 credentials: reads from ./corpus/documents/, writes to ./corpus/extracted/
 
-Already-extracted files are automatically skipped (tracked in index.jsonl).
+Extraction progress is tracked in the database (extracted_at column).
 
 Options
   --batch, -b <n>         Limit to n documents (default: all)
@@ -61,6 +61,7 @@ Options
   --help, -h              Show this help
 
 Environment Variables
+  DATABASE_URL            PostgreSQL connection string (required)
   STORAGE_PATH            Local storage path (default: ./corpus)
   CLOUDFLARE_ACCOUNT_ID   Cloudflare account ID (enables R2)
   R2_ACCESS_KEY_ID        R2 access key
@@ -85,8 +86,18 @@ export async function runExtract(args: string[]) {
 
   const flags = parseFlags(args);
   const envConfig = loadExtractorConfig();
+
+  // Validate database URL
+  if (!envConfig.database.url) {
+    console.error("Error: DATABASE_URL environment variable is required");
+    process.exit(1);
+  }
+
   const useCloud = hasCloudflareCredentials(envConfig);
 
+  // Create database client
+  const db = await createDb(envConfig.database.url);
+
   // Create storage based on credentials
   const storage = useCloud
     ? createR2Storage({
@@ -98,10 +109,11 @@ export async function runExtract(args: string[]) {
     : createLocalStorage(envConfig.storage.localPath);
 
   const config: ExtractConfig = {
+    db,
     storage,
     inputPrefix: envConfig.extract.inputPrefix,
     outputPrefix: envConfig.extract.outputPrefix,
-    batchSize: flags.batchSize ?? Infinity,
+    batchSize: flags.batchSize ?? 1000000,
     workers: flags.workers ?? envConfig.extract.workers,
   };
 
@@ -113,13 +125,15 @@ export async function runExtract(args: string[]) {
   console.log(`Input:   ${config.inputPrefix}/`);
   console.log(`Output:  ${config.outputPrefix}/`);
   console.log(`Workers: ${config.workers}`);
-  console.log(`Batch:   ${config.batchSize === Infinity ? "all" : config.batchSize}`);
+  console.log(`Batch:   ${config.batchSize >= 1000000 ? "all" : config.batchSize}`);
   console.log("");
 
   try {
     await processDirectory(config, flags.verbose);
   } catch (err) {
     console.error("Fatal error:", err);
     process.exit(1);
+  } finally {
+    await db.close();
   }
 }