fix(edge): meter and size-cap knowledge embeddings in embed-chunks and search-knowledge
This commit is contained in:
Yun Chan
parent
b35676c75c
commit
1afaea7214
12 changed files with 1111 additions and 270 deletions
250
server/supabase/functions/_shared/embedding-quota.ts
Normal file
250
server/supabase/functions/_shared/embedding-quota.ts
Normal file
|
|
@ -0,0 +1,250 @@
|
|||
// server/supabase/functions/_shared/embedding-quota.ts
|
||||
// Cost guard for the paid OpenAI embedding calls made by embed-chunks and
|
||||
// search-knowledge.
|
||||
//
|
||||
// Before this module both functions only checked the JWT: any signed-in account
|
||||
// could insert an unbounded number of knowledge_chunks rows (RLS only checks
|
||||
// document ownership, `content` is unbounded text) and have every one of them
|
||||
// embedded, repeatedly, and search-knowledge embedded a 4,000-char query per
|
||||
// call with no counter. stt/llm/realtime all go through quota.ts; embeddings
|
||||
// did not.
|
||||
//
|
||||
// Policy (pure): size caps + per-tier unit budget + batch planning.
|
||||
// Port: EmbeddingUsageStore (tier lookup + atomic counter).
|
||||
// Adapter: createSupabaseEmbeddingUsageStore (subscriptions + daily_usage
|
||||
// via the existing service-role `increment_daily_usage` RPC).
|
||||
//
|
||||
// TODO(plan-catalog): the per-tier budget below belongs in
|
||||
// packages/core/src/plan-catalog.ts `PLAN_QUOTA` (feature `knowledge_embedding`)
|
||||
// and should be read from core-contract.generated.ts once it is added there.
|
||||
|
||||
import type { createClient } from '@supabase/supabase-js'
|
||||
import type { PlanQuotaPeriod, PlanQuotaTier } from './core-contract.generated.ts'
|
||||
|
||||
/** daily_usage.feature key for embedding spend. */
|
||||
export const EMBEDDING_QUOTA_FEATURE = 'knowledge_embedding'
|
||||
|
||||
export const EMBEDDING_LIMITS = Object.freeze({
|
||||
/** Longest chunk accepted for embedding (the OpenAI per-input cap is ~8k tokens). */
|
||||
maxChunkChars: 8_000,
|
||||
/** Most chunks one document may have. Desktop chunking stops at 2,000 chunks. */
|
||||
maxChunksPerDocument: 2_000,
|
||||
/** Longest search query accepted. */
|
||||
maxQueryChars: 4_000,
|
||||
/** Inputs per provider request. */
|
||||
maxBatchInputs: 100,
|
||||
/** Characters per provider request (keeps a batch under the per-request token cap). */
|
||||
maxBatchChars: 100_000,
|
||||
/** One quota unit = up to this many characters of embedded text. */
|
||||
charsPerUnit: 1_000,
|
||||
})
|
||||
|
||||
export interface EmbeddingQuotaPolicy {
|
||||
/** Units per period. -1 = unlimited, 0 = not available. */
|
||||
limit: number
|
||||
period: PlanQuotaPeriod
|
||||
}
|
||||
|
||||
/**
|
||||
* Units (1 unit = up to 1,000 chars) per tier. Cloud knowledge is offered on
|
||||
* every tier (web/mobile have no paywall for it), so free gets a small weekly
|
||||
* budget instead of none: 5,000 units ≈ 5M chars ≈ 20 maximum-size mobile
|
||||
* documents a week.
|
||||
*/
|
||||
export const EMBEDDING_QUOTA: Readonly<Record<PlanQuotaTier, EmbeddingQuotaPolicy>> = Object.freeze({
|
||||
free: { limit: 5_000, period: 'weekly' },
|
||||
pro: { limit: 50_000, period: 'daily' },
|
||||
pro_plus: { limit: 150_000, period: 'daily' },
|
||||
team: { limit: 300_000, period: 'daily' },
|
||||
enterprise: { limit: -1, period: 'daily' },
|
||||
})
|
||||
|
||||
const QUOTA_TIERS: readonly PlanQuotaTier[] = ['free', 'pro', 'pro_plus', 'team', 'enterprise']
|
||||
|
||||
export function normalizeQuotaTier(value: unknown): PlanQuotaTier {
|
||||
return typeof value === 'string' && (QUOTA_TIERS as readonly string[]).includes(value)
|
||||
? value as PlanQuotaTier
|
||||
: 'free'
|
||||
}
|
||||
|
||||
export function embeddingPolicyFor(tier: PlanQuotaTier): EmbeddingQuotaPolicy {
|
||||
return EMBEDDING_QUOTA[tier] ?? { limit: 0, period: 'daily' }
|
||||
}
|
||||
|
||||
/** Units one text costs: every started 1,000 characters, at least 1. */
|
||||
export function embeddingUnitsFor(text: string): number {
|
||||
return Math.max(1, Math.ceil(text.length / EMBEDDING_LIMITS.charsPerUnit))
|
||||
}
|
||||
|
||||
export function embeddingUnits(texts: readonly string[]): number {
|
||||
return texts.reduce((sum, text) => sum + embeddingUnitsFor(text), 0)
|
||||
}
|
||||
|
||||
/** Index of the first chunk over `maxChunkChars`, or -1. */
|
||||
export function firstOversizedChunk(chunks: readonly { content: string }[]): number {
|
||||
return chunks.findIndex((chunk) => chunk.content.length > EMBEDDING_LIMITS.maxChunkChars)
|
||||
}
|
||||
|
||||
/** Split chunks into provider requests bounded by input count and characters. */
|
||||
export function planEmbeddingBatches<T extends { content: string }>(chunks: readonly T[]): T[][] {
|
||||
const batches: T[][] = []
|
||||
let current: T[] = []
|
||||
let chars = 0
|
||||
for (const chunk of chunks) {
|
||||
const size = chunk.content.length
|
||||
if (
|
||||
current.length > 0
|
||||
&& (current.length >= EMBEDDING_LIMITS.maxBatchInputs || chars + size > EMBEDDING_LIMITS.maxBatchChars)
|
||||
) {
|
||||
batches.push(current)
|
||||
current = []
|
||||
chars = 0
|
||||
}
|
||||
current.push(chunk)
|
||||
chars += size
|
||||
}
|
||||
if (current.length > 0) batches.push(current)
|
||||
return batches
|
||||
}
|
||||
|
||||
function isoDate(date: Date): string {
|
||||
return date.toISOString().slice(0, 10)
|
||||
}
|
||||
|
||||
/** First date (inclusive) of the usage window before today, or null for a daily period. */
|
||||
export function usageWindowStart(period: PlanQuotaPeriod, now: Date): string | null {
|
||||
if (period !== 'weekly') return null
|
||||
const start = new Date(now.getTime())
|
||||
start.setUTCDate(start.getUTCDate() - 6)
|
||||
return isoDate(start)
|
||||
}
|
||||
|
||||
// ── Port ────────────────────────────────────────────────────────────────
|
||||
|
||||
export interface EmbeddingUsageStore {
|
||||
readTier(userId: string): Promise<PlanQuotaTier>
|
||||
/** Atomically add `amount` (may be negative) to today's counter; returns the new count. */
|
||||
increment(userId: string, feature: string, amount: number): Promise<number>
|
||||
/** Sum of counters with fromDate <= date < beforeDate. */
|
||||
sumBetween(userId: string, feature: string, fromDate: string, beforeDate: string): Promise<number>
|
||||
}
|
||||
|
||||
export interface EmbeddingReservation {
|
||||
allowed: boolean
|
||||
/** Units held by this reservation (0 when denied or nothing was requested). */
|
||||
units: number
|
||||
/** Usage in the current window including this reservation when allowed. */
|
||||
current: number
|
||||
limit: number
|
||||
period: PlanQuotaPeriod
|
||||
tier: PlanQuotaTier
|
||||
/** UTC date the units were counted on. */
|
||||
date: string
|
||||
}
|
||||
|
||||
/**
|
||||
* Reserve `units` before spending. The counter is incremented first (atomic in
|
||||
* the database) and rolled back when the window total then exceeds the limit,
|
||||
* so concurrent requests can never admit more than the limit between them.
|
||||
* Throws when the store fails: callers must fail closed.
|
||||
*/
|
||||
export async function reserveEmbeddingUnits(
|
||||
store: EmbeddingUsageStore,
|
||||
userId: string,
|
||||
units: number,
|
||||
now: Date = new Date(),
|
||||
): Promise<EmbeddingReservation> {
|
||||
const tier = await store.readTier(userId)
|
||||
const policy = embeddingPolicyFor(tier)
|
||||
const date = isoDate(now)
|
||||
const base = { limit: policy.limit, period: policy.period, tier, date }
|
||||
|
||||
if (units <= 0) return { ...base, allowed: true, units: 0, current: 0 }
|
||||
if (policy.limit === 0) return { ...base, allowed: false, units: 0, current: 0 }
|
||||
|
||||
const today = await store.increment(userId, EMBEDDING_QUOTA_FEATURE, units)
|
||||
if (policy.limit === -1) return { ...base, allowed: true, units, current: today }
|
||||
|
||||
const windowStart = usageWindowStart(policy.period, now)
|
||||
const earlier = windowStart
|
||||
? await store.sumBetween(userId, EMBEDDING_QUOTA_FEATURE, windowStart, date)
|
||||
: 0
|
||||
const used = earlier + today
|
||||
if (used > policy.limit) {
|
||||
await store.increment(userId, EMBEDDING_QUOTA_FEATURE, -units)
|
||||
return { ...base, allowed: false, units: 0, current: used - units }
|
||||
}
|
||||
return { ...base, allowed: true, units, current: used }
|
||||
}
|
||||
|
||||
/**
|
||||
* Give back units that were reserved but not spent. Best effort: a failure
|
||||
* leaves the user slightly over-counted, never under-counted. A refund after
|
||||
* the UTC day rolled over is skipped so it cannot credit the new day.
|
||||
*/
|
||||
export async function refundEmbeddingUnits(
|
||||
store: EmbeddingUsageStore,
|
||||
userId: string,
|
||||
reservation: EmbeddingReservation,
|
||||
units: number,
|
||||
now: Date = new Date(),
|
||||
): Promise<void> {
|
||||
const amount = Math.min(units, reservation.units)
|
||||
if (amount <= 0 || isoDate(now) !== reservation.date) return
|
||||
try {
|
||||
await store.increment(userId, EMBEDDING_QUOTA_FEATURE, -amount)
|
||||
} catch {
|
||||
// over-counting is the safe direction
|
||||
}
|
||||
}
|
||||
|
||||
export function quotaExceededBody(reservation: EmbeddingReservation, requested: number): Record<string, unknown> {
|
||||
return {
|
||||
error: 'quota_exceeded',
|
||||
feature: EMBEDDING_QUOTA_FEATURE,
|
||||
requested,
|
||||
current: reservation.current,
|
||||
limit: reservation.limit,
|
||||
period: reservation.period,
|
||||
tier: reservation.tier,
|
||||
}
|
||||
}
|
||||
|
||||
// ── Supabase adapter ────────────────────────────────────────────────────
|
||||
|
||||
type SupabaseClient = ReturnType<typeof createClient>
|
||||
|
||||
export function createSupabaseEmbeddingUsageStore(client: SupabaseClient): EmbeddingUsageStore {
|
||||
return {
|
||||
async readTier(userId) {
|
||||
const { data, error } = await client
|
||||
.from('subscriptions')
|
||||
.select('tier')
|
||||
.eq('user_id', userId)
|
||||
.maybeSingle()
|
||||
if (error) throw new Error('Failed to read subscription tier.')
|
||||
return normalizeQuotaTier((data as { tier?: unknown } | null)?.tier)
|
||||
},
|
||||
async increment(userId, feature, amount) {
|
||||
const { data, error } = await client.rpc('increment_daily_usage', {
|
||||
p_user_id: userId,
|
||||
p_feature: feature,
|
||||
p_amount: amount,
|
||||
})
|
||||
if (error || typeof data !== 'number') throw new Error('Failed to update embedding usage.')
|
||||
return data
|
||||
},
|
||||
async sumBetween(userId, feature, fromDate, beforeDate) {
|
||||
const { data, error } = await client
|
||||
.from('daily_usage')
|
||||
.select('count')
|
||||
.eq('user_id', userId)
|
||||
.eq('feature', feature)
|
||||
.gte('date', fromDate)
|
||||
.lt('date', beforeDate)
|
||||
if (error) throw new Error('Failed to read embedding usage.')
|
||||
const rows = (data ?? []) as Array<{ count: number | null }>
|
||||
return rows.reduce((sum, row) => sum + (row.count ?? 0), 0)
|
||||
},
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue