@opensaas/stack-rag 0.1.6

package/src/runtime/batch.ts

@@ -0,0 +1,424 @@
+/**
+ * Batch processing utilities with rate limiting and progress tracking
+ */
+
+import type { EmbeddingProvider } from '../providers/types.js'
+import type { StoredEmbedding } from '../config/types.js'
+import { generateEmbeddings } from './embeddings.js'
+
+export interface BatchProcessOptions {
+  /**
+   * Embedding provider to use
+   */
+  provider: EmbeddingProvider
+
+  /**
+   * Array of texts to process
+   */
+  texts: string[]
+
+  /**
+   * Batch size for processing
+   * @default 10
+   */
+  batchSize?: number
+
+  /**
+   * Rate limit in requests per minute
+   * @default 100
+   */
+  rateLimit?: number
+
+  /**
+   * Progress callback called after each batch
+   */
+  onProgress?: (progress: BatchProgress) => void
+
+  /**
+   * Error callback called when a batch fails
+   * If not provided, errors will be thrown
+   */
+  onError?: (error: BatchError) => void
+
+  /**
+   * Number of retries for failed batches
+   * @default 3
+   */
+  maxRetries?: number
+
+  /**
+   * Initial retry delay in milliseconds
+   * @default 1000
+   */
+  retryDelay?: number
+
+  /**
+   * Whether to include source hash in metadata
+   * @default true
+   */
+  includeSourceHash?: boolean
+}
+
+export interface BatchProgress {
+  /**
+   * Number of items processed so far
+   */
+  processed: number
+
+  /**
+   * Total number of items to process
+   */
+  total: number
+
+  /**
+   * Number of items that failed
+   */
+  failed: number
+
+  /**
+   * Percentage completed (0-100)
+   */
+  percentage: number
+
+  /**
+   * Current batch number (1-indexed)
+   */
+  currentBatch: number
+
+  /**
+   * Total number of batches
+   */
+  totalBatches: number
+}
+
+export interface BatchError {
+  /**
+   * Batch number that failed
+   */
+  batchNumber: number
+
+  /**
+   * Items in the failed batch
+   */
+  items: string[]
+
+  /**
+   * Error that occurred
+   */
+  error: Error
+
+  /**
+   * Number of retry attempts made
+   */
+  retries: number
+}
+
+export interface BatchProcessResult {
+  /**
+   * Successfully generated embeddings
+   */
+  embeddings: StoredEmbedding[]
+
+  /**
+   * Texts that failed to process
+   */
+  failed: Array<{ text: string; error: Error }>
+
+  /**
+   * Total processing statistics
+   */
+  stats: {
+    total: number
+    successful: number
+    failed: number
+    duration: number
+  }
+}
+
+/**
+ * Process embeddings in batches with rate limiting and retry logic
+ *
+ * @example
+ * ```typescript
+ * const result = await batchProcess({
+ *   provider: createEmbeddingProvider({ type: 'openai', apiKey: '...' }),
+ *   texts: largeTextArray,
+ *   batchSize: 10,
+ *   rateLimit: 60, // 60 requests per minute
+ *   onProgress: (progress) => {
+ *     console.log(`Progress: ${progress.percentage}%`)
+ *   },
+ * })
+ *
+ * console.log(`Successfully processed: ${result.stats.successful}`)
+ * console.log(`Failed: ${result.stats.failed}`)
+ * ```
+ */
+export async function batchProcess(options: BatchProcessOptions): Promise<BatchProcessResult> {
+  const {
+    provider,
+    texts,
+    batchSize = 10,
+    rateLimit = 100,
+    onProgress,
+    onError,
+    maxRetries = 3,
+    retryDelay = 1000,
+    includeSourceHash = true,
+  } = options
+
+  const startTime = Date.now()
+  const totalBatches = Math.ceil(texts.length / batchSize)
+  const embeddings: StoredEmbedding[] = []
+  const failed: Array<{ text: string; error: Error }> = []
+
+  // Calculate delay between batches to respect rate limit
+  const delayBetweenBatches = calculateBatchDelay(rateLimit)
+
+  for (let i = 0; i < texts.length; i += batchSize) {
+    const batchNumber = Math.floor(i / batchSize) + 1
+    const batch = texts.slice(i, i + batchSize)
+
+    try {
+      // Process batch with retry logic
+      const batchEmbeddings = await retryWithBackoff(
+        async () =>
+          generateEmbeddings({
+            provider,
+            texts: batch,
+            includeSourceHash,
+            batchSize: batch.length,
+          }),
+        maxRetries,
+        retryDelay,
+      )
+
+      embeddings.push(...batchEmbeddings)
+
+      // Report progress
+      if (onProgress) {
+        const processed = Math.min(i + batchSize, texts.length)
+        onProgress({
+          processed,
+          total: texts.length,
+          failed: failed.length,
+          percentage: Math.round((processed / texts.length) * 100),
+          currentBatch: batchNumber,
+          totalBatches,
+        })
+      }
+
+      // Rate limiting: wait before next batch (except for last batch)
+      if (batchNumber < totalBatches && delayBetweenBatches > 0) {
+        await sleep(delayBetweenBatches)
+      }
+    } catch (error) {
+      const batchError: BatchError = {
+        batchNumber,
+        items: batch,
+        error: error instanceof Error ? error : new Error(String(error)),
+        retries: maxRetries,
+      }
+
+      if (onError) {
+        onError(batchError)
+      } else {
+        // If no error handler, throw the error
+        throw error
+      }
+
+      // Add all items in batch to failed list
+      for (const text of batch) {
+        failed.push({
+          text,
+          error: batchError.error,
+        })
+      }
+    }
+  }
+
+  const duration = Date.now() - startTime
+
+  return {
+    embeddings,
+    failed,
+    stats: {
+      total: texts.length,
+      successful: embeddings.length,
+      failed: failed.length,
+      duration,
+    },
+  }
+}
+
+/**
+ * Rate limiter for controlling API request rate
+ */
+export class RateLimiter {
+  private queue: Array<() => void> = []
+  private requestTimestamps: number[] = []
+  private readonly requestsPerMinute: number
+
+  constructor(requestsPerMinute: number) {
+    this.requestsPerMinute = requestsPerMinute
+  }
+
+  /**
+   * Wait until rate limit allows next request
+   */
+  async waitForSlot(): Promise<void> {
+    return new Promise((resolve) => {
+      const tryAcquire = () => {
+        const now = Date.now()
+        const oneMinuteAgo = now - 60_000
+
+        // Remove timestamps older than 1 minute
+        this.requestTimestamps = this.requestTimestamps.filter((t) => t > oneMinuteAgo)
+
+        // Check if we can make a request
+        if (this.requestTimestamps.length < this.requestsPerMinute) {
+          this.requestTimestamps.push(now)
+          resolve()
+          return
+        }
+
+        // Calculate wait time until oldest request expires
+        const oldestTimestamp = this.requestTimestamps[0]
+        const waitTime = oldestTimestamp + 60_000 - now
+
+        // Try again after wait time
+        setTimeout(tryAcquire, Math.max(waitTime, 100))
+      }
+
+      tryAcquire()
+    })
+  }
+
+  /**
+   * Execute a function with rate limiting
+   */
+  async execute<T>(fn: () => Promise<T>): Promise<T> {
+    await this.waitForSlot()
+    return fn()
+  }
+}
+
+/**
+ * Calculate delay between batches to respect rate limit
+ */
+function calculateBatchDelay(requestsPerMinute: number): number {
+  // Each batch counts as one request
+  const requestsPerSecond = requestsPerMinute / 60
+  const delayPerRequest = 1000 / requestsPerSecond
+
+  // Return delay in milliseconds
+  return Math.ceil(delayPerRequest)
+}
+
+/**
+ * Sleep for specified milliseconds
+ */
+function sleep(ms: number): Promise<void> {
+  return new Promise((resolve) => setTimeout(resolve, ms))
+}
+
+/**
+ * Retry a function with exponential backoff
+ */
+async function retryWithBackoff<T>(
+  fn: () => Promise<T>,
+  maxRetries: number,
+  initialDelay: number,
+): Promise<T> {
+  let lastError: Error | undefined
+
+  for (let attempt = 0; attempt <= maxRetries; attempt++) {
+    try {
+      return await fn()
+    } catch (error) {
+      lastError = error instanceof Error ? error : new Error(String(error))
+
+      // Don't retry on last attempt
+      if (attempt === maxRetries) {
+        break
+      }
+
+      // Exponential backoff: delay * 2^attempt
+      const delay = initialDelay * Math.pow(2, attempt)
+      await sleep(delay)
+    }
+  }
+
+  throw lastError
+}
+
+/**
+ * Queue for processing items with concurrency control
+ */
+export class ProcessingQueue<T, R> {
+  private queue: Array<{ item: T; resolve: (value: R) => void; reject: (error: Error) => void }> =
+    []
+  private processing = 0
+  private readonly concurrency: number
+  private readonly processor: (item: T) => Promise<R>
+
+  constructor(processor: (item: T) => Promise<R>, concurrency: number = 1) {
+    this.processor = processor
+    this.concurrency = concurrency
+  }
+
+  /**
+   * Add an item to the queue
+   */
+  async add(item: T): Promise<R> {
+    return new Promise((resolve, reject) => {
+      this.queue.push({ item, resolve, reject })
+      this.processNext()
+    })
+  }
+
+  /**
+   * Add multiple items to the queue
+   */
+  async addBatch(items: T[]): Promise<R[]> {
+    return Promise.all(items.map((item) => this.add(item)))
+  }
+
+  /**
+   * Process next item in queue
+   */
+  private async processNext(): Promise<void> {
+    if (this.processing >= this.concurrency || this.queue.length === 0) {
+      return
+    }
+
+    this.processing++
+    const next = this.queue.shift()
+
+    if (next) {
+      try {
+        const result = await this.processor(next.item)
+        next.resolve(result)
+      } catch (error) {
+        next.reject(error instanceof Error ? error : new Error(String(error)))
+      } finally {
+        this.processing--
+        this.processNext()
+      }
+    }
+  }
+
+  /**
+   * Get current queue size
+   */
+  get size(): number {
+    return this.queue.length
+  }
+
+  /**
+   * Get number of items currently being processed
+   */
+  get activeCount(): number {
+    return this.processing
+  }
+}

package/src/runtime/chunking.test.ts

@@ -0,0 +1,258 @@
+import { describe, it, expect } from 'vitest'
+import { chunkText, estimateTokenCount, mergeSmallChunks } from './chunking.js'
+
+describe('chunkText', () => {
+  describe('recursive strategy', () => {
+    it('should chunk text at paragraph boundaries', () => {
+      const text = 'Paragraph 1\n\nParagraph 2\n\nParagraph 3'
+      const chunks = chunkText(text, {
+        strategy: 'recursive',
+        chunkSize: 15,
+        chunkOverlap: 0,
+      })
+
+      expect(chunks).toHaveLength(3)
+      expect(chunks[0].text).toBe('Paragraph 1\n\n')
+      expect(chunks[1].text).toBe('Paragraph 2\n\n')
+      expect(chunks[2].text).toBe('Paragraph 3')
+    })
+
+    it('should handle overlap between chunks', () => {
+      const text = 'First chunk here. Second chunk here. Third chunk here.'
+      const chunks = chunkText(text, {
+        strategy: 'recursive',
+        chunkSize: 20,
+        chunkOverlap: 5,
+      })
+
+      expect(chunks.length).toBeGreaterThan(1)
+
+      // Check that chunks have some overlap
+      for (let i = 1; i < chunks.length; i++) {
+        const prevChunk = chunks[i - 1].text
+        const currChunk = chunks[i].text
+        const overlap = prevChunk.slice(-5)
+        // Overlap might not be exact due to sentence boundaries
+        expect(currChunk).toContain(overlap.trim().split(' ')[0])
+      }
+    })
+
+    it('should respect chunk size limits', () => {
+      const text = 'A'.repeat(1000)
+      const chunks = chunkText(text, {
+        strategy: 'recursive',
+        chunkSize: 100,
+        chunkOverlap: 0,
+      })
+
+      for (const chunk of chunks) {
+        expect(chunk.text.length).toBeLessThanOrEqual(100)
+      }
+    })
+
+    it('should track chunk positions correctly', () => {
+      const text = 'Start. Middle. End.'
+      const chunks = chunkText(text, {
+        strategy: 'recursive',
+        chunkSize: 100,
+        chunkOverlap: 0,
+      })
+
+      expect(chunks[0].start).toBe(0)
+      expect(chunks[0].end).toBe(text.length)
+      expect(chunks[0].text).toBe(text)
+    })
+  })
+
+  describe('sentence strategy', () => {
+    it('should preserve sentence boundaries', () => {
+      const text = 'First sentence. Second sentence. Third sentence.'
+      const chunks = chunkText(text, {
+        strategy: 'sentence',
+        chunkSize: 20,
+        chunkOverlap: 0,
+      })
+
+      expect(chunks.length).toBeGreaterThan(1)
+
+      // Each chunk should end with sentence punctuation
+      for (const chunk of chunks) {
+        expect(chunk.text.trim()).toMatch(/[.!?]$/)
+      }
+    })
+
+    it('should handle text with no sentences', () => {
+      const text = 'No sentence markers here'
+      const chunks = chunkText(text, {
+        strategy: 'sentence',
+        chunkSize: 10,
+        chunkOverlap: 0,
+      })
+
+      expect(chunks).toHaveLength(1)
+      expect(chunks[0].text).toBe(text)
+    })
+  })
+
+  describe('sliding-window strategy', () => {
+    it('should create overlapping fixed-size chunks', () => {
+      const text = 'A'.repeat(100)
+      const chunks = chunkText(text, {
+        strategy: 'sliding-window',
+        chunkSize: 30,
+        chunkOverlap: 10,
+      })
+
+      expect(chunks.length).toBeGreaterThan(1)
+
+      // Check fixed size (except possibly last chunk)
+      for (let i = 0; i < chunks.length - 1; i++) {
+        expect(chunks[i].text.length).toBe(30)
+      }
+
+      // Check overlap
+      for (let i = 1; i < chunks.length; i++) {
+        const prevChunk = chunks[i - 1]
+        const currChunk = chunks[i]
+        expect(currChunk.start).toBe(prevChunk.start + 20) // step = chunkSize - overlap
+      }
+    })
+
+    it('should skip empty chunks', () => {
+      const text = '   ' // Just whitespace
+      const chunks = chunkText(text, {
+        strategy: 'sliding-window',
+        chunkSize: 10,
+        chunkOverlap: 0,
+      })
+
+      expect(chunks).toHaveLength(0)
+    })
+  })
+
+  describe('token-aware strategy', () => {
+    it('should estimate token limits', () => {
+      const text = 'A'.repeat(400) // ~100 tokens at 4 chars/token
+      const chunks = chunkText(text, {
+        strategy: 'token-aware',
+        tokenLimit: 50,
+        chunkOverlap: 0,
+      })
+
+      expect(chunks.length).toBeGreaterThanOrEqual(2)
+
+      // Each chunk should be roughly under token limit * 4 chars
+      for (const chunk of chunks) {
+        expect(chunk.text.length).toBeLessThanOrEqual(50 * 4)
+      }
+    })
+  })
+
+  describe('edge cases', () => {
+    it('should handle empty text', () => {
+      const chunks = chunkText('', { chunkSize: 100 })
+      expect(chunks).toHaveLength(0)
+    })
+
+    it('should handle text smaller than chunk size', () => {
+      const text = 'Small text'
+      const chunks = chunkText(text, { chunkSize: 1000 })
+
+      expect(chunks).toHaveLength(1)
+      expect(chunks[0].text).toBe(text)
+      expect(chunks[0].start).toBe(0)
+      expect(chunks[0].end).toBe(text.length)
+    })
+
+    it('should throw error if overlap >= chunk size', () => {
+      expect(() => {
+        chunkText('text', { chunkSize: 10, chunkOverlap: 10 })
+      }).toThrow('chunkOverlap must be less than chunkSize')
+    })
+
+    it('should assign correct chunk indexes', () => {
+      const text = 'A'.repeat(300)
+      const chunks = chunkText(text, { chunkSize: 100, chunkOverlap: 0 })
+
+      chunks.forEach((chunk, i) => {
+        expect(chunk.index).toBe(i)
+      })
+    })
+  })
+})
+
+describe('estimateTokenCount', () => {
+  it('should estimate token count', () => {
+    const text = 'Hello world'
+    const count = estimateTokenCount(text)
+
+    // "Hello world" is 11 chars / 4 = ~3 tokens
+    expect(count).toBe(3)
+  })
+
+  it('should handle empty text', () => {
+    expect(estimateTokenCount('')).toBe(0)
+  })
+
+  it('should handle long text', () => {
+    const text = 'A'.repeat(1000)
+    const count = estimateTokenCount(text)
+
+    expect(count).toBe(250) // 1000 / 4
+  })
+})
+
+describe('mergeSmallChunks', () => {
+  it('should merge chunks below minimum size', () => {
+    const chunks = [
+      { text: 'A', start: 0, end: 1, index: 0 },
+      { text: 'B', start: 1, end: 2, index: 1 },
+      { text: 'C', start: 2, end: 3, index: 2 },
+    ]
+
+    const merged = mergeSmallChunks(chunks, 2)
+
+    expect(merged.length).toBeLessThan(chunks.length)
+    expect(merged[0].text.length).toBeGreaterThanOrEqual(2)
+  })
+
+  it('should not merge chunks already above minimum size', () => {
+    const chunks = [
+      { text: 'AAA', start: 0, end: 3, index: 0 },
+      { text: 'BBB', start: 3, end: 6, index: 1 },
+    ]
+
+    const merged = mergeSmallChunks(chunks, 2)
+
+    expect(merged).toHaveLength(2)
+  })
+
+  it('should handle empty array', () => {
+    const merged = mergeSmallChunks([], 10)
+    expect(merged).toHaveLength(0)
+  })
+
+  it('should update chunk indexes after merge', () => {
+    const chunks = [
+      { text: 'A', start: 0, end: 1, index: 0 },
+      { text: 'B', start: 1, end: 2, index: 1 },
+    ]
+
+    const merged = mergeSmallChunks(chunks, 5)
+
+    merged.forEach((chunk, i) => {
+      expect(chunk.index).toBe(i)
+    })
+  })
+
+  it('should merge metadata from merged chunks', () => {
+    const chunks = [
+      { text: 'A', start: 0, end: 1, index: 0, metadata: { foo: 1 } },
+      { text: 'B', start: 1, end: 2, index: 1, metadata: { bar: 2 } },
+    ]
+
+    const merged = mergeSmallChunks(chunks, 5)
+
+    expect(merged[0].metadata).toEqual({ foo: 1, bar: 2 })
+  })
+})