// 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 /** Exact row count, or null when the database returned no count. */ countChunks(documentId: string, pendingOnly: boolean): Promise listPendingChunks(documentId: string, limit: number): Promise /** Returns false when the write failed. */ saveEmbedding(documentId: string, chunkId: string, embedding: number[]): Promise /** 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 /** Set the indexed flag. Returns false when the write failed or matched no row. */ markIndexed(documentId: string, userId: string): Promise } 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 { const outcome = await provider.embed(batch.map((chunk) => chunk.content)) if (!outcome.ok || outcome.data.length !== batch.length) return 0 const received = new Set() 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 { const now = () => deps.now?.() ?? new Date() return async (req: Request): Promise => { 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) } } }