diff --git a/Makefile b/Makefile index a588c50..1da4651 100644 --- a/Makefile +++ b/Makefile @@ -34,7 +34,7 @@ DATA = sql/$(EXTENSION)--$(EXTVERSION).sql \ sql/$(EXTENSION)--1.0-beta3--1.0.sql # Test configuration for pg_regress -REGRESS = setup chunking multibyte_chunking hybrid_chunking queue delete_truncate delete_truncate_pk pk_type_session max_retries vectorization multi_column maintenance edge_cases providers worker cleanup embedding pk_types stale_embeddings hybrid_test count_tokens vectorizer_status per_table_model +REGRESS = setup chunking multibyte_chunking hybrid_chunking queue delete_truncate delete_truncate_pk pk_type_session max_retries vectorization multi_column maintenance edge_cases providers worker cleanup embedding pk_types stale_embeddings hybrid_test count_tokens vectorizer_status per_table_model model_drift REGRESS_OPTS = --inputdir=test --outputdir=test # Documentation files (if any) diff --git a/docs/api_reference.md b/docs/api_reference.md index bf44d0d..6e6fe7a 100644 --- a/docs/api_reference.md +++ b/docs/api_reference.md @@ -208,6 +208,88 @@ SELECT pgedge_vectorizer.set_embedding_model( ); ``` +### embedding_model_status() + +Report which provider and model each vectorizer's chunks were actually embedded +with, and where that disagrees with what it would use now. + +```sql +SELECT * FROM pgedge_vectorizer.embedding_model_status( + source_table REGCLASS DEFAULT NULL, + source_column NAME DEFAULT NULL +); +``` + +**Parameters:** both optional, narrowing the result to one source table or one +column of it. With neither, every registered vectorizer is reported. + +Columns: + +- `source_table`, `source_column`, `chunk_table`: The vectorizer, as registered +- `effective_provider`, `effective_model`: What it would use now, inheritance + resolved +- `chunks_embedded`: Chunks with a vector. Every count below is a subset of + this one; a chunk with no vector has no model to disagree about and is + excluded throughout +- `chunks_current`: Embedded by the effective provider and model +- `chunks_other_model`: Embedded by something else. Vectors from two models are + not comparable, so these rows are effectively invisible to search +- `chunks_model_unknown`: Embedded before the extension recorded this, which is + every row on an installation that has just upgraded. Reported apart from a + mismatch because they may well be current +- `embedded_models`: The distinct `provider/model` pairs actually present, + ordered + +Each row scans a chunk table, so this costs considerably more than the queue +views. A chunk table that has been dropped, or that the caller cannot read, +gives NULL counts rather than failing the whole result set. + +### reembed() + +Re-embed a vectorizer's chunks with the provider and model it would use now. + +```sql +SELECT pgedge_vectorizer.reembed( + source_table REGCLASS, + source_column NAME, + embedding_dimension INT DEFAULT NULL +); +``` + +**Parameters:** + +- `source_table`, `source_column`: The vectorizer to repair +- `embedding_dimension`: Dimension of the effective model. When NULL (the + default) the provider is probed for it, which is a real request + +Returns: `BIGINT` - The number of chunks queued + +Clears and requeues every chunk not known to have been produced by the +effective provider and model, which includes chunks with nothing recorded: +a row that cannot be shown to be current is treated as needing doing again, so +the first call on a freshly upgraded installation re-embeds the whole table. +Chunks already current are left alone. + +If the effective model is a different width from the chunk table's vector +column, that distinction cannot hold: the column is altered and every chunk is +requeued, since a column cannot carry two widths. A notice says so. + +Chunk rows, token counts, sparse embeddings and the BM25 statistics are +untouched either way, because none of them depends on the embedding model. + +The vectorizer's pending queue rows are replaced, since anything queued was +queued before this decision was made, except rows queued for sparse work +alone: those carry no embedding to redo and are left where they are. + +Unlike `set_embedding_model()`, there is no confirmation flag: this function +does what its name says. It does spend money against a metered provider, and +raises a notice with the count for that reason. + +This is the supported repair for a vectorizer that drifted because +`pgedge_vectorizer.model` changed under it. `set_embedding_model()` will not do +it, because an inheriting vectorizer's effective model already is the new one, +so from that function's point of view nothing has changed. + ### retry_failed() Retry failed queue items. diff --git a/docs/best_practices.md b/docs/best_practices.md index 14c4c2b..4e2a68f 100644 --- a/docs/best_practices.md +++ b/docs/best_practices.md @@ -34,6 +34,13 @@ counts survive untouched. smaller size from providers that support shortening. - Pin the model rather than inheriting it wherever the embeddings matter, since an inheriting vectorizer follows `pgedge_vectorizer.model` with no guard. +- Treat a change to `pgedge_vectorizer.model` as a data migration rather than a + configuration change, because for every inheriting vectorizer that is what it + is. The sequence is: change the setting, run `embedding_model_status()` to see + which tables are now a mixture, and `reembed()` each of them, having budgeted + for embedding all of it again. Until that is done those tables hold vectors + from two models, and similarity between them is noise, so the affected rows + are effectively invisible to search rather than merely out of date. **Performance** diff --git a/docs/changelog.md b/docs/changelog.md index fb6f3bc..cb7bb7e 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -53,6 +53,17 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). - `generate_embedding()` and `detect_embedding_dimension()` accept an optional provider and model, so a query can be embedded with the same model as the chunks it will be compared against. +- Chunk tables now record the provider and model that produced each embedding, + and `pgedge_vectorizer.embedding_model_status()` reports where that disagrees + with what the vectorizer would use now + ([#75](https://github.com/pgEdge/pgedge-vectorizer/issues/75)). A vectorizer + that inherits follows `pgedge_vectorizer.model` as it changes, so a chunk + table can end up holding vectors from two models with nothing reporting it; + where the widths match the existing dimension check cannot see it either. + `pgedge_vectorizer.reembed()` repairs it, redoing the chunks that are not + known to be current. Rows embedded before this release have nothing recorded + and are counted separately, but `reembed()` treats them as needing redoing, + so the first call on an upgraded installation re-embeds the whole table. ### Changed diff --git a/docs/configuration.md b/docs/configuration.md index d3ecdeb..be96c8e 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -97,6 +97,57 @@ is not limited to changes of dimension. setting existed. Pin the model on any vectorizer whose embeddings matter. + Where it has already happened, it is at least visible and repairable: + `embedding_model_status()` reports what each chunk table is a mixture + of, and `reembed()` redoes the chunks that are not current. See + [Seeing what a table was embedded with](#seeing-what-a-table-was-embedded-with). + +### Seeing what a table was embedded with + +Every chunk records the provider and model that produced its vector, so a +disagreement between that and what the vectorizer would use now is visible +rather than something you discover through poor search results: + +```sql +SELECT * FROM pgedge_vectorizer.embedding_model_status('articles'::regclass); +``` + +``` +source_table | articles +source_column | body +effective_provider | openai +effective_model | text-embedding-3-large +chunks_embedded | 12043 +chunks_current | 9945 +chunks_other_model | 2098 +chunks_model_unknown | 0 +embedded_models | {openai/text-embedding-3-large,openai/text-embedding-3-small} +``` + +`chunks_other_model` is the count that matters: those rows hold vectors from a +different model, and similarity between two models' vectors is meaningless, so +they are effectively invisible to search rather than merely stale. +`chunks_model_unknown` counts rows embedded before this was recorded, which is +every row on an installation that has just upgraded; they may well be current, +so they are reported separately rather than assumed wrong. + +To repair it: + +```sql +SELECT pgedge_vectorizer.reembed('articles'::regclass, 'body'); +``` + +That clears and requeues everything not known to have come from the provider +and model the vectorizer would use now, which includes the unknown rows, since +a row that cannot be shown to be current is one that needs doing again. Rows +that are already current are left alone, unless the new model is a different +width, in which case the column has to be altered and every chunk goes with it. +Chunks, token counts, sparse embeddings and the BM25 statistics are untouched +throughout, because none of them depends on the embedding model. + +Both functions scan the chunk table, so give them a source table rather than +running them across every vectorizer out of habit. + ## Worker Settings These settings control the background workers that process the embedding queue, including concurrency, batch sizes, and retry behavior. diff --git a/sql/pgedge_vectorizer--1.1--1.2.sql b/sql/pgedge_vectorizer--1.1--1.2.sql index fd0d752..ab3b364 100644 --- a/sql/pgedge_vectorizer--1.1--1.2.sql +++ b/sql/pgedge_vectorizer--1.1--1.2.sql @@ -22,6 +22,62 @@ COMMENT ON COLUMN pgedge_vectorizer.vectorizers.provider IS COMMENT ON COLUMN pgedge_vectorizer.vectorizers.model IS 'Embedding model for this vectorizer; NULL inherits pgedge_vectorizer.model'; +--------------------------------------------------------------------------- +-- Provenance columns on the chunk tables that already exist +-- +-- enable_vectorization() adds these to a chunk table it finds without them, +-- but nothing re-runs it on upgrade, and the worker writes both columns in +-- the same statement as the embedding: an existing installation would fail +-- every embedding write until someone happened to re-enable the vectorizer. +-- So the upgrade alters what the registry knows about, here and now. +--------------------------------------------------------------------------- + +DO $$ +DECLARE + v RECORD; + chunk_oid OID; + n_found INT; +BEGIN + FOR v IN SELECT r.chunk_table FROM pgedge_vectorizer.vectorizers r LOOP + /* + * The registry stores the chunk table's name as one identifier, with + * a dot in it for a schema-qualified source, and records nothing + * about the schema it was created in. to_regclass() would resolve it + * through whatever search_path the upgrade happens to run with, and + * return NULL without complaint for a table outside it, leaving the + * columns unadded and every later embedding write failing. So look + * through the catalogue instead of the search_path. + */ + SELECT count(*), min(c.oid) INTO n_found, chunk_oid + FROM pg_class c + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE c.relname = v.chunk_table + AND c.relkind = 'r' + AND n.nspname NOT IN ('pg_catalog', 'information_schema'); + + -- A chunk table dropped from under the registry is not this script's + -- problem to fix, and must not stop the upgrade. + CONTINUE WHEN n_found = 0; + + /* + * Two schemas holding a relation of that name leaves no way to tell + * which one the registry means, and altering the wrong table would be + * worse than altering neither. + */ + IF n_found > 1 THEN + RAISE WARNING 'chunk table % exists in more than one schema; add embedding_provider and embedding_model to it by hand', v.chunk_table; + CONTINUE; + END IF; + + EXECUTE format( + 'ALTER TABLE %s + ADD COLUMN IF NOT EXISTS embedding_provider TEXT, + ADD COLUMN IF NOT EXISTS embedding_model TEXT', + chunk_oid::REGCLASS); + END LOOP; +END; +$$; + --------------------------------------------------------------------------- -- Approximate token counter, shared with the C chunking code -- @@ -170,18 +226,38 @@ BEGIN token_count INT, embedding vector(%s), sparse_embedding sparsevec(65536), + embedding_provider TEXT, + embedding_model TEXT, created_at TIMESTAMPTZ DEFAULT NOW(), updated_at TIMESTAMPTZ DEFAULT NOW(), UNIQUE(source_id, chunk_index) )', chunk_table, pk_col_type, embedding_dimension); - -- Add sparse columns to pre-existing chunk tables (upgrade path). - -- These are no-ops for freshly created tables (columns exist already). + -- Add sparse and provenance columns to pre-existing chunk tables (upgrade + -- path). These are no-ops for freshly created tables (columns exist + -- already). EXECUTE format(' ALTER TABLE %I ADD COLUMN IF NOT EXISTS sparse_embedding sparsevec(65536)', chunk_table); + -- Guarded rather than ADD COLUMN IF NOT EXISTS, which would print two + -- "already exists, skipping" notices on every call for a table that was + -- just created with them, which is every call but an upgrade. + IF NOT EXISTS ( + SELECT 1 + FROM pg_attribute a + WHERE a.attrelid = to_regclass(quote_ident(chunk_table)) + AND a.attname = 'embedding_model' + AND NOT a.attisdropped + ) THEN + EXECUTE format(' + ALTER TABLE %I + ADD COLUMN embedding_provider TEXT, + ADD COLUMN embedding_model TEXT', + chunk_table); + END IF; + -- Create vector index for similarity search EXECUTE format(' CREATE INDEX IF NOT EXISTS %I ON %I @@ -1100,13 +1176,13 @@ BEGIN -- Compare effective values, not stored ones: moving a table from an -- explicit 'openai' to NULL whilst the GUC also says 'openai' changes -- nothing, and must not cost a re-embed. - old_provider := COALESCE(v_row.provider, + old_provider := COALESCE(NULLIF(v_row.provider, ''), current_setting('pgedge_vectorizer.provider')); - old_model := COALESCE(v_row.model, + old_model := COALESCE(NULLIF(v_row.model, ''), current_setting('pgedge_vectorizer.model')); - new_provider := COALESCE(set_embedding_model.provider, + new_provider := COALESCE(NULLIF(set_embedding_model.provider, ''), current_setting('pgedge_vectorizer.provider')); - new_model := COALESCE(set_embedding_model.model, + new_model := COALESCE(NULLIF(set_embedding_model.model, ''), current_setting('pgedge_vectorizer.model')); IF old_provider = new_provider AND old_model = new_model THEN @@ -1152,10 +1228,12 @@ BEGIN * and then fail every embedding the worker tried to write, which is the * failure this function exists to prevent. */ + -- The probe asks about the effective values rather than the raw + -- arguments: an empty string means inherit everywhere else, and would + -- otherwise reach the provider as a model name of ''. new_dim := COALESCE( set_embedding_model.embedding_dimension, - pgedge_vectorizer.detect_embedding_dimension( - set_embedding_model.provider, set_embedding_model.model)); + pgedge_vectorizer.detect_embedding_dimension(new_provider, new_model)); SELECT a.atttypmod INTO current_dim FROM pg_attribute a @@ -1164,7 +1242,9 @@ BEGIN IF chunk_count > 0 THEN -- NULL first: a vector column cannot change width with values in it. - EXECUTE format('UPDATE %s SET embedding = NULL ' + -- A cleared embedding has no model, so its provenance goes with it. + EXECUTE format('UPDATE %s SET embedding = NULL, ' + 'embedding_provider = NULL, embedding_model = NULL ' 'WHERE embedding IS NOT NULL', chunk_oid::REGCLASS); -- Anything already queued was queued against the old model. @@ -1209,3 +1289,429 @@ COMMENT ON FUNCTION pgedge_vectorizer.set_embedding_model IS 'force_reembed is true, in which case every embedding is cleared and every ' 'chunk requeued. Returns the number of chunks requeued'; +--------------------------------------------------------------------------- +-- embedding_model_status(): what each vectorizer's chunks were embedded with +-- +-- A vectorizer with NULL provider and model inherits the GUCs, and inheritance +-- resolves when the work runs rather than being copied at creation. Changing +-- pgedge_vectorizer.model therefore re-points every inheriting vectorizer at +-- once, leaving a chunk table holding vectors from the old model beside new +-- ones from the new. Similarity between two models' vectors is noise, so +-- search degrades quietly; where the widths match, as they do between +-- text-embedding-3-small and text-embedding-ada-002, nothing catches it at all. +-- +-- set_embedding_model() guards the per-vectorizer path and cannot guard this +-- one: the extension does not own that GUC and cannot intercept every way it +-- changes. So the chunk table records what produced each vector, and this +-- reports where that disagrees with what the vectorizer would use now. It +-- diagnoses rather than prevents, but it catches drift from any cause, +-- including a setting changed months ago by someone since departed. +-- +-- Each row costs a scan of one chunk table, so the arguments narrow it. +--------------------------------------------------------------------------- + +CREATE FUNCTION pgedge_vectorizer.embedding_model_status( + p_source_table REGCLASS DEFAULT NULL, + p_source_column NAME DEFAULT NULL +) RETURNS TABLE ( + source_table TEXT, + source_column NAME, + chunk_table TEXT, + effective_provider TEXT, + effective_model TEXT, + chunks_embedded BIGINT, + chunks_current BIGINT, + chunks_other_model BIGINT, + chunks_model_unknown BIGINT, + embedded_models TEXT[] +) AS $$ +DECLARE + v RECORD; + chunk_oid OID; +BEGIN + FOR v IN + SELECT r.source_table, r.source_column, r.chunk_table, + COALESCE(NULLIF(r.provider, ''), + current_setting('pgedge_vectorizer.provider')) + AS eff_provider, + COALESCE(NULLIF(r.model, ''), + current_setting('pgedge_vectorizer.model')) + AS eff_model + FROM pgedge_vectorizer.vectorizers r + WHERE (p_source_table IS NULL + OR r.source_table = p_source_table::TEXT) + AND (p_source_column IS NULL OR r.source_column = p_source_column) + ORDER BY r.source_table, r.source_column + LOOP + source_table := v.source_table; + source_column := v.source_column; + chunk_table := v.chunk_table; + effective_provider := v.eff_provider; + effective_model := v.eff_model; + + chunks_embedded := NULL; + chunks_current := NULL; + chunks_other_model := NULL; + chunks_model_unknown := NULL; + embedded_models := NULL; + + /* + * A chunk table that has been dropped, or that the caller cannot + * read, leaves this row's counts NULL rather than failing the whole + * result set. The name is one identifier with a dot in it for a + * schema-qualified source, so it is quoted rather than parsed. + */ + chunk_oid := to_regclass(quote_ident(v.chunk_table)); + IF chunk_oid IS NOT NULL + AND has_table_privilege(chunk_oid, 'SELECT') THEN + /* + * Every count is over embedded rows only: a chunk with no vector + * has no model to disagree about, and counting it as drifted + * would confuse work still to do with work done wrongly. + * + * A row with no recorded model is reported apart from a mismatch + * rather than lumped in with it. It predates these columns, so it + * may well be current; the report says what is there, and leaves + * the pessimistic reading to reembed(), which has to act. + */ + EXECUTE format( + 'SELECT count(*) FILTER (WHERE embedding IS NOT NULL), + count(*) FILTER (WHERE embedding IS NOT NULL + AND embedding_provider = %L + AND embedding_model = %L), + count(*) FILTER (WHERE embedding IS NOT NULL + AND embedding_model IS NOT NULL + AND (embedding_provider + IS DISTINCT FROM %L + OR embedding_model + IS DISTINCT FROM %L)), + count(*) FILTER (WHERE embedding IS NOT NULL + AND embedding_model IS NULL), + (SELECT array_agg(pair ORDER BY pair) + FROM (SELECT DISTINCT + embedding_provider || ''/'' || + embedding_model AS pair + FROM %s + WHERE embedding IS NOT NULL + AND embedding_model IS NOT NULL) d) + FROM %s', + v.eff_provider, v.eff_model, v.eff_provider, v.eff_model, + chunk_oid::REGCLASS, chunk_oid::REGCLASS) + INTO chunks_embedded, chunks_current, chunks_other_model, + chunks_model_unknown, embedded_models; + END IF; + + RETURN NEXT; + END LOOP; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.embedding_model_status IS +'Report, per vectorizer, how many embedded chunks were produced by the ' +'provider and model it would use now, how many by something else, and how ' +'many predate the columns that record it. Scans the chunk tables'; + +--------------------------------------------------------------------------- +-- reembed(): redo the embeddings that are not known to be current +-- +-- The report above leaves a user with a number and nothing to do about it. +-- set_embedding_model(..., force_reembed => true) is not the answer, because +-- an inheriting vectorizer's effective model already is the new one, so that +-- function sees no change and takes its no-op branch. +-- +-- No confirmation flag. Unlike set_embedding_model(), whose re-embed is a +-- surprising consequence of a settings change, this function does what its +-- name says; it raises a notice with the count, because the cost lands on a +-- metered provider. +--------------------------------------------------------------------------- + +CREATE FUNCTION pgedge_vectorizer.reembed( + source_table REGCLASS, + source_column NAME, + embedding_dimension INT DEFAULT NULL +) RETURNS BIGINT AS $$ +DECLARE + v_row RECORD; + chunk_oid OID; + eff_provider TEXT; + eff_model TEXT; + new_dim INT; + current_dim INT; + width_change BOOLEAN; + requeued BIGINT := 0; +BEGIN + SELECT r.* INTO v_row + FROM pgedge_vectorizer.vectorizers r + WHERE r.source_table = reembed.source_table::TEXT + AND r.source_column = reembed.source_column; + + IF NOT FOUND THEN + RAISE EXCEPTION 'no vectorizer registered for %.%', + reembed.source_table::TEXT, reembed.source_column; + END IF; + + eff_provider := COALESCE(NULLIF(v_row.provider, ''), + current_setting('pgedge_vectorizer.provider')); + eff_model := COALESCE(NULLIF(v_row.model, ''), + current_setting('pgedge_vectorizer.model')); + + chunk_oid := to_regclass(quote_ident(v_row.chunk_table)); + IF chunk_oid IS NULL THEN + RAISE EXCEPTION 'chunk table % for %.% no longer exists', + v_row.chunk_table, + reembed.source_table::TEXT, reembed.source_column; + END IF; + + new_dim := COALESCE( + reembed.embedding_dimension, + pgedge_vectorizer.detect_embedding_dimension(eff_provider, eff_model)); + + SELECT a.atttypmod INTO current_dim + FROM pg_attribute a + WHERE a.attrelid = chunk_oid + AND a.attname = 'embedding'; + + width_change := new_dim IS DISTINCT FROM current_dim; + + IF width_change THEN + /* + * A column cannot hold two widths, so a change of width takes every + * row with it whether or not it had drifted. Clearing has to come + * first: a vector column cannot be altered with values in it. + */ + EXECUTE format('UPDATE %s SET embedding = NULL, ' + 'embedding_provider = NULL, embedding_model = NULL ' + 'WHERE embedding IS NOT NULL', chunk_oid::REGCLASS); + + EXECUTE format('ALTER TABLE %s ALTER COLUMN embedding ' + 'TYPE vector(%s)', chunk_oid::REGCLASS, new_dim); + + RAISE NOTICE 'Embedding dimension changed from % to %, so every chunk ' + 'is being re-embedded', current_dim, new_dim; + ELSE + /* + * Same width, so rows already produced by this provider and model are + * left exactly as they are. Everything else goes, including rows with + * nothing recorded: those predate the columns and cannot be shown to + * be current, and the safe reading of a row that cannot be proved + * current is that it needs doing again. On an installation freshly + * upgraded to 1.2 that is every row, which the documentation says. + */ + EXECUTE format( + 'UPDATE %s SET embedding = NULL, ' + ' embedding_provider = NULL, embedding_model = NULL ' + ' WHERE embedding IS NOT NULL ' + ' AND (embedding_model IS NULL ' + ' OR embedding_provider IS DISTINCT FROM %L ' + ' OR embedding_model IS DISTINCT FROM %L)', + chunk_oid::REGCLASS, eff_provider, eff_model); + END IF; + + /* + * Anything already queued for embedding was queued before this decision + * was made. Sparse-only rows are left alone: they carry no dense work to + * redo, the probe only marks a chunk sparse-only whilst it still has an + * embedding, and deleting them would strand the sparse vector as NULL + * until someone thought to run reprocess_chunks(). + */ + DELETE FROM pgedge_vectorizer.queue q + WHERE q.chunk_table = v_row.chunk_table + AND NOT COALESCE((q.metadata->>'sparse_only')::BOOLEAN, FALSE); + + /* + * Whatever now has no embedding needs one, which after the clearing above + * is exactly the set chosen, plus any chunk that was never embedded in + * the first place and would have been picked up by reprocess_chunks() + * anyway. + */ + EXECUTE format( + 'INSERT INTO pgedge_vectorizer.queue ' + ' (chunk_id, chunk_table, content, max_attempts) ' + 'SELECT id, %L, content, %s FROM %s WHERE embedding IS NULL', + v_row.chunk_table, + current_setting('pgedge_vectorizer.max_retries')::INT, + chunk_oid::REGCLASS); + + GET DIAGNOSTICS requeued = ROW_COUNT; + + RAISE NOTICE 'Queued % chunks to be embedded with %/%', + requeued, eff_provider, eff_model; + + RETURN requeued; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.reembed IS +'Re-embed a vectorizer''s chunks with the provider and model it would use now, ' +'leaving alone any already produced by them. A change of embedding dimension ' +'takes every chunk with it. Returns the number queued'; + +--------------------------------------------------------------------------- +-- hybrid_search(): embed the query with the vectorizer's own model +-- +-- The query vector was generated from the GUCs, which was right whilst that +-- was the only place a model could come from. Now that a vectorizer can pin +-- its own, a query embedded by one model would be compared against chunks +-- embedded by another: meaningless distances where the widths match, and an +-- outright error where they do not. Not redefined by the 1.1 script, so it is +-- replaced here in full. +--------------------------------------------------------------------------- + +CREATE OR REPLACE FUNCTION pgedge_vectorizer.hybrid_search( + p_source_table REGCLASS, + p_query TEXT, + p_limit INT DEFAULT 10, + p_alpha FLOAT8 DEFAULT 0.7, + p_rrf_k INT DEFAULT 60, + p_source_column NAME DEFAULT NULL +) +RETURNS TABLE ( + source_id TEXT, + chunk TEXT, + dense_rank INT, + sparse_rank INT, + rrf_score FLOAT8 +) +LANGUAGE plpgsql AS $$ +DECLARE + v_chunk_table TEXT; + v_provider TEXT; + v_model TEXT; + v_query_dense vector; + v_query_sparse sparsevec; +BEGIN + IF COALESCE(current_setting('pgedge_vectorizer.enable_hybrid', true), 'false')::boolean IS NOT TRUE THEN + RAISE EXCEPTION + 'Hybrid search is disabled. Set pgedge_vectorizer.enable_hybrid = true and allow workers to populate sparse_embedding.'; + END IF; + + -- Look up the chunk table from the vectorizers registry. + -- When p_source_column is provided, use the exact mapping. + -- When NULL, raise an exception if the table has more than one + -- vectorized column to avoid silently returning results from the + -- wrong chunk table. + IF p_source_column IS NOT NULL THEN + SELECT vz.chunk_table, vz.provider, vz.model + INTO v_chunk_table, v_provider, v_model + FROM pgedge_vectorizer.vectorizers vz + WHERE vz.source_table = p_source_table::TEXT + AND vz.source_column = p_source_column; + ELSE + SELECT vz.chunk_table, vz.provider, vz.model + INTO v_chunk_table, v_provider, v_model + FROM pgedge_vectorizer.vectorizers vz + WHERE vz.source_table = p_source_table::TEXT + LIMIT 1; + + IF v_chunk_table IS NOT NULL AND + (SELECT count(*) FROM pgedge_vectorizer.vectorizers + WHERE source_table = p_source_table::TEXT) > 1 + THEN + RAISE EXCEPTION + 'Table % has multiple vectorized columns. ' + 'Pass p_source_column to disambiguate.', + p_source_table; + END IF; + END IF; + + IF v_chunk_table IS NULL THEN + RAISE EXCEPTION + 'No vectorizer found for table %. ' + 'Call pgedge_vectorizer.enable_vectorization() first.', + p_source_table; + END IF; + + /* + * Embed the query with this vectorizer's own provider and model rather + * than the GUCs. A query embedded by one model and compared against chunks + * embedded by another gives meaningless distances, and where the widths + * differ it fails outright. NULL passes straight through and means + * inherit, so a vectorizer that has pinned nothing behaves as before. + */ + v_query_dense := pgedge_vectorizer.generate_embedding(p_query, + v_provider, v_model); + + -- Generate sparse BM25 query vector + v_query_sparse := pgedge_vectorizer.bm25_query_vector( + p_query, v_chunk_table); + + -- Run both ranked lists and merge with Reciprocal Rank Fusion. + -- Join on chunk id (not source_id) to avoid mixing unrelated chunks + -- from the same document. source_id is cast to TEXT to support + -- arbitrary PK types (BIGINT, UUID, VARCHAR, etc.). + RETURN QUERY EXECUTE format($sql$ + WITH dense_candidates AS ( + SELECT + id, + source_id::text AS source_id, + content AS chunk, + embedding <=> %L::vector AS dist + FROM %I + WHERE embedding IS NOT NULL + ORDER BY dist + LIMIT %s * 3 + ), + dense AS ( + SELECT + id, + source_id, + chunk, + ROW_NUMBER() OVER (ORDER BY dist) AS rnk + FROM dense_candidates + ), + sparse_candidates AS ( + SELECT + id, + source_id::text AS source_id, + content AS chunk, + sparse_embedding <#> %L::sparsevec AS dist + FROM %I + WHERE sparse_embedding IS NOT NULL + ORDER BY dist ASC + LIMIT %s * 3 + ), + sparse AS ( + SELECT + id, + source_id, + chunk, + ROW_NUMBER() OVER (ORDER BY dist ASC) AS rnk + FROM sparse_candidates + ), + merged AS ( + SELECT + COALESCE(d.source_id, s.source_id) AS source_id, + COALESCE(d.chunk, s.chunk) AS chunk, + COALESCE(d.rnk, 9999)::INT AS dense_rank, + COALESCE(s.rnk, 9999)::INT AS sparse_rank, + ( + %s::float8 / (%s + COALESCE(d.rnk, 9999)) + + (1.0 - %s::float8) / (%s + COALESCE(s.rnk, 9999)) + ) AS rrf_score + FROM dense d + FULL OUTER JOIN sparse s USING (id) + ) + SELECT + source_id, + chunk, + dense_rank, + sparse_rank, + rrf_score + FROM merged + ORDER BY rrf_score DESC + LIMIT %s + $sql$, + v_query_dense, v_chunk_table, p_limit, + v_query_sparse, v_chunk_table, p_limit, + p_alpha, p_rrf_k, + p_alpha, p_rrf_k, + p_limit + ); +END; +$$; + +COMMENT ON FUNCTION pgedge_vectorizer.hybrid_search IS +'Hybrid BM25 + dense vector search using Reciprocal Rank Fusion. + p_alpha controls the weight of dense results (0 = pure sparse, 1 = pure dense). + p_rrf_k is the RRF rank smoothing constant (default 60). + Requires pgedge_vectorizer.enable_hybrid = true in postgresql.conf.'; diff --git a/sql/pgedge_vectorizer--1.2.sql b/sql/pgedge_vectorizer--1.2.sql index 81551ac..d4af8ab 100644 --- a/sql/pgedge_vectorizer--1.2.sql +++ b/sql/pgedge_vectorizer--1.2.sql @@ -315,18 +315,38 @@ BEGIN token_count INT, embedding vector(%s), sparse_embedding sparsevec(65536), + embedding_provider TEXT, + embedding_model TEXT, created_at TIMESTAMPTZ DEFAULT NOW(), updated_at TIMESTAMPTZ DEFAULT NOW(), UNIQUE(source_id, chunk_index) )', chunk_table, pk_col_type, embedding_dimension); - -- Add sparse columns to pre-existing chunk tables (upgrade path). - -- These are no-ops for freshly created tables (columns exist already). + -- Add sparse and provenance columns to pre-existing chunk tables (upgrade + -- path). These are no-ops for freshly created tables (columns exist + -- already). EXECUTE format(' ALTER TABLE %I ADD COLUMN IF NOT EXISTS sparse_embedding sparsevec(65536)', chunk_table); + -- Guarded rather than ADD COLUMN IF NOT EXISTS, which would print two + -- "already exists, skipping" notices on every call for a table that was + -- just created with them, which is every call but an upgrade. + IF NOT EXISTS ( + SELECT 1 + FROM pg_attribute a + WHERE a.attrelid = to_regclass(quote_ident(chunk_table)) + AND a.attname = 'embedding_model' + AND NOT a.attisdropped + ) THEN + EXECUTE format(' + ALTER TABLE %I + ADD COLUMN embedding_provider TEXT, + ADD COLUMN embedding_model TEXT', + chunk_table); + END IF; + -- Create vector index for similarity search EXECUTE format(' CREATE INDEX IF NOT EXISTS %I ON %I @@ -704,13 +724,13 @@ BEGIN -- Compare effective values, not stored ones: moving a table from an -- explicit 'openai' to NULL whilst the GUC also says 'openai' changes -- nothing, and must not cost a re-embed. - old_provider := COALESCE(v_row.provider, + old_provider := COALESCE(NULLIF(v_row.provider, ''), current_setting('pgedge_vectorizer.provider')); - old_model := COALESCE(v_row.model, + old_model := COALESCE(NULLIF(v_row.model, ''), current_setting('pgedge_vectorizer.model')); - new_provider := COALESCE(set_embedding_model.provider, + new_provider := COALESCE(NULLIF(set_embedding_model.provider, ''), current_setting('pgedge_vectorizer.provider')); - new_model := COALESCE(set_embedding_model.model, + new_model := COALESCE(NULLIF(set_embedding_model.model, ''), current_setting('pgedge_vectorizer.model')); IF old_provider = new_provider AND old_model = new_model THEN @@ -756,10 +776,12 @@ BEGIN * and then fail every embedding the worker tried to write, which is the * failure this function exists to prevent. */ + -- The probe asks about the effective values rather than the raw + -- arguments: an empty string means inherit everywhere else, and would + -- otherwise reach the provider as a model name of ''. new_dim := COALESCE( set_embedding_model.embedding_dimension, - pgedge_vectorizer.detect_embedding_dimension( - set_embedding_model.provider, set_embedding_model.model)); + pgedge_vectorizer.detect_embedding_dimension(new_provider, new_model)); SELECT a.atttypmod INTO current_dim FROM pg_attribute a @@ -768,7 +790,9 @@ BEGIN IF chunk_count > 0 THEN -- NULL first: a vector column cannot change width with values in it. - EXECUTE format('UPDATE %s SET embedding = NULL ' + -- A cleared embedding has no model, so its provenance goes with it. + EXECUTE format('UPDATE %s SET embedding = NULL, ' + 'embedding_provider = NULL, embedding_model = NULL ' 'WHERE embedding IS NOT NULL', chunk_oid::REGCLASS); -- Anything already queued was queued against the old model. @@ -813,6 +837,263 @@ COMMENT ON FUNCTION pgedge_vectorizer.set_embedding_model IS 'force_reembed is true, in which case every embedding is cleared and every ' 'chunk requeued. Returns the number of chunks requeued'; +--------------------------------------------------------------------------- +-- embedding_model_status(): what each vectorizer's chunks were embedded with +-- +-- A vectorizer with NULL provider and model inherits the GUCs, and inheritance +-- resolves when the work runs rather than being copied at creation. Changing +-- pgedge_vectorizer.model therefore re-points every inheriting vectorizer at +-- once, leaving a chunk table holding vectors from the old model beside new +-- ones from the new. Similarity between two models' vectors is noise, so +-- search degrades quietly; where the widths match, as they do between +-- text-embedding-3-small and text-embedding-ada-002, nothing catches it at all. +-- +-- set_embedding_model() guards the per-vectorizer path and cannot guard this +-- one: the extension does not own that GUC and cannot intercept every way it +-- changes. So the chunk table records what produced each vector, and this +-- reports where that disagrees with what the vectorizer would use now. It +-- diagnoses rather than prevents, but it catches drift from any cause, +-- including a setting changed months ago by someone since departed. +-- +-- Each row costs a scan of one chunk table, so the arguments narrow it. +--------------------------------------------------------------------------- + +CREATE FUNCTION pgedge_vectorizer.embedding_model_status( + p_source_table REGCLASS DEFAULT NULL, + p_source_column NAME DEFAULT NULL +) RETURNS TABLE ( + source_table TEXT, + source_column NAME, + chunk_table TEXT, + effective_provider TEXT, + effective_model TEXT, + chunks_embedded BIGINT, + chunks_current BIGINT, + chunks_other_model BIGINT, + chunks_model_unknown BIGINT, + embedded_models TEXT[] +) AS $$ +DECLARE + v RECORD; + chunk_oid OID; +BEGIN + FOR v IN + SELECT r.source_table, r.source_column, r.chunk_table, + COALESCE(NULLIF(r.provider, ''), + current_setting('pgedge_vectorizer.provider')) + AS eff_provider, + COALESCE(NULLIF(r.model, ''), + current_setting('pgedge_vectorizer.model')) + AS eff_model + FROM pgedge_vectorizer.vectorizers r + WHERE (p_source_table IS NULL + OR r.source_table = p_source_table::TEXT) + AND (p_source_column IS NULL OR r.source_column = p_source_column) + ORDER BY r.source_table, r.source_column + LOOP + source_table := v.source_table; + source_column := v.source_column; + chunk_table := v.chunk_table; + effective_provider := v.eff_provider; + effective_model := v.eff_model; + + chunks_embedded := NULL; + chunks_current := NULL; + chunks_other_model := NULL; + chunks_model_unknown := NULL; + embedded_models := NULL; + + /* + * A chunk table that has been dropped, or that the caller cannot + * read, leaves this row's counts NULL rather than failing the whole + * result set. The name is one identifier with a dot in it for a + * schema-qualified source, so it is quoted rather than parsed. + */ + chunk_oid := to_regclass(quote_ident(v.chunk_table)); + IF chunk_oid IS NOT NULL + AND has_table_privilege(chunk_oid, 'SELECT') THEN + /* + * Every count is over embedded rows only: a chunk with no vector + * has no model to disagree about, and counting it as drifted + * would confuse work still to do with work done wrongly. + * + * A row with no recorded model is reported apart from a mismatch + * rather than lumped in with it. It predates these columns, so it + * may well be current; the report says what is there, and leaves + * the pessimistic reading to reembed(), which has to act. + */ + EXECUTE format( + 'SELECT count(*) FILTER (WHERE embedding IS NOT NULL), + count(*) FILTER (WHERE embedding IS NOT NULL + AND embedding_provider = %L + AND embedding_model = %L), + count(*) FILTER (WHERE embedding IS NOT NULL + AND embedding_model IS NOT NULL + AND (embedding_provider + IS DISTINCT FROM %L + OR embedding_model + IS DISTINCT FROM %L)), + count(*) FILTER (WHERE embedding IS NOT NULL + AND embedding_model IS NULL), + (SELECT array_agg(pair ORDER BY pair) + FROM (SELECT DISTINCT + embedding_provider || ''/'' || + embedding_model AS pair + FROM %s + WHERE embedding IS NOT NULL + AND embedding_model IS NOT NULL) d) + FROM %s', + v.eff_provider, v.eff_model, v.eff_provider, v.eff_model, + chunk_oid::REGCLASS, chunk_oid::REGCLASS) + INTO chunks_embedded, chunks_current, chunks_other_model, + chunks_model_unknown, embedded_models; + END IF; + + RETURN NEXT; + END LOOP; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.embedding_model_status IS +'Report, per vectorizer, how many embedded chunks were produced by the ' +'provider and model it would use now, how many by something else, and how ' +'many predate the columns that record it. Scans the chunk tables'; + +--------------------------------------------------------------------------- +-- reembed(): redo the embeddings that are not known to be current +-- +-- The report above leaves a user with a number and nothing to do about it. +-- set_embedding_model(..., force_reembed => true) is not the answer, because +-- an inheriting vectorizer's effective model already is the new one, so that +-- function sees no change and takes its no-op branch. +-- +-- No confirmation flag. Unlike set_embedding_model(), whose re-embed is a +-- surprising consequence of a settings change, this function does what its +-- name says; it raises a notice with the count, because the cost lands on a +-- metered provider. +--------------------------------------------------------------------------- + +CREATE FUNCTION pgedge_vectorizer.reembed( + source_table REGCLASS, + source_column NAME, + embedding_dimension INT DEFAULT NULL +) RETURNS BIGINT AS $$ +DECLARE + v_row RECORD; + chunk_oid OID; + eff_provider TEXT; + eff_model TEXT; + new_dim INT; + current_dim INT; + width_change BOOLEAN; + requeued BIGINT := 0; +BEGIN + SELECT r.* INTO v_row + FROM pgedge_vectorizer.vectorizers r + WHERE r.source_table = reembed.source_table::TEXT + AND r.source_column = reembed.source_column; + + IF NOT FOUND THEN + RAISE EXCEPTION 'no vectorizer registered for %.%', + reembed.source_table::TEXT, reembed.source_column; + END IF; + + eff_provider := COALESCE(NULLIF(v_row.provider, ''), + current_setting('pgedge_vectorizer.provider')); + eff_model := COALESCE(NULLIF(v_row.model, ''), + current_setting('pgedge_vectorizer.model')); + + chunk_oid := to_regclass(quote_ident(v_row.chunk_table)); + IF chunk_oid IS NULL THEN + RAISE EXCEPTION 'chunk table % for %.% no longer exists', + v_row.chunk_table, + reembed.source_table::TEXT, reembed.source_column; + END IF; + + new_dim := COALESCE( + reembed.embedding_dimension, + pgedge_vectorizer.detect_embedding_dimension(eff_provider, eff_model)); + + SELECT a.atttypmod INTO current_dim + FROM pg_attribute a + WHERE a.attrelid = chunk_oid + AND a.attname = 'embedding'; + + width_change := new_dim IS DISTINCT FROM current_dim; + + IF width_change THEN + /* + * A column cannot hold two widths, so a change of width takes every + * row with it whether or not it had drifted. Clearing has to come + * first: a vector column cannot be altered with values in it. + */ + EXECUTE format('UPDATE %s SET embedding = NULL, ' + 'embedding_provider = NULL, embedding_model = NULL ' + 'WHERE embedding IS NOT NULL', chunk_oid::REGCLASS); + + EXECUTE format('ALTER TABLE %s ALTER COLUMN embedding ' + 'TYPE vector(%s)', chunk_oid::REGCLASS, new_dim); + + RAISE NOTICE 'Embedding dimension changed from % to %, so every chunk ' + 'is being re-embedded', current_dim, new_dim; + ELSE + /* + * Same width, so rows already produced by this provider and model are + * left exactly as they are. Everything else goes, including rows with + * nothing recorded: those predate the columns and cannot be shown to + * be current, and the safe reading of a row that cannot be proved + * current is that it needs doing again. On an installation freshly + * upgraded to 1.2 that is every row, which the documentation says. + */ + EXECUTE format( + 'UPDATE %s SET embedding = NULL, ' + ' embedding_provider = NULL, embedding_model = NULL ' + ' WHERE embedding IS NOT NULL ' + ' AND (embedding_model IS NULL ' + ' OR embedding_provider IS DISTINCT FROM %L ' + ' OR embedding_model IS DISTINCT FROM %L)', + chunk_oid::REGCLASS, eff_provider, eff_model); + END IF; + + /* + * Anything already queued for embedding was queued before this decision + * was made. Sparse-only rows are left alone: they carry no dense work to + * redo, the probe only marks a chunk sparse-only whilst it still has an + * embedding, and deleting them would strand the sparse vector as NULL + * until someone thought to run reprocess_chunks(). + */ + DELETE FROM pgedge_vectorizer.queue q + WHERE q.chunk_table = v_row.chunk_table + AND NOT COALESCE((q.metadata->>'sparse_only')::BOOLEAN, FALSE); + + /* + * Whatever now has no embedding needs one, which after the clearing above + * is exactly the set chosen, plus any chunk that was never embedded in + * the first place and would have been picked up by reprocess_chunks() + * anyway. + */ + EXECUTE format( + 'INSERT INTO pgedge_vectorizer.queue ' + ' (chunk_id, chunk_table, content, max_attempts) ' + 'SELECT id, %L, content, %s FROM %s WHERE embedding IS NULL', + v_row.chunk_table, + current_setting('pgedge_vectorizer.max_retries')::INT, + chunk_oid::REGCLASS); + + GET DIAGNOSTICS requeued = ROW_COUNT; + + RAISE NOTICE 'Queued % chunks to be embedded with %/%', + requeued, eff_provider, eff_model; + + RETURN requeued; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.reembed IS +'Re-embed a vectorizer''s chunks with the provider and model it would use now, ' +'leaving alone any already produced by them. A change of embedding dimension ' +'takes every chunk with it. Returns the number queued'; + -- Recreate the DELETE and TRUNCATE cleanup triggers for every registered -- vectorizer. -- @@ -1688,6 +1969,8 @@ RETURNS TABLE ( LANGUAGE plpgsql AS $$ DECLARE v_chunk_table TEXT; + v_provider TEXT; + v_model TEXT; v_query_dense vector; v_query_sparse sparsevec; BEGIN @@ -1702,12 +1985,14 @@ BEGIN -- vectorized column to avoid silently returning results from the -- wrong chunk table. IF p_source_column IS NOT NULL THEN - SELECT vz.chunk_table INTO v_chunk_table + SELECT vz.chunk_table, vz.provider, vz.model + INTO v_chunk_table, v_provider, v_model FROM pgedge_vectorizer.vectorizers vz WHERE vz.source_table = p_source_table::TEXT AND vz.source_column = p_source_column; ELSE - SELECT vz.chunk_table INTO v_chunk_table + SELECT vz.chunk_table, vz.provider, vz.model + INTO v_chunk_table, v_provider, v_model FROM pgedge_vectorizer.vectorizers vz WHERE vz.source_table = p_source_table::TEXT LIMIT 1; @@ -1730,8 +2015,15 @@ BEGIN p_source_table; END IF; - -- Generate dense query vector via the existing C function - v_query_dense := pgedge_vectorizer.generate_embedding(p_query); + /* + * Embed the query with this vectorizer's own provider and model rather + * than the GUCs. A query embedded by one model and compared against chunks + * embedded by another gives meaningless distances, and where the widths + * differ it fails outright. NULL passes straight through and means + * inherit, so a vectorizer that has pinned nothing behaves as before. + */ + v_query_dense := pgedge_vectorizer.generate_embedding(p_query, + v_provider, v_model); -- Generate sparse BM25 query vector v_query_sparse := pgedge_vectorizer.bm25_query_vector( diff --git a/src/embed.c b/src/embed.c index a6ddb98..0b3c8e2 100644 --- a/src/embed.c +++ b/src/embed.c @@ -46,12 +46,26 @@ resolve_provider(const char *name) return provider; } -/* As resolve_provider(), for the model name. */ +/* + * As resolve_provider(), for the model name. + * + * An unset model is a configuration error rather than something to send: the + * providers interpolate this straight into their request bodies, so an empty + * GUC would put "model":"" on the wire and a NULL one would be dereferenced + * while escaping it. resolve_provider() already refuses an unset provider for + * the same reason. + */ static const char * resolve_model(const char *model) { - return (model != NULL && model[0] != '\0') - ? model : pgedge_vectorizer_model; + const char *use; + + use = (model != NULL && model[0] != '\0') ? model : pgedge_vectorizer_model; + + if (use == NULL || use[0] == '\0') + elog(ERROR, "pgedge_vectorizer.model is not set"); + + return use; } /* Read an optional text argument as a cstring, or NULL if it was not given. */ diff --git a/src/worker.c b/src/worker.c index 01df934..c2be578 100644 --- a/src/worker.c +++ b/src/worker.c @@ -249,7 +249,8 @@ static bool queue_item_record_failure(void); static void process_queue_batch(const char *dbname); static void cleanup_completed_items(const char *dbname); static void update_embedding(int64 chunk_id, const char *chunk_table, - const float *embedding, int dim); + const float *embedding, int dim, + const char *provider, const char *model); static char *trim_whitespace(char *str); static int parse_database_list(char ***names); static int worker_quantum_secs(void); @@ -1649,6 +1650,14 @@ process_queue_batch(const char *dbname) * FOR UPDATE OF q, not a bare FOR UPDATE: the registry rows are not * being changed, and locking the nullable side of a left join is * rejected outright. + * + * LATERAL ... LIMIT 1 rather than a plain join, because only + * (source_table, source_column) is unique in the registry: two + * vectorizers pointed at one chunk table by an explicit chunk_table_name + * would otherwise match a queue row twice, and the item would be embedded + * twice and counted twice into the BM25 corpus statistics. That + * configuration is already incoherent, but it must not corrupt the + * statistics of a table that is merely nearby. */ ret = SPI_execute(psprintf( "SELECT q.id, q.chunk_id, q.chunk_table, q.content, q.attempts, " @@ -1856,6 +1865,15 @@ process_queue_batch(const char *dbname) * provider that fails before reaching the network records * nothing of its own. */ + /* + * A provider that cannot be resolved or initialised is a + * fault of the configuration rather than of any item in + * the request, so it is raised rather than charged: see + * 005_batch_failure_backoff.pl, which exists to keep a + * mistyped provider name from retiring the whole queue + * one blameless row at a time. The outer handler backs + * the worker off instead. + */ provider = get_embedding_provider(providers[batch_start]); if (provider == NULL) elog(ERROR, "embedding provider \"%s\" is not available", @@ -1889,11 +1907,17 @@ process_queue_batch(const char *dbname) bool isnull_dim; Datum val_dim; + /* + * The chunk table's generated name is one identifier + * with a dot in it for a schema-qualified source, so it + * is quoted rather than left for regclass to parse as a + * qualified reference. Same rule as update_embedding(). + */ ret_dim = SPI_execute(psprintf( "SELECT atttypmod FROM pg_attribute " - "WHERE attrelid = '%s'::regclass " + "WHERE attrelid = %s::regclass " "AND attname = 'embedding'", - chunk_tables[idx0]), + quote_literal_cstr(quote_identifier(chunk_tables[idx0]))), true, 1); if (ret_dim == SPI_OK_SELECT && SPI_processed == 1) @@ -1955,7 +1979,9 @@ process_queue_batch(const char *dbname) queue_item_begin(queue_ids[idx], attempts[idx], max_attempts[idx]); if (!sparse_only[idx]) - update_embedding(chunk_ids[idx], chunk_tables[idx], embeddings[i], dim); + update_embedding(chunk_ids[idx], chunk_tables[idx], + embeddings[i], dim, + providers[idx], models[idx]); else if (!pgedge_vectorizer_enable_hybrid) elog(ERROR, "cannot process sparse-only queue item while pgedge_vectorizer.enable_hybrid is disabled"); @@ -2236,7 +2262,9 @@ process_queue_batch(const char *dbname) * Update a chunk table with the generated embedding */ static void -update_embedding(int64 chunk_id, const char *chunk_table, const float *embedding, int dim) +update_embedding(int64 chunk_id, const char *chunk_table, + const float *embedding, int dim, + const char *provider, const char *model) { StringInfoData vector_str; int ret; @@ -2252,10 +2280,22 @@ update_embedding(int64 chunk_id, const char *chunk_table, const float *embedding } appendStringInfoChar(&vector_str, ']'); - /* Update the chunk table */ + /* + * Record what produced the vector in the same statement that writes it, + * so the two cannot disagree. Nothing else writes these columns. + * + * quote_identifier() because the chunk table's generated name is one + * identifier with a dot in it for a schema-qualified source, not a + * schema-qualified reference, which is how every other statement in this + * file spells it. + */ ret = SPI_execute(psprintf( - "UPDATE %s SET embedding = '%s'::vector WHERE id = %ld", - chunk_table, vector_str.data, chunk_id), + "UPDATE %s SET embedding = '%s'::vector, " + "embedding_provider = %s, embedding_model = %s " + "WHERE id = %ld", + quote_identifier(chunk_table), vector_str.data, + quote_literal_cstr(provider), quote_literal_cstr(model), + chunk_id), false, 0); if (ret != SPI_OK_UPDATE) diff --git a/test/expected/model_drift.out b/test/expected/model_drift.out new file mode 100644 index 0000000..687efea --- /dev/null +++ b/test/expected/model_drift.out @@ -0,0 +1,326 @@ +-- model_drift test +-- +-- A vectorizer that inherits follows pgedge_vectorizer.model as it changes, so +-- a chunk table can end up holding vectors from two models with nothing +-- reporting it. Each chunk now records what produced its vector, +-- embedding_model_status() reports where that disagrees with what the +-- vectorizer would use now, and reembed() repairs it. +-- +-- Nothing here reaches a provider: embeddings and their provenance are written +-- by hand, as the vectorizer_status tests do, and every call passes a +-- dimension explicitly. +SET pgedge_vectorizer.provider = 'openai'; +SET pgedge_vectorizer.model = 'text-embedding-3-small'; +CREATE TABLE drift_docs (id BIGSERIAL PRIMARY KEY, body TEXT); +INSERT INTO drift_docs (body) +VALUES ('The first document, long enough to be worth embedding.'), + ('The second document, on an unrelated subject.'), + ('The third document, which predates the new columns.'), + ('The fourth document, never embedded at all.'); +SELECT pgedge_vectorizer.enable_vectorization( + 'drift_docs'::regclass, 'body', 'token_based', 100, 10, 1536); +NOTICE: Using primary key column: id (bigint) +NOTICE: column "sparse_embedding" of relation "drift_docs_body_chunks" already exists, skipping +NOTICE: Vectorization enabled: drift_docs -> drift_docs_body_chunks +NOTICE: Strategy: token_based, chunk_size: 100, overlap: 10 +NOTICE: Processing existing rows... +NOTICE: Processed 4 existing rows + enable_vectorization +---------------------- + +(1 row) + +--------------------------------------------------------------------------- +-- The chunk table records what produced each vector +--------------------------------------------------------------------------- +SELECT a.attname, format_type(a.atttypid, a.atttypmod) AS type + FROM pg_attribute a + WHERE a.attrelid = 'drift_docs_body_chunks'::regclass + AND a.attname IN ('embedding_provider', 'embedding_model') + ORDER BY a.attname; + attname | type +--------------------+------ + embedding_model | text + embedding_provider | text +(2 rows) + +--------------------------------------------------------------------------- +-- Four states, one per chunk: current, drifted, unknown, unembedded +--------------------------------------------------------------------------- +-- Chunk 1 was produced by the model the vectorizer would use now. +UPDATE drift_docs_body_chunks + SET embedding = array_fill(0.1::real, ARRAY[1536])::vector, + embedding_provider = 'openai', + embedding_model = 'text-embedding-3-small', + sparse_embedding = '{1:0.5}/65536'::sparsevec + WHERE source_id = 1; +-- Chunk 2 came from a different model of the same width, which is the case +-- nothing catches today: the vectors are the right shape and meaningless. +UPDATE drift_docs_body_chunks + SET embedding = array_fill(0.2::real, ARRAY[1536])::vector, + embedding_provider = 'openai', + embedding_model = 'text-embedding-ada-002', + sparse_embedding = '{1:0.5}/65536'::sparsevec + WHERE source_id = 2; +-- Chunk 3 was embedded before the columns existed, so nothing is recorded. +UPDATE drift_docs_body_chunks + SET embedding = array_fill(0.3::real, ARRAY[1536])::vector, + sparse_embedding = '{1:0.5}/65536'::sparsevec + WHERE source_id = 3; +-- Chunk 4 is left unembedded. +SELECT source_id, embedding IS NOT NULL AS embedded, + embedding_provider, embedding_model + FROM drift_docs_body_chunks + ORDER BY source_id; + source_id | embedded | embedding_provider | embedding_model +-----------+----------+--------------------+------------------------ + 1 | t | openai | text-embedding-3-small + 2 | t | openai | text-embedding-ada-002 + 3 | t | | + 4 | f | | +(4 rows) + +--------------------------------------------------------------------------- +-- The report +--------------------------------------------------------------------------- +SELECT source_table, source_column, effective_provider, effective_model, + chunks_embedded, chunks_current, chunks_other_model, + chunks_model_unknown, embedded_models + FROM pgedge_vectorizer.embedding_model_status('drift_docs'::regclass); + source_table | source_column | effective_provider | effective_model | chunks_embedded | chunks_current | chunks_other_model | chunks_model_unknown | embedded_models +--------------+---------------+--------------------+------------------------+-----------------+----------------+--------------------+----------------------+--------------------------------------------------------------- + drift_docs | body | openai | text-embedding-3-small | 3 | 1 | 1 | 1 | {openai/text-embedding-3-small,openai/text-embedding-ada-002} +(1 row) + +-- Narrowing by column, and a column that is not vectorized returns nothing +-- rather than erroring. +SELECT count(*) AS rows_for_body + FROM pgedge_vectorizer.embedding_model_status('drift_docs'::regclass, 'body'); + rows_for_body +--------------- + 1 +(1 row) + +SELECT count(*) AS rows_for_unknown_column + FROM pgedge_vectorizer.embedding_model_status('drift_docs'::regclass, 'title'); + rows_for_unknown_column +------------------------- + 0 +(1 row) + +--------------------------------------------------------------------------- +-- Narrowing does not depend on the caller being able to resolve every other +-- vectorizer's source table +-- +-- The registry stores source_table as text, and resolving each row's name +-- back to a regclass in order to compare it raises for a schema the caller +-- has no USAGE on, so narrowing to the one table they can read failed whilst +-- the unnarrowed call worked. The comparison is made as text, as every other +-- lookup in the extension makes it. +--------------------------------------------------------------------------- +CREATE SCHEMA drift_private; +CREATE TABLE drift_private.secret (id BIGSERIAL PRIMARY KEY, body TEXT); +SELECT pgedge_vectorizer.enable_vectorization( + 'drift_private.secret'::regclass, 'body', 'token_based', 100, 10, 1536); +NOTICE: Using primary key column: id (bigint) +NOTICE: column "sparse_embedding" of relation "drift_private.secret_body_chunks" already exists, skipping +NOTICE: Vectorization enabled: drift_private.secret -> drift_private.secret_body_chunks +NOTICE: Strategy: token_based, chunk_size: 100, overlap: 10 +NOTICE: Processing existing rows... +NOTICE: Processed 0 existing rows + enable_vectorization +---------------------- + +(1 row) + +CREATE ROLE drift_reader_role; +GRANT USAGE ON SCHEMA pgedge_vectorizer TO drift_reader_role; +GRANT SELECT ON pgedge_vectorizer.vectorizers TO drift_reader_role; +GRANT SELECT ON drift_docs, drift_docs_body_chunks TO drift_reader_role; +SET ROLE drift_reader_role; +SELECT source_table, source_column, chunks_embedded + FROM pgedge_vectorizer.embedding_model_status('drift_docs'::regclass); + source_table | source_column | chunks_embedded +--------------+---------------+----------------- + drift_docs | body | 3 +(1 row) + +RESET ROLE; +SELECT pgedge_vectorizer.disable_vectorization( + 'drift_private.secret'::regclass, 'body', TRUE); +NOTICE: Vectorization disabled and chunk table dropped: drift_private.secret_body_chunks + disable_vectorization +----------------------- + +(1 row) + +REVOKE SELECT ON drift_docs, drift_docs_body_chunks FROM drift_reader_role; +REVOKE SELECT ON pgedge_vectorizer.vectorizers FROM drift_reader_role; +REVOKE USAGE ON SCHEMA pgedge_vectorizer FROM drift_reader_role; +DROP ROLE drift_reader_role; +DROP TABLE drift_private.secret; +DROP SCHEMA drift_private; +--------------------------------------------------------------------------- +-- reembed() at a matching width leaves the current row alone +--------------------------------------------------------------------------- +-- Chunk 1 has sparse work still outstanding, which has nothing to do with +-- the embedding model: clearing the queue must not take it with it. +INSERT INTO pgedge_vectorizer.queue (chunk_id, chunk_table, content, metadata) +SELECT id, 'drift_docs_body_chunks', content, + jsonb_build_object('sparse_only', true) + FROM drift_docs_body_chunks WHERE source_id = 1; +SELECT pgedge_vectorizer.reembed( + 'drift_docs'::regclass, 'body', embedding_dimension => 1536) AS queued; +NOTICE: Queued 3 chunks to be embedded with openai/text-embedding-3-small + queued +-------- + 3 +(1 row) + +SELECT count(*) AS sparse_only_rows_kept + FROM pgedge_vectorizer.queue + WHERE chunk_table = 'drift_docs_body_chunks' + AND (metadata->>'sparse_only')::BOOLEAN; + sparse_only_rows_kept +----------------------- + 1 +(1 row) + +-- Out of the way again, so the counts below are about the dense requeue. +DELETE FROM pgedge_vectorizer.queue + WHERE chunk_table = 'drift_docs_body_chunks' + AND (metadata->>'sparse_only')::BOOLEAN; +-- Chunk 1 keeps its embedding and its provenance; 2 and 3 lose both; 4 was +-- never embedded and is queued alongside them. +SELECT source_id, embedding IS NOT NULL AS embedded, + embedding_provider, embedding_model + FROM drift_docs_body_chunks + ORDER BY source_id; + source_id | embedded | embedding_provider | embedding_model +-----------+----------+--------------------+------------------------ + 1 | t | openai | text-embedding-3-small + 2 | f | | + 3 | f | | + 4 | f | | +(4 rows) + +-- Nothing that does not depend on the embedding model was disturbed. +SELECT count(*) AS chunks, + count(token_count) AS token_counts_kept, + count(sparse_embedding) AS sparse_kept + FROM drift_docs_body_chunks; + chunks | token_counts_kept | sparse_kept +--------+-------------------+------------- + 4 | 4 | 3 +(1 row) + +SELECT count(*) AS queued_rows + FROM pgedge_vectorizer.queue + WHERE chunk_table = 'drift_docs_body_chunks' AND status = 'pending'; + queued_rows +------------- + 3 +(1 row) + +-- The report now shows one current row and nothing drifted. +SELECT chunks_embedded, chunks_current, chunks_other_model, + chunks_model_unknown, embedded_models + FROM pgedge_vectorizer.embedding_model_status('drift_docs'::regclass); + chunks_embedded | chunks_current | chunks_other_model | chunks_model_unknown | embedded_models +-----------------+----------------+--------------------+----------------------+--------------------------------- + 1 | 1 | 0 | 0 | {openai/text-embedding-3-small} +(1 row) + +--------------------------------------------------------------------------- +-- A vectorizer with nothing to redo queues nothing +--------------------------------------------------------------------------- +-- Give the three cleared chunks their embeddings back, as the worker would. +UPDATE drift_docs_body_chunks + SET embedding = array_fill(0.1::real, ARRAY[1536])::vector, + embedding_provider = 'openai', + embedding_model = 'text-embedding-3-small' + WHERE embedding IS NULL; +SELECT pgedge_vectorizer.reembed( + 'drift_docs'::regclass, 'body', embedding_dimension => 1536) AS queued; +NOTICE: Queued 0 chunks to be embedded with openai/text-embedding-3-small + queued +-------- + 0 +(1 row) + +SELECT count(*) AS still_embedded FROM drift_docs_body_chunks + WHERE embedding IS NOT NULL; + still_embedded +---------------- + 4 +(1 row) + +--------------------------------------------------------------------------- +-- A change of width takes every chunk with it, drifted or not +--------------------------------------------------------------------------- +SELECT pgedge_vectorizer.reembed( + 'drift_docs'::regclass, 'body', embedding_dimension => 768) AS queued; +NOTICE: Embedding dimension changed from 1536 to 768, so every chunk is being re-embedded +NOTICE: Queued 4 chunks to be embedded with openai/text-embedding-3-small + queued +-------- + 4 +(1 row) + +SELECT format_type(a.atttypid, a.atttypmod) AS embedding_type + FROM pg_attribute a + WHERE a.attrelid = 'drift_docs_body_chunks'::regclass + AND a.attname = 'embedding'; + embedding_type +---------------- + vector(768) +(1 row) + +SELECT count(*) AS chunks, + count(embedding) AS embedded, + count(token_count) AS token_counts_kept, + count(sparse_embedding) AS sparse_kept + FROM drift_docs_body_chunks; + chunks | embedded | token_counts_kept | sparse_kept +--------+----------+-------------------+------------- + 4 | 0 | 4 | 3 +(1 row) + +--------------------------------------------------------------------------- +-- A table with no vectorizer is an error +--------------------------------------------------------------------------- +-- The sentinel is raised outside the handler, so that a reembed() which +-- wrongly succeeded fails the test rather than being caught by it. +DO $$ +DECLARE + raised TEXT; +BEGIN + BEGIN + PERFORM pgedge_vectorizer.reembed('drift_docs'::regclass, 'nosuchcolumn'); + EXCEPTION WHEN OTHERS THEN + raised := SQLERRM; + END; + + IF raised IS NULL THEN + RAISE EXCEPTION 'expected an error, got none'; + END IF; + + RAISE NOTICE '%', raised; +END; +$$; +NOTICE: no vectorizer registered for drift_docs.nosuchcolumn +--------------------------------------------------------------------------- +-- Cleanup +--------------------------------------------------------------------------- +SELECT pgedge_vectorizer.disable_vectorization('drift_docs'::regclass, + 'body', TRUE); +NOTICE: Vectorization disabled and chunk table dropped: drift_docs_body_chunks + disable_vectorization +----------------------- + +(1 row) + +DROP TABLE drift_docs; +DELETE FROM pgedge_vectorizer.queue; +RESET pgedge_vectorizer.provider; +RESET pgedge_vectorizer.model; diff --git a/test/sql/model_drift.sql b/test/sql/model_drift.sql new file mode 100644 index 0000000..8fd1fe6 --- /dev/null +++ b/test/sql/model_drift.sql @@ -0,0 +1,237 @@ +-- model_drift test +-- +-- A vectorizer that inherits follows pgedge_vectorizer.model as it changes, so +-- a chunk table can end up holding vectors from two models with nothing +-- reporting it. Each chunk now records what produced its vector, +-- embedding_model_status() reports where that disagrees with what the +-- vectorizer would use now, and reembed() repairs it. +-- +-- Nothing here reaches a provider: embeddings and their provenance are written +-- by hand, as the vectorizer_status tests do, and every call passes a +-- dimension explicitly. + +SET pgedge_vectorizer.provider = 'openai'; +SET pgedge_vectorizer.model = 'text-embedding-3-small'; + +CREATE TABLE drift_docs (id BIGSERIAL PRIMARY KEY, body TEXT); + +INSERT INTO drift_docs (body) +VALUES ('The first document, long enough to be worth embedding.'), + ('The second document, on an unrelated subject.'), + ('The third document, which predates the new columns.'), + ('The fourth document, never embedded at all.'); + +SELECT pgedge_vectorizer.enable_vectorization( + 'drift_docs'::regclass, 'body', 'token_based', 100, 10, 1536); + +--------------------------------------------------------------------------- +-- The chunk table records what produced each vector +--------------------------------------------------------------------------- + +SELECT a.attname, format_type(a.atttypid, a.atttypmod) AS type + FROM pg_attribute a + WHERE a.attrelid = 'drift_docs_body_chunks'::regclass + AND a.attname IN ('embedding_provider', 'embedding_model') + ORDER BY a.attname; + +--------------------------------------------------------------------------- +-- Four states, one per chunk: current, drifted, unknown, unembedded +--------------------------------------------------------------------------- + +-- Chunk 1 was produced by the model the vectorizer would use now. +UPDATE drift_docs_body_chunks + SET embedding = array_fill(0.1::real, ARRAY[1536])::vector, + embedding_provider = 'openai', + embedding_model = 'text-embedding-3-small', + sparse_embedding = '{1:0.5}/65536'::sparsevec + WHERE source_id = 1; + +-- Chunk 2 came from a different model of the same width, which is the case +-- nothing catches today: the vectors are the right shape and meaningless. +UPDATE drift_docs_body_chunks + SET embedding = array_fill(0.2::real, ARRAY[1536])::vector, + embedding_provider = 'openai', + embedding_model = 'text-embedding-ada-002', + sparse_embedding = '{1:0.5}/65536'::sparsevec + WHERE source_id = 2; + +-- Chunk 3 was embedded before the columns existed, so nothing is recorded. +UPDATE drift_docs_body_chunks + SET embedding = array_fill(0.3::real, ARRAY[1536])::vector, + sparse_embedding = '{1:0.5}/65536'::sparsevec + WHERE source_id = 3; + +-- Chunk 4 is left unembedded. + +SELECT source_id, embedding IS NOT NULL AS embedded, + embedding_provider, embedding_model + FROM drift_docs_body_chunks + ORDER BY source_id; + +--------------------------------------------------------------------------- +-- The report +--------------------------------------------------------------------------- + +SELECT source_table, source_column, effective_provider, effective_model, + chunks_embedded, chunks_current, chunks_other_model, + chunks_model_unknown, embedded_models + FROM pgedge_vectorizer.embedding_model_status('drift_docs'::regclass); + +-- Narrowing by column, and a column that is not vectorized returns nothing +-- rather than erroring. +SELECT count(*) AS rows_for_body + FROM pgedge_vectorizer.embedding_model_status('drift_docs'::regclass, 'body'); + +SELECT count(*) AS rows_for_unknown_column + FROM pgedge_vectorizer.embedding_model_status('drift_docs'::regclass, 'title'); + +--------------------------------------------------------------------------- +-- Narrowing does not depend on the caller being able to resolve every other +-- vectorizer's source table +-- +-- The registry stores source_table as text, and resolving each row's name +-- back to a regclass in order to compare it raises for a schema the caller +-- has no USAGE on, so narrowing to the one table they can read failed whilst +-- the unnarrowed call worked. The comparison is made as text, as every other +-- lookup in the extension makes it. +--------------------------------------------------------------------------- + +CREATE SCHEMA drift_private; +CREATE TABLE drift_private.secret (id BIGSERIAL PRIMARY KEY, body TEXT); +SELECT pgedge_vectorizer.enable_vectorization( + 'drift_private.secret'::regclass, 'body', 'token_based', 100, 10, 1536); + +CREATE ROLE drift_reader_role; +GRANT USAGE ON SCHEMA pgedge_vectorizer TO drift_reader_role; +GRANT SELECT ON pgedge_vectorizer.vectorizers TO drift_reader_role; +GRANT SELECT ON drift_docs, drift_docs_body_chunks TO drift_reader_role; + +SET ROLE drift_reader_role; +SELECT source_table, source_column, chunks_embedded + FROM pgedge_vectorizer.embedding_model_status('drift_docs'::regclass); +RESET ROLE; + +SELECT pgedge_vectorizer.disable_vectorization( + 'drift_private.secret'::regclass, 'body', TRUE); +REVOKE SELECT ON drift_docs, drift_docs_body_chunks FROM drift_reader_role; +REVOKE SELECT ON pgedge_vectorizer.vectorizers FROM drift_reader_role; +REVOKE USAGE ON SCHEMA pgedge_vectorizer FROM drift_reader_role; +DROP ROLE drift_reader_role; +DROP TABLE drift_private.secret; +DROP SCHEMA drift_private; + +--------------------------------------------------------------------------- +-- reembed() at a matching width leaves the current row alone +--------------------------------------------------------------------------- + +-- Chunk 1 has sparse work still outstanding, which has nothing to do with +-- the embedding model: clearing the queue must not take it with it. +INSERT INTO pgedge_vectorizer.queue (chunk_id, chunk_table, content, metadata) +SELECT id, 'drift_docs_body_chunks', content, + jsonb_build_object('sparse_only', true) + FROM drift_docs_body_chunks WHERE source_id = 1; + +SELECT pgedge_vectorizer.reembed( + 'drift_docs'::regclass, 'body', embedding_dimension => 1536) AS queued; + +SELECT count(*) AS sparse_only_rows_kept + FROM pgedge_vectorizer.queue + WHERE chunk_table = 'drift_docs_body_chunks' + AND (metadata->>'sparse_only')::BOOLEAN; + +-- Out of the way again, so the counts below are about the dense requeue. +DELETE FROM pgedge_vectorizer.queue + WHERE chunk_table = 'drift_docs_body_chunks' + AND (metadata->>'sparse_only')::BOOLEAN; + +-- Chunk 1 keeps its embedding and its provenance; 2 and 3 lose both; 4 was +-- never embedded and is queued alongside them. +SELECT source_id, embedding IS NOT NULL AS embedded, + embedding_provider, embedding_model + FROM drift_docs_body_chunks + ORDER BY source_id; + +-- Nothing that does not depend on the embedding model was disturbed. +SELECT count(*) AS chunks, + count(token_count) AS token_counts_kept, + count(sparse_embedding) AS sparse_kept + FROM drift_docs_body_chunks; + +SELECT count(*) AS queued_rows + FROM pgedge_vectorizer.queue + WHERE chunk_table = 'drift_docs_body_chunks' AND status = 'pending'; + +-- The report now shows one current row and nothing drifted. +SELECT chunks_embedded, chunks_current, chunks_other_model, + chunks_model_unknown, embedded_models + FROM pgedge_vectorizer.embedding_model_status('drift_docs'::regclass); + +--------------------------------------------------------------------------- +-- A vectorizer with nothing to redo queues nothing +--------------------------------------------------------------------------- + +-- Give the three cleared chunks their embeddings back, as the worker would. +UPDATE drift_docs_body_chunks + SET embedding = array_fill(0.1::real, ARRAY[1536])::vector, + embedding_provider = 'openai', + embedding_model = 'text-embedding-3-small' + WHERE embedding IS NULL; + +SELECT pgedge_vectorizer.reembed( + 'drift_docs'::regclass, 'body', embedding_dimension => 1536) AS queued; + +SELECT count(*) AS still_embedded FROM drift_docs_body_chunks + WHERE embedding IS NOT NULL; + +--------------------------------------------------------------------------- +-- A change of width takes every chunk with it, drifted or not +--------------------------------------------------------------------------- + +SELECT pgedge_vectorizer.reembed( + 'drift_docs'::regclass, 'body', embedding_dimension => 768) AS queued; + +SELECT format_type(a.atttypid, a.atttypmod) AS embedding_type + FROM pg_attribute a + WHERE a.attrelid = 'drift_docs_body_chunks'::regclass + AND a.attname = 'embedding'; + +SELECT count(*) AS chunks, + count(embedding) AS embedded, + count(token_count) AS token_counts_kept, + count(sparse_embedding) AS sparse_kept + FROM drift_docs_body_chunks; + +--------------------------------------------------------------------------- +-- A table with no vectorizer is an error +--------------------------------------------------------------------------- + +-- The sentinel is raised outside the handler, so that a reembed() which +-- wrongly succeeded fails the test rather than being caught by it. +DO $$ +DECLARE + raised TEXT; +BEGIN + BEGIN + PERFORM pgedge_vectorizer.reembed('drift_docs'::regclass, 'nosuchcolumn'); + EXCEPTION WHEN OTHERS THEN + raised := SQLERRM; + END; + + IF raised IS NULL THEN + RAISE EXCEPTION 'expected an error, got none'; + END IF; + + RAISE NOTICE '%', raised; +END; +$$; + +--------------------------------------------------------------------------- +-- Cleanup +--------------------------------------------------------------------------- + +SELECT pgedge_vectorizer.disable_vectorization('drift_docs'::regclass, + 'body', TRUE); +DROP TABLE drift_docs; +DELETE FROM pgedge_vectorizer.queue; +RESET pgedge_vectorizer.provider; +RESET pgedge_vectorizer.model; diff --git a/test/t/012_upgrade_1_1_to_1_2.pl b/test/t/012_upgrade_1_1_to_1_2.pl index 4adc070..8985dc4 100644 --- a/test/t/012_upgrade_1_1_to_1_2.pl +++ b/test/t/012_upgrade_1_1_to_1_2.pl @@ -58,6 +58,20 @@ q(SELECT pgedge_vectorizer.enable_vectorization('docs', 'body', 'token_based', 100, 10, 1536))); +# A second vectorizer whose chunk table ends up in a schema the upgrade will +# not have on its search_path. The migration adds the provenance columns to +# the chunk tables the registry knows about, and the registry stores only a +# bare name, so resolving that name through the search_path would silently +# miss this one and leave the worker failing every embedding write against it. +$node->safe_psql('upgraded', q( +CREATE SCHEMA tucked_away; +SET search_path = tucked_away, public; +CREATE TABLE tucked_away.notes (id BIGSERIAL PRIMARY KEY, body TEXT); +INSERT INTO tucked_away.notes (body) VALUES ('Out of the way.'); +SELECT pgedge_vectorizer.enable_vectorization('tucked_away.notes', 'body', + 'token_based', 100, 10, 1536); +)); + $node->safe_psql('upgraded', "ALTER EXTENSION pgedge_vectorizer UPDATE TO '1.2'"); @@ -113,18 +127,52 @@ $node->safe_psql('fresh', $views), 'an upgraded install has the same views as a fresh one'); +# Chunk tables too, which is a separate trap: enable_vectorization() adds +# columns to a chunk table it finds without them, but nothing re-runs it on +# upgrade, so anything the worker writes has to be added by the upgrade script +# itself. Compare a chunk table created at 1.1 and upgraded against one created +# fresh at 1.2. +$node->safe_psql('fresh', q( +CREATE TABLE docs (id BIGSERIAL PRIMARY KEY, body TEXT); +INSERT INTO docs (body) VALUES ('Written on a fresh 1.2 install.'); +)); +$node->safe_psql('fresh', + q(SELECT pgedge_vectorizer.enable_vectorization('docs', 'body', + 'token_based', 100, 10, 1536))); + +my $chunk_columns = q( + SELECT string_agg(a.attname || ' ' || format_type(a.atttypid, a.atttypmod), + E'\n' ORDER BY a.attname) + FROM pg_attribute a + WHERE a.attrelid = 'docs_body_chunks'::regclass + AND a.attnum > 0 + AND NOT a.attisdropped +); + +is($node->safe_psql('upgraded', $chunk_columns), + $node->safe_psql('fresh', $chunk_columns), + 'an upgraded chunk table has the same columns as a freshly created one'); + # The data that was there before the upgrade is still there, and the new # columns default to inheriting. is($node->safe_psql('upgraded', q(SELECT source_table || ' ' || COALESCE(provider, 'NULL') || ' ' || COALESCE(model, 'NULL') - FROM pgedge_vectorizer.vectorizers)), + FROM pgedge_vectorizer.vectorizers + WHERE source_table = 'docs')), 'docs NULL NULL', 'a vectorizer registered before the upgrade survives it, inheriting'); is($node->safe_psql('upgraded', 'SELECT count(*) FROM docs_body_chunks'), '1', 'the chunks written before the upgrade survive it'); +is($node->safe_psql('upgraded', q( + SELECT count(*) FROM pg_attribute + WHERE attrelid = 'tucked_away.notes_body_chunks'::regclass + AND attname IN ('embedding_provider', 'embedding_model'))), + '2', + 'the upgrade reaches a chunk table outside its own search_path'); + $node->stop; done_testing(); diff --git a/test/t/013_embedding_model_recorded.pl b/test/t/013_embedding_model_recorded.pl new file mode 100644 index 0000000..e68de6e --- /dev/null +++ b/test/t/013_embedding_model_recorded.pl @@ -0,0 +1,222 @@ +# Copyright (c) 2025 - 2026, pgEdge, Inc. +# +# Verify that the worker records the provider and model it actually embedded +# with, per chunk. +# +# The regression tests write embedding_provider and embedding_model by hand, +# because a real embedding needs a provider, so they check what +# embedding_model_status() and reembed() make of those columns but never that +# anything fills them in. That is the half worth a running worker: the columns +# are written in the same statement as the vector, in update_embedding(), and +# the value has to be the one that request actually carried rather than +# whatever the GUC happens to say. +# +# Two vectorizers on different models make the difference visible. Recording +# the GUC instead of the per-request value would put the same model on both. + +use strict; +use warnings; + +# See the comment in 001_worker_coverage.pl about loading these at compile time. +use PostgreSQL::Test::Cluster; +use PostgreSQL::Test::Utils; +use Test::More; + +use IO::Socket::INET; + +my $dbname = 'model_recorded'; +my $rows = 3; + +# The socket is created before the fork, so the port is known without guessing +# one or waiting for the child to report it. +my $listener = IO::Socket::INET->new( + LocalAddr => '127.0.0.1', + LocalPort => 0, + Proto => 'tcp', + Listen => 16, + ReuseAddr => 1) or die "could not listen: $!"; + +my $port = $listener->sockport; +my $tempdir = PostgreSQL::Test::Utils::tempdir; + +my $server_pid = fork(); +die "fork failed: $!" unless defined $server_pid; + +if ($server_pid == 0) +{ + fake_provider($listener); + exit 0; +} + +$listener->close; + +# Never checked by the socket above, but the provider will not start without it. +my $keyfile = "$tempdir/api_key"; +open my $kf, '>', $keyfile or die "could not write $keyfile: $!"; +print $kf "not-a-real-key\n"; +close $kf; +chmod 0600, $keyfile; + +my $node = PostgreSQL::Test::Cluster->new('vectorizer_model_recorded'); +$node->init; +$node->append_conf( + 'postgresql.conf', qq( +shared_preload_libraries = 'pgedge_vectorizer' +pgedge_vectorizer.worker_poll_interval = 500 +pgedge_vectorizer.provider = 'voyage' +pgedge_vectorizer.model = 'inherited-model' +pgedge_vectorizer.api_url = 'http://127.0.0.1:$port/v1' +pgedge_vectorizer.api_key_file = '$keyfile' +pgedge_vectorizer.batch_size = 25 +max_worker_processes = 16 +)); + +$node->start; + +$node->safe_psql('postgres', "CREATE DATABASE $dbname"); +$node->safe_psql($dbname, 'CREATE EXTENSION vector'); +$node->safe_psql($dbname, 'CREATE EXTENSION pgedge_vectorizer'); + +$node->safe_psql($dbname, + 'CREATE TABLE inherits (id BIGSERIAL PRIMARY KEY, body TEXT)'); +$node->safe_psql($dbname, + 'CREATE TABLE pinned (id BIGSERIAL PRIMARY KEY, body TEXT)'); + +$node->safe_psql($dbname, + q(SELECT pgedge_vectorizer.enable_vectorization('inherits', 'body', + embedding_dimension => 3))); +$node->safe_psql($dbname, + q(SELECT pgedge_vectorizer.enable_vectorization('pinned', 'body', + embedding_dimension => 3, + model => 'pinned-model'))); + +# Name the database only once it is ready to be serviced, so that no worker can +# arrive before the extension exists and take its five second backoff instead. +$node->append_conf('postgresql.conf', + "pgedge_vectorizer.databases = '$dbname'\n"); +$node->reload; + +$node->safe_psql($dbname, qq( +BEGIN; +INSERT INTO inherits (body) + SELECT 'inherited chunk ' || g FROM generate_series(1, $rows) g; +INSERT INTO pinned (body) + SELECT 'pinned chunk ' || g FROM generate_series(1, $rows) g; +COMMIT; +)); + +my $wanted = $rows * 2; +my $deadline = time() + 30; +my $completed = 0; + +while (time() < $deadline) +{ + $completed = $node->safe_psql($dbname, + "SELECT count(*) FROM pgedge_vectorizer.queue WHERE status = 'completed'"); + + last if $completed == $wanted; + + sleep 1; +} + +is($completed, $wanted, 'both vectorizers drain'); + +# What each table's rows say produced them. +is($node->safe_psql($dbname, + q(SELECT string_agg(DISTINCT embedding_provider || '/' || + embedding_model, ' ') + FROM inherits_body_chunks WHERE embedding IS NOT NULL)), + 'voyage/inherited-model', + 'the inheriting table records the GUC model it was embedded with'); + +is($node->safe_psql($dbname, + q(SELECT string_agg(DISTINCT embedding_provider || '/' || + embedding_model, ' ') + FROM pinned_body_chunks WHERE embedding IS NOT NULL)), + 'voyage/pinned-model', + 'the pinned table records its own model, not the GUC'); + +# Every row, not merely some: a missed write would show as a NULL here. +is($node->safe_psql($dbname, + q(SELECT count(*) FROM inherits_body_chunks + WHERE embedding IS NOT NULL AND embedding_model IS NULL) + ), + '0', 'no embedded chunk is left without a recorded model'); + +# And what the report makes of it: everything current, nothing drifted, and +# nothing unknown, since every row was written by this worker. +is($node->safe_psql($dbname, + q(SELECT string_agg(source_table || ' ' || chunks_embedded || ' ' || + chunks_current || ' ' || chunks_other_model || + ' ' || chunks_model_unknown, E'\n' + ORDER BY source_table) + FROM pgedge_vectorizer.embedding_model_status())), + "inherits $rows $rows 0 0\npinned $rows $rows 0 0", + 'the report sees every chunk as current'); + +# Nothing to redo, so reembed() queues nothing and disturbs nothing. +is($node->safe_psql($dbname, + q(SELECT pgedge_vectorizer.reembed('pinned', 'body', + embedding_dimension => 3))), + '0', 'reembed() queues nothing when every chunk is already current'); + +$node->stop; + +kill 'TERM', $server_pid; +waitpid $server_pid, 0; + +done_testing(); + +# Answer every request with an embedding per input. The model each request +# names does not matter here; what is being checked is what the worker writes +# to the chunk table afterwards. +sub fake_provider +{ + my ($socket) = @_; + + while (my $conn = $socket->accept()) + { + my $headers = ''; + my $body = ''; + my $length; + my $inputs = 0; + my $texts; + my $payload; + + $conn->autoflush(1); + + while (my $line = <$conn>) + { + $headers .= $line; + last if $line =~ /^\r?\n\z/; + } + + # read() can come back short of Content-Length on a socket, and a + # partial body would count the wrong number of inputs. + ($length) = $headers =~ /^content-length:\s*(\d+)/im; + while ($length && length($body) < $length) + { + my $chunk = ''; + my $got = read $conn, $chunk, $length - length($body); + + die 'truncated request body' unless $got; + $body .= $chunk; + } + + # {"input":["...","..."],"model":"..."}. The chunk text is ours and + # has no quotes in it, so counting quoted strings counts the items. + ($texts) = $body =~ /"input"\s*:\s*\[(.*?)\]/s; + $inputs++ while defined $texts && $texts =~ /"(?:[^"\\]|\\.)*"/g; + + $payload = '{"data":[' + . join(',', ('{"embedding":[0.1,0.2,0.3]}') x $inputs) + . ']}'; + $conn->print("HTTP/1.1 200 OK\r\n" + . "Content-Type: application/json\r\n" + . "Content-Length: " . length($payload) . "\r\n" + . "Connection: close\r\n\r\n" + . $payload); + + $conn->close; + } +}