spine-framework-cortex 0.2.20 → 0.2.22

package/functions/custom_kb-chunker.ts

@@ -0,0 +1,576 @@
+/**
+ * Adaptive Article Chunker for KB Embeddings
+ *
+ * Strategy 5 (Recursive/Hierarchical) as the framework:
+ *   - Strategy 1 (Heading-Based) for structured docs
+ *   - Strategy 3 (Heading + Size Guards) for oversized sections
+ *   - Strategy 4 (Paragraph Grouping) for unstructured prose
+ *   - Strategy 2 (Fixed Window) deliberately excluded
+ *
+ * Hard rules:
+ *   - Never split a code block or table
+ *   - Never split mid-paragraph
+ *   - Merge tiny chunks (<100 tokens) into neighbor
+ *   - Prefix every chunk with context
+ */
+
+// ---------------------------------------------------------------------------
+// Types
+// ---------------------------------------------------------------------------
+
+export interface Chunk {
+  /** The text content to embed */
+  content: string
+  /** Section path for heading-based chunks, e.g. "Fields > data_type reference" */
+  sectionPath: string | null
+  /** 0-based index within the article */
+  chunkIndex: number
+  /** Total chunks for this article */
+  chunkTotal: number
+}
+
+export interface ChunkerOptions {
+  /** Article title — used as prefix context for every chunk */
+  articleTitle: string
+  /** Max tokens per chunk before sub-splitting. Default 800. */
+  maxTokens?: number
+  /** Min tokens per chunk before merging with neighbor. Default 100. */
+  minTokens?: number
+  /** Token count below which the entire article is a single chunk. Default 600. */
+  singleChunkThreshold?: number
+}
+
+// ---------------------------------------------------------------------------
+// Token estimation
+// ---------------------------------------------------------------------------
+
+/**
+ * Rough token estimate: ~4 chars per token for English text.
+ * Good enough for chunking decisions — actual tokenization happens at OpenAI.
+ */
+export function estimateTokens(text: string): number {
+  return Math.ceil(text.length / 4)
+}
+
+// ---------------------------------------------------------------------------
+// HTML → plain text
+// ---------------------------------------------------------------------------
+
+/**
+ * Strip HTML tags to plain text, preserving structure via newlines.
+ * Headings, paragraphs, divs get newlines. Inline tags are stripped.
+ */
+export function htmlToPlainText(html: string): string {
+  let text = html
+  // Normalize line breaks
+  text = text.replace(/\r\n?/g, '\n')
+  // Block-level elements → newlines
+  text = text.replace(/<\/(p|div|li|tr|blockquote)>/gi, '\n')
+  text = text.replace(/<br\s*\/?>/gi, '\n')
+  text = text.replace(/<\/(h[1-6])>/gi, '\n')
+  // Strip remaining tags
+  text = text.replace(/<[^>]+>/g, '')
+  // Decode common entities
+  text = text.replace(/&amp;/g, '&')
+  text = text.replace(/&lt;/g, '<')
+  text = text.replace(/&gt;/g, '>')
+  text = text.replace(/&quot;/g, '"')
+  text = text.replace(/&#39;/g, "'")
+  text = text.replace(/&nbsp;/g, ' ')
+  // Collapse excessive newlines
+  text = text.replace(/\n{3,}/g, '\n\n')
+  return text.trim()
+}
+
+// ---------------------------------------------------------------------------
+// Detect content format
+// ---------------------------------------------------------------------------
+
+type ContentFormat = 'markdown' | 'html'
+
+function detectFormat(content: string): ContentFormat {
+  // If it has HTML block tags, treat as HTML
+  if (/<(p|div|h[1-6]|ul|ol|table|pre)\b/i.test(content)) return 'html'
+  return 'markdown'
+}
+
+// ---------------------------------------------------------------------------
+// Heading extraction
+// ---------------------------------------------------------------------------
+
+interface Section {
+  /** Heading level: 1-6 for explicit headings, 0 for preamble */
+  level: number
+  /** Heading text (empty for preamble) */
+  heading: string
+  /** Body content below this heading (not including sub-sections) */
+  body: string
+}
+
+/**
+ * Parse markdown into a flat list of sections by heading.
+ * Each section includes its heading text and the body up to the next heading.
+ *
+ * @param splitLevel - Only split on headings at this level or shallower.
+ *   Deeper headings (e.g. ### when splitLevel=2) are kept as body content.
+ *   Default 6 = split on all headings.
+ */
+function parseMarkdownSections(content: string, splitLevel: number = 6): Section[] {
+  const lines = content.split('\n')
+  const sections: Section[] = []
+  let currentLevel = 0
+  let currentHeading = ''
+  let bodyLines: string[] = []
+
+  for (const line of lines) {
+    const headingMatch = line.match(/^(#{1,6})\s+(.+)$/)
+    if (headingMatch && headingMatch[1].length <= splitLevel) {
+      // Flush previous section
+      if (bodyLines.length > 0 || currentHeading) {
+        sections.push({
+          level: currentLevel,
+          heading: currentHeading,
+          body: bodyLines.join('\n').trim()
+        })
+      }
+      currentLevel = headingMatch[1].length
+      currentHeading = headingMatch[2]
+      bodyLines = []
+    } else {
+      bodyLines.push(line)
+    }
+  }
+
+  // Flush final section
+  if (bodyLines.length > 0 || currentHeading) {
+    sections.push({
+      level: currentLevel,
+      heading: currentHeading,
+      body: bodyLines.join('\n').trim()
+    })
+  }
+
+  return sections
+}
+
+/**
+ * Parse HTML into sections by h2/h3/h4 tags.
+ * Converts to plain text first, then extracts heading structure.
+ */
+function parseHtmlSections(html: string, splitLevel: number = 6): Section[] {
+  // Insert markdown-style headings before stripping, so we can parse them
+  let marked = html.replace(/<h([1-6])[^>]*>(.*?)<\/h\1>/gi, (_, level, text) => {
+    const hashes = '#'.repeat(parseInt(level))
+    // Strip any HTML inside the heading text
+    const cleanText = text.replace(/<[^>]+>/g, '').trim()
+    return `\n${hashes} ${cleanText}\n`
+  })
+  // Convert the rest to plain text
+  marked = htmlToPlainText(marked)
+  return parseMarkdownSections(marked, splitLevel)
+}
+
+// ---------------------------------------------------------------------------
+// Atomic block detection
+// ---------------------------------------------------------------------------
+
+/**
+ * Check if a line is the start of a fenced code block.
+ */
+function isCodeFenceStart(line: string): boolean {
+  return /^```/.test(line.trim())
+}
+
+/**
+ * Check if a line is inside a markdown table (starts with |).
+ */
+function isTableRow(line: string): boolean {
+  return /^\|/.test(line.trim())
+}
+
+// ---------------------------------------------------------------------------
+// Paragraph-based splitting (Strategy 4)
+// ---------------------------------------------------------------------------
+
+/**
+ * Split text into paragraphs, keeping code blocks and tables as atomic units.
+ * Returns an array of paragraph strings.
+ */
+function splitIntoParagraphs(text: string): string[] {
+  const lines = text.split('\n')
+  const paragraphs: string[] = []
+  let current: string[] = []
+  let inCodeBlock = false
+
+  for (let i = 0; i < lines.length; i++) {
+    const line = lines[i]
+
+    // Track code fence boundaries
+    if (isCodeFenceStart(line)) {
+      if (inCodeBlock) {
+        // End of code block — include closing fence, flush as atomic unit
+        current.push(line)
+        paragraphs.push(current.join('\n'))
+        current = []
+        inCodeBlock = false
+        continue
+      } else {
+        // Start of code block — flush what we have, start atomic unit
+        if (current.length > 0 && current.some(l => l.trim())) {
+          paragraphs.push(current.join('\n').trim())
+        }
+        current = [line]
+        inCodeBlock = true
+        continue
+      }
+    }
+
+    if (inCodeBlock) {
+      current.push(line)
+      continue
+    }
+
+    // Table rows are grouped together as atomic
+    if (isTableRow(line)) {
+      // If previous content wasn't a table, flush it
+      if (current.length > 0 && !isTableRow(current[current.length - 1])) {
+        if (current.some(l => l.trim())) {
+          paragraphs.push(current.join('\n').trim())
+        }
+        current = []
+      }
+      current.push(line)
+      continue
+    }
+
+    // If we were in a table and hit a non-table line, flush the table
+    if (current.length > 0 && isTableRow(current[current.length - 1]) && !isTableRow(line)) {
+      paragraphs.push(current.join('\n').trim())
+      current = []
+    }
+
+    // Empty line = paragraph boundary
+    if (line.trim() === '') {
+      if (current.length > 0 && current.some(l => l.trim())) {
+        paragraphs.push(current.join('\n').trim())
+        current = []
+      }
+      continue
+    }
+
+    current.push(line)
+  }
+
+  // Flush remaining
+  if (inCodeBlock && current.length > 0) {
+    // Unclosed code block — flush as-is
+    paragraphs.push(current.join('\n').trim())
+  } else if (current.length > 0 && current.some(l => l.trim())) {
+    paragraphs.push(current.join('\n').trim())
+  }
+
+  return paragraphs.filter(p => p.length > 0)
+}
+
+/**
+ * Group paragraphs into chunks of approximately maxTokens.
+ * Never splits a paragraph — it's the atomic unit.
+ */
+function groupParagraphs(paragraphs: string[], maxTokens: number): string[] {
+  const chunks: string[] = []
+  let current: string[] = []
+  let currentTokens = 0
+
+  for (const para of paragraphs) {
+    const paraTokens = estimateTokens(para)
+
+    // If a single paragraph exceeds max, it goes as its own chunk (atomic — don't break it)
+    if (paraTokens > maxTokens && current.length === 0) {
+      chunks.push(para)
+      continue
+    }
+
+    // Would adding this paragraph exceed the limit?
+    if (currentTokens + paraTokens > maxTokens && current.length > 0) {
+      chunks.push(current.join('\n\n'))
+      current = [para]
+      currentTokens = paraTokens
+    } else {
+      current.push(para)
+      currentTokens += paraTokens
+    }
+  }
+
+  if (current.length > 0) {
+    chunks.push(current.join('\n\n'))
+  }
+
+  return chunks
+}
+
+// ---------------------------------------------------------------------------
+// Section-based splitting (Strategy 1 + 3)
+// ---------------------------------------------------------------------------
+
+/**
+ * Build section path from nested heading context.
+ * E.g. "fields" → "Fields" or "fields > data_type reference" → "Fields > data_type reference"
+ */
+function buildSectionPath(headingStack: string[]): string {
+  return headingStack.filter(Boolean).join(' > ')
+}
+
+/**
+ * Recursively split sections that are too large.
+ * Strategy 3: sub-split on sub-headings if available, else paragraph split (Strategy 4).
+ */
+function splitOversizedSection(
+  body: string,
+  currentLevel: number,
+  maxTokens: number
+): string[] {
+  // Try to find sub-headings at the next level
+  const subHeadingPattern = new RegExp(`^${'#'.repeat(currentLevel + 1)}\\s+`, 'm')
+
+  if (subHeadingPattern.test(body)) {
+    // Has sub-headings — split on them (recurse Strategy 1)
+    const subSections = parseMarkdownSections(body).filter(s =>
+      s.level > currentLevel || s.level === 0
+    )
+
+    // Re-parse keeping proper section boundaries
+    const subLines = body.split('\n')
+    const subChunks: string[] = []
+    let currentChunk: string[] = []
+
+    for (const line of subLines) {
+      const subMatch = line.match(new RegExp(`^(#{${currentLevel + 1},6})\\s+(.+)$`))
+      if (subMatch && currentChunk.length > 0) {
+        const chunkText = currentChunk.join('\n').trim()
+        if (chunkText) {
+          if (estimateTokens(chunkText) > maxTokens) {
+            // Still too big — go deeper
+            subChunks.push(...splitOversizedSection(chunkText, currentLevel + 1, maxTokens))
+          } else {
+            subChunks.push(chunkText)
+          }
+        }
+        currentChunk = [line]
+      } else {
+        currentChunk.push(line)
+      }
+    }
+    if (currentChunk.length > 0) {
+      const chunkText = currentChunk.join('\n').trim()
+      if (chunkText) {
+        if (estimateTokens(chunkText) > maxTokens) {
+          subChunks.push(...splitOversizedSection(chunkText, currentLevel + 1, maxTokens))
+        } else {
+          subChunks.push(chunkText)
+        }
+      }
+    }
+    return subChunks
+  }
+
+  // No sub-headings — fall back to paragraph grouping (Strategy 4)
+  const paragraphs = splitIntoParagraphs(body)
+  return groupParagraphs(paragraphs, maxTokens)
+}
+
+// ---------------------------------------------------------------------------
+// Merge tiny chunks
+// ---------------------------------------------------------------------------
+
+/**
+ * Merge chunks that are below the minimum token threshold into their neighbor.
+ * Prefers merging with the next chunk; if last, merges with previous.
+ */
+function mergeTinyChunks(chunks: string[], minTokens: number): string[] {
+  if (chunks.length <= 1) return chunks
+
+  const result: string[] = []
+
+  for (let i = 0; i < chunks.length; i++) {
+    const tokens = estimateTokens(chunks[i])
+
+    if (tokens < minTokens && result.length > 0) {
+      // Merge with previous
+      result[result.length - 1] += '\n\n' + chunks[i]
+    } else if (tokens < minTokens && i < chunks.length - 1) {
+      // Merge with next
+      chunks[i + 1] = chunks[i] + '\n\n' + chunks[i + 1]
+    } else {
+      result.push(chunks[i])
+    }
+  }
+
+  return result
+}
+
+// ---------------------------------------------------------------------------
+// Main chunker
+// ---------------------------------------------------------------------------
+
+/**
+ * Chunk an article for embedding using the adaptive recursive strategy.
+ *
+ * Decision tree:
+ *   1. ≤ singleChunkThreshold tokens → single chunk, no splitting
+ *   2. Has headings → heading-based split (Strategy 1)
+ *      - Oversized sections → sub-split (Strategy 3 → recurse or Strategy 4)
+ *   3. No headings → paragraph grouping (Strategy 4)
+ *   4. Merge tiny chunks
+ *   5. Prefix every chunk with context
+ */
+export function chunkArticle(content: string, options: ChunkerOptions): Chunk[] {
+  const {
+    articleTitle,
+    maxTokens = 800,
+    minTokens = 100,
+    singleChunkThreshold = 600,
+  } = options
+
+  if (!content || content.trim().length === 0) {
+    return [{
+      content: articleTitle,
+      sectionPath: null,
+      chunkIndex: 0,
+      chunkTotal: 1,
+    }]
+  }
+
+  // Detect format and normalize to plain text for token counting / splitting
+  const format = detectFormat(content)
+  let plainContent: string
+  let sections: Section[]
+
+  // Primary split at ## level — ### and deeper stay as body content
+  // Sub-splitting (Strategy 3) will split on ### when a section is too large
+  const primarySplitLevel = 2
+
+  if (format === 'html') {
+    sections = parseHtmlSections(content, primarySplitLevel)
+    plainContent = htmlToPlainText(content)
+  } else {
+    sections = parseMarkdownSections(content, primarySplitLevel)
+    plainContent = content
+  }
+
+  const totalTokens = estimateTokens(plainContent)
+
+  // ── Step 1: Small article → single chunk ──────────────────────────
+  if (totalTokens <= singleChunkThreshold) {
+    const prefix = articleTitle
+    return [{
+      content: `${prefix}\n\n${plainContent}`,
+      sectionPath: null,
+      chunkIndex: 0,
+      chunkTotal: 1,
+    }]
+  }
+
+  // ── Step 2/3: Check for headings ──────────────────────────────────
+  // h1 is treated as preamble/title — only ## and deeper count as section headings
+  const hasHeadings = sections.some(s => s.level >= 2)
+
+  let rawChunks: { text: string; sectionPath: string | null }[]
+
+  if (hasHeadings) {
+    // Strategy 1: heading-based split
+    rawChunks = []
+    // Track headings by level for proper nesting: level → heading text
+    const headingByLevel = new Map<number, string>()
+
+    for (const section of sections) {
+      if ((section.level === 0 && !section.heading) || section.level === 1) {
+        // Preamble or h1 title — treat as intro context, not a section to chunk
+        if (section.body.trim()) {
+          rawChunks.push({
+            text: section.body,
+            sectionPath: null,
+          })
+        }
+        continue
+      }
+
+      // Clear this level and all deeper levels, then set current heading
+      for (const lvl of headingByLevel.keys()) {
+        if (lvl >= section.level) headingByLevel.delete(lvl)
+      }
+      headingByLevel.set(section.level, section.heading)
+
+      // Build path from sorted levels: ## Parent > ### Child
+      const sortedLevels = [...headingByLevel.keys()].sort((a, b) => a - b)
+      const sectionPath = sortedLevels.map(l => headingByLevel.get(l)!).join(' > ')
+      const fullSection = section.heading + '\n' + section.body
+      const sectionTokens = estimateTokens(fullSection)
+
+      if (sectionTokens > maxTokens) {
+        // Strategy 3: sub-split oversized section
+        const subChunks = splitOversizedSection(fullSection, section.level, maxTokens)
+        for (let i = 0; i < subChunks.length; i++) {
+          rawChunks.push({
+            text: subChunks[i],
+            sectionPath: subChunks.length > 1 ? `${sectionPath} (${i + 1}/${subChunks.length})` : sectionPath,
+          })
+        }
+      } else {
+        rawChunks.push({
+          text: fullSection,
+          sectionPath,
+        })
+      }
+    }
+  } else {
+    // Strategy 4: paragraph grouping for unstructured content
+    const paragraphs = splitIntoParagraphs(plainContent)
+    const grouped = groupParagraphs(paragraphs, maxTokens)
+    rawChunks = grouped.map(text => ({
+      text,
+      sectionPath: null,
+    }))
+  }
+
+  // ── Step 4: Merge tiny chunks ─────────────────────────────────────
+  // We need to merge while preserving sectionPath, so we work at the rawChunks level
+  const mergedChunks: { text: string; sectionPath: string | null }[] = []
+
+  for (let i = 0; i < rawChunks.length; i++) {
+    const tokens = estimateTokens(rawChunks[i].text)
+
+    if (tokens < minTokens && mergedChunks.length > 0) {
+      // Merge with previous chunk
+      const prev = mergedChunks[mergedChunks.length - 1]
+      prev.text += '\n\n' + rawChunks[i].text
+      // Keep the previous section path (the primary one)
+    } else if (tokens < minTokens && i < rawChunks.length - 1) {
+      // Merge with next chunk
+      rawChunks[i + 1].text = rawChunks[i].text + '\n\n' + rawChunks[i + 1].text
+      // Next chunk keeps its section path
+    } else {
+      mergedChunks.push({ ...rawChunks[i] })
+    }
+  }
+
+  // ── Step 5: Prefix and build final chunks ─────────────────────────
+  const total = mergedChunks.length
+
+  return mergedChunks.map((chunk, index) => {
+    let prefix: string
+    if (chunk.sectionPath) {
+      prefix = `${articleTitle} > ${chunk.sectionPath}`
+    } else if (total > 1) {
+      prefix = `${articleTitle} (chunk ${index + 1} of ${total})`
+    } else {
+      prefix = articleTitle
+    }
+
+    return {
+      content: `${prefix}\n\n${chunk.text}`,
+      sectionPath: chunk.sectionPath,
+      chunkIndex: index,
+      chunkTotal: total,
+    }
+  })
+}