fix(ai-kb): address code-review findings

- Security: REVOKE ALL FROM PUBLIC + GRANT EXECUTE (authenticated,
  service_role) on the match RPCs so anon can't read a KB by supplying
  an account_id (mirrors migrations 018/025).
- Ingest: store chunks even when embedding fails (insert first, rethrow
  after) so a doc stays lexically searchable; the "lexical still works"
  warning is now truthful.
- Retrieval: short-circuit on an empty KB (cheap COUNT) to avoid a paid
  query embedding + 2 RPCs on every draft/auto-reply.
- Migration: HNSW instead of IVFFlat (robust on small/growing tables);
  language-neutral 'simple' FTS config (was 'english'); semantic RPC
  takes text and casts ::vector(1536) to remove PostgREST bind ambiguity.
- Embeddings key: distinguish corrupt-from-absent (loadEmbeddingsKey
  returns {key,corrupt}); log on decrypt failure; reindex now reports
  failure instead of false success; POST/PATCH warn on keyword-only save.
- embedTexts: fail loud when the provider omits result indices (was
  silently misaligning chunks with vectors).

Tests: +3 (empty-KB guard, embed-failure-still-indexes, missing-index).
587 pass; typecheck/lint/build green.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Arnas Donauskas 2026-07-02 17:20:38 +03:00
parent ce1fa4705d
commit 020dc2ae11
9 changed files with 187 additions and 30 deletions

View File

@ -77,7 +77,10 @@ export async function PATCH(request: Request, { params }: Params) {
if (!updated) return NextResponse.json({ error: 'Not found' }, { status: 404 }) if (!updated) return NextResponse.json({ error: 'Not found' }, { status: 404 })
if (content !== undefined) { if (content !== undefined) {
const embeddingsApiKey = await loadEmbeddingsKey(supabase, accountId) const { key: embeddingsApiKey, corrupt } = await loadEmbeddingsKey(
supabase,
accountId,
)
try { try {
await ingestDocument(supabase, accountId, { embeddingsApiKey }, id, content) await ingestDocument(supabase, accountId, { embeddingsApiKey }, id, content)
} catch (err) { } catch (err) {
@ -91,6 +94,13 @@ export async function PATCH(request: Request, { params }: Params) {
{ status: 200 }, { status: 200 },
) )
} }
if (corrupt) {
return NextResponse.json({
success: true,
warning:
'Updated with keyword search only — your embeddings key could not be decrypted (check ENCRYPTION_KEY, then re-enter the key).',
})
}
} }
return NextResponse.json({ success: true }) return NextResponse.json({ success: true })

View File

@ -31,7 +31,24 @@ export async function POST() {
) )
} }
const embeddingsApiKey = await loadEmbeddingsKey(supabase, accountId) const { key: embeddingsApiKey, corrupt } = await loadEmbeddingsKey(
supabase,
accountId,
)
// The whole point of Reindex is usually to backfill embeddings — so
// if a key is configured but can't be decrypted, don't quietly do a
// lexical-only pass and report success. Stop and tell the admin.
if (corrupt) {
return NextResponse.json(
{
success: false,
reindexed: 0,
error:
'Your embeddings key could not be decrypted (check ENCRYPTION_KEY, then re-enter the key in Settings → AI Assistant). Nothing was reindexed.',
},
{ status: 200 },
)
}
let reindexed = 0 let reindexed = 0
for (const doc of docs ?? []) { for (const doc of docs ?? []) {

View File

@ -70,12 +70,20 @@ export async function POST(request: Request) {
) )
} }
const embeddingsApiKey = await loadEmbeddingsKey(supabase, accountId) const { key: embeddingsApiKey, corrupt } = await loadEmbeddingsKey(
supabase,
accountId,
)
try { try {
await ingestDocument(supabase, accountId, { embeddingsApiKey }, doc.id, content) await ingestDocument(
supabase,
accountId,
{ embeddingsApiKey },
doc.id,
content,
)
} catch (err) { } catch (err) {
const message = const message = err instanceof AiError ? err.message : 'indexing failed'
err instanceof AiError ? err.message : 'indexing failed'
console.error('[ai/knowledge POST] ingest error:', err) console.error('[ai/knowledge POST] ingest error:', err)
return NextResponse.json( return NextResponse.json(
{ {
@ -87,6 +95,14 @@ export async function POST(request: Request) {
) )
} }
if (corrupt) {
return NextResponse.json({
success: true,
id: doc.id,
warning:
'Saved with keyword search only — your embeddings key could not be decrypted (check ENCRYPTION_KEY, then re-enter the key).',
})
}
return NextResponse.json({ success: true, id: doc.id }) return NextResponse.json({ success: true, id: doc.id })
} catch (err) { } catch (err) {
return toErrorResponse(err) return toErrorResponse(err)

View File

@ -55,6 +55,11 @@ export async function loadAiConfig(
try { try {
embeddingsApiKey = decrypt(row.embeddings_api_key) embeddingsApiKey = decrypt(row.embeddings_api_key)
} catch { } catch {
// Not silent — a rotated/mismatched ENCRYPTION_KEY here means
// semantic search quietly stops working, so leave a breadcrumb.
console.error(
`[ai config] embeddings key for account ${accountId} could not be decrypted — check ENCRYPTION_KEY; semantic search is disabled until it is re-entered.`,
)
embeddingsApiKey = null embeddingsApiKey = null
} }
} }
@ -75,22 +80,29 @@ export async function loadAiConfig(
* Load + decrypt just the embeddings key, independent of `is_active`. * Load + decrypt just the embeddings key, independent of `is_active`.
* Used by the knowledge-base ingest routes so the KB gets embedded (and * Used by the knowledge-base ingest routes so the KB gets embedded (and
* semantic search works) whenever an embeddings key is present, even if * semantic search works) whenever an embeddings key is present, even if
* the assistant's master switch is currently off. Returns null when * the assistant's master switch is currently off.
* there's no key or it can't be decrypted. *
* Returns `{ key, corrupt }`: `key` is null when there's no key OR it
* can't be decrypted; `corrupt` distinguishes those cases so callers can
* warn ("a key is set but unusable") rather than silently indexing
* lexical-only and reporting success.
*/ */
export async function loadEmbeddingsKey( export async function loadEmbeddingsKey(
db: SupabaseClient, db: SupabaseClient,
accountId: string, accountId: string,
): Promise<string | null> { ): Promise<{ key: string | null; corrupt: boolean }> {
const { data, error } = await db const { data, error } = await db
.from('ai_configs') .from('ai_configs')
.select('embeddings_api_key') .select('embeddings_api_key')
.eq('account_id', accountId) .eq('account_id', accountId)
.maybeSingle() .maybeSingle()
if (error || !data?.embeddings_api_key) return null if (error || !data?.embeddings_api_key) return { key: null, corrupt: false }
try { try {
return decrypt(data.embeddings_api_key) return { key: decrypt(data.embeddings_api_key), corrupt: false }
} catch { } catch {
return null console.error(
`[ai config] embeddings key for account ${accountId} could not be decrypted — check ENCRYPTION_KEY.`,
)
return { key: null, corrupt: true }
} }
} }

View File

@ -85,6 +85,18 @@ describe('embedTexts', () => {
}) })
}) })
it('throws when the provider omits result indices', async () => {
vi.stubGlobal(
'fetch',
vi.fn().mockResolvedValue({
ok: true,
status: 200,
json: async () => ({ data: [{ embedding: [0.1] }, { embedding: [0.2] }] }),
} as unknown as Response),
)
await expect(embedTexts('sk-x', ['a', 'b'])).rejects.toBeInstanceOf(AiError)
})
it('throws on a malformed response (count mismatch)', async () => { it('throws on a malformed response (count mismatch)', async () => {
vi.stubGlobal( vi.stubGlobal(
'fetch', 'fetch',

View File

@ -77,8 +77,15 @@ export async function embedTexts(
} }
// Sort by index so order matches the input batch regardless of how // Sort by index so order matches the input batch regardless of how
// the provider returns them. // the provider returns them. Require a real numeric index — defaulting
const ordered = [...rows].sort((a, b) => (a.index ?? 0) - (b.index ?? 0)) // a missing one to 0 would silently misalign chunks with their
// vectors (chunk N gets chunk M's embedding), so fail loud instead.
if (rows.some((r) => typeof r.index !== 'number')) {
throw new AiError('Embeddings response was missing result indices.', {
code: 'embeddings_malformed',
})
}
const ordered = [...rows].sort((a, b) => a.index! - b.index!)
for (const r of ordered) { for (const r of ordered) {
if (!Array.isArray(r.embedding)) { if (!Array.isArray(r.embedding)) {
throw new AiError('Embeddings response missing a vector.', { throw new AiError('Embeddings response missing a vector.', {

View File

@ -12,6 +12,7 @@ import { retrieveKnowledge, ingestDocument } from './knowledge'
interface FakeState { interface FakeState {
semantic: { id: string; content: string }[] semantic: { id: string; content: string }[]
fts: { id: string; content: string }[] fts: { id: string; content: string }[]
chunkCount: number
rpcCalls: string[] rpcCalls: string[]
inserted: Record<string, unknown>[] | null inserted: Record<string, unknown>[] | null
deletedFor: string | null deletedFor: string | null
@ -21,6 +22,7 @@ function makeDb() {
const state: FakeState = { const state: FakeState = {
semantic: [], semantic: [],
fts: [], fts: [],
chunkCount: 5, // account has a non-empty KB by default
rpcCalls: [], rpcCalls: [],
inserted: null, inserted: null,
deletedFor: null, deletedFor: null,
@ -35,6 +37,10 @@ function makeDb() {
return Promise.resolve({ data: null, error: null }) return Promise.resolve({ data: null, error: null })
}, },
from: () => ({ from: () => ({
// retrieveKnowledge's empty-KB count guard.
select: () => ({
eq: () => Promise.resolve({ count: state.chunkCount, error: null }),
}),
delete: () => ({ delete: () => ({
eq: (_col: string, val: string) => { eq: (_col: string, val: string) => {
state.deletedFor = val state.deletedFor = val
@ -64,6 +70,15 @@ describe('retrieveKnowledge', () => {
expect(state.rpcCalls).toEqual([]) expect(state.rpcCalls).toEqual([])
}) })
it('short-circuits (no embed, no RPC) when the KB is empty', async () => {
const { db, state } = makeDb()
state.chunkCount = 0
const out = await retrieveKnowledge(db, 'acct', { embeddingsApiKey: 'sk-x' }, 'q')
expect(out).toEqual([])
expect(h.embedTexts).not.toHaveBeenCalled()
expect(state.rpcCalls).toEqual([])
})
it('uses lexical FTS only when there is no embeddings key', async () => { it('uses lexical FTS only when there is no embeddings key', async () => {
const { db, state } = makeDb() const { db, state } = makeDb()
state.fts = [{ id: 'f1', content: 'F1' }] state.fts = [{ id: 'f1', content: 'F1' }]
@ -131,4 +146,15 @@ describe('ingestDocument', () => {
expect(state.inserted).toBeNull() expect(state.inserted).toBeNull()
expect(h.embedTexts).not.toHaveBeenCalled() expect(h.embedTexts).not.toHaveBeenCalled()
}) })
it('still stores lexical chunks when embedding fails, then rethrows', async () => {
const { db, state } = makeDb()
h.embedTexts.mockRejectedValueOnce(new Error('rate limited'))
await expect(
ingestDocument(db, 'acct', { embeddingsApiKey: 'sk-x' }, 'doc-1', 'hello world'),
).rejects.toThrow('rate limited')
// Chunks were inserted (lexical search works) despite the embed failure…
expect(state.inserted).toHaveLength(1)
expect(state.inserted![0].embedding).toBeNull()
})
}) })

View File

@ -33,11 +33,6 @@ export async function ingestDocument(
): Promise<void> { ): Promise<void> {
const chunks = chunkText(content) const chunks = chunkText(content)
let embeddings: number[][] | null = null
if (chunks.length > 0 && config.embeddingsApiKey) {
embeddings = await embedTexts(config.embeddingsApiKey, chunks)
}
// Replace, don't append — re-ingest must be idempotent. // Replace, don't append — re-ingest must be idempotent.
const { error: delErr } = await db const { error: delErr } = await db
.from('ai_knowledge_chunks') .from('ai_knowledge_chunks')
@ -47,6 +42,22 @@ export async function ingestDocument(
if (chunks.length === 0) return if (chunks.length === 0) return
// Embed if a key is set, but DON'T let an embedding failure stop the
// chunks from being stored: a failed embed must still leave the
// document searchable lexically. We record the error and rethrow it
// AFTER inserting (embedding-less) rows, so the route can warn
// "semantic indexing failed" — which is now truthful, because lexical
// search really does still work.
let embeddings: number[][] | null = null
let embedError: unknown = null
if (config.embeddingsApiKey) {
try {
embeddings = await embedTexts(config.embeddingsApiKey, chunks)
} catch (err) {
embedError = err
}
}
const rows = chunks.map((content, i) => ({ const rows = chunks.map((content, i) => ({
document_id: documentId, document_id: documentId,
account_id: accountId, account_id: accountId,
@ -57,6 +68,8 @@ export async function ingestDocument(
const { error: insErr } = await db.from('ai_knowledge_chunks').insert(rows) const { error: insErr } = await db.from('ai_knowledge_chunks').insert(rows)
if (insErr) throw insErr if (insErr) throw insErr
if (embedError) throw embedError
} }
/** /**
@ -78,6 +91,20 @@ export async function retrieveKnowledge(
const query = queryText.trim() const query = queryText.trim()
if (!query || k <= 0) return [] if (!query || k <= 0) return []
// Skip everything when the account has no knowledge base — otherwise
// every draft / auto-reply would pay for a query embedding + two RPCs
// just to get []. One cheap indexed COUNT (head, no rows) instead of a
// paid embeddings call on the hot path.
try {
const { count, error } = await db
.from('ai_knowledge_chunks')
.select('id', { count: 'exact', head: true })
.eq('account_id', accountId)
if (error || !count) return []
} catch {
return []
}
const picked = new Map<string, string>() // id → content, preserves order const picked = new Map<string, string>() // id → content, preserves order
// Semantic path. // Semantic path.

View File

@ -96,7 +96,14 @@ CREATE TABLE IF NOT EXISTS ai_knowledge_chunks (
account_id uuid NOT NULL REFERENCES accounts(id) ON DELETE CASCADE, account_id uuid NOT NULL REFERENCES accounts(id) ON DELETE CASCADE,
chunk_index integer NOT NULL DEFAULT 0, chunk_index integer NOT NULL DEFAULT 0,
content text NOT NULL, content text NOT NULL,
fts tsvector GENERATED ALWAYS AS (to_tsvector('english', content)) STORED, -- Language-neutral FTS config: wacrm is used in many languages
-- (its markets include BR / LATAM / IN), and this lexical path is the
-- fallback for accounts without an embeddings key. `'simple'` tokenizes
-- + lowercases without English-only stemming/stopwords, so it degrades
-- gracefully in any language. (Per-account language config is a
-- follow-up; accounts wanting paraphrase/morphology matching add an
-- embeddings key for the semantic path.)
fts tsvector GENERATED ALWAYS AS (to_tsvector('simple', content)) STORED,
embedding vector(1536), embedding vector(1536),
created_at timestamptz NOT NULL DEFAULT now() created_at timestamptz NOT NULL DEFAULT now()
); );
@ -109,9 +116,14 @@ CREATE INDEX IF NOT EXISTS ai_knowledge_chunks_fts_idx
ON ai_knowledge_chunks USING gin (fts); ON ai_knowledge_chunks USING gin (fts);
-- Cosine-distance ANN index for the semantic path. Rows with a NULL -- Cosine-distance ANN index for the semantic path. Rows with a NULL
-- embedding (lexical-only accounts) are simply absent from it. -- embedding (lexical-only accounts) are simply absent from it.
--
-- HNSW (not IVFFlat): per-account knowledge bases start empty and grow
-- incrementally, and IVFFlat must be trained on existing rows — built
-- against an empty/tiny table its centroids are meaningless and recall
-- is poor until it's large and REINDEXed. HNSW needs no training and is
-- accurate from the first row.
CREATE INDEX IF NOT EXISTS ai_knowledge_chunks_embedding_idx CREATE INDEX IF NOT EXISTS ai_knowledge_chunks_embedding_idx
ON ai_knowledge_chunks USING ivfflat (embedding vector_cosine_ops) ON ai_knowledge_chunks USING hnsw (embedding vector_cosine_ops);
WITH (lists = 100);
ALTER TABLE ai_knowledge_chunks ENABLE ROW LEVEL SECURITY; ALTER TABLE ai_knowledge_chunks ENABLE ROW LEVEL SECURITY;
@ -138,7 +150,8 @@ CREATE POLICY ai_knowledge_chunks_delete ON ai_knowledge_chunks FOR DELETE
-- ============================================================ -- ============================================================
-- Lexical: full-text rank. `plainto_tsquery` turns a raw customer -- Lexical: full-text rank. `plainto_tsquery` turns a raw customer
-- message into a query safely (no operator injection). -- message into a query safely (no operator injection). Uses the same
-- language-neutral `'simple'` config as the stored `fts` column.
CREATE OR REPLACE FUNCTION public.match_ai_knowledge_fts( CREATE OR REPLACE FUNCTION public.match_ai_knowledge_fts(
p_account_id uuid, p_account_id uuid,
p_query text, p_query text,
@ -147,28 +160,45 @@ CREATE OR REPLACE FUNCTION public.match_ai_knowledge_fts(
RETURNS TABLE (id uuid, content text, rank real) AS $$ RETURNS TABLE (id uuid, content text, rank real) AS $$
SELECT c.id, SELECT c.id,
c.content, c.content,
ts_rank(c.fts, plainto_tsquery('english', p_query)) AS rank ts_rank(c.fts, plainto_tsquery('simple', p_query)) AS rank
FROM ai_knowledge_chunks c FROM ai_knowledge_chunks c
WHERE c.account_id = p_account_id WHERE c.account_id = p_account_id
AND c.fts @@ plainto_tsquery('english', p_query) AND c.fts @@ plainto_tsquery('simple', p_query)
ORDER BY rank DESC ORDER BY rank DESC
LIMIT GREATEST(p_match_count, 0); LIMIT GREATEST(p_match_count, 0);
$$ LANGUAGE sql STABLE SECURITY DEFINER SET search_path = public; $$ LANGUAGE sql STABLE SECURITY DEFINER SET search_path = public;
-- Semantic: cosine distance against the query embedding. Only rows -- Semantic: cosine distance against the query embedding. Only rows
-- that actually have an embedding participate. -- that actually have an embedding participate.
--
-- `p_query_embedding` is declared `text` (not `vector`) and cast inside:
-- the caller sends the canonical pgvector literal `[0.1,0.2,...]` as a
-- plain string, so there's no ambiguity in how PostgREST binds a JSON
-- value to a `vector` parameter. Casting a literal to a constant vector
-- still lets the HNSW index serve the `<=>` order-by.
CREATE OR REPLACE FUNCTION public.match_ai_knowledge_semantic( CREATE OR REPLACE FUNCTION public.match_ai_knowledge_semantic(
p_account_id uuid, p_account_id uuid,
p_query_embedding vector(1536), p_query_embedding text,
p_match_count integer p_match_count integer
) )
RETURNS TABLE (id uuid, content text, distance real) AS $$ RETURNS TABLE (id uuid, content text, distance real) AS $$
SELECT c.id, SELECT c.id,
c.content, c.content,
(c.embedding <=> p_query_embedding) AS distance (c.embedding <=> p_query_embedding::vector(1536)) AS distance
FROM ai_knowledge_chunks c FROM ai_knowledge_chunks c
WHERE c.account_id = p_account_id WHERE c.account_id = p_account_id
AND c.embedding IS NOT NULL AND c.embedding IS NOT NULL
ORDER BY c.embedding <=> p_query_embedding ORDER BY c.embedding <=> p_query_embedding::vector(1536)
LIMIT GREATEST(p_match_count, 0); LIMIT GREATEST(p_match_count, 0);
$$ LANGUAGE sql STABLE SECURITY DEFINER SET search_path = public; $$ LANGUAGE sql STABLE SECURITY DEFINER SET search_path = public;
-- Lock down EXECUTE (mirrors migrations 018 / 025). These are
-- SECURITY DEFINER and would otherwise default to PUBLIC — i.e. the
-- anon role — which, since the function bypasses RLS and only gates on
-- the passed account_id, would let an unauthenticated caller read any
-- account's knowledge base. The draft path calls them as `authenticated`
-- and the auto-reply bot as `service_role`.
REVOKE ALL ON FUNCTION public.match_ai_knowledge_fts(uuid, text, integer) FROM PUBLIC;
GRANT EXECUTE ON FUNCTION public.match_ai_knowledge_fts(uuid, text, integer) TO authenticated, service_role;
REVOKE ALL ON FUNCTION public.match_ai_knowledge_semantic(uuid, text, integer) FROM PUBLIC;
GRANT EXECUTE ON FUNCTION public.match_ai_knowledge_semantic(uuid, text, integer) TO authenticated, service_role;