Momento/memento-note/lib/ai/services/chunk-indexing.service.ts

/**
 * ChunkIndexingService — Indexation incrémentale d'une note en fragments.
 *
 * Inspiré d'AppFlowy flowy-ai/src/embeddings/scheduler.rs.
 *
 * Fonctionnement :
 *   1. Récupère les fragmentIds existants en DB pour la note
 *   2. Chunk le contenu actuel (via chunkNoteContent)
 *   3. Compare les hash :
 *      - Inchangés (hash identique) → skip
 *      - Nouveaux (hash absent) → embed + insert
 *      - Supprimés (hash en DB mais plus dans le contenu) → delete
 *   4. Embed les nouveaux fragments via une queue à concurrence limitée
 *   5. Retry avec backoff exponentiel en cas d'erreur API
 *
 * Garanties :
 *   - Pas de re-embed des fragments inchangés (économie API)
 *   - Pas de fragments orphelins (stale supprimés)
 *   - Pas de race condition (verrou par noteId)
 */

import { createHash } from 'crypto'
import PQueue from 'p-queue'
import { prisma } from '@/lib/prisma'
import { embeddingService } from './embedding.service'
import {
  chunkNoteContent,
  type NoteChunk,
} from '@/lib/text/note-chunking'
import {
  prepareNoteTextForEmbedding,
} from '@/lib/text/plain-text'
import { Prisma } from '@prisma/client'

const EMBEDDING_CONCURRENCY = 4
const MAX_RETRIES = 3
const RETRY_BASE_DELAY_MS = 1000

const embeddingQueue = new PQueue({ concurrency: EMBEDDING_CONCURRENCY })

const noteLocks = new Map<string, Promise<void>>()

export interface IndexResult {
  noteId: string
  totalFragments: number
  newFragments: number
  skipped: number
  deleted: number
  durationMs: number
}

export class ChunkIndexingService {
  /**
   * Indexe une note en fragments. Ne re-embed que les fragments modifiés.
   */
  async indexNote(
    noteId: string,
    title: string | null | undefined,
    content: string,
  ): Promise<IndexResult> {
    while (noteLocks.has(noteId)) {
      await noteLocks.get(noteId)
    }

    const task = this.doIndexNote(noteId, title, content)
    noteLocks.set(noteId, task.then(() => {}, () => {}))

    try {
      return await task
    } finally {
      noteLocks.delete(noteId)
    }
  }

  private async doIndexNote(
    noteId: string,
    title: string | null | undefined,
    content: string,
  ): Promise<IndexResult> {
    const start = Date.now()
    const plain = prepareNoteTextForEmbedding(title, content)
    // Index even very short notes (title-only or one-liners)
    let newChunks = chunkNoteContent(noteId, plain)
    if (newChunks.length === 0 && plain.trim().length > 0) {
      // Force a single fragment when paragraph split yields nothing
      const fragmentId = createHash('sha256')
        .update(`${noteId}::${plain}`)
        .digest('hex')
        .slice(0, 32)
      newChunks = [{ fragmentId, content: plain, chunkIndex: 0, charCount: plain.length }]
    }
    const newFragmentIds = new Set(newChunks.map((c) => c.fragmentId))

    const existing = await prisma.noteEmbeddingChunk.findMany({
      where: { noteId },
      select: { fragmentId: true },
    })
    const existingIds = new Set<string>(existing.map((e) => e.fragmentId))

    const toDelete = [...existingIds].filter((id) => !newFragmentIds.has(id))
    const toEmbed = newChunks.filter((c) => !existingIds.has(c.fragmentId))
    const skipped = newChunks.filter((c) => existingIds.has(c.fragmentId))

    if (toDelete.length > 0) {
      await prisma.noteEmbeddingChunk.deleteMany({
        where: {
          noteId,
          fragmentId: { in: toDelete },
        },
      })
    }

    const embeddedChunks = await this.embedChunks(toEmbed)

    if (embeddedChunks.length > 0) {
      await this.upsertChunks(noteId, embeddedChunks)
    }

    return {
      noteId,
      totalFragments: newChunks.length,
      newFragments: embeddedChunks.length,
      skipped: skipped.length,
      deleted: toDelete.length,
      durationMs: Date.now() - start,
    }
  }

  /**
   * Embed une liste de fragments avec queue à concurrence limitée + retry.
   */
  private async embedChunks(chunks: NoteChunk[]): Promise<
    Array<NoteChunk & { embedding: number[] }>
  > {
    if (chunks.length === 0) return []

    const results = await Promise.all(
      chunks.map((chunk) =>
        embeddingQueue.add(() => this.embedWithRetry(chunk)),
      ),
    )

    return results.filter(
      (r): r is NoteChunk & { embedding: number[] } => r !== null,
    )
  }

  private async embedWithRetry(
    chunk: NoteChunk,
  ): Promise<(NoteChunk & { embedding: number[] }) | null> {
    let lastError: Error | null = null

    for (let attempt = 0; attempt < MAX_RETRIES; attempt++) {
      try {
        const embedding = await embeddingService.embedText(chunk.content)
        return { ...chunk, embedding }
      } catch (err: any) {
        lastError = err
        if (attempt < MAX_RETRIES - 1) {
          const delay = RETRY_BASE_DELAY_MS * Math.pow(2, attempt)
          console.warn(
            `[ChunkIndexing] Retry ${attempt + 1}/${MAX_RETRIES} for fragment ${chunk.fragmentId} after ${delay}ms: ${err.message}`,
          )
          await new Promise((resolve) => setTimeout(resolve, delay))
        }
      }
    }

    console.error(
      `[ChunkIndexing] Failed to embed fragment ${chunk.fragmentId} after ${MAX_RETRIES} attempts: ${lastError?.message}`,
    )
    return null
  }

  /**
   * Upsert transactionnel des fragments embeddés.
   * Utilise $executeRawUnsafe pour la colonne vector(1536).
   */
  private async upsertChunks(
    noteId: string,
    chunks: Array<NoteChunk & { embedding: number[] }>,
  ): Promise<void> {
    for (const chunk of chunks) {
      const vecStr = `[${chunk.embedding.join(',')}]`
      await prisma.$executeRaw`
        INSERT INTO "NoteEmbeddingChunk" (
          "id", "noteId", "fragmentId", "chunkIndex",
          "content", "charCount", "embedding", "embeddingModel",
          "createdAt", "updatedAt"
        )
        VALUES (
          gen_random_uuid(), ${noteId}, ${chunk.fragmentId}, ${chunk.chunkIndex},
          ${chunk.content}, ${chunk.charCount},
          ${vecStr}::vector, 'text-embedding-3-small',
          now(), now()
        )
        ON CONFLICT ("noteId", "fragmentId")
        DO UPDATE SET
          "content" = EXCLUDED."content",
          "charCount" = EXCLUDED."charCount",
          "embedding" = EXCLUDED."embedding",
          "chunkIndex" = EXCLUDED."chunkIndex",
          "updatedAt" = now()
      `
    }
  }

  /**
   * Supprime tous les fragments d'une note.
   */
  async deleteNoteChunks(noteId: string): Promise<void> {
    await prisma.noteEmbeddingChunk.deleteMany({ where: { noteId } })
  }

  /**
   * Vérifie si une note a déjà des fragments indexés.
   */
  async hasChunks(noteId: string): Promise<boolean> {
    const count = await prisma.noteEmbeddingChunk.count({
      where: { noteId },
      take: 1,
    })
    return count > 0
  }
}

export const chunkIndexingService = new ChunkIndexingService()