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

import { corsHeaders, handleCorsPreflightRequest } from '../_shared/cors.ts'
import { requireUser, authErrorResponse, type AuthError } from '../_shared/auth.ts'
import { createServiceRoleClient } from '../_shared/quota.ts'
import { readProviderKey } from '../_shared/provider-key.ts'

const EMBEDDING_DIMENSIONS = 1536
const PROVIDER_TIMEOUT_MS = 45_000
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

interface OpenAIEmbeddingResponse {
  data?: Array<{ embedding?: unknown; index?: unknown }>
}

function json(status: number, body: Record<string, unknown>): Response {
  return new Response(JSON.stringify(body), {
    status,
    headers: { ...corsHeaders, 'Content-Type': 'application/json' },
  })
}

function validEmbedding(value: unknown): value is number[] {
  return Array.isArray(value)
    && value.length === EMBEDDING_DIMENSIONS
    && value.every((entry) => typeof entry === 'number' && Number.isFinite(entry))
}

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

  try {
    const user = await requireUser(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 openaiKey = readProviderKey('OPENAI_API_KEY')
    if (!openaiKey) return json(503, { error: 'embedding_provider_unavailable' })

    const serviceClient = createServiceRoleClient()
    const documentResult = await serviceClient
      .from('knowledge_documents')
      .select('id,user_id')
      .eq('id', body.document_id)
      .maybeSingle()
    if (documentResult.error) return json(500, { error: 'knowledge_storage_failed' })
    if (!documentResult.data || documentResult.data.user_id !== user.id) {
      return json(404, { error: 'knowledge_document_not_found' })
    }

    const totalResult = await serviceClient
      .from('knowledge_chunks')
      .select('id', { count: 'exact', head: true })
      .eq('document_id', body.document_id)
    if (totalResult.error) return json(500, { error: 'knowledge_storage_failed' })
    const total = totalResult.count ?? 0
    if (total === 0) {
      await serviceClient
        .from('knowledge_documents')
        .update({ indexed: false, indexed_at: null })
        .eq('id', body.document_id)
      return json(409, { error: 'knowledge_document_empty' })
    }

    const pendingResult = await serviceClient
      .from('knowledge_chunks')
      .select('id,content')
      .eq('document_id', body.document_id)
      .is('embedding', null)
      .order('chunk_index', { ascending: true })
    if (pendingResult.error) return json(500, { error: 'knowledge_storage_failed' })

    const chunks = (pendingResult.data ?? []) as Array<{ id: string; content: string }>
    let embedded = 0
    let failed = 0

    for (let offset = 0; offset < chunks.length; offset += 100) {
      const batch = chunks.slice(offset, offset + 100)
      try {
        const response = await fetch('https://api.openai.com/v1/embeddings', {
          method: 'POST',
          headers: {
            Authorization: `Bearer ${openaiKey}`,
            'Content-Type': 'application/json',
          },
          body: JSON.stringify({
            model: 'text-embedding-3-small',
            input: batch.map((chunk) => chunk.content),
            dimensions: EMBEDDING_DIMENSIONS,
          }),
          signal: AbortSignal.timeout(PROVIDER_TIMEOUT_MS),
        })
        if (!response.ok) {
          failed += batch.length
          continue
        }

        const payload = await response.json().catch(() => null) as OpenAIEmbeddingResponse | null
        if (!payload || !Array.isArray(payload.data) || payload.data.length !== batch.length) {
          failed += batch.length
          continue
        }

        const received = new Set<number>()
        let batchSuccess = 0
        for (const item of payload.data) {
          if (
            typeof item.index !== 'number'
            || !Number.isInteger(item.index)
            || item.index < 0
            || item.index >= batch.length
            || received.has(item.index)
            || !validEmbedding(item.embedding)
          ) {
            continue
          }
          received.add(item.index)
          const update = await serviceClient
            .from('knowledge_chunks')
            .update({ embedding: item.embedding })
            .eq('id', batch[item.index].id)
            .eq('document_id', body.document_id)
          if (!update.error) {
            embedded += 1
            batchSuccess += 1
          }
        }
        failed += batch.length - batchSuccess
      } catch {
        failed += batch.length
      }
    }

    const remainingResult = await serviceClient
      .from('knowledge_chunks')
      .select('id', { count: 'exact', head: true })
      .eq('document_id', body.document_id)
      .is('embedding', null)
    if (remainingResult.error) return json(500, { error: 'knowledge_storage_failed' })

    const remaining = remainingResult.count ?? total
    if (failed > 0 || remaining > 0) {
      const rollback = await serviceClient
        .from('knowledge_documents')
        .update({ indexed: false, indexed_at: null })
        .eq('id', body.document_id)
        .eq('user_id', user.id)
      if (rollback.error) return json(500, { error: 'knowledge_storage_failed' })
      return json(502, {
        error: 'embedding_failed',
        embedded,
        remaining,
      })
    }

    const indexed = await serviceClient
      .from('knowledge_documents')
      .update({ indexed: true, indexed_at: new Date().toISOString() })
      .eq('id', body.document_id)
      .eq('user_id', user.id)
      .select('id')
      .maybeSingle()
    if (indexed.error || !indexed.data) return json(500, { error: 'knowledge_storage_failed' })

    return json(200, { embedded, total, indexed: true })
  } catch (error) {
    if (
      error
      && typeof error === 'object'
      && 'status' in error
      && (error.status === 401 || error.status === 403)
      && 'message' in error
      && typeof error.message === 'string'
    ) {
      return authErrorResponse(error as AuthError, corsHeaders)
    }
    return json(500, { error: 'internal_error' })
  }
})