@twelvehart/supermemory-runtime 1.0.0-next.0

package/src/workers/chunking.worker.ts

@@ -0,0 +1,242 @@
+/**
+ * Chunking Worker
+ *
+ * BullMQ worker that processes documents from extraction queue,
+ * chunks them using appropriate strategies, stores chunks in database,
+ * and chains to embedding queue.
+ *
+ * Part of TASK-008: Content Processing Pipeline
+ */
+
+import { Job, Worker, Queue } from 'bullmq'
+import { v4 as uuidv4 } from 'uuid'
+import { eq } from 'drizzle-orm'
+import { chunks } from '../db/schema/chunks.schema.js'
+import { memories } from '../db/schema/memories.schema.js'
+import { chunkContent, detectContentType } from '../services/chunking/index.js'
+import { workerDb as db } from '../db/worker-connection.js'
+import { getLogger } from '../utils/logger.js'
+import { NotFoundError, ErrorCode } from '../utils/errors.js'
+
+const logger = getLogger('ChunkingWorker')
+
+// Job data interfaces
+export interface ChunkingJobData {
+  documentId: string
+  memoryId: string
+  content: string
+  contentType?: 'markdown' | 'code' | 'text'
+  chunkSize?: number
+  overlap?: number
+}
+
+export interface ChunkingJobResult {
+  documentId: string
+  memoryId: string
+  chunkCount: number
+  chunkIds: string[]
+  contentType: 'markdown' | 'code' | 'text'
+  totalTokens: number
+}
+
+export interface EmbeddingJobData {
+  documentId: string
+  memoryId: string
+  chunkIds: string[]
+}
+
+// Queue configuration
+const REDIS_URL = process.env.REDIS_URL || 'redis://localhost:6379'
+const QUEUE_NAME = 'chunking'
+const EMBEDDING_QUEUE_NAME = 'embedding'
+const CONCURRENCY = parseInt(process.env.BULLMQ_CONCURRENCY_CHUNKING || '3', 10)
+
+// Retry configuration
+const MAX_ATTEMPTS = 3
+const BACKOFF_DELAY = 2000 // 2 seconds
+
+/**
+ * Process a chunking job
+ */
+async function processChunkingJob(job: Job<ChunkingJobData>): Promise<ChunkingJobResult> {
+  const { documentId, memoryId, content, contentType, chunkSize, overlap } = job.data
+
+  try {
+    // Update progress: starting
+    await job.updateProgress(0)
+    await job.log(`Starting chunking for document ${documentId}`)
+
+    // Detect content type if not provided
+    const detectedType = contentType || detectContentType(content)
+    await job.log(`Detected content type: ${detectedType}`)
+
+    // Update progress: content type detected
+    await job.updateProgress(20)
+
+    // Chunk the content using appropriate strategy
+    const contentChunks = chunkContent(content, memoryId, {
+      chunkSize,
+      overlap,
+      contentType: detectedType,
+    })
+
+    await job.log(`Generated ${contentChunks.length} chunks`)
+    await job.updateProgress(50)
+
+    // Verify memory exists
+    const memory = await db.query.memories.findFirst({
+      where: eq(memories.id, memoryId),
+    })
+
+    if (!memory) {
+      throw new NotFoundError('Memory', memoryId, ErrorCode.MEMORY_NOT_FOUND)
+    }
+
+    // Store chunks in database
+    const chunkIds: string[] = []
+    const totalTokens = contentChunks.reduce((sum, chunk) => sum + chunk.tokenCount, 0)
+
+    for (let i = 0; i < contentChunks.length; i++) {
+      const chunk = contentChunks[i]
+      if (!chunk) continue
+      const chunkId = uuidv4()
+
+      await db.insert(chunks).values({
+        id: chunkId,
+        memoryId: memoryId,
+        content: chunk.content,
+        chunkIndex: i,
+        startOffset: chunk.metadata.startOffset,
+        endOffset: chunk.metadata.endOffset,
+        tokenCount: chunk.tokenCount,
+        metadata: {
+          contentType: chunk.metadata.contentType,
+          language: chunk.metadata.language,
+          heading: chunk.metadata.heading,
+          position: chunk.metadata.position,
+        },
+      })
+
+      chunkIds.push(chunkId)
+
+      // Update progress per chunk
+      const progress = 50 + Math.floor(((i + 1) / contentChunks.length) * 40)
+      await job.updateProgress(progress)
+    }
+
+    await job.log(`Stored ${chunkIds.length} chunks in database`)
+    await job.updateProgress(90)
+
+    // Chain to embedding queue
+    const embeddingQueue = new Queue<EmbeddingJobData>(EMBEDDING_QUEUE_NAME, {
+      connection: {
+        host: new URL(REDIS_URL).hostname,
+        port: parseInt(new URL(REDIS_URL).port || '6379', 10),
+      },
+    })
+
+    await embeddingQueue.add(
+      'embed',
+      {
+        documentId,
+        memoryId,
+        chunkIds,
+      },
+      {
+        priority: 5, // Medium priority
+        attempts: MAX_ATTEMPTS,
+        backoff: {
+          type: 'exponential',
+          delay: BACKOFF_DELAY,
+        },
+      }
+    )
+
+    await job.log(`Chained to embedding queue with ${chunkIds.length} chunks`)
+    await job.updateProgress(100)
+
+    return {
+      documentId,
+      memoryId,
+      chunkCount: contentChunks.length,
+      chunkIds,
+      contentType: detectedType,
+      totalTokens,
+    }
+  } catch (error) {
+    const errorMessage = error instanceof Error ? error.message : String(error)
+    await job.log(`Error: ${errorMessage}`)
+    throw error
+  }
+}
+
+/**
+ * Create and start chunking worker
+ */
+export function createChunkingWorker(): Worker<ChunkingJobData, ChunkingJobResult> {
+  const worker = new Worker<ChunkingJobData, ChunkingJobResult>(QUEUE_NAME, async (job) => processChunkingJob(job), {
+    connection: {
+      host: new URL(REDIS_URL).hostname,
+      port: parseInt(new URL(REDIS_URL).port || '6379', 10),
+    },
+    concurrency: CONCURRENCY,
+    autorun: true,
+    removeOnComplete: {
+      count: 100, // Keep last 100 completed jobs
+    },
+    removeOnFail: {
+      count: 500, // Keep last 500 failed jobs for debugging
+    },
+  })
+
+  // Event handlers
+  worker.on('completed', (job, result) => {
+    logger.info('Job completed', { jobId: job.id, chunkCount: result.chunkCount })
+  })
+
+  worker.on('failed', (job, error) => {
+    logger.error('Job failed', { jobId: job?.id, error: error.message })
+  })
+
+  worker.on('error', (error) => {
+    logger.error('Worker error', { error: error.message })
+  })
+
+  worker.on('stalled', (jobId) => {
+    logger.warn('Job stalled', { jobId })
+  })
+
+  logger.info('Worker started', { concurrency: CONCURRENCY })
+
+  return worker
+}
+
+/**
+ * Create chunking queue (for adding jobs)
+ */
+export function createChunkingQueue(): Queue<ChunkingJobData> {
+  return new Queue<ChunkingJobData>(QUEUE_NAME, {
+    connection: {
+      host: new URL(REDIS_URL).hostname,
+      port: parseInt(new URL(REDIS_URL).port || '6379', 10),
+    },
+    defaultJobOptions: {
+      attempts: MAX_ATTEMPTS,
+      backoff: {
+        type: 'exponential',
+        delay: BACKOFF_DELAY,
+      },
+      removeOnComplete: 100,
+      removeOnFail: 500,
+    },
+  })
+}
+
+/**
+ * Graceful shutdown
+ */
+export async function shutdownChunkingWorker(worker: Worker): Promise<void> {
+  logger.info('Shutting down...')
+  await worker.close()
+  logger.info('Shutdown complete')
+}

package/src/workers/embedding.worker.ts

@@ -0,0 +1,358 @@
+/**
+ * Embedding Worker - Generates embeddings for chunks in batches
+ *
+ * Responsibilities:
+ * - Receive chunks from chunking queue
+ * - Group into batches of 100 (OpenAI API limit)
+ * - Generate embeddings using EmbeddingService
+ * - Store in vector_embeddings staging table via PgVectorStore
+ * - Chain to indexing queue with embedding IDs
+ * - Track cost and progress per batch
+ * - Rate limiting: 3500 RPM (58 req/sec)
+ */
+
+import { Queue, Worker, Job } from 'bullmq'
+import pLimit from 'p-limit'
+import { getEmbeddingService } from '../services/embedding.service.js'
+import { createPgVectorStore } from '../services/vectorstore/pgvector.js'
+import type { PgVectorStore } from '../services/vectorstore/pgvector.js'
+import type { VectorEntry } from '../services/vectorstore/types.js'
+import { getLogger } from '../utils/logger.js'
+import { DatabaseError, EmbeddingError, ErrorCode } from '../utils/errors.js'
+
+const logger = getLogger('EmbeddingWorker')
+
+/**
+ * Job data structure for embedding worker
+ */
+export interface EmbeddingJobData {
+  /** Document ID for tracking */
+  documentId: string
+  /** Chunks to embed */
+  chunks: Array<{
+    id: string
+    content: string
+    metadata?: Record<string, unknown>
+  }>
+  /** Optional: Override default batch size */
+  batchSize?: number
+  /** Optional: Processing queue ID for status updates */
+  processingQueueId?: string
+}
+
+/**
+ * Job result structure
+ */
+export interface EmbeddingJobResult {
+  /** Total number of embeddings generated */
+  embeddingCount: number
+  /** Total cost in USD */
+  costUsd: number
+  /** Number of batches processed */
+  batchesProcessed: number
+  /** Embedding IDs for chaining to indexing queue */
+  embeddingIds: string[]
+  /** Processing time in milliseconds */
+  processingTimeMs: number
+}
+
+/**
+ * Embedding cost constants
+ * Based on OpenAI text-embedding-3-small pricing: $0.0001 per 1K tokens
+ */
+const COST_PER_1K_TOKENS = 0.0001
+const AVG_TOKENS_PER_CHAR = 0.25 // Rough estimate: 4 chars = 1 token
+
+/**
+ * Rate limiting constants
+ * 3500 RPM = 58.33 requests per second
+ * Conservative limit: 58 concurrent requests
+ */
+const MAX_CONCURRENT_REQUESTS = 58
+
+/**
+ * Default batch size for OpenAI API
+ */
+const DEFAULT_BATCH_SIZE = 100
+
+/**
+ * Estimate token count from text length
+ */
+function estimateTokens(text: string): number {
+  return Math.ceil(text.length * AVG_TOKENS_PER_CHAR)
+}
+
+/**
+ * Calculate cost based on token count
+ */
+function calculateCost(tokens: number): number {
+  return (tokens / 1000) * COST_PER_1K_TOKENS
+}
+
+/**
+ * Group chunks into batches
+ */
+function createBatches<T>(items: T[], batchSize: number): T[][] {
+  const batches: T[][] = []
+  for (let i = 0; i < items.length; i += batchSize) {
+    batches.push(items.slice(i, i + batchSize))
+  }
+  return batches
+}
+
+/**
+ * Embedding Worker class
+ */
+export class EmbeddingWorker {
+  private worker: Worker<EmbeddingJobData, EmbeddingJobResult> | null = null
+  private readonly queueName: string
+  private readonly connectionString: string
+  private vectorStore: PgVectorStore | null = null
+  private rateLimiter = pLimit(MAX_CONCURRENT_REQUESTS)
+
+  constructor(queueName: string = 'embedding', connectionString?: string) {
+    this.queueName = queueName
+    this.connectionString = connectionString || process.env.DATABASE_URL || 'postgresql://localhost:5432/supermemory'
+  }
+
+  /**
+   * Initialize the worker
+   */
+  async initialize(): Promise<void> {
+    // Initialize vector store
+    const embeddingService = getEmbeddingService()
+    const dimensions = embeddingService.getDimensions()
+
+    this.vectorStore = createPgVectorStore(this.connectionString, dimensions, {
+      batchSize: DEFAULT_BATCH_SIZE,
+      hnswConfig: { M: 16, efConstruction: 64 },
+      metric: 'cosine',
+    })
+
+    await this.vectorStore.initialize()
+
+    // Create worker
+    this.worker = new Worker<EmbeddingJobData, EmbeddingJobResult>(this.queueName, this.processJob.bind(this), {
+      connection: {
+        host: 'localhost',
+        port: 6379,
+      },
+      concurrency: 1, // Process one job at a time to control rate limiting globally
+      removeOnComplete: { count: 100 }, // Keep last 100 completed jobs
+      removeOnFail: { count: 500 }, // Keep last 500 failed jobs
+    })
+
+    // Error handling
+    this.worker.on('error', (error) => {
+      logger.error('Worker error', { error: error.message })
+    })
+
+    this.worker.on('failed', (job, error) => {
+      logger.error('Job failed', { jobId: job?.id, error: error.message })
+    })
+
+    logger.info('Worker initialized', { queueName: this.queueName })
+  }
+
+  /**
+   * Process embedding job
+   */
+  private async processJob(job: Job<EmbeddingJobData, EmbeddingJobResult>): Promise<EmbeddingJobResult> {
+    const startTime = Date.now()
+    const { documentId, chunks, batchSize = DEFAULT_BATCH_SIZE } = job.data
+
+    logger.info('Processing job', { jobId: job.id, documentId, chunkCount: chunks.length })
+
+    if (!this.vectorStore) {
+      throw new DatabaseError('Vector store not initialized', 'embedding', {
+        code: ErrorCode.DATABASE_NOT_INITIALIZED,
+      })
+    }
+
+    // Filter out empty chunks
+    const validChunks = chunks.filter((chunk) => chunk.content && chunk.content.trim().length > 0)
+
+    if (validChunks.length === 0) {
+      logger.warn('No valid chunks to process', { jobId: job.id })
+      return {
+        embeddingCount: 0,
+        costUsd: 0,
+        batchesProcessed: 0,
+        embeddingIds: [],
+        processingTimeMs: Date.now() - startTime,
+      }
+    }
+
+    // Create batches
+    const batches = createBatches(validChunks, batchSize)
+    logger.info('Created batches', { batchCount: batches.length, batchSize })
+
+    const embeddingService = getEmbeddingService()
+    const embeddingIds: string[] = []
+    let totalCost = 0
+    let totalTokens = 0
+
+    // Process batches with rate limiting
+    for (let i = 0; i < batches.length; i++) {
+      const batch = batches[i]
+      if (!batch) continue
+
+      const batchProgress = Math.round(((i + 1) / batches.length) * 100)
+      await job.updateProgress(batchProgress)
+      logger.info('Processing batch', {
+        batchNum: i + 1,
+        totalBatches: batches.length,
+        progress: batchProgress,
+      })
+
+      // Extract texts from batch
+      const texts = batch.map((chunk) => chunk.content)
+
+      // Estimate tokens and cost
+      const batchTokens = texts.reduce((sum, text) => sum + estimateTokens(text), 0)
+      const batchCost = calculateCost(batchTokens)
+      totalTokens += batchTokens
+      totalCost += batchCost
+
+      // Generate embeddings with rate limiting
+      const embeddings = await this.rateLimiter(async () => {
+        try {
+          return await embeddingService.batchEmbed(texts)
+        } catch (error) {
+          logger.error('Batch failed, retrying', { batchNum: i + 1, error })
+          // Retry once after exponential backoff
+          await new Promise((resolve) => setTimeout(resolve, 1000 * Math.pow(2, i)))
+          return await embeddingService.batchEmbed(texts)
+        }
+      })
+
+      // Store embeddings in vector store
+      const vectorEntries: VectorEntry[] = batch.map((chunk, idx) => {
+        const embedding = embeddings[idx]
+        if (!embedding || embedding.length === 0) {
+          throw new EmbeddingError(`Empty embedding for chunk ${chunk.id}`, undefined, {
+            chunkId: chunk.id,
+            batchIndex: idx,
+          })
+        }
+
+        return {
+          id: chunk.id,
+          embedding,
+          metadata: {
+            ...chunk.metadata,
+            documentId,
+            chunkId: chunk.id,
+            createdAt: new Date().toISOString(),
+          },
+        }
+      })
+
+      // Add to vector store in batch
+      const batchResult = await this.vectorStore.addBatch(vectorEntries, {
+        overwrite: false,
+        namespace: 'memories',
+      })
+
+      if (batchResult.failed > 0) {
+        logger.warn('Batch had failures', {
+          batchNum: i + 1,
+          failures: batchResult.failed,
+          errors: batchResult.errors,
+        })
+      }
+
+      // Collect embedding IDs
+      embeddingIds.push(...vectorEntries.map((entry) => entry.id))
+
+      logger.info('Batch complete', {
+        batchNum: i + 1,
+        totalBatches: batches.length,
+        embeddingCount: vectorEntries.length,
+        tokens: batchTokens,
+        cost: batchCost.toFixed(6),
+      })
+    }
+
+    const processingTimeMs = Date.now() - startTime
+
+    logger.info('Job complete', {
+      jobId: job.id,
+      embeddingCount: embeddingIds.length,
+      tokens: totalTokens,
+      cost: totalCost.toFixed(6),
+      processingTimeMs,
+    })
+
+    // Chain to indexing queue (if configured)
+    await this.chainToIndexingQueue(documentId, embeddingIds)
+
+    return {
+      embeddingCount: embeddingIds.length,
+      costUsd: totalCost,
+      batchesProcessed: batches.length,
+      embeddingIds,
+      processingTimeMs,
+    }
+  }
+
+  /**
+   * Chain to indexing queue with embedding IDs
+   */
+  private async chainToIndexingQueue(documentId: string, embeddingIds: string[]): Promise<void> {
+    try {
+      const indexingQueue = new Queue('indexing', {
+        connection: {
+          host: 'localhost',
+          port: 6379,
+        },
+      })
+
+      await indexingQueue.add(
+        'index',
+        {
+          documentId,
+          embeddingIds,
+        },
+        {
+          attempts: 3,
+          backoff: {
+            type: 'exponential',
+            delay: 2000,
+          },
+        }
+      )
+
+      logger.info('Chained to indexing queue', { documentId })
+    } catch (error) {
+      logger.error('Failed to chain to indexing queue', { documentId, error })
+      // Don't throw - embedding job succeeded
+    }
+  }
+
+  /**
+   * Close the worker and cleanup resources
+   */
+  async close(): Promise<void> {
+    if (this.worker) {
+      await this.worker.close()
+      this.worker = null
+    }
+
+    if (this.vectorStore) {
+      await this.vectorStore.close()
+      this.vectorStore = null
+    }
+
+    logger.info('Worker closed')
+  }
+}
+
+/**
+ * Create and initialize an embedding worker
+ */
+export async function createEmbeddingWorker(queueName?: string, connectionString?: string): Promise<EmbeddingWorker> {
+  const worker = new EmbeddingWorker(queueName, connectionString)
+  await worker.initialize()
+  return worker
+}