@lota-sdk/core 0.1.5

package/src/services/document-chunk.service.ts

@@ -0,0 +1,141 @@
+import { createHash } from 'node:crypto'
+
+import { chunkMarkdownDocument, chunkPagedDocument, chunkPlainTextDocument } from '../document/org-document-chunking'
+import type { ParsedDocumentChunk } from '../document/org-document-chunking'
+import { createDefaultEmbeddings } from '../embeddings/provider'
+
+type DocumentChunkEmbeddings = {
+  embedDocuments(documents: string[]): Promise<number[][]>
+  embedQuery(query: string): Promise<number[]>
+}
+
+function createDocumentChunkEmbeddings(): DocumentChunkEmbeddings {
+  const embeddings = createDefaultEmbeddings()
+
+  return {
+    embedDocuments: async (documents) => await embeddings.embedDocuments(documents),
+    embedQuery: async (query) => await embeddings.embedQuery(query),
+  }
+}
+
+export interface VersionedDocumentChunkRecordShape {
+  chunkKey: string
+  chunkIndex: number
+  sectionPath?: string | null
+  pageStart?: number | null
+  pageEnd?: number | null
+  contentHash: string
+  sourceVersionKey: string
+  archivedAt?: string | number | Date | null
+}
+
+export class DocumentChunkService {
+  constructor(private readonly embeddings: DocumentChunkEmbeddings = createDocumentChunkEmbeddings()) {}
+
+  buildChunks(params: {
+    source: string
+    renderMode: 'markdown' | 'text' | 'pdf'
+    text: string
+    sectionPath?: string
+    pages?: Array<{ pageNumber: number; text: string }>
+  }): ParsedDocumentChunk[] {
+    if (
+      params.source === 'indexedOutcome' ||
+      params.source === 'websiteIntelligence' ||
+      params.renderMode === 'markdown'
+    ) {
+      return chunkMarkdownDocument({ text: params.text, baseSectionPath: params.sectionPath })
+    }
+
+    if (params.renderMode === 'pdf' && params.pages && params.pages.length > 0) {
+      return chunkPagedDocument({ pages: params.pages })
+    }
+
+    return chunkPlainTextDocument({ text: params.text, sectionPath: params.sectionPath })
+  }
+
+  hashContent(content: string): string {
+    return createHash('sha256').update(content).digest('hex')
+  }
+
+  estimateTokenCount(content: string): number {
+    return Math.max(1, Math.ceil(content.length / 4))
+  }
+
+  async embedQuery(query: string): Promise<number[]> {
+    return await this.embeddings.embedQuery(query)
+  }
+
+  async syncVersionedChunks<TRecord, TPayload>(params: {
+    sourceVersionKey: string
+    chunks: ParsedDocumentChunk[]
+    loadExisting: () => Promise<TRecord[]>
+    archive: (rows: TRecord[]) => Promise<void>
+    create: (payload: TPayload) => Promise<void>
+    update: (row: TRecord, payload: TPayload) => Promise<void>
+    buildPayload: (params: {
+      chunk: ParsedDocumentChunk
+      embedding: number[]
+      contentHash: string
+      tokenEstimate: number
+    }) => TPayload
+    selectShape: (row: TRecord) => VersionedDocumentChunkRecordShape
+  }): Promise<void> {
+    const existingRows = await params.loadExisting()
+    const existingByChunkKey = new Map(
+      existingRows
+        .filter((row) => params.selectShape(row).sourceVersionKey === params.sourceVersionKey)
+        .map((row) => [params.selectShape(row).chunkKey, row]),
+    )
+    const embeddings = await this.embeddings.embedDocuments(params.chunks.map((chunk) => chunk.content))
+    const seenChunkKeys = new Set<string>()
+    const staleVersionRows = existingRows.filter(
+      (row) => params.selectShape(row).sourceVersionKey !== params.sourceVersionKey,
+    )
+
+    await params.archive(staleVersionRows)
+
+    for (const [index, chunk] of params.chunks.entries()) {
+      const contentHash = this.hashContent(chunk.content)
+      const existingRow = existingByChunkKey.get(chunk.chunkKey)
+      const payload = params.buildPayload({
+        chunk,
+        embedding: embeddings[index] ?? [],
+        contentHash,
+        tokenEstimate: this.estimateTokenCount(chunk.content),
+      })
+
+      seenChunkKeys.add(chunk.chunkKey)
+
+      if (!existingRow) {
+        await params.create(payload)
+        continue
+      }
+
+      const current = params.selectShape(existingRow)
+      const hasChanged =
+        current.contentHash !== contentHash ||
+        current.chunkIndex !== chunk.chunkIndex ||
+        (current.sectionPath ?? null) !== (chunk.sectionPath ?? null) ||
+        (current.pageStart ?? null) !== (chunk.pageStart ?? null) ||
+        (current.pageEnd ?? null) !== (chunk.pageEnd ?? null)
+
+      if (!hasChanged) {
+        continue
+      }
+
+      await params.update(existingRow, payload)
+    }
+
+    const removedCurrentVersionRows = existingRows.filter((row) => {
+      const current = params.selectShape(row)
+      return current.sourceVersionKey === params.sourceVersionKey && !seenChunkKeys.has(current.chunkKey)
+    })
+
+    await params.archive(removedCurrentVersionRows)
+  }
+}
+
+export const documentChunkService = new DocumentChunkService()
+
+export type { ParsedDocumentChunk }