d3ro-voice/server/supabase/functions/embed-chunks/handler.ts

// server/supabase/functions/embed-chunks/handler.ts
// Use case: embed a document's pending knowledge chunks.
// IO is injected through ports (auth, knowledge store, embedding provider,
// usage store) so the size and quota guards are testable without Supabase or
// OpenAI. index.ts wires the real adapters.

import { handleCorsPreflightRequest } from '../_shared/cors.ts'
import { errorResponse, jsonResponse as json } from '../_shared/json-response.ts'
import { type EmbeddingProvider, isEmbedding } from '../_shared/openai-embeddings.ts'
import {
  EMBEDDING_LIMITS,
  type EmbeddingReservation,
  type EmbeddingUsageStore,
  embeddingUnits,
  firstOversizedChunk,
  planEmbeddingBatches,
  quotaExceededBody,
  refundEmbeddingUnits,
  reserveEmbeddingUnits,
} from '../_shared/embedding-quota.ts'

const UUID_PATTERN = /^[0-9a-f]{8}-[0-9a-f]{4}-[1-5][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/i

export interface KnowledgeChunkRow {
  id: string
  content: string
}

/** Thrown by a KnowledgeIndexStore when the database call fails. */
export class KnowledgeStorageError extends Error {
  constructor() {
    super('knowledge_storage_failed')
  }
}

/** Port over knowledge_documents / knowledge_chunks. Read methods throw KnowledgeStorageError. */
export interface KnowledgeIndexStore {
  /** Owner of the document, or null when it does not exist. */
  documentOwner(documentId: string): Promise<string | null>
  /** Exact row count, or null when the database returned no count. */
  countChunks(documentId: string, pendingOnly: boolean): Promise<number | null>
  listPendingChunks(documentId: string, limit: number): Promise<KnowledgeChunkRow[]>
  /** Returns false when the write failed. */
  saveEmbedding(documentId: string, chunkId: string, embedding: number[]): Promise<boolean>
  /** Clear the indexed flag. With a userId the update is scoped to that owner. Returns false on failure. */
  clearIndexed(documentId: string, userId: string | null): Promise<boolean>
  /** Set the indexed flag. Returns false when the write failed or matched no row. */
  markIndexed(documentId: string, userId: string): Promise<boolean>
}

export interface EmbedChunksDeps {
  authenticate(req: Request): Promise<{ id: string }>
  /** null when the provider key is not configured. */
  embeddingProvider(): EmbeddingProvider | null
  knowledgeStore(): KnowledgeIndexStore
  usageStore(): EmbeddingUsageStore
  now?(): Date
}

async function embedBatch(
  provider: EmbeddingProvider,
  store: KnowledgeIndexStore,
  documentId: string,
  batch: KnowledgeChunkRow[],
): Promise<number> {
  const outcome = await provider.embed(batch.map((chunk) => chunk.content))
  if (!outcome.ok || outcome.data.length !== batch.length) return 0

  const received = new Set<number>()
  let saved = 0
  for (const item of outcome.data) {
    if (
      typeof item.index !== 'number'
      || !Number.isInteger(item.index)
      || item.index < 0
      || item.index >= batch.length
      || received.has(item.index)
      || !isEmbedding(item.embedding)
    ) {
      continue
    }
    received.add(item.index)
    if (await store.saveEmbedding(documentId, batch[item.index].id, item.embedding)) saved += 1
  }
  return saved
}

export function createEmbedChunksHandler(deps: EmbedChunksDeps): (req: Request) => Promise<Response> {
  const now = () => deps.now?.() ?? new Date()

  return async (req: Request): Promise<Response> => {
    const preflight = handleCorsPreflightRequest(req)
    if (preflight) return preflight
    if (req.method !== 'POST') return json(405, { error: 'method_not_allowed' })

    try {
      const user = await deps.authenticate(req)
      const body = await req.json().catch(() => null) as { document_id?: unknown } | null
      if (!body || typeof body.document_id !== 'string' || !UUID_PATTERN.test(body.document_id)) {
        return json(400, { error: 'invalid_document_id' })
      }
      const documentId = body.document_id

      const provider = deps.embeddingProvider()
      if (!provider) return json(503, { error: 'embedding_provider_unavailable' })

      const store = deps.knowledgeStore()
      const owner = await store.documentOwner(documentId)
      if (owner !== user.id) return json(404, { error: 'knowledge_document_not_found' })

      const total = await store.countChunks(documentId, false) ?? 0
      if (total === 0) {
        await store.clearIndexed(documentId, null)
        return json(409, { error: 'knowledge_document_empty' })
      }
      if (total > EMBEDDING_LIMITS.maxChunksPerDocument) {
        return json(413, {
          error: 'knowledge_document_too_large',
          total,
          max_chunks: EMBEDDING_LIMITS.maxChunksPerDocument,
        })
      }

      const chunks = await store.listPendingChunks(documentId, EMBEDDING_LIMITS.maxChunksPerDocument)
      const oversized = firstOversizedChunk(chunks)
      if (oversized !== -1) {
        return json(413, {
          error: 'knowledge_chunk_too_large',
          chunk_id: chunks[oversized].id,
          max_chars: EMBEDDING_LIMITS.maxChunkChars,
        })
      }

      // Reserve the whole run before the first paid call; unspent units are refunded.
      const usage = deps.usageStore()
      const units = embeddingUnits(chunks.map((chunk) => chunk.content))
      let reservation: EmbeddingReservation
      try {
        reservation = await reserveEmbeddingUnits(usage, user.id, units, now())
      } catch {
        return json(503, { error: 'quota_unavailable' })
      }
      if (!reservation.allowed) return json(429, quotaExceededBody(reservation, units))

      let embedded = 0
      let failed = 0
      let unspentUnits = 0
      for (const batch of planEmbeddingBatches(chunks)) {
        let saved = 0
        try {
          saved = await embedBatch(provider, store, documentId, batch)
        } catch {
          saved = 0
        }
        embedded += saved
        failed += batch.length - saved
        if (saved === 0) unspentUnits += embeddingUnits(batch.map((chunk) => chunk.content))
      }
      await refundEmbeddingUnits(usage, user.id, reservation, unspentUnits, now())

      const remaining = await store.countChunks(documentId, true) ?? total
      if (failed > 0 || remaining > 0) {
        if (!await store.clearIndexed(documentId, user.id)) {
          return json(500, { error: 'knowledge_storage_failed' })
        }
        return json(502, { error: 'embedding_failed', embedded, remaining })
      }

      if (!await store.markIndexed(documentId, user.id)) {
        return json(500, { error: 'knowledge_storage_failed' })
      }
      return json(200, { embedded, total, indexed: true })
    } catch (error) {
      if (error instanceof KnowledgeStorageError) return json(500, { error: 'knowledge_storage_failed' })
      return errorResponse(error)
    }
  }
}