d3ro-voice/server/supabase/tests/knowledge-chunks-vector-search.integration.sql

\set ON_ERROR_STOP on

-- Regression for 20260929030000_knowledge_chunks_vector_index.sql.
-- match_knowledge_chunks must return the caller's exact top-k chunks (own
-- documents + documents of teams they belong to), even when other tenants own
-- almost every row of knowledge_chunks and the planner is pushed towards
-- index scans. Before the fix an IVFFlat index trained on an empty table was
-- scanned across all tenants with probes = 1 and filtered by tenant only
-- afterwards, so the caller usually got 0-1 rows although exact matches
-- existed.

BEGIN;

CREATE OR REPLACE FUNCTION pg_temp.assert_true(condition boolean, message text)
RETURNS void
LANGUAGE plpgsql
AS $$
BEGIN
  IF condition IS NOT TRUE THEN
    RAISE EXCEPTION 'assertion_failed: %', message;
  END IF;
END;
$$;

CREATE OR REPLACE FUNCTION pg_temp.act_as(uid uuid)
RETURNS void
LANGUAGE sql
AS $$
  SELECT set_config(
    'request.jwt.claims',
    json_build_object('sub', uid, 'role', 'authenticated')::text,
    true
  );
$$;

-- Uniform random direction in [-1, 1]^1536 (deterministic after setseed).
CREATE OR REPLACE FUNCTION pg_temp.rand_vec()
RETURNS public.vector
LANGUAGE sql
VOLATILE
AS $$
  SELECT array_agg(random() * 2 - 1 ORDER BY i)::public.vector
  FROM generate_series(1, 1536) AS g(i);
$$;

-- base plus a small random perturbation (still very similar to base).
CREATE OR REPLACE FUNCTION pg_temp.near_vec(base public.vector, eps double precision)
RETURNS public.vector
LANGUAGE sql
VOLATILE
AS $$
  SELECT array_agg(x + eps * (random() * 2 - 1) ORDER BY i)::public.vector
  FROM unnest(base::real[]) WITH ORDINALITY AS t(x, i);
$$;

-- Ids returned by match_knowledge_chunks, in result order.
CREATE OR REPLACE FUNCTION pg_temp.match_ids(q public.vector, k integer, threshold double precision)
RETURNS uuid[]
LANGUAGE sql
VOLATILE
AS $$
  SELECT coalesce(array_agg(m.id ORDER BY m.ord), ARRAY[]::uuid[])
  FROM public.match_knowledge_chunks(q, k, threshold)
    WITH ORDINALITY AS m(id, document_id, chunk_index, content, similarity, ord);
$$;

-- ── 0) the broken IVFFlat index is gone ──────────────────────────────────

SELECT pg_temp.assert_true(
  NOT EXISTS (
    SELECT 1 FROM pg_indexes
    WHERE schemaname = 'public'
      AND tablename = 'knowledge_chunks'
      AND indexdef ILIKE '%ivfflat%'
  ),
  'knowledge_chunks has no IVFFlat index trained on an empty table'
);

-- ── fixtures (as postgres) ───────────────────────────────────────────────

SELECT setseed(0.4242);

INSERT INTO auth.users (
  id, aud, role, email, encrypted_password, email_confirmed_at,
  raw_app_meta_data, raw_user_meta_data, created_at, updated_at
) VALUES
  (
    '39000000-0000-4000-8000-00000000000a', 'authenticated', 'authenticated',
    'rag-searcher@example.invalid', crypt('fixture-password', gen_salt('bf')), now(),
    '{"provider":"email","providers":["email"]}'::jsonb, '{"name":"Searcher"}'::jsonb, now(), now()
  ),
  (
    '39000000-0000-4000-8000-00000000000b', 'authenticated', 'authenticated',
    'rag-other-tenant@example.invalid', crypt('fixture-password', gen_salt('bf')), now(),
    '{"provider":"email","providers":["email"]}'::jsonb, '{"name":"Other"}'::jsonb, now(), now()
  ),
  (
    '39000000-0000-4000-8000-00000000000c', 'authenticated', 'authenticated',
    'rag-teammate@example.invalid', crypt('fixture-password', gen_salt('bf')), now(),
    '{"provider":"email","providers":["email"]}'::jsonb, '{"name":"Teammate"}'::jsonb, now(), now()
  );

-- Team T: searcher + teammate. Team T2: other tenant only.
INSERT INTO public.teams (id, name, owner_id) VALUES
  ('39100000-0000-4000-8000-00000000000a', 'RAG Team',   '39000000-0000-4000-8000-00000000000c'),
  ('39100000-0000-4000-8000-00000000000b', 'Other Team', '39000000-0000-4000-8000-00000000000b');
INSERT INTO public.team_members (team_id, user_id, role) VALUES
  ('39100000-0000-4000-8000-00000000000a', '39000000-0000-4000-8000-00000000000c', 'owner'),
  ('39100000-0000-4000-8000-00000000000a', '39000000-0000-4000-8000-00000000000a', 'member'),
  ('39100000-0000-4000-8000-00000000000b', '39000000-0000-4000-8000-00000000000b', 'owner');

INSERT INTO public.knowledge_documents (id, user_id, team_id, title) VALUES
  -- visible to the searcher
  ('39200000-0000-4000-8000-00000000000a', '39000000-0000-4000-8000-00000000000a', NULL,
   'Searcher personal doc'),
  ('39200000-0000-4000-8000-00000000000c', '39000000-0000-4000-8000-00000000000c',
   '39100000-0000-4000-8000-00000000000a', 'Teammate shared doc'),
  -- invisible to the searcher
  ('39200000-0000-4000-8000-00000000000b', '39000000-0000-4000-8000-00000000000b', NULL,
   'Other tenant personal doc'),
  ('39200000-0000-4000-8000-0000000000bb', '39000000-0000-4000-8000-00000000000b',
   '39100000-0000-4000-8000-00000000000b', 'Other tenant team doc');

CREATE TEMP TABLE fixture_query ON COMMIT DROP AS
SELECT pg_temp.rand_vec() AS q;
GRANT SELECT ON fixture_query TO authenticated;

-- Searcher: 60 random chunks + chunk 60 identical to the query.
INSERT INTO public.knowledge_chunks (document_id, chunk_index, content, embedding)
SELECT '39200000-0000-4000-8000-00000000000a', g, 'searcher chunk ' || g, pg_temp.rand_vec()
FROM generate_series(0, 59) AS g;
INSERT INTO public.knowledge_chunks (id, document_id, chunk_index, content, embedding)
SELECT '39300000-0000-4000-8000-00000000000a', '39200000-0000-4000-8000-00000000000a', 60,
       'searcher exact match', q
FROM fixture_query;

-- Team doc: 4 random chunks + one close to the query.
INSERT INTO public.knowledge_chunks (document_id, chunk_index, content, embedding)
SELECT '39200000-0000-4000-8000-00000000000c', g, 'team chunk ' || g, pg_temp.rand_vec()
FROM generate_series(0, 3) AS g;
INSERT INTO public.knowledge_chunks (id, document_id, chunk_index, content, embedding)
SELECT '39300000-0000-4000-8000-00000000000c', '39200000-0000-4000-8000-00000000000c', 4,
       'team near match', pg_temp.near_vec(q, 0.05)
FROM fixture_query;

-- Other tenant owns almost the whole table, including an exact match of its
-- own and near matches in a team the searcher is not in.
INSERT INTO public.knowledge_chunks (document_id, chunk_index, content, embedding)
SELECT '39200000-0000-4000-8000-00000000000b', g, 'other chunk ' || g, pg_temp.rand_vec()
FROM generate_series(0, 1999) AS g;
INSERT INTO public.knowledge_chunks (id, document_id, chunk_index, content, embedding)
SELECT '39300000-0000-4000-8000-00000000000b', '39200000-0000-4000-8000-00000000000b', 2000,
       'other exact match', q
FROM fixture_query;
INSERT INTO public.knowledge_chunks (document_id, chunk_index, content, embedding)
SELECT '39200000-0000-4000-8000-0000000000bb', g, 'other team near ' || g, pg_temp.near_vec(q, 0.02)
FROM fixture_query, generate_series(0, 9) AS g;
-- A chunk without an embedding must never be returned.
INSERT INTO public.knowledge_chunks (document_id, chunk_index, content)
VALUES ('39200000-0000-4000-8000-00000000000a', 61, 'searcher unembedded');

ANALYZE public.knowledge_chunks;
ANALYZE public.knowledge_documents;

-- Exact brute-force ranking over the searcher's visible chunks.
CREATE TEMP TABLE fixture_expected ON COMMIT DROP AS
SELECT
  (SELECT array_agg(r.id ORDER BY r.distance, r.id) FROM (
     SELECT kc.id, kc.embedding <=> f.q AS distance
     FROM public.knowledge_chunks kc, fixture_query f
     WHERE kc.embedding IS NOT NULL
       AND kc.document_id IN ('39200000-0000-4000-8000-00000000000a',
                              '39200000-0000-4000-8000-00000000000c')
     ORDER BY 2, 1
     LIMIT 5
   ) r) AS top5,
  (SELECT array_agg(r.id ORDER BY r.distance, r.id) FROM (
     SELECT kc.id, kc.embedding <=> f.q AS distance
     FROM public.knowledge_chunks kc, fixture_query f
     WHERE kc.embedding IS NOT NULL
       AND kc.document_id IN ('39200000-0000-4000-8000-00000000000a',
                              '39200000-0000-4000-8000-00000000000c')
     ORDER BY 2, 1
     LIMIT 20
   ) r) AS top20;
GRANT SELECT ON fixture_expected TO authenticated;

SELECT pg_temp.assert_true(
  (SELECT top5[1:2] FROM fixture_expected)
    = ARRAY['39300000-0000-4000-8000-00000000000a',
            '39300000-0000-4000-8000-00000000000c']::uuid[],
  'fixture sanity: exact match then team near match rank first'
);

-- ── 1) searcher gets the exact top-k of their visible chunks ─────────────

CREATE OR REPLACE FUNCTION pg_temp.assert_searcher_results(label text)
RETURNS void
LANGUAGE plpgsql
AS $$
DECLARE
  q public.vector := (SELECT f.q FROM fixture_query f);
  got uuid[];
BEGIN
  got := pg_temp.match_ids(q, 5, -1);
  PERFORM pg_temp.assert_true(
    got = (SELECT top5 FROM fixture_expected),
    format('%s: top-5 equals exact ranking (got %s)', label, got)
  );

  got := pg_temp.match_ids(q, 20, -1);
  PERFORM pg_temp.assert_true(
    got = (SELECT top20 FROM fixture_expected),
    format('%s: top-20 equals exact ranking (got %s rows)', label, cardinality(got))
  );

  -- default threshold (0.5) as used by search-knowledge: only the two
  -- planted matches qualify; the other tenant's exact/near matches never leak.
  got := pg_temp.match_ids(q, 20, 0.5);
  PERFORM pg_temp.assert_true(
    got = ARRAY['39300000-0000-4000-8000-00000000000a',
                '39300000-0000-4000-8000-00000000000c']::uuid[],
    format('%s: threshold keeps only the caller''s planted matches (got %s)', label, got)
  );

  PERFORM pg_temp.assert_true(
    NOT EXISTS (
      SELECT 1 FROM public.match_knowledge_chunks(q, 50, -1) m
      WHERE m.document_id NOT IN ('39200000-0000-4000-8000-00000000000a',
                                  '39200000-0000-4000-8000-00000000000c')
         OR m.content = 'searcher unembedded'
    ),
    format('%s: only visible, embedded chunks are returned', label)
  );

  -- match_count is clamped: 66 visible embedded chunks, cap 50.
  PERFORM pg_temp.assert_true(
    cardinality(pg_temp.match_ids(q, 1000, -1)) = 50,
    format('%s: match_count is capped at 50', label)
  );
  PERFORM pg_temp.assert_true(
    cardinality(pg_temp.match_ids(q, NULL, -1)) = 5,
    format('%s: NULL match_count falls back to 5 instead of unlimited', label)
  );
  PERFORM pg_temp.assert_true(
    cardinality(pg_temp.match_ids(q, 0, -1)) = 0
      AND cardinality(pg_temp.match_ids(q, -3, -1)) = 0,
    format('%s: non-positive match_count returns nothing', label)
  );
END;
$$;

SET LOCAL ROLE authenticated;
SELECT pg_temp.act_as('39000000-0000-4000-8000-00000000000a');
-- Push the planner towards index-ordered scans, the worst case for recall:
-- with seq scans and explicit sorts penalised, an ORDER BY distance LIMIT k
-- query is planned as an ANN index scan whenever one is available (what
-- happens in production once knowledge_chunks is large).
SET LOCAL enable_seqscan = off;
SET LOCAL enable_sort = off;

SELECT pg_temp.assert_searcher_results('no ANN index');

-- ── 2) still exact when an ANN index exists and is attractive ────────────
-- A properly trained IVFFlat index scanned with probes = 1 visits ~2% of all
-- rows across tenants. The function must not let the planner use it for the
-- ranking (the pre-fix shape did).

RESET ROLE;
CREATE INDEX knowledge_chunks_embedding_probe_test
  ON public.knowledge_chunks
  USING ivfflat (embedding public.vector_cosine_ops)
  WITH (lists = 50);
SET LOCAL ivfflat.probes = 1;

SET LOCAL ROLE authenticated;
SELECT pg_temp.act_as('39000000-0000-4000-8000-00000000000a');

SELECT pg_temp.assert_searcher_results('trained IVFFlat index present');

-- ── 3) the other tenant sees only their own chunks ───────────────────────

SELECT pg_temp.act_as('39000000-0000-4000-8000-00000000000b');

SELECT pg_temp.assert_true(
  (SELECT pg_temp.match_ids(f.q, 1, 0.5) FROM fixture_query f)
    = ARRAY['39300000-0000-4000-8000-00000000000b']::uuid[],
  'other tenant gets their own exact match first'
);
SELECT pg_temp.assert_true(
  NOT EXISTS (
    SELECT 1 FROM fixture_query f, public.match_knowledge_chunks(f.q, 50, -1) m
    WHERE m.document_id NOT IN ('39200000-0000-4000-8000-00000000000b',
                                '39200000-0000-4000-8000-0000000000bb')
  ),
  'other tenant never sees the searcher''s or team T''s chunks'
);

-- ── 4) anonymous callers get nothing ─────────────────────────────────────

SELECT set_config('request.jwt.claims', '', true);
SELECT pg_temp.assert_true(
  (SELECT cardinality(pg_temp.match_ids(f.q, 50, -1)) FROM fixture_query f) = 0,
  'caller without a user id sees no chunks'
);

ROLLBACK;