diff --git a/Makefile b/Makefile index 0fb1cb3..8f22346 100644 --- a/Makefile +++ b/Makefile @@ -2,7 +2,7 @@ # Supports PostgreSQL 14-17 EXTENSION = pgedge_vectorizer -EXTVERSION = 1.1 +EXTVERSION = 1.2 # Extension module and data files MODULE_big = $(EXTENSION) @@ -24,7 +24,9 @@ OBJS = src/pgedge_vectorizer.o \ DATA = sql/$(EXTENSION)--$(EXTVERSION).sql \ sql/$(EXTENSION)--1.0.sql \ + sql/$(EXTENSION)--1.1.sql \ sql/$(EXTENSION)--1.0--1.1.sql \ + sql/$(EXTENSION)--1.1--1.2.sql \ sql/$(EXTENSION)--1.0-beta2.sql \ sql/$(EXTENSION)--1.0-beta3.sql \ sql/$(EXTENSION)--1.0-beta1--1.0-beta2.sql \ @@ -32,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 +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 per_table_model REGRESS_OPTS = --inputdir=test --outputdir=test # Documentation files (if any) diff --git a/docs/api_reference.md b/docs/api_reference.md index 076b00d..3c1fd56 100644 --- a/docs/api_reference.md +++ b/docs/api_reference.md @@ -15,7 +15,9 @@ SELECT pgedge_vectorizer.enable_vectorization( chunk_overlap INT DEFAULT NULL, embedding_dimension INT DEFAULT NULL, chunk_table_name TEXT DEFAULT NULL, - source_pk NAME DEFAULT NULL + source_pk NAME DEFAULT NULL, + provider TEXT DEFAULT NULL, + model TEXT DEFAULT NULL ); ``` @@ -29,6 +31,8 @@ SELECT pgedge_vectorizer.enable_vectorization( - `embedding_dimension`: Vector dimension. When NULL (the default), the dimension is auto-detected by making a probe call to the configured embedding provider/model. Can be set explicitly to override auto-detection. - `chunk_table_name`: Custom chunk table name (default: `{table}_{column}_chunks`) - `source_pk`: Primary key column to use as the document identifier in the chunk table. When NULL (the default), the primary key column name and type are auto-detected from the table's primary key index via `pg_index`. Set explicitly to use a specific column (e.g., `'external_id'`). +- `provider`: Embedding provider for this vectorizer. When NULL (the default), `pgedge_vectorizer.provider` is used, and continues to be used as it changes. +- `model`: Embedding model for this vectorizer. When NULL (the default), `pgedge_vectorizer.model` is used, and continues to be used as it changes. Where `embedding_dimension` is not given, the probe asks about this model rather than the configured one. **Primary Key Handling:** @@ -107,15 +111,24 @@ Generate an embedding vector from query text. ```sql SELECT pgedge_vectorizer.generate_embedding( - query_text TEXT + query_text TEXT, + provider TEXT DEFAULT NULL, + model TEXT DEFAULT NULL ); ``` **Parameters:** - `query_text`: Text to generate an embedding for +- `provider`: Provider to use. NULL (the default) uses `pgedge_vectorizer.provider`. +- `model`: Model to use. NULL (the default) uses `pgedge_vectorizer.model`. -Returns: `vector` - The embedding vector using the configured provider +Returns: `vector` - The embedding vector + +A query embedding must come from the same model as the embeddings it is +compared against, so name the model explicitly when searching a chunk table +whose vectorizer pins one. Vectors from two models are not comparable, and +nothing will report an error if you mix them. **Example:** @@ -135,15 +148,65 @@ LIMIT 5; ### detect_embedding_dimension() -Detect the embedding dimension of the currently configured provider/model. +Detect the embedding dimension of a provider and model. ```sql -SELECT pgedge_vectorizer.detect_embedding_dimension(); +SELECT pgedge_vectorizer.detect_embedding_dimension( + provider TEXT DEFAULT NULL, + model TEXT DEFAULT NULL +); ``` +**Parameters:** + +- `provider`: Provider to probe. NULL (the default) uses `pgedge_vectorizer.provider`. +- `model`: Model to probe. NULL (the default) uses `pgedge_vectorizer.model`. + Returns: `INT` - The number of dimensions in the embedding vector -This function generates a probe embedding using the configured provider and model, and returns the dimension of the resulting vector. It is called automatically by `enable_vectorization()` when `embedding_dimension` is not specified. +This function generates a probe embedding and returns the dimension of the result, which means a real request to the provider. It is called automatically by `enable_vectorization()` and `set_embedding_model()` when `embedding_dimension` is not specified. + +### set_embedding_model() + +Change the embedding provider and model for one vectorizer. + +```sql +SELECT pgedge_vectorizer.set_embedding_model( + source_table REGCLASS, + source_column NAME, + model TEXT, + provider TEXT DEFAULT NULL, + embedding_dimension INT DEFAULT NULL, + force_reembed BOOLEAN DEFAULT FALSE +); +``` + +**Parameters:** + +- `source_table`, `source_column`: The vectorizer to change +- `model`: Model to use. NULL means inherit `pgedge_vectorizer.model`. +- `provider`: Provider to use. NULL means inherit `pgedge_vectorizer.provider`. +- `embedding_dimension`: Dimension of the new model. When NULL (the default), the new provider and model are probed for it, which is a real request. The chunk table's vector column is altered to match whether or not the vectorizer has any chunks yet, since a column left at the old width would fail every embedding written afterwards. +- `force_reembed`: Whether to clear the existing embeddings and requeue every chunk. Required to change a vectorizer that has any chunks. + +Returns: `BIGINT` - The number of chunks requeued, which is zero unless the re-embed ran + +Both columns are written to exactly what you pass, NULL included, so this is also how a vectorizer goes back to inheriting the GUCs. Where the effective provider and model do not actually change, nothing is requeued. + +Changing a vectorizer that has chunks raises an error unless `force_reembed` is true. With it, every `embedding` is set to NULL, the column's dimension is altered if the new model differs, the vectorizer's queue rows are cleared and every chunk is requeued, all in one transaction. Chunk rows, their token counts, their sparse embeddings and the BM25 statistics are left alone, because none of them depends on the embedding model. + +The refusal triggers on the model changing rather than on the dimension changing. See [Best Practices](best_practices.md) for why, and for what a re-embed costs. + +**Example:** + +```sql +-- Move one table to a local model, re-embedding what is already there +SELECT pgedge_vectorizer.set_embedding_model( + 'articles'::regclass, 'body', 'nomic-embed-text', + provider => 'ollama', + force_reembed => true +); +``` ### retry_failed() @@ -328,6 +391,25 @@ SELECT pgedge_vectorizer.bm25_tokenize(query TEXT); Returns: `TEXT[]` -- Array of distinct non-stopword terms. +### count_tokens() + +Approximate the number of tokens in a piece of text. This is the same estimate +the chunking engine uses when it decides where a chunk ends, and it is what +gets stored in the `token_count` column of a chunk table, so it is useful for +working out why a given piece of text chunked the way it did. + +```sql +SELECT pgedge_vectorizer.count_tokens(content TEXT); +``` + +Returns: `INT` -- The estimated token count, or `NULL` for `NULL` input. + +The estimate counts UTF-8 characters and divides by four, rounding up, which +is a reasonable rule of thumb for English prose but no more than that: text +that tokenises unusually, such as code, dense punctuation or languages other +than English, will be some way out. Do not use it where an exact count +matters, such as checking a payload against a provider's hard token limit. + ### show_config() Display all pgedge_vectorizer configuration settings. diff --git a/docs/best_practices.md b/docs/best_practices.md index f8cde5c..14c4c2b 100644 --- a/docs/best_practices.md +++ b/docs/best_practices.md @@ -10,6 +10,31 @@ Proper chunking is essential for effective vector search because it balances sem - An overlap of 10-20% (50-100 tokens) provides good context between adjacent chunks. - Use a token-based strategy for general purpose content and the markdown strategy for structured documents. +**Changing an embedding model** + +Changing the model for a table that already has embeddings is not free, and +`set_embedding_model()` refuses to do it silently for that reason. Passing +`force_reembed => true` clears every embedding and requeues every chunk, so the +whole table is embedded again: against a metered provider that is a bill, and +if the dimension changes it is also a rewrite of the chunk table. The chunks +themselves are not rebuilt, because neither chunk boundaries nor the BM25 +statistics depend on the embedding model, so the sparse embeddings and token +counts survive untouched. + +- The refusal keys on the model changing, not on the dimension changing. Two + models of the same width, such as `text-embedding-3-small` and + `text-embedding-ada-002`, both produce 1536 values, so swapping one for the + other would leave the old vectors in place, correctly shaped and meaningless + beside the new ones. Similarity between two models' vectors is noise, and + nothing else in the system would report a problem, which makes it the more + dangerous of the two cases. +- A model wider than 2000 dimensions cannot be used, because the HNSW index + that `enable_vectorization()` creates does not support one. That rules out + `text-embedding-3-large` at its full 3072, though it can be requested at a + 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. + **Performance** Optimizing performance ensures efficient resource utilization and faster embedding generation. These settings help minimize API costs while maintaining responsive processing speeds. @@ -34,14 +59,12 @@ Effective data management ensures clean operations and provides flexibility when - Use the `reprocess_chunks()` function to queue existing chunks that are missing embeddings. - Use the `recreate_chunks()` function for a complete chunk regeneration, which deletes all existing chunks first. - Each column gets independent chunk tables and triggers, so you can disable them selectively as needed. -- Settle on an embedding model before enabling vectorization, because the - model's dimension is fixed into the chunk table when the table is - created. -- Rebuild the vectorizer with `disable_vectorization(..., - drop_chunk_table => TRUE)` and then `enable_vectorization()` after - changing to a model of a different dimension, repeating this for every - vectorized column. Neither `recreate_chunks()` nor a - `disable_vectorization()` that keeps the chunk table alters the column, - so neither resolves the mismatch. +- Change a vectorizer's model with `set_embedding_model()`, which alters the + chunk table's vector column for you where the new model is a different + width. Dropping the chunk table and enabling vectorization again also works, + but throws away chunk rows, sparse embeddings and BM25 statistics that were + never wrong, and has to be repeated for every vectorized column. + `recreate_chunks()` is not an alternative: it rebuilds the chunks and leaves + the column exactly as it was. - Budget for the provider cost of re-embedding an entire table before changing the model on a populated one. diff --git a/docs/changelog.md b/docs/changelog.md index 384a4e4..88b25d2 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -6,6 +6,47 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). ## [Unreleased] +### Fixed + +- The `token_count` recorded for each chunk is now computed the same way + everywhere. The chunking code in C rounds its four-characters-per-token + estimate up, whilst the plpgsql paths that actually write the column + open-coded the same estimate as `length(chunk_text) / 4`, which truncates, + so the two disagreed by a token on most chunks. Because `token_count` feeds + the BM25 document-length normalisation, hybrid search scored chunks written + by the trigger slightly differently from chunks written by the C chunker. + Both now call the new `count_tokens()` function, so there is one definition + of the rule. + +### Added + +- `pgedge_vectorizer.count_tokens(text)`, which exposes the chunking engine's + token estimate so you can see why a piece of text chunked the way it did. +- A per-vectorizer embedding provider and model + ([#27](https://github.com/pgEdge/pgedge-vectorizer/issues/27)). Each + vectorizer may now name its own, through new `provider` and `model` + parameters on `enable_vectorization()` or through the new + `set_embedding_model()`, so one table can be embedded locally whilst another + goes to a hosted provider. Both default to `pgedge_vectorizer.provider` and + `pgedge_vectorizer.model`, so an existing installation is unaffected. + `set_embedding_model()` refuses to change a vectorizer that already has + embeddings unless `force_reembed` is passed, because vectors from two models + are not comparable and mixing them degrades search without failing. +- `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. + +### Changed + +- `generate_embedding(NULL)` now raises an error rather than returning NULL. + The function always meant to reject a NULL query, and said so in its own + code, but was declared `STRICT`, which returned NULL before that check could + run. It can no longer be `STRICT`, because a NULL provider or model has to + reach the function to mean "use the GUC". +- The chunk tables that `disable_vectorization()` drops are now processed in a + defined order, so a disable that covers several columns reports them the same + way twice. + ## [1.1] - 2026-08-28 ### Changed diff --git a/docs/configuration.md b/docs/configuration.md index 5b595e0..d3ecdeb 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -28,6 +28,75 @@ These settings configure the connection to your embedding provider, including th [Troubleshooting](troubleshooting.md) document describes how to recover. +### Per-vectorizer provider and model + +The settings above are the defaults for the whole database, which is the right +thing when every table wants the same embeddings, and the wrong thing when they +do not: a table of short product titles and a table of long technical documents +are rarely well served by one model, and you may want one table embedded +locally through Ollama whilst another goes to a hosted provider. A vectorizer +can therefore name its own provider and model, and falls back to the settings +above where it does not. + +Pin them when the vectorizer is created: + +```sql +SELECT pgedge_vectorizer.enable_vectorization( + 'articles'::regclass, 'body', + provider => 'ollama', + model => 'nomic-embed-text' +); +``` + +Or change them afterwards with `set_embedding_model()`, which takes the model +first because that is the argument you usually want: + +```sql +SELECT pgedge_vectorizer.set_embedding_model( + 'articles'::regclass, 'body', 'nomic-embed-text', provider => 'ollama'); +``` + +Both settings live in `pgedge_vectorizer.vectorizers` as nullable columns, +where NULL means inherit. Inheritance is resolved when the work runs rather +than copied at creation, so a vectorizer that inherits follows the GUC as the +GUC changes. + +`set_embedding_model()` always writes both columns to exactly what you pass, +and both default to NULL, so the shortest call resets both to inheriting: + +```sql +-- Back to inheriting the provider and the model +SELECT pgedge_vectorizer.set_embedding_model('articles'::regclass, 'body', NULL); +``` + +That cuts both ways: to change only the model whilst keeping a pinned +provider, name the provider again, or it reverts to inheriting alongside the +model. + +```sql +-- Keep the pinned provider, change only the model +SELECT pgedge_vectorizer.set_embedding_model( + 'articles'::regclass, 'body', 'mxbai-embed-large', provider => 'ollama'); +``` + +Either call needs `force_reembed => true` if the vectorizer already has +embeddings and the effective model actually moves, as below. + +`set_embedding_model()` refuses to change a vectorizer that already has +embeddings unless you pass `force_reembed => true`, which clears every +embedding and requeues every chunk. See +[Best Practices](best_practices.md) for what that costs and why the refusal +is not limited to changes of dimension. + +!!! warning "Changing the GUC still moves every inheriting vectorizer" + + The refusal above protects a vectorizer that has pinned its model. A + vectorizer that inherits has not, so changing + `pgedge_vectorizer.model` globally re-points every inheriting table at + once, with no guard and no re-embed, exactly as it did before this + setting existed. Pin the model on any vectorizer whose embeddings + matter. + ## Worker Settings These settings control the background workers that process the embedding queue, including concurrency, batch sizes, and retry behavior. diff --git a/pgedge_vectorizer.control b/pgedge_vectorizer.control index 27245d6..e4f4598 100644 --- a/pgedge_vectorizer.control +++ b/pgedge_vectorizer.control @@ -1,6 +1,6 @@ # pgedge_vectorizer extension comment = 'Asynchronous text chunking and vectorization for PostgreSQL' -default_version = '1.1' +default_version = '1.2' module_pathname = '$libdir/pgedge_vectorizer' relocatable = false requires = 'vector' diff --git a/sql/pgedge_vectorizer--1.1--1.2.sql b/sql/pgedge_vectorizer--1.1--1.2.sql new file mode 100644 index 0000000..001ab6c --- /dev/null +++ b/sql/pgedge_vectorizer--1.1--1.2.sql @@ -0,0 +1,997 @@ +-- pgedge_vectorizer extension +-- Version 1.2 +-- +-- Asynchronous text chunking and vectorization for PostgreSQL + +-- complain if script is sourced in psql, rather than via CREATE EXTENSION +\echo Use "ALTER EXTENSION pgedge_vectorizer UPDATE TO '1.2'" to load this file. \quit + +--------------------------------------------------------------------------- +-- Per-vectorizer provider and model +-- +-- NULL in either column means "inherit the GUC at the time the work runs", +-- so an existing installation carries on behaving exactly as it did. +--------------------------------------------------------------------------- + +ALTER TABLE pgedge_vectorizer.vectorizers + ADD COLUMN IF NOT EXISTS provider TEXT, + ADD COLUMN IF NOT EXISTS model TEXT; + +COMMENT ON COLUMN pgedge_vectorizer.vectorizers.provider IS +'Embedding provider for this vectorizer; NULL inherits pgedge_vectorizer.provider'; +COMMENT ON COLUMN pgedge_vectorizer.vectorizers.model IS +'Embedding model for this vectorizer; NULL inherits pgedge_vectorizer.model'; + +--------------------------------------------------------------------------- +-- Approximate token counter, shared with the C chunking code +-- +-- The chunking engine in C has always sized chunks with this estimate, but +-- the plpgsql paths that write the token_count column open-coded it as +-- length(chunk_text) / 4, which truncates where the C code rounds up. The two +-- therefore disagreed by a token on most chunks, and since token_count feeds +-- the BM25 document-length normalisation, hybrid search scored chunks written +-- by the trigger slightly differently from those written by the C chunker. +-- Exposing the C function and calling it from plpgsql leaves one definition +-- of the rule. +--------------------------------------------------------------------------- + +CREATE FUNCTION pgedge_vectorizer.count_tokens( + content TEXT +) RETURNS INT +AS 'MODULE_PATHNAME', 'pgedge_vectorizer_count_tokens' +LANGUAGE C STABLE STRICT; + +COMMENT ON FUNCTION pgedge_vectorizer.count_tokens IS +'Approximate the token count of the given text (UTF-8 characters divided by ' +'four, rounded up). This is the same estimate the chunking engine uses, and ' +'is what gets stored in the token_count column of a chunk table'; + +--------------------------------------------------------------------------- +-- Redefine the three functions that wrote token_count themselves, so that +-- they call count_tokens() instead. Existing rows keep whatever count they +-- were written with; the values are an approximation either way, and a +-- rewrite of every chunk table is not worth a one-token correction. +--------------------------------------------------------------------------- + +--------------------------------------------------------------------------- +-- enable_vectorization() gains provider and model +-- +-- The two new parameters are defaulted, which means CREATE OR REPLACE would +-- define a second function rather than replace the eight-argument one, +-- leaving both in place: a call passing eight arguments could then reach the +-- old body, which knows nothing about the registry's new columns, and even +-- COMMENT ON FUNCTION becomes ambiguous. Drop the old signature first. +--------------------------------------------------------------------------- + +DROP FUNCTION IF EXISTS pgedge_vectorizer.enable_vectorization( + REGCLASS, NAME, TEXT, INT, INT, INT, TEXT, NAME); + +CREATE OR REPLACE FUNCTION pgedge_vectorizer.enable_vectorization( + source_table REGCLASS, + source_column NAME, + chunk_strategy TEXT DEFAULT NULL, + chunk_size INT DEFAULT NULL, + chunk_overlap INT DEFAULT NULL, + embedding_dimension INT DEFAULT NULL, + chunk_table_name TEXT DEFAULT NULL, + source_pk NAME DEFAULT NULL, + provider TEXT DEFAULT NULL, + model TEXT DEFAULT NULL +) RETURNS VOID AS $$ +DECLARE + chunk_table TEXT; + trigger_name TEXT; + actual_strategy TEXT; + actual_chunk_size INT; + actual_chunk_overlap INT; + pk_col_type TEXT; + pk_count INT; +BEGIN + -- Use defaults from GUC if not provided + actual_strategy := COALESCE(chunk_strategy, + current_setting('pgedge_vectorizer.default_chunk_strategy')); + actual_chunk_size := COALESCE(chunk_size, + current_setting('pgedge_vectorizer.default_chunk_size')::INT); + actual_chunk_overlap := COALESCE(chunk_overlap, + current_setting('pgedge_vectorizer.default_chunk_overlap')::INT); + + -- Auto-detect embedding dimension from configured model if not specified + IF embedding_dimension IS NULL THEN + -- Probe the model this vectorizer will actually use, which is not + -- necessarily the one the GUCs name. + embedding_dimension := pgedge_vectorizer.detect_embedding_dimension( + enable_vectorization.provider, enable_vectorization.model); + RAISE NOTICE 'Auto-detected embedding dimension: %', embedding_dimension; + END IF; + + -- Detect PK column count to reject composite PKs + SELECT count(*) + INTO pk_count + FROM pg_index i + JOIN pg_attribute a ON a.attrelid = i.indrelid + AND a.attnum = ANY(i.indkey) + WHERE i.indrelid = source_table + AND i.indisprimary; + + IF pk_count = 0 AND source_pk IS NULL THEN + RAISE EXCEPTION 'Table % has no primary key. Use the source_pk parameter to specify the column to use as document identifier.', + source_table; + END IF; + + IF pk_count > 1 AND source_pk IS NULL THEN + RAISE EXCEPTION 'Table % has a composite primary key (% columns), which is not supported by auto-detection. Use the source_pk parameter to specify a single column.', + source_table, pk_count; + END IF; + + -- Auto-detect PK column name and type if source_pk not specified + IF source_pk IS NULL THEN + SELECT a.attname, format_type(a.atttypid, a.atttypmod) + INTO source_pk, pk_col_type + FROM pg_index i + JOIN pg_attribute a ON a.attrelid = i.indrelid + AND a.attnum = ANY(i.indkey) + WHERE i.indrelid = source_table + AND i.indisprimary; + ELSE + -- User specified a column; look up its type + SELECT format_type(a.atttypid, a.atttypmod) + INTO pk_col_type + FROM pg_attribute a + WHERE a.attrelid = source_table + AND a.attname = source_pk + AND NOT a.attisdropped; + + IF pk_col_type IS NULL THEN + RAISE EXCEPTION 'Column "%" does not exist on table %', + source_pk, source_table; + END IF; + END IF; + + RAISE NOTICE 'Using primary key column: % (%)', source_pk, pk_col_type; + + -- Determine chunk table name. + -- Include source schema in the generated identifier text to avoid + -- collisions when two schemas have the same relname. + chunk_table := COALESCE(chunk_table_name, + source_table::TEXT || '_' || source_column || '_chunks'); + + -- Create chunks table + -- Note: pk_col_type uses %s (not %I) because format_type() returns + -- canonical SQL type names (e.g. "character varying(26)") that would + -- be incorrectly double-quoted by %I. This value is system-controlled. + EXECUTE format(' + CREATE TABLE IF NOT EXISTS %I ( + id BIGSERIAL PRIMARY KEY, + source_id %s NOT NULL, + chunk_index INT NOT NULL, + content TEXT NOT NULL, + token_count INT, + embedding vector(%s), + sparse_embedding sparsevec(65536), + 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). + EXECUTE format(' + ALTER TABLE %I + ADD COLUMN IF NOT EXISTS sparse_embedding sparsevec(65536)', + chunk_table); + + -- Create vector index for similarity search + EXECUTE format(' + CREATE INDEX IF NOT EXISTS %I ON %I + USING hnsw (embedding vector_cosine_ops)', + chunk_table || '_embedding_idx', chunk_table); + + -- Create index on source_id for joins + EXECUTE format(' + CREATE INDEX IF NOT EXISTS %I ON %I (source_id)', + chunk_table || '_source_id_idx', chunk_table); + + -- Create HNSW index on sparse_embedding for fast sparse search + EXECUTE format(' + CREATE INDEX IF NOT EXISTS %I ON %I + USING hnsw (sparse_embedding sparsevec_ip_ops) + WHERE sparse_embedding IS NOT NULL', + chunk_table || '_sparse_idx', chunk_table); + + -- Create BM25 IDF statistics table for this chunk table. + -- Only doc_freq is stored; the IDF weight is computed on read. + EXECUTE format(' + CREATE TABLE IF NOT EXISTS %I ( + term TEXT PRIMARY KEY, + doc_freq INT NOT NULL DEFAULT 1, + updated_at TIMESTAMPTZ DEFAULT now() + )', chunk_table || '_idf_stats'); + + -- Register in vectorizers table for hybrid_search() lookups. + -- Use EXECUTE...USING to avoid PL/pgSQL variable/column ambiguity. + EXECUTE + 'INSERT INTO pgedge_vectorizer.vectorizers + (source_table, source_column, chunk_table, source_pk, pk_type, + provider, model) + VALUES ($1, $2, $3, $4, $5, $6, $7) + ON CONFLICT (source_table, source_column) + DO UPDATE SET chunk_table = EXCLUDED.chunk_table, + source_pk = EXCLUDED.source_pk, + pk_type = EXCLUDED.pk_type, + provider = EXCLUDED.provider, + model = EXCLUDED.model' + USING source_table::TEXT, source_column, chunk_table, source_pk, pk_col_type, + enable_vectorization.provider, enable_vectorization.model; + + -- Create trigger to chunk and queue on insert/update + trigger_name := source_table::TEXT || '_' || source_column || '_vectorization_trigger'; + + EXECUTE format(' + CREATE OR REPLACE TRIGGER %I + AFTER INSERT OR UPDATE ON %s + FOR EACH ROW + EXECUTE FUNCTION pgedge_vectorizer.vectorization_trigger(%L, %L, %L, %L, %L, %L, %L)', + trigger_name, source_table, + source_column, chunk_table, actual_strategy, + actual_chunk_size, actual_chunk_overlap, source_pk, pk_col_type); + + -- Clean up derived data when source rows are deleted. Statement-level with + -- a transition table so that bulk deletes do not degenerate into per-row + -- work. + EXECUTE format(' + CREATE OR REPLACE TRIGGER %I + AFTER DELETE ON %s + REFERENCING OLD TABLE AS old_rows + FOR EACH STATEMENT + EXECUTE FUNCTION pgedge_vectorizer.vectorization_delete_trigger(%L, %L, %L, %L)', + pgedge_vectorizer.cleanup_trigger_name( + source_table::TEXT, source_column, '_vectorization_delete_trigger'), + source_table, + source_column, chunk_table, source_pk, pk_col_type); + + -- Clean up when the whole source table is truncated. + EXECUTE format(' + CREATE OR REPLACE TRIGGER %I + AFTER TRUNCATE ON %s + FOR EACH STATEMENT + EXECUTE FUNCTION pgedge_vectorizer.vectorization_truncate_trigger(%L)', + pgedge_vectorizer.cleanup_trigger_name( + source_table::TEXT, source_column, '_vectorization_truncate_trigger'), + source_table, chunk_table); + + RAISE NOTICE 'Vectorization enabled: % -> %', source_table, chunk_table; + RAISE NOTICE 'Strategy: %, chunk_size: %, overlap: %', + actual_strategy, actual_chunk_size, actual_chunk_overlap; + + -- Process existing rows + DECLARE + row_record RECORD; + doc_content TEXT; + chunks TEXT[]; + chunk_text TEXT; + i INT; + chunk_id BIGINT; + needs_embedding BOOLEAN; + needs_sparse BOOLEAN; + rows_processed INT := 0; + BEGIN + RAISE NOTICE 'Processing existing rows...'; + + -- pk_val is cast to text here so that row_record.pk_val is always the + -- same type across every call to this function within a session, + -- regardless of the source table's actual primary key type. PL/pgSQL + -- fixes the parameter type of a RECORD field the first time a dynamic + -- EXECUTE ... USING statement evaluates it, and reusing that same + -- statement later with a differently-typed record field fails with + -- "type of parameter N does not match that when preparing the plan" + -- (issue #39). Casting at the source, rather than at each USING site, + -- is required: PostgreSQL still binds the RECORD field's own runtime + -- type before any cast written into the later query text is applied. + FOR row_record IN EXECUTE format('SELECT %I::text as pk_val, %I as content FROM %s WHERE %I IS NOT NULL AND %I != ''''', + source_pk, source_column, source_table, source_column, source_column) + LOOP + doc_content := row_record.content; + + -- Chunk the document + chunks := pgedge_vectorizer.chunk_text(doc_content, actual_strategy, actual_chunk_size, actual_chunk_overlap); + + -- Insert chunks and queue for embedding + FOR i IN 1..array_length(chunks, 1) LOOP + chunk_text := chunks[i]; + + -- Insert or update chunk (only clear embedding if content changed). + -- pk_col_type uses %s: value from format_type() is system-controlled + -- (see the comment where the chunk table is created, above). + -- $1::%s casts pk_val, now always text, back to the source + -- table's actual primary key type. + EXECUTE format(' + INSERT INTO %I (source_id, chunk_index, content, token_count) + VALUES ($1::%s, $2, $3, $4) + ON CONFLICT (source_id, chunk_index) + DO UPDATE SET content = EXCLUDED.content, + token_count = EXCLUDED.token_count, + embedding = CASE + WHEN %I.content = EXCLUDED.content THEN %I.embedding + ELSE NULL + END, + sparse_embedding = CASE + WHEN %I.content = EXCLUDED.content THEN %I.sparse_embedding + ELSE NULL + END, + updated_at = NOW() + RETURNING id, + (embedding IS NULL) AS needs_embedding, + (sparse_embedding IS NULL) AS needs_sparse', + chunk_table, pk_col_type, chunk_table, chunk_table, chunk_table, chunk_table) + USING row_record.pk_val, i, chunk_text, + pgedge_vectorizer.count_tokens(chunk_text) + INTO chunk_id, needs_embedding, needs_sparse; + + -- Queue if dense or sparse work is needed. + IF needs_embedding OR needs_sparse THEN + INSERT INTO pgedge_vectorizer.queue (chunk_id, chunk_table, content, metadata, max_attempts) + VALUES ( + chunk_id, + chunk_table, + chunk_text, + CASE + WHEN NOT needs_embedding AND needs_sparse + THEN jsonb_build_object('sparse_only', true) + ELSE NULL + END, + current_setting('pgedge_vectorizer.max_retries')::INT + ); + END IF; + END LOOP; + + -- Remove queue entries for stale high-index chunks before deleting them. + -- Only targets 'pending'/'failed'; 'processing' items are left for the + -- worker to handle gracefully via its SPI_processed == 0 check. + -- pk_col_type uses %s: value from format_type() is system-controlled + EXECUTE format( + 'DELETE FROM pgedge_vectorizer.queue + WHERE chunk_table = %L + AND chunk_id IN ( + SELECT id FROM %I WHERE source_id = $1::%s AND chunk_index > $2 + ) + AND status IN (''pending'', ''failed'')', + chunk_table, chunk_table, pk_col_type) + USING row_record.pk_val, COALESCE(array_length(chunks, 1), 0); + + -- Remove any stale chunks beyond the new chunk count + -- pk_col_type uses %s: value from format_type() is system-controlled + EXECUTE format('DELETE FROM %I WHERE source_id = $1::%s AND chunk_index > $2', + chunk_table, pk_col_type) + USING row_record.pk_val, COALESCE(array_length(chunks, 1), 0); + + rows_processed := rows_processed + 1; + END LOOP; + + RAISE NOTICE 'Processed % existing rows', rows_processed; + END; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.enable_vectorization IS +'Enable automatic chunking and vectorization for a table column'; + +CREATE OR REPLACE FUNCTION pgedge_vectorizer.vectorization_trigger() +RETURNS TRIGGER AS $$ +DECLARE + content_col TEXT; + chunk_table TEXT; + strategy TEXT; + chunk_sz INT; + overlap INT; + pk_col TEXT; + pk_type TEXT; + doc_content TEXT; + chunks TEXT[]; + chunk_text TEXT; + i INT; + chunk_id BIGINT; + source_id_val TEXT; + deleted_chunks_count INT := 0; +BEGIN + -- Extract trigger arguments + content_col := TG_ARGV[0]; + chunk_table := TG_ARGV[1]; + strategy := TG_ARGV[2]; + chunk_sz := TG_ARGV[3]::INT; + overlap := TG_ARGV[4]::INT; + pk_col := COALESCE(TG_ARGV[5], 'id'); + pk_type := COALESCE(TG_ARGV[6], 'bigint'); + + -- Get source document ID + EXECUTE format('SELECT ($1).%I', pk_col) USING NEW INTO source_id_val; + + -- Get document content + EXECUTE format('SELECT $1.%I', content_col) USING NEW INTO doc_content; + + -- Trim whitespace for empty check + IF doc_content IS NOT NULL THEN + doc_content := trim(doc_content); + END IF; + + -- Skip if content unchanged (on UPDATE) + IF TG_OP = 'UPDATE' THEN + DECLARE + old_content TEXT; + BEGIN + EXECUTE format('SELECT $1.%I', content_col) USING OLD INTO old_content; + IF old_content IS NOT NULL THEN + old_content := trim(old_content); + END IF; + IF doc_content = old_content OR (doc_content IS NULL AND old_content IS NULL) THEN + RETURN NEW; + END IF; + END; + END IF; + + -- On UPDATE, decrement IDF stats for the old document's terms before + -- deleting the old chunks. This prevents doc_freq from drifting upward + -- when the worker later re-increments stats for the new chunks. + IF TG_OP = 'UPDATE' THEN + DECLARE + old_terms TEXT[]; + old_content_for_idf TEXT; + BEGIN + EXECUTE format('SELECT $1.%I', content_col) USING OLD INTO old_content_for_idf; + IF old_content_for_idf IS NOT NULL THEN + old_content_for_idf := trim(old_content_for_idf); + END IF; + IF old_content_for_idf IS NOT NULL AND old_content_for_idf <> '' THEN + old_terms := pgedge_vectorizer.bm25_tokenize(old_content_for_idf); + EXECUTE format( + 'SELECT count(*)::int FROM %I WHERE source_id = $1::%s', + chunk_table, pk_type + ) + INTO deleted_chunks_count + USING source_id_val; + + PERFORM pgedge_vectorizer.bm25_decrement_idf_stats( + chunk_table, old_terms, deleted_chunks_count); + END IF; + END; + END IF; + + -- Delete queue entries for this document's chunks before deleting the chunks. + -- Prevents orphaned queue entries that waste embedding API calls on deleted chunks. + -- Only targets 'pending'/'failed'; 'processing' items are left for the + -- worker to handle gracefully via its SPI_processed == 0 check. + EXECUTE format( + 'DELETE FROM pgedge_vectorizer.queue + WHERE chunk_table = %L + AND chunk_id IN (SELECT id FROM %I WHERE source_id = $1::%s) + AND status IN (''pending'', ''failed'')', + chunk_table, chunk_table, pk_type) + USING source_id_val; + + -- Delete existing chunks for this document + -- pk_type uses %s: value from format_type() is system-controlled (see enable_vectorization) + EXECUTE format('DELETE FROM %I WHERE source_id = $1::%s', chunk_table, pk_type) + USING source_id_val; + + -- Skip if content is NULL or empty (after deleting old chunks) + IF doc_content IS NULL OR doc_content = '' THEN + RETURN NEW; + END IF; + + -- Chunk the document + chunks := pgedge_vectorizer.chunk_text(doc_content, strategy, chunk_sz, overlap); + + -- Insert chunks and queue for embedding + FOR i IN 1..array_length(chunks, 1) LOOP + chunk_text := chunks[i]; + + -- Insert chunk + EXECUTE format(' + INSERT INTO %I (source_id, chunk_index, content, token_count) + VALUES ($1::%s, $2, $3, $4) + RETURNING id', chunk_table, pk_type) + USING source_id_val, i, chunk_text, + pgedge_vectorizer.count_tokens(chunk_text) + INTO chunk_id; + + -- Queue for embedding + INSERT INTO pgedge_vectorizer.queue (chunk_id, chunk_table, content, max_attempts) + VALUES (chunk_id, chunk_table, chunk_text, + current_setting('pgedge_vectorizer.max_retries')::INT); + END LOOP; + + -- Notify workers (they will pick up work via polling and SKIP LOCKED) + PERFORM pg_notify('pgedge_vectorizer_queue', source_id_val::TEXT); + + RETURN NEW; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.vectorization_trigger IS +'Trigger function that chunks text and queues for vectorization'; + +CREATE OR REPLACE FUNCTION pgedge_vectorizer.recreate_chunks( + source_table_name REGCLASS, + source_column_name NAME +) RETURNS INT AS $$ +DECLARE + chunk_table_name TEXT; + rows_affected INT := 0; + trigger_name TEXT; + trigger_exists BOOLEAN; +BEGIN + -- Prefer authoritative mapping from vectorizers registry. + SELECT v.chunk_table + INTO chunk_table_name + FROM pgedge_vectorizer.vectorizers v + WHERE v.source_table = source_table_name::TEXT + AND v.source_column = source_column_name; + + -- Fallback to legacy default naming if no registry row exists. + IF chunk_table_name IS NULL THEN + chunk_table_name := source_table_name::TEXT || '_' || source_column_name || '_chunks'; + END IF; + + -- Verify chunk table exists + IF to_regclass(chunk_table_name) IS NULL THEN + RAISE EXCEPTION 'Chunk table % does not exist. Use enable_vectorization() first.', chunk_table_name; + END IF; + + -- Verify trigger exists + trigger_name := source_table_name::TEXT || '_' || source_column_name || '_vectorization_trigger'; + SELECT EXISTS ( + SELECT 1 FROM pg_trigger t + JOIN pg_class c ON t.tgrelid = c.oid + WHERE c.oid = source_table_name + AND t.tgname = trigger_name + ) INTO trigger_exists; + + IF NOT trigger_exists THEN + RAISE EXCEPTION 'Vectorization trigger % does not exist. Use enable_vectorization() first.', trigger_name; + END IF; + + RAISE NOTICE 'Recreating chunks for %.% -> %', source_table_name, source_column_name, chunk_table_name; + + -- Delete all existing chunks and reset IDF stats. + -- Truncating _idf_stats is safe here because recreate_chunks rebuilds + -- all chunks from scratch; the worker will repopulate IDF stats as it + -- processes the newly queued chunks. + EXECUTE format('DELETE FROM %I', chunk_table_name); + EXECUTE format('TRUNCATE TABLE %I', chunk_table_name || '_idf_stats'); + GET DIAGNOSTICS rows_affected = ROW_COUNT; + RAISE NOTICE 'Deleted % existing chunks', rows_affected; + + -- Delete all queue items for this chunk table (with retry logic) + BEGIN + -- Try to delete with a lock timeout + SET LOCAL lock_timeout = '5s'; + DELETE FROM pgedge_vectorizer.queue WHERE chunk_table = chunk_table_name; + RAISE NOTICE 'Cleared queue for %', chunk_table_name; + EXCEPTION WHEN lock_not_available OR deadlock_detected THEN + -- If we can't get the lock, just mark them for cleanup + RAISE WARNING 'Could not clear queue due to concurrent access, continuing anyway'; + END; + + -- Manually process each row to bypass trigger's unchanged-content optimization + DECLARE + row_record RECORD; + doc_content TEXT; + chunks TEXT[]; + chunk_text TEXT; + i INT; + chunk_id BIGINT; + rows_processed INT := 0; + actual_strategy TEXT; + actual_chunk_size INT; + actual_chunk_overlap INT; + pk_col TEXT; + pk_type TEXT; + BEGIN + -- Get chunking configuration from trigger arguments + -- In PostgreSQL 17+, tgargs is bytea and needs to be decoded + DECLARE + tgargs_array TEXT[]; + BEGIN + SELECT string_to_array(encode(t.tgargs, 'escape'), E'\\000') + INTO tgargs_array + FROM pg_trigger t + JOIN pg_class c ON t.tgrelid = c.oid + WHERE c.oid = source_table_name + AND t.tgname = trigger_name; + + -- Arguments: 1=content_col, 2=chunk_table, 3=strategy, 4=size, 5=overlap, 6=pk_col, 7=pk_type + actual_strategy := tgargs_array[3]; + actual_chunk_size := tgargs_array[4]::INT; + actual_chunk_overlap := tgargs_array[5]::INT; + pk_col := COALESCE(tgargs_array[6], 'id'); + pk_type := COALESCE(tgargs_array[7], 'bigint'); + END; + + RAISE NOTICE 'Re-chunking with strategy=%, size=%, overlap=%', + actual_strategy, actual_chunk_size, actual_chunk_overlap; + + -- pk_val is cast to text so that row_record.pk_val is always the same + -- type across calls in a session, whatever the source table's actual + -- primary key type. See the identical comment in enable_vectorization() + -- for why: PL/pgSQL fixes a RECORD field's parameter type the first + -- time a dynamic EXECUTE ... USING statement evaluates it, and this + -- statement's own "$1::%s" cast below does not protect it, because + -- that cast is applied after PostgreSQL has already bound the record + -- field's raw runtime type (issue #39). + FOR row_record IN EXECUTE format( + 'SELECT %I::text as pk_val, %I as content FROM %s WHERE %I IS NOT NULL AND %I != ''''', + pk_col, source_column_name, source_table_name, source_column_name, source_column_name + ) + LOOP + doc_content := row_record.content; + + -- Chunk the document + chunks := pgedge_vectorizer.chunk_text(doc_content, actual_strategy, actual_chunk_size, actual_chunk_overlap); + + -- Insert chunks and queue for embedding + FOR i IN 1..array_length(chunks, 1) LOOP + chunk_text := chunks[i]; + + -- Insert chunk + -- pk_type uses %s: value from format_type() is system-controlled (see enable_vectorization) + EXECUTE format(' + INSERT INTO %I (source_id, chunk_index, content, token_count) + VALUES ($1::%s, $2, $3, $4) + RETURNING id', chunk_table_name, pk_type) + USING row_record.pk_val, i, chunk_text, + pgedge_vectorizer.count_tokens(chunk_text) + INTO chunk_id; + + -- Queue for embedding + INSERT INTO pgedge_vectorizer.queue (chunk_id, chunk_table, content, max_attempts) + VALUES (chunk_id, chunk_table_name, chunk_text, + current_setting('pgedge_vectorizer.max_retries')::INT); + END LOOP; + + rows_processed := rows_processed + 1; + END LOOP; + + RAISE NOTICE 'Processed % rows', rows_processed; + RETURN rows_processed; + END; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.recreate_chunks IS +'Delete all chunks and recreate from source table (complete rebuild)'; + +--------------------------------------------------------------------------- +-- Provider and model may now be named per call +-- +-- Adding defaulted parameters creates a new function rather than replacing +-- the old one, and the two would then be ambiguous for a caller passing only +-- the arguments they share, so the old forms are dropped first. Neither is +-- STRICT any more: NULL has to reach the C, where it means "fall back to the +-- GUC", and a STRICT function would return NULL before getting there. +--------------------------------------------------------------------------- + +DROP FUNCTION IF EXISTS pgedge_vectorizer.generate_embedding(TEXT); +DROP FUNCTION IF EXISTS pgedge_vectorizer.detect_embedding_dimension(); + +CREATE FUNCTION pgedge_vectorizer.generate_embedding( + query_text TEXT, + provider TEXT DEFAULT NULL, + model TEXT DEFAULT NULL +) RETURNS vector +AS 'MODULE_PATHNAME', 'pgedge_vectorizer_generate_embedding' +LANGUAGE C STABLE; + +COMMENT ON FUNCTION pgedge_vectorizer.generate_embedding IS +'Generate an embedding vector from query text. The provider and model ' +'default to pgedge_vectorizer.provider and pgedge_vectorizer.model'; + +-- Embedding dimension detection function +CREATE FUNCTION pgedge_vectorizer.detect_embedding_dimension( + provider TEXT DEFAULT NULL, + model TEXT DEFAULT NULL +) RETURNS INT +AS 'MODULE_PATHNAME', 'pgedge_vectorizer_detect_embedding_dimension' +LANGUAGE C; + +COMMENT ON FUNCTION pgedge_vectorizer.detect_embedding_dimension IS +'Detect the embedding dimension of the given provider and model, defaulting ' +'to pgedge_vectorizer.provider and pgedge_vectorizer.model'; + +--------------------------------------------------------------------------- +-- disable_vectorization(): drop the chunk tables in a defined order +-- +-- The array of chunk tables to drop was collected with no ORDER BY, so the +-- notices a multi-column disable emits came out in whatever order the scan +-- happened to return, which changed when the registry gained columns. Order +-- by the column name so that the same disable says the same thing twice. +--------------------------------------------------------------------------- + +CREATE OR REPLACE FUNCTION pgedge_vectorizer.disable_vectorization( + source_table REGCLASS, + source_column NAME DEFAULT NULL, + drop_chunk_table BOOLEAN DEFAULT FALSE +) RETURNS VOID AS $$ +DECLARE + trigger_name TEXT; + chunk_table TEXT; + trigger_rec RECORD; + chunk_tables_to_drop TEXT[]; + ct TEXT; +BEGIN + -- If column specified, drop that specific trigger + IF source_column IS NOT NULL THEN + trigger_name := source_table::TEXT || '_' || source_column || '_vectorization_trigger'; + + -- Look up the authoritative chunk table name from the registry so that + -- custom chunk_table_name values (passed to enable_vectorization) are + -- honored; fall back to the default convention only when not registered. + -- Use EXECUTE...USING to avoid variable/column name ambiguity for + -- source_table and source_column (same pattern as the DELETE below). + EXECUTE + 'SELECT v.chunk_table FROM pgedge_vectorizer.vectorizers v + WHERE v.source_table = $1 AND v.source_column = $2' + INTO chunk_table + USING source_table::TEXT, source_column; + + IF chunk_table IS NULL THEN + chunk_table := source_table::TEXT || '_' || source_column || '_chunks'; + END IF; + + -- Drop triggers + EXECUTE format('DROP TRIGGER IF EXISTS %I ON %s', trigger_name, source_table); + EXECUTE format('DROP TRIGGER IF EXISTS %I ON %s', + pgedge_vectorizer.cleanup_trigger_name( + source_table::TEXT, source_column, '_vectorization_delete_trigger'), + source_table); + EXECUTE format('DROP TRIGGER IF EXISTS %I ON %s', + pgedge_vectorizer.cleanup_trigger_name( + source_table::TEXT, source_column, '_vectorization_truncate_trigger'), + source_table); + + -- Remove orphaned queue items for this chunk table + EXECUTE format('DELETE FROM pgedge_vectorizer.queue WHERE chunk_table = %L AND status IN (''pending'', ''processing'')', chunk_table); + + -- Remove from vectorizers registry. + -- Use EXECUTE...USING to avoid PL/pgSQL variable/column + -- name ambiguity for source_table and source_column. + EXECUTE + 'DELETE FROM pgedge_vectorizer.vectorizers + WHERE source_table = $1 AND source_column = $2' + USING source_table::TEXT, source_column; + + -- Optionally drop chunk table and IDF stats table + IF drop_chunk_table THEN + EXECUTE format('DROP TABLE IF EXISTS %I CASCADE', + chunk_table || '_idf_stats'); + EXECUTE format('DROP TABLE IF EXISTS %I CASCADE', chunk_table); + RAISE NOTICE 'Vectorization disabled and chunk table dropped: %', chunk_table; + ELSE + RAISE NOTICE 'Vectorization disabled (chunk table preserved): %', chunk_table; + END IF; + ELSE + -- Drop all vectorization triggers for this table + -- Find vectorization triggers by their trigger function rather than by + -- name pattern. Cleanup trigger names are shortened when the table and + -- column are long, so a shortened name need not begin with the source + -- table text and a LIKE pattern anchored on it would miss them, + -- silently leaving cleanup triggers behind. + FOR trigger_rec IN + SELECT t.tgname + FROM pg_trigger t + WHERE t.tgrelid = source_table + AND NOT t.tgisinternal + AND t.tgfoid IN ( + SELECT p.oid + FROM pg_proc p + JOIN pg_namespace n ON n.oid = p.pronamespace + WHERE n.nspname = 'pgedge_vectorizer' + AND p.proname IN ('vectorization_trigger', + 'vectorization_delete_trigger', + 'vectorization_truncate_trigger')) + LOOP + EXECUTE format('DROP TRIGGER IF EXISTS %I ON %s', trigger_rec.tgname, source_table); + RAISE NOTICE 'Dropped trigger: %', trigger_rec.tgname; + END LOOP; + + -- Collect chunk table names before deleting registry entries. + -- Use EXECUTE...USING to avoid PL/pgSQL variable/column + -- name ambiguity for source_table. + EXECUTE + 'SELECT ARRAY( + SELECT v.chunk_table + FROM pgedge_vectorizer.vectorizers v + WHERE v.source_table = $1 + ORDER BY v.source_column + )' + INTO chunk_tables_to_drop + USING source_table::TEXT; + + -- Remove orphaned queue items for exact chunk tables from registry. + DELETE FROM pgedge_vectorizer.queue q + WHERE q.chunk_table = ANY(COALESCE(chunk_tables_to_drop, '{}')) + AND q.status IN ('pending', 'processing'); + + -- Remove all vectorizer registry entries for this source table + EXECUTE + 'DELETE FROM pgedge_vectorizer.vectorizers WHERE source_table = $1' + USING source_table::TEXT; + + -- Optionally drop all chunk tables and their IDF stats tables + IF drop_chunk_table THEN + FOREACH ct IN ARRAY COALESCE(chunk_tables_to_drop, '{}') LOOP + EXECUTE format('DROP TABLE IF EXISTS %I CASCADE', ct || '_idf_stats'); + EXECUTE format('DROP TABLE IF EXISTS %I CASCADE', ct); + RAISE NOTICE 'Vectorization disabled and chunk table dropped: %', ct; + END LOOP; + END IF; + END IF; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.disable_vectorization IS +'Disable automatic vectorization for a table'; + +--------------------------------------------------------------------------- +-- set_embedding_model(): change a vectorizer's provider and model +-- +-- Both columns are written to exactly what was passed, NULL included, so +-- reverting a table to the global default is a call with a NULL model rather +-- than a separate function, and there is no hidden "leave it alone" state. +-- +-- Changing the model on a populated vectorizer is refused unless the caller +-- asks for the re-embed, and the refusal keys on the model rather than on the +-- dimension. A dimension change is the loud failure and the worker already +-- catches it before writing anything. The quiet one is a change that keeps the +-- same width: text-embedding-3-small and text-embedding-ada-002 are both 1536, +-- so swapping them would leave the old vectors in place, correctly shaped and +-- meaningless beside the new ones, with nothing reporting a problem. +-- +-- The re-embed leaves the chunks themselves alone. Chunking does not depend on +-- the embedding model, since count_tokens() ignores the model it is given, and +-- BM25 is lexical, so the chunk rows, their token counts and their sparse +-- embeddings are all still correct. Only the dense embeddings are wrong, which +-- is why this does not go near recreate_chunks(). +--------------------------------------------------------------------------- + +CREATE FUNCTION pgedge_vectorizer.set_embedding_model( + source_table REGCLASS, + source_column NAME, + model TEXT, + provider TEXT DEFAULT NULL, + embedding_dimension INT DEFAULT NULL, + force_reembed BOOLEAN DEFAULT FALSE +) RETURNS BIGINT AS $$ +DECLARE + v_row RECORD; + chunk_oid OID; + old_provider TEXT; + old_model TEXT; + new_provider TEXT; + new_model TEXT; + chunk_count BIGINT; + new_dim INT; + current_dim INT; + requeued BIGINT := 0; +BEGIN + SELECT r.* INTO v_row + FROM pgedge_vectorizer.vectorizers r + WHERE r.source_table = set_embedding_model.source_table::TEXT + AND r.source_column = set_embedding_model.source_column; + + IF NOT FOUND THEN + RAISE EXCEPTION 'no vectorizer registered for %.%', + set_embedding_model.source_table::TEXT, + set_embedding_model.source_column; + END IF; + + -- 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, + current_setting('pgedge_vectorizer.provider')); + old_model := COALESCE(v_row.model, + current_setting('pgedge_vectorizer.model')); + new_provider := COALESCE(set_embedding_model.provider, + current_setting('pgedge_vectorizer.provider')); + new_model := COALESCE(set_embedding_model.model, + current_setting('pgedge_vectorizer.model')); + + IF old_provider = new_provider AND old_model = new_model THEN + UPDATE pgedge_vectorizer.vectorizers r + SET provider = set_embedding_model.provider, + model = set_embedding_model.model + WHERE r.id = v_row.id; + + RAISE NOTICE 'Effective provider and model unchanged (%/%)', + new_provider, new_model; + RETURN 0; + END IF; + + -- The chunk table's name is one identifier, dot included, so it is quoted + -- rather than parsed as schema.relation. + 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, + set_embedding_model.source_table::TEXT, + set_embedding_model.source_column; + END IF; + + EXECUTE format('SELECT count(*) FROM %s', chunk_oid::REGCLASS) + INTO chunk_count; + + IF chunk_count > 0 AND NOT force_reembed THEN + RAISE EXCEPTION + 'changing the embedding model for %.% would leave % chunks ' + 'embedded with %/% whilst everything after uses %/%', + set_embedding_model.source_table::TEXT, + set_embedding_model.source_column, chunk_count, + old_provider, old_model, new_provider, new_model + USING HINT = 'Pass force_reembed => true to clear every embedding ' + 'and requeue the chunks. Vectors from two models are ' + 'not comparable, so leaving the old ones in place ' + 'would quietly degrade search rather than fail.'; + END IF; + + /* + * The column has to be rewidened whether or not there are chunks. An + * empty vectorizer left at its old width would accept the change happily + * and then fail every embedding the worker tried to write, which is the + * failure this function exists to prevent. + */ + new_dim := COALESCE( + set_embedding_model.embedding_dimension, + pgedge_vectorizer.detect_embedding_dimension( + set_embedding_model.provider, set_embedding_model.model)); + + SELECT a.atttypmod INTO current_dim + FROM pg_attribute a + WHERE a.attrelid = chunk_oid + AND a.attname = 'embedding'; + + IF chunk_count > 0 THEN + -- NULL first: a vector column cannot change width with values in it. + EXECUTE format('UPDATE %s SET embedding = NULL ' + 'WHERE embedding IS NOT NULL', chunk_oid::REGCLASS); + + -- Anything already queued was queued against the old model. + DELETE FROM pgedge_vectorizer.queue q + WHERE q.chunk_table = v_row.chunk_table; + END IF; + + IF new_dim IS DISTINCT FROM current_dim THEN + EXECUTE format('ALTER TABLE %s ALTER COLUMN embedding ' + 'TYPE vector(%s)', chunk_oid::REGCLASS, new_dim); + RAISE NOTICE 'Embedding dimension changed from % to %', + current_dim, new_dim; + END IF; + + UPDATE pgedge_vectorizer.vectorizers r + SET provider = set_embedding_model.provider, + model = set_embedding_model.model + WHERE r.id = v_row.id; + + IF chunk_count > 0 THEN + EXECUTE format( + 'INSERT INTO pgedge_vectorizer.queue ' + ' (chunk_id, chunk_table, content, max_attempts) ' + 'SELECT id, %L, content, %s FROM %s', + v_row.chunk_table, + current_setting('pgedge_vectorizer.max_retries')::INT, + chunk_oid::REGCLASS); + + GET DIAGNOSTICS requeued = ROW_COUNT; + + RAISE NOTICE 'Requeued % chunks for re-embedding with %/%', + requeued, new_provider, new_model; + END IF; + + RETURN requeued; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.set_embedding_model IS +'Set the embedding provider and model for one vectorizer, NULL meaning ' +'inherit the GUC. Refuses to change a populated vectorizer unless ' +'force_reembed is true, in which case every embedding is cleared and every ' +'chunk requeued. Returns the number of chunks requeued'; + diff --git a/sql/pgedge_vectorizer--1.2.sql b/sql/pgedge_vectorizer--1.2.sql new file mode 100644 index 0000000..52f4c6f --- /dev/null +++ b/sql/pgedge_vectorizer--1.2.sql @@ -0,0 +1,1665 @@ +-- pgedge_vectorizer extension +-- Version 1.2 +-- +-- Asynchronous text chunking and vectorization for PostgreSQL + +-- complain if script is sourced in psql, rather than via CREATE EXTENSION +\echo Use "CREATE EXTENSION pgedge_vectorizer" to load this file. \quit + +--------------------------------------------------------------------------- +-- Create schema +--------------------------------------------------------------------------- + +CREATE SCHEMA IF NOT EXISTS pgedge_vectorizer; + +--------------------------------------------------------------------------- +-- Vectorizers registry +-- Tracks which chunk tables have been created for source tables. +-- Used by hybrid_search() to resolve chunk table names. +--------------------------------------------------------------------------- + +CREATE TABLE pgedge_vectorizer.vectorizers ( + id BIGSERIAL PRIMARY KEY, + source_table TEXT NOT NULL, + source_column NAME NOT NULL, + chunk_table TEXT NOT NULL, + source_pk NAME, + pk_type TEXT, + provider TEXT, + model TEXT, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + UNIQUE (source_table, source_column) +); + +COMMENT ON TABLE pgedge_vectorizer.vectorizers IS +'Registry of active vectorizer configurations (source table → chunk table)'; + +COMMENT ON COLUMN pgedge_vectorizer.vectorizers.provider IS +'Embedding provider for this vectorizer; NULL inherits pgedge_vectorizer.provider'; +COMMENT ON COLUMN pgedge_vectorizer.vectorizers.model IS +'Embedding model for this vectorizer; NULL inherits pgedge_vectorizer.model'; + +--------------------------------------------------------------------------- +-- Queue table for async embedding generation +--------------------------------------------------------------------------- + +CREATE TABLE pgedge_vectorizer.queue ( + id BIGSERIAL PRIMARY KEY, + chunk_id BIGINT NOT NULL, -- ID of the chunk in the chunk table + chunk_table TEXT NOT NULL, -- Name of the chunk table + content TEXT NOT NULL, -- Text content to embed + status TEXT NOT NULL DEFAULT 'pending' + CHECK (status IN ('pending', 'processing', 'completed', 'failed')), + attempts INT NOT NULL DEFAULT 0, + max_attempts INT NOT NULL DEFAULT 3, + error_message TEXT, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + processing_started_at TIMESTAMPTZ, + processed_at TIMESTAMPTZ, + next_retry_at TIMESTAMPTZ, + -- Times the provider asked for a slower rate. Counted apart from + -- attempts, which max_attempts bounds, so throttling cannot spend an + -- item's retries. + rate_limit_deferrals INT NOT NULL DEFAULT 0, + metadata JSONB +); + +-- Indexes for efficient queue processing +CREATE INDEX idx_queue_status ON pgedge_vectorizer.queue(status, next_retry_at) + WHERE status IN ('pending', 'failed'); + +CREATE INDEX idx_queue_chunk ON pgedge_vectorizer.queue(chunk_table, chunk_id); + +CREATE INDEX idx_queue_created_at ON pgedge_vectorizer.queue(created_at) + WHERE status = 'pending'; + +--------------------------------------------------------------------------- +-- C function declarations +--------------------------------------------------------------------------- + +-- Chunking function +CREATE FUNCTION pgedge_vectorizer.chunk_text( + content TEXT, + strategy TEXT DEFAULT NULL, + chunk_size INT DEFAULT NULL, + overlap INT DEFAULT NULL +) RETURNS TEXT[] +AS 'MODULE_PATHNAME', 'pgedge_vectorizer_chunk_text_sql' +LANGUAGE C IMMUTABLE STRICT; + +COMMENT ON FUNCTION pgedge_vectorizer.chunk_text IS +'Split text into chunks according to the specified strategy'; + +-- Embedding generation function +CREATE FUNCTION pgedge_vectorizer.generate_embedding( + query_text TEXT, + provider TEXT DEFAULT NULL, + model TEXT DEFAULT NULL +) RETURNS vector +AS 'MODULE_PATHNAME', 'pgedge_vectorizer_generate_embedding' +LANGUAGE C STABLE; + +COMMENT ON FUNCTION pgedge_vectorizer.generate_embedding IS +'Generate an embedding vector from query text. The provider and model ' +'default to pgedge_vectorizer.provider and pgedge_vectorizer.model'; + +-- Embedding dimension detection function +CREATE FUNCTION pgedge_vectorizer.detect_embedding_dimension( + provider TEXT DEFAULT NULL, + model TEXT DEFAULT NULL +) RETURNS INT +AS 'MODULE_PATHNAME', 'pgedge_vectorizer_detect_embedding_dimension' +LANGUAGE C; + +COMMENT ON FUNCTION pgedge_vectorizer.detect_embedding_dimension IS +'Detect the embedding dimension of the given provider and model, defaulting ' +'to pgedge_vectorizer.provider and pgedge_vectorizer.model'; + +-- BM25 query vector function +-- Tokenizes the query and computes a sparse vector using current IDF stats. +-- Used by hybrid_search() for the query-side sparse representation. +CREATE FUNCTION pgedge_vectorizer.bm25_query_vector( + query TEXT, + chunk_table TEXT +) RETURNS sparsevec +AS 'MODULE_PATHNAME', 'pgedge_vectorizer_bm25_query_vector' +LANGUAGE C STRICT; + +COMMENT ON FUNCTION pgedge_vectorizer.bm25_query_vector IS +'Compute a BM25 sparse vector for a query text using IDF stats from the given chunk table'; + +-- BM25 average document length helper +CREATE FUNCTION pgedge_vectorizer.bm25_avg_doc_len( + chunk_table TEXT +) RETURNS FLOAT8 +AS 'MODULE_PATHNAME', 'pgedge_vectorizer_bm25_avg_doc_len' +LANGUAGE C STRICT; + +COMMENT ON FUNCTION pgedge_vectorizer.bm25_avg_doc_len IS +'Return the average document length (in tokens) for the given chunk table'; + +-- BM25 tokenizer (exposed for testing) +CREATE FUNCTION pgedge_vectorizer.bm25_tokenize( + query TEXT +) RETURNS TEXT[] +AS 'MODULE_PATHNAME', 'pgedge_vectorizer_bm25_tokenize' +LANGUAGE C STRICT; + +COMMENT ON FUNCTION pgedge_vectorizer.bm25_tokenize IS +'Tokenize text and return the non-stopword terms (useful for testing)'; + +-- Approximate token counter, shared with the C chunking code +CREATE FUNCTION pgedge_vectorizer.count_tokens( + content TEXT +) RETURNS INT +AS 'MODULE_PATHNAME', 'pgedge_vectorizer_count_tokens' +LANGUAGE C STABLE STRICT; + +COMMENT ON FUNCTION pgedge_vectorizer.count_tokens IS +'Approximate the token count of the given text (UTF-8 characters divided by ' +'four, rounded up). This is the same estimate the chunking engine uses, and ' +'is what gets stored in the token_count column of a chunk table'; + +--------------------------------------------------------------------------- +-- SQL Functions +--------------------------------------------------------------------------- + +-- Build a cleanup trigger name that stays unique within PostgreSQL's 63-byte +-- identifier limit. +-- +-- Two problems have to be solved at once. The names differ only in their final +-- word, so plain truncation would make the delete and truncate names identical +-- for a long enough table and column, and the second CREATE OR REPLACE TRIGGER +-- would silently replace the first. Shortening the readable part instead is not +-- sufficient either, because two columns on a long-named table would then +-- shorten to the same string. +-- +-- So the readable prefix is shortened to fit and a digest of the exact table and +-- column is appended, which keeps the name unique per vectorized column however +-- much of the readable part had to go. +-- +-- Note that a shortened name no longer begins with the source table text, so +-- teardown must not look these up by name pattern. disable_vectorization() +-- finds them by trigger function instead. +CREATE FUNCTION pgedge_vectorizer.cleanup_trigger_name( + p_source_table TEXT, + p_source_column TEXT, + p_suffix TEXT +) RETURNS TEXT AS $$ +DECLARE + digest TEXT; + prefix TEXT; + room INT; +BEGIN + digest := substr(md5(p_source_table || '.' || p_source_column), 1, 8); + + -- Budget in bytes, not characters: the limit is 63 bytes, so a multibyte + -- name measured in characters would overflow and be truncated by the + -- server, cutting into the digest. One byte for the separator. + room := GREATEST(63 - octet_length(p_suffix) - octet_length(digest) - 1, 0); + + -- No more than `room` characters can fit in `room` bytes, so start there + -- and drop whole characters until the prefix is within the byte budget. + prefix := left(p_source_table || '_' || p_source_column, room); + WHILE octet_length(prefix) > room LOOP + prefix := left(prefix, -1); + END LOOP; + + RETURN prefix || '_' || digest || p_suffix; +END; +$$ LANGUAGE plpgsql IMMUTABLE; + +COMMENT ON FUNCTION pgedge_vectorizer.cleanup_trigger_name IS +'Build a cleanup trigger name that remains unique per vectorized column within the identifier length limit'; + +-- Enable vectorization for a table/column +CREATE FUNCTION pgedge_vectorizer.enable_vectorization( + source_table REGCLASS, + source_column NAME, + chunk_strategy TEXT DEFAULT NULL, + chunk_size INT DEFAULT NULL, + chunk_overlap INT DEFAULT NULL, + embedding_dimension INT DEFAULT NULL, + chunk_table_name TEXT DEFAULT NULL, + source_pk NAME DEFAULT NULL, + provider TEXT DEFAULT NULL, + model TEXT DEFAULT NULL +) RETURNS VOID AS $$ +DECLARE + chunk_table TEXT; + trigger_name TEXT; + actual_strategy TEXT; + actual_chunk_size INT; + actual_chunk_overlap INT; + pk_col_type TEXT; + pk_count INT; +BEGIN + -- Use defaults from GUC if not provided + actual_strategy := COALESCE(chunk_strategy, + current_setting('pgedge_vectorizer.default_chunk_strategy')); + actual_chunk_size := COALESCE(chunk_size, + current_setting('pgedge_vectorizer.default_chunk_size')::INT); + actual_chunk_overlap := COALESCE(chunk_overlap, + current_setting('pgedge_vectorizer.default_chunk_overlap')::INT); + + -- Auto-detect embedding dimension from configured model if not specified + IF embedding_dimension IS NULL THEN + -- Probe the model this vectorizer will actually use, which is not + -- necessarily the one the GUCs name. + embedding_dimension := pgedge_vectorizer.detect_embedding_dimension( + enable_vectorization.provider, enable_vectorization.model); + RAISE NOTICE 'Auto-detected embedding dimension: %', embedding_dimension; + END IF; + + -- Detect PK column count to reject composite PKs + SELECT count(*) + INTO pk_count + FROM pg_index i + JOIN pg_attribute a ON a.attrelid = i.indrelid + AND a.attnum = ANY(i.indkey) + WHERE i.indrelid = source_table + AND i.indisprimary; + + IF pk_count = 0 AND source_pk IS NULL THEN + RAISE EXCEPTION 'Table % has no primary key. Use the source_pk parameter to specify the column to use as document identifier.', + source_table; + END IF; + + IF pk_count > 1 AND source_pk IS NULL THEN + RAISE EXCEPTION 'Table % has a composite primary key (% columns), which is not supported by auto-detection. Use the source_pk parameter to specify a single column.', + source_table, pk_count; + END IF; + + -- Auto-detect PK column name and type if source_pk not specified + IF source_pk IS NULL THEN + SELECT a.attname, format_type(a.atttypid, a.atttypmod) + INTO source_pk, pk_col_type + FROM pg_index i + JOIN pg_attribute a ON a.attrelid = i.indrelid + AND a.attnum = ANY(i.indkey) + WHERE i.indrelid = source_table + AND i.indisprimary; + ELSE + -- User specified a column; look up its type + SELECT format_type(a.atttypid, a.atttypmod) + INTO pk_col_type + FROM pg_attribute a + WHERE a.attrelid = source_table + AND a.attname = source_pk + AND NOT a.attisdropped; + + IF pk_col_type IS NULL THEN + RAISE EXCEPTION 'Column "%" does not exist on table %', + source_pk, source_table; + END IF; + END IF; + + RAISE NOTICE 'Using primary key column: % (%)', source_pk, pk_col_type; + + -- Determine chunk table name. + -- Include source schema in the generated identifier text to avoid + -- collisions when two schemas have the same relname. + chunk_table := COALESCE(chunk_table_name, + source_table::TEXT || '_' || source_column || '_chunks'); + + -- Create chunks table + -- Note: pk_col_type uses %s (not %I) because format_type() returns + -- canonical SQL type names (e.g. "character varying(26)") that would + -- be incorrectly double-quoted by %I. This value is system-controlled. + EXECUTE format(' + CREATE TABLE IF NOT EXISTS %I ( + id BIGSERIAL PRIMARY KEY, + source_id %s NOT NULL, + chunk_index INT NOT NULL, + content TEXT NOT NULL, + token_count INT, + embedding vector(%s), + sparse_embedding sparsevec(65536), + 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). + EXECUTE format(' + ALTER TABLE %I + ADD COLUMN IF NOT EXISTS sparse_embedding sparsevec(65536)', + chunk_table); + + -- Create vector index for similarity search + EXECUTE format(' + CREATE INDEX IF NOT EXISTS %I ON %I + USING hnsw (embedding vector_cosine_ops)', + chunk_table || '_embedding_idx', chunk_table); + + -- Create index on source_id for joins + EXECUTE format(' + CREATE INDEX IF NOT EXISTS %I ON %I (source_id)', + chunk_table || '_source_id_idx', chunk_table); + + -- Create HNSW index on sparse_embedding for fast sparse search + EXECUTE format(' + CREATE INDEX IF NOT EXISTS %I ON %I + USING hnsw (sparse_embedding sparsevec_ip_ops) + WHERE sparse_embedding IS NOT NULL', + chunk_table || '_sparse_idx', chunk_table); + + -- Create BM25 IDF statistics table for this chunk table. + -- Only doc_freq is stored; the IDF weight is computed on read. + EXECUTE format(' + CREATE TABLE IF NOT EXISTS %I ( + term TEXT PRIMARY KEY, + doc_freq INT NOT NULL DEFAULT 1, + updated_at TIMESTAMPTZ DEFAULT now() + )', chunk_table || '_idf_stats'); + + -- Register in vectorizers table for hybrid_search() lookups. + -- Use EXECUTE...USING to avoid PL/pgSQL variable/column ambiguity. + EXECUTE + 'INSERT INTO pgedge_vectorizer.vectorizers + (source_table, source_column, chunk_table, source_pk, pk_type, + provider, model) + VALUES ($1, $2, $3, $4, $5, $6, $7) + ON CONFLICT (source_table, source_column) + DO UPDATE SET chunk_table = EXCLUDED.chunk_table, + source_pk = EXCLUDED.source_pk, + pk_type = EXCLUDED.pk_type, + provider = EXCLUDED.provider, + model = EXCLUDED.model' + USING source_table::TEXT, source_column, chunk_table, source_pk, pk_col_type, + enable_vectorization.provider, enable_vectorization.model; + + -- Create trigger to chunk and queue on insert/update + trigger_name := source_table::TEXT || '_' || source_column || '_vectorization_trigger'; + + EXECUTE format(' + CREATE OR REPLACE TRIGGER %I + AFTER INSERT OR UPDATE ON %s + FOR EACH ROW + EXECUTE FUNCTION pgedge_vectorizer.vectorization_trigger(%L, %L, %L, %L, %L, %L, %L)', + trigger_name, source_table, + source_column, chunk_table, actual_strategy, + actual_chunk_size, actual_chunk_overlap, source_pk, pk_col_type); + + -- Clean up derived data when source rows are deleted. Statement-level with + -- a transition table so that bulk deletes do not degenerate into per-row + -- work. + EXECUTE format(' + CREATE OR REPLACE TRIGGER %I + AFTER DELETE ON %s + REFERENCING OLD TABLE AS old_rows + FOR EACH STATEMENT + EXECUTE FUNCTION pgedge_vectorizer.vectorization_delete_trigger(%L, %L, %L, %L)', + pgedge_vectorizer.cleanup_trigger_name( + source_table::TEXT, source_column, '_vectorization_delete_trigger'), + source_table, + source_column, chunk_table, source_pk, pk_col_type); + + -- Clean up when the whole source table is truncated. + EXECUTE format(' + CREATE OR REPLACE TRIGGER %I + AFTER TRUNCATE ON %s + FOR EACH STATEMENT + EXECUTE FUNCTION pgedge_vectorizer.vectorization_truncate_trigger(%L)', + pgedge_vectorizer.cleanup_trigger_name( + source_table::TEXT, source_column, '_vectorization_truncate_trigger'), + source_table, chunk_table); + + RAISE NOTICE 'Vectorization enabled: % -> %', source_table, chunk_table; + RAISE NOTICE 'Strategy: %, chunk_size: %, overlap: %', + actual_strategy, actual_chunk_size, actual_chunk_overlap; + + -- Process existing rows + DECLARE + row_record RECORD; + doc_content TEXT; + chunks TEXT[]; + chunk_text TEXT; + i INT; + chunk_id BIGINT; + needs_embedding BOOLEAN; + needs_sparse BOOLEAN; + rows_processed INT := 0; + BEGIN + RAISE NOTICE 'Processing existing rows...'; + + -- pk_val is cast to text here so that row_record.pk_val is always the + -- same type across every call to this function within a session, + -- regardless of the source table's actual primary key type. PL/pgSQL + -- fixes the parameter type of a RECORD field the first time a dynamic + -- EXECUTE ... USING statement evaluates it, and reusing that same + -- statement later with a differently-typed record field fails with + -- "type of parameter N does not match that when preparing the plan" + -- (issue #39). Casting at the source, rather than at each USING site, + -- is required: PostgreSQL still binds the RECORD field's own runtime + -- type before any cast written into the later query text is applied. + FOR row_record IN EXECUTE format('SELECT %I::text as pk_val, %I as content FROM %s WHERE %I IS NOT NULL AND %I != ''''', + source_pk, source_column, source_table, source_column, source_column) + LOOP + doc_content := row_record.content; + + -- Chunk the document + chunks := pgedge_vectorizer.chunk_text(doc_content, actual_strategy, actual_chunk_size, actual_chunk_overlap); + + -- Insert chunks and queue for embedding + FOR i IN 1..array_length(chunks, 1) LOOP + chunk_text := chunks[i]; + + -- Insert or update chunk (only clear embedding if content changed). + -- pk_col_type uses %s: value from format_type() is system-controlled + -- (see the comment where the chunk table is created, above). + -- $1::%s casts pk_val, now always text, back to the source + -- table's actual primary key type. + EXECUTE format(' + INSERT INTO %I (source_id, chunk_index, content, token_count) + VALUES ($1::%s, $2, $3, $4) + ON CONFLICT (source_id, chunk_index) + DO UPDATE SET content = EXCLUDED.content, + token_count = EXCLUDED.token_count, + embedding = CASE + WHEN %I.content = EXCLUDED.content THEN %I.embedding + ELSE NULL + END, + sparse_embedding = CASE + WHEN %I.content = EXCLUDED.content THEN %I.sparse_embedding + ELSE NULL + END, + updated_at = NOW() + RETURNING id, + (embedding IS NULL) AS needs_embedding, + (sparse_embedding IS NULL) AS needs_sparse', + chunk_table, pk_col_type, chunk_table, chunk_table, chunk_table, chunk_table) + USING row_record.pk_val, i, chunk_text, + pgedge_vectorizer.count_tokens(chunk_text) + INTO chunk_id, needs_embedding, needs_sparse; + + -- Queue if dense or sparse work is needed. + IF needs_embedding OR needs_sparse THEN + INSERT INTO pgedge_vectorizer.queue (chunk_id, chunk_table, content, metadata, max_attempts) + VALUES ( + chunk_id, + chunk_table, + chunk_text, + CASE + WHEN NOT needs_embedding AND needs_sparse + THEN jsonb_build_object('sparse_only', true) + ELSE NULL + END, + current_setting('pgedge_vectorizer.max_retries')::INT + ); + END IF; + END LOOP; + + -- Remove queue entries for stale high-index chunks before deleting them. + -- Only targets 'pending'/'failed'; 'processing' items are left for the + -- worker to handle gracefully via its SPI_processed == 0 check. + -- pk_col_type uses %s: value from format_type() is system-controlled + EXECUTE format( + 'DELETE FROM pgedge_vectorizer.queue + WHERE chunk_table = %L + AND chunk_id IN ( + SELECT id FROM %I WHERE source_id = $1::%s AND chunk_index > $2 + ) + AND status IN (''pending'', ''failed'')', + chunk_table, chunk_table, pk_col_type) + USING row_record.pk_val, COALESCE(array_length(chunks, 1), 0); + + -- Remove any stale chunks beyond the new chunk count + -- pk_col_type uses %s: value from format_type() is system-controlled + EXECUTE format('DELETE FROM %I WHERE source_id = $1::%s AND chunk_index > $2', + chunk_table, pk_col_type) + USING row_record.pk_val, COALESCE(array_length(chunks, 1), 0); + + rows_processed := rows_processed + 1; + END LOOP; + + RAISE NOTICE 'Processed % existing rows', rows_processed; + END; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.enable_vectorization IS +'Enable automatic chunking and vectorization for a table column'; + +-- Disable vectorization for a table column +CREATE FUNCTION pgedge_vectorizer.disable_vectorization( + source_table REGCLASS, + source_column NAME DEFAULT NULL, + drop_chunk_table BOOLEAN DEFAULT FALSE +) RETURNS VOID AS $$ +DECLARE + trigger_name TEXT; + chunk_table TEXT; + trigger_rec RECORD; + chunk_tables_to_drop TEXT[]; + ct TEXT; +BEGIN + -- If column specified, drop that specific trigger + IF source_column IS NOT NULL THEN + trigger_name := source_table::TEXT || '_' || source_column || '_vectorization_trigger'; + + -- Look up the authoritative chunk table name from the registry so that + -- custom chunk_table_name values (passed to enable_vectorization) are + -- honored; fall back to the default convention only when not registered. + -- Use EXECUTE...USING to avoid variable/column name ambiguity for + -- source_table and source_column (same pattern as the DELETE below). + EXECUTE + 'SELECT v.chunk_table FROM pgedge_vectorizer.vectorizers v + WHERE v.source_table = $1 AND v.source_column = $2' + INTO chunk_table + USING source_table::TEXT, source_column; + + IF chunk_table IS NULL THEN + chunk_table := source_table::TEXT || '_' || source_column || '_chunks'; + END IF; + + -- Drop triggers + EXECUTE format('DROP TRIGGER IF EXISTS %I ON %s', trigger_name, source_table); + EXECUTE format('DROP TRIGGER IF EXISTS %I ON %s', + pgedge_vectorizer.cleanup_trigger_name( + source_table::TEXT, source_column, '_vectorization_delete_trigger'), + source_table); + EXECUTE format('DROP TRIGGER IF EXISTS %I ON %s', + pgedge_vectorizer.cleanup_trigger_name( + source_table::TEXT, source_column, '_vectorization_truncate_trigger'), + source_table); + + -- Remove orphaned queue items for this chunk table + EXECUTE format('DELETE FROM pgedge_vectorizer.queue WHERE chunk_table = %L AND status IN (''pending'', ''processing'')', chunk_table); + + -- Remove from vectorizers registry. + -- Use EXECUTE...USING to avoid PL/pgSQL variable/column + -- name ambiguity for source_table and source_column. + EXECUTE + 'DELETE FROM pgedge_vectorizer.vectorizers + WHERE source_table = $1 AND source_column = $2' + USING source_table::TEXT, source_column; + + -- Optionally drop chunk table and IDF stats table + IF drop_chunk_table THEN + EXECUTE format('DROP TABLE IF EXISTS %I CASCADE', + chunk_table || '_idf_stats'); + EXECUTE format('DROP TABLE IF EXISTS %I CASCADE', chunk_table); + RAISE NOTICE 'Vectorization disabled and chunk table dropped: %', chunk_table; + ELSE + RAISE NOTICE 'Vectorization disabled (chunk table preserved): %', chunk_table; + END IF; + ELSE + -- Drop all vectorization triggers for this table + -- Find vectorization triggers by their trigger function rather than by + -- name pattern. Cleanup trigger names are shortened when the table and + -- column are long, so a shortened name need not begin with the source + -- table text and a LIKE pattern anchored on it would miss them, + -- silently leaving cleanup triggers behind. + FOR trigger_rec IN + SELECT t.tgname + FROM pg_trigger t + WHERE t.tgrelid = source_table + AND NOT t.tgisinternal + AND t.tgfoid IN ( + SELECT p.oid + FROM pg_proc p + JOIN pg_namespace n ON n.oid = p.pronamespace + WHERE n.nspname = 'pgedge_vectorizer' + AND p.proname IN ('vectorization_trigger', + 'vectorization_delete_trigger', + 'vectorization_truncate_trigger')) + LOOP + EXECUTE format('DROP TRIGGER IF EXISTS %I ON %s', trigger_rec.tgname, source_table); + RAISE NOTICE 'Dropped trigger: %', trigger_rec.tgname; + END LOOP; + + -- Collect chunk table names before deleting registry entries. + -- Use EXECUTE...USING to avoid PL/pgSQL variable/column + -- name ambiguity for source_table. + EXECUTE + 'SELECT ARRAY( + SELECT v.chunk_table + FROM pgedge_vectorizer.vectorizers v + WHERE v.source_table = $1 + ORDER BY v.source_column + )' + INTO chunk_tables_to_drop + USING source_table::TEXT; + + -- Remove orphaned queue items for exact chunk tables from registry. + DELETE FROM pgedge_vectorizer.queue q + WHERE q.chunk_table = ANY(COALESCE(chunk_tables_to_drop, '{}')) + AND q.status IN ('pending', 'processing'); + + -- Remove all vectorizer registry entries for this source table + EXECUTE + 'DELETE FROM pgedge_vectorizer.vectorizers WHERE source_table = $1' + USING source_table::TEXT; + + -- Optionally drop all chunk tables and their IDF stats tables + IF drop_chunk_table THEN + FOREACH ct IN ARRAY COALESCE(chunk_tables_to_drop, '{}') LOOP + EXECUTE format('DROP TABLE IF EXISTS %I CASCADE', ct || '_idf_stats'); + EXECUTE format('DROP TABLE IF EXISTS %I CASCADE', ct); + RAISE NOTICE 'Vectorization disabled and chunk table dropped: %', ct; + END LOOP; + END IF; + END IF; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.disable_vectorization IS +'Disable automatic vectorization for a table'; + +--------------------------------------------------------------------------- +-- set_embedding_model(): change a vectorizer's provider and model +-- +-- Both columns are written to exactly what was passed, NULL included, so +-- reverting a table to the global default is a call with a NULL model rather +-- than a separate function, and there is no hidden "leave it alone" state. +-- +-- Changing the model on a populated vectorizer is refused unless the caller +-- asks for the re-embed, and the refusal keys on the model rather than on the +-- dimension. A dimension change is the loud failure and the worker already +-- catches it before writing anything. The quiet one is a change that keeps the +-- same width: text-embedding-3-small and text-embedding-ada-002 are both 1536, +-- so swapping them would leave the old vectors in place, correctly shaped and +-- meaningless beside the new ones, with nothing reporting a problem. +-- +-- The re-embed leaves the chunks themselves alone. Chunking does not depend on +-- the embedding model, since count_tokens() ignores the model it is given, and +-- BM25 is lexical, so the chunk rows, their token counts and their sparse +-- embeddings are all still correct. Only the dense embeddings are wrong, which +-- is why this does not go near recreate_chunks(). +--------------------------------------------------------------------------- + +CREATE FUNCTION pgedge_vectorizer.set_embedding_model( + source_table REGCLASS, + source_column NAME, + model TEXT, + provider TEXT DEFAULT NULL, + embedding_dimension INT DEFAULT NULL, + force_reembed BOOLEAN DEFAULT FALSE +) RETURNS BIGINT AS $$ +DECLARE + v_row RECORD; + chunk_oid OID; + old_provider TEXT; + old_model TEXT; + new_provider TEXT; + new_model TEXT; + chunk_count BIGINT; + new_dim INT; + current_dim INT; + requeued BIGINT := 0; +BEGIN + SELECT r.* INTO v_row + FROM pgedge_vectorizer.vectorizers r + WHERE r.source_table = set_embedding_model.source_table::TEXT + AND r.source_column = set_embedding_model.source_column; + + IF NOT FOUND THEN + RAISE EXCEPTION 'no vectorizer registered for %.%', + set_embedding_model.source_table::TEXT, + set_embedding_model.source_column; + END IF; + + -- 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, + current_setting('pgedge_vectorizer.provider')); + old_model := COALESCE(v_row.model, + current_setting('pgedge_vectorizer.model')); + new_provider := COALESCE(set_embedding_model.provider, + current_setting('pgedge_vectorizer.provider')); + new_model := COALESCE(set_embedding_model.model, + current_setting('pgedge_vectorizer.model')); + + IF old_provider = new_provider AND old_model = new_model THEN + UPDATE pgedge_vectorizer.vectorizers r + SET provider = set_embedding_model.provider, + model = set_embedding_model.model + WHERE r.id = v_row.id; + + RAISE NOTICE 'Effective provider and model unchanged (%/%)', + new_provider, new_model; + RETURN 0; + END IF; + + -- The chunk table's name is one identifier, dot included, so it is quoted + -- rather than parsed as schema.relation. + 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, + set_embedding_model.source_table::TEXT, + set_embedding_model.source_column; + END IF; + + EXECUTE format('SELECT count(*) FROM %s', chunk_oid::REGCLASS) + INTO chunk_count; + + IF chunk_count > 0 AND NOT force_reembed THEN + RAISE EXCEPTION + 'changing the embedding model for %.% would leave % chunks ' + 'embedded with %/% whilst everything after uses %/%', + set_embedding_model.source_table::TEXT, + set_embedding_model.source_column, chunk_count, + old_provider, old_model, new_provider, new_model + USING HINT = 'Pass force_reembed => true to clear every embedding ' + 'and requeue the chunks. Vectors from two models are ' + 'not comparable, so leaving the old ones in place ' + 'would quietly degrade search rather than fail.'; + END IF; + + /* + * The column has to be rewidened whether or not there are chunks. An + * empty vectorizer left at its old width would accept the change happily + * and then fail every embedding the worker tried to write, which is the + * failure this function exists to prevent. + */ + new_dim := COALESCE( + set_embedding_model.embedding_dimension, + pgedge_vectorizer.detect_embedding_dimension( + set_embedding_model.provider, set_embedding_model.model)); + + SELECT a.atttypmod INTO current_dim + FROM pg_attribute a + WHERE a.attrelid = chunk_oid + AND a.attname = 'embedding'; + + IF chunk_count > 0 THEN + -- NULL first: a vector column cannot change width with values in it. + EXECUTE format('UPDATE %s SET embedding = NULL ' + 'WHERE embedding IS NOT NULL', chunk_oid::REGCLASS); + + -- Anything already queued was queued against the old model. + DELETE FROM pgedge_vectorizer.queue q + WHERE q.chunk_table = v_row.chunk_table; + END IF; + + IF new_dim IS DISTINCT FROM current_dim THEN + EXECUTE format('ALTER TABLE %s ALTER COLUMN embedding ' + 'TYPE vector(%s)', chunk_oid::REGCLASS, new_dim); + RAISE NOTICE 'Embedding dimension changed from % to %', + current_dim, new_dim; + END IF; + + UPDATE pgedge_vectorizer.vectorizers r + SET provider = set_embedding_model.provider, + model = set_embedding_model.model + WHERE r.id = v_row.id; + + IF chunk_count > 0 THEN + EXECUTE format( + 'INSERT INTO pgedge_vectorizer.queue ' + ' (chunk_id, chunk_table, content, max_attempts) ' + 'SELECT id, %L, content, %s FROM %s', + v_row.chunk_table, + current_setting('pgedge_vectorizer.max_retries')::INT, + chunk_oid::REGCLASS); + + GET DIAGNOSTICS requeued = ROW_COUNT; + + RAISE NOTICE 'Requeued % chunks for re-embedding with %/%', + requeued, new_provider, new_model; + END IF; + + RETURN requeued; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.set_embedding_model IS +'Set the embedding provider and model for one vectorizer, NULL meaning ' +'inherit the GUC. Refuses to change a populated vectorizer unless ' +'force_reembed is true, in which case every embedding is cleared and every ' +'chunk requeued. Returns the number of chunks requeued'; + +-- Recreate the DELETE and TRUNCATE cleanup triggers for every registered +-- vectorizer. +-- +-- Upgrade scripts only run on a version change, so an installation of an +-- unreleased build that already has vectorized tables would otherwise have no +-- supported way to acquire the new triggers. This is also the repair route if +-- triggers are ever dropped by hand. +-- +-- Returns the number of vectorizers whose triggers were recreated. Entries +-- whose primary key is not recorded are skipped with a warning rather than +-- guessed at, because enable_vectorization() accepts an explicit source_pk that +-- re-detection could get wrong. +CREATE FUNCTION pgedge_vectorizer.refresh_triggers() +RETURNS INT AS $$ +DECLARE + v RECORD; + refreshed INT := 0; +BEGIN + FOR v IN + SELECT source_table, source_column, chunk_table, source_pk, pk_type + FROM pgedge_vectorizer.vectorizers + LOOP + IF to_regclass(v.source_table) IS NULL THEN + RAISE WARNING 'Skipping %: source table no longer exists', v.source_table; + CONTINUE; + END IF; + + IF v.source_pk IS NULL OR v.pk_type IS NULL THEN + RAISE WARNING 'Skipping %.%: primary key not recorded, re-run enable_vectorization() for this column', + v.source_table, v.source_column; + CONTINUE; + END IF; + + EXECUTE format(' + CREATE OR REPLACE TRIGGER %I + AFTER DELETE ON %s + REFERENCING OLD TABLE AS old_rows + FOR EACH STATEMENT + EXECUTE FUNCTION pgedge_vectorizer.vectorization_delete_trigger(%L, %L, %L, %L)', + pgedge_vectorizer.cleanup_trigger_name( + v.source_table, v.source_column, '_vectorization_delete_trigger'), + v.source_table, + v.source_column, v.chunk_table, v.source_pk, v.pk_type); + + EXECUTE format(' + CREATE OR REPLACE TRIGGER %I + AFTER TRUNCATE ON %s + FOR EACH STATEMENT + EXECUTE FUNCTION pgedge_vectorizer.vectorization_truncate_trigger(%L)', + pgedge_vectorizer.cleanup_trigger_name( + v.source_table, v.source_column, '_vectorization_truncate_trigger'), + v.source_table, v.chunk_table); + + refreshed := refreshed + 1; + END LOOP; + + RETURN refreshed; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.refresh_triggers IS +'Recreate the DELETE and TRUNCATE cleanup triggers for all registered vectorizers; returns the number refreshed'; + +-- BM25 IDF stats decrement helper +-- Called by vectorization_trigger before deleting old chunks on UPDATE so that +-- the doc_freq counts for the old document's terms are removed before the +-- worker re-increments them for the new chunks. This prevents permanent +-- overcount of IDF stats when documents are updated. +CREATE OR REPLACE FUNCTION pgedge_vectorizer.bm25_decrement_idf_stats( + p_chunk_table TEXT, + p_terms TEXT[], + p_deleted_chunks_count INT DEFAULT 1 +) RETURNS VOID AS $$ +BEGIN + IF p_terms IS NULL OR array_length(p_terms, 1) IS NULL THEN + RETURN; + END IF; + + IF p_deleted_chunks_count IS NULL OR p_deleted_chunks_count <= 0 THEN + RETURN; + END IF; + + EXECUTE format( + 'UPDATE %I SET' + ' doc_freq = GREATEST(doc_freq - $1, 0),' + ' updated_at = now()' + ' WHERE term = ANY($2)', + p_chunk_table || '_idf_stats') + USING p_deleted_chunks_count, p_terms; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.bm25_decrement_idf_stats IS +'Decrement doc_freq in the _idf_stats table for a set of terms before their source chunks are deleted'; + +-- Statement-level trigger function cleaning up after DELETE on a source table. +-- +-- Statement-level rather than row-level so that a bulk delete does not run three +-- statements plus a tokenisation per row: with the old_rows transition table the +-- queue and chunk deletions are one set-based statement each. Only the BM25 +-- decrement needs a loop, because bm25_decrement_idf_stats() takes one +-- document's terms at a time. +CREATE FUNCTION pgedge_vectorizer.vectorization_delete_trigger() +RETURNS TRIGGER AS $$ +DECLARE + content_col TEXT; + chunk_table TEXT; + pk_col TEXT; + pk_type TEXT; + old_row RECORD; + old_terms TEXT[]; + chunk_count INT; +BEGIN + content_col := TG_ARGV[0]; + chunk_table := TG_ARGV[1]; + pk_col := COALESCE(TG_ARGV[2], 'id'); + pk_type := COALESCE(TG_ARGV[3], 'bigint'); + + -- Decrement BM25 document frequencies first: the loop counts each + -- document's chunks, which is only possible while they still exist. + FOR old_row IN EXECUTE + format('SELECT %I::text AS pk_value, %I::text AS content FROM old_rows', + pk_col, content_col) + LOOP + IF old_row.content IS NULL OR trim(old_row.content) = '' THEN + CONTINUE; + END IF; + + EXECUTE format( + 'SELECT count(*)::int FROM %I WHERE source_id = $1::%s', + chunk_table, pk_type) + INTO chunk_count + USING old_row.pk_value; + + IF chunk_count > 0 THEN + old_terms := pgedge_vectorizer.bm25_tokenize(trim(old_row.content)); + PERFORM pgedge_vectorizer.bm25_decrement_idf_stats( + chunk_table, old_terms, chunk_count); + END IF; + END LOOP; + + -- Remove queue entries for the doomed chunks, so the worker does not spend + -- embedding API calls on them. 'processing' rows are left alone, matching + -- the INSERT/UPDATE path: the worker copes with the chunk having gone. + EXECUTE format( + 'DELETE FROM pgedge_vectorizer.queue + WHERE chunk_table = %L + AND status IN (''pending'', ''failed'') + AND chunk_id IN ( + SELECT c.id FROM %I c + WHERE c.source_id IN (SELECT o.%I::%s FROM old_rows o))', + chunk_table, chunk_table, pk_col, pk_type); + + -- Finally the chunks themselves, which takes the embeddings with them. + EXECUTE format( + 'DELETE FROM %I WHERE source_id IN (SELECT o.%I::%s FROM old_rows o)', + chunk_table, pk_col, pk_type); + + RETURN NULL; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.vectorization_delete_trigger IS +'Statement-level AFTER DELETE trigger removing chunks, queue entries and BM25 statistics for deleted source rows'; + +-- Statement-level trigger function cleaning up after TRUNCATE on a source table. +-- +-- Truncating the source orphans every chunk, so there is no per-row work and no +-- transition table (which TRUNCATE triggers cannot have in any case). Resetting +-- _idf_stats wholesale is right because the corpus becomes empty, and mirrors +-- what recreate_chunks() does when rebuilding from scratch. +CREATE FUNCTION pgedge_vectorizer.vectorization_truncate_trigger() +RETURNS TRIGGER AS $$ +DECLARE + chunk_table TEXT; +BEGIN + chunk_table := TG_ARGV[0]; + + EXECUTE format( + 'DELETE FROM pgedge_vectorizer.queue + WHERE chunk_table = %L AND status IN (''pending'', ''failed'')', + chunk_table); + + EXECUTE format('TRUNCATE TABLE %I', chunk_table); + + IF to_regclass(chunk_table || '_idf_stats') IS NOT NULL THEN + EXECUTE format('TRUNCATE TABLE %I', chunk_table || '_idf_stats'); + END IF; + + RETURN NULL; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.vectorization_truncate_trigger IS +'Statement-level AFTER TRUNCATE trigger emptying the chunk table, its queue entries and its BM25 statistics'; + +-- Trigger function for vectorization +CREATE FUNCTION pgedge_vectorizer.vectorization_trigger() +RETURNS TRIGGER AS $$ +DECLARE + content_col TEXT; + chunk_table TEXT; + strategy TEXT; + chunk_sz INT; + overlap INT; + pk_col TEXT; + pk_type TEXT; + doc_content TEXT; + chunks TEXT[]; + chunk_text TEXT; + i INT; + chunk_id BIGINT; + source_id_val TEXT; + deleted_chunks_count INT := 0; +BEGIN + -- Extract trigger arguments + content_col := TG_ARGV[0]; + chunk_table := TG_ARGV[1]; + strategy := TG_ARGV[2]; + chunk_sz := TG_ARGV[3]::INT; + overlap := TG_ARGV[4]::INT; + pk_col := COALESCE(TG_ARGV[5], 'id'); + pk_type := COALESCE(TG_ARGV[6], 'bigint'); + + -- Get source document ID + EXECUTE format('SELECT ($1).%I', pk_col) USING NEW INTO source_id_val; + + -- Get document content + EXECUTE format('SELECT $1.%I', content_col) USING NEW INTO doc_content; + + -- Trim whitespace for empty check + IF doc_content IS NOT NULL THEN + doc_content := trim(doc_content); + END IF; + + -- Skip if content unchanged (on UPDATE) + IF TG_OP = 'UPDATE' THEN + DECLARE + old_content TEXT; + BEGIN + EXECUTE format('SELECT $1.%I', content_col) USING OLD INTO old_content; + IF old_content IS NOT NULL THEN + old_content := trim(old_content); + END IF; + IF doc_content = old_content OR (doc_content IS NULL AND old_content IS NULL) THEN + RETURN NEW; + END IF; + END; + END IF; + + -- On UPDATE, decrement IDF stats for the old document's terms before + -- deleting the old chunks. This prevents doc_freq from drifting upward + -- when the worker later re-increments stats for the new chunks. + IF TG_OP = 'UPDATE' THEN + DECLARE + old_terms TEXT[]; + old_content_for_idf TEXT; + BEGIN + EXECUTE format('SELECT $1.%I', content_col) USING OLD INTO old_content_for_idf; + IF old_content_for_idf IS NOT NULL THEN + old_content_for_idf := trim(old_content_for_idf); + END IF; + IF old_content_for_idf IS NOT NULL AND old_content_for_idf <> '' THEN + old_terms := pgedge_vectorizer.bm25_tokenize(old_content_for_idf); + EXECUTE format( + 'SELECT count(*)::int FROM %I WHERE source_id = $1::%s', + chunk_table, pk_type + ) + INTO deleted_chunks_count + USING source_id_val; + + PERFORM pgedge_vectorizer.bm25_decrement_idf_stats( + chunk_table, old_terms, deleted_chunks_count); + END IF; + END; + END IF; + + -- Delete queue entries for this document's chunks before deleting the chunks. + -- Prevents orphaned queue entries that waste embedding API calls on deleted chunks. + -- Only targets 'pending'/'failed'; 'processing' items are left for the + -- worker to handle gracefully via its SPI_processed == 0 check. + EXECUTE format( + 'DELETE FROM pgedge_vectorizer.queue + WHERE chunk_table = %L + AND chunk_id IN (SELECT id FROM %I WHERE source_id = $1::%s) + AND status IN (''pending'', ''failed'')', + chunk_table, chunk_table, pk_type) + USING source_id_val; + + -- Delete existing chunks for this document + -- pk_type uses %s: value from format_type() is system-controlled (see enable_vectorization) + EXECUTE format('DELETE FROM %I WHERE source_id = $1::%s', chunk_table, pk_type) + USING source_id_val; + + -- Skip if content is NULL or empty (after deleting old chunks) + IF doc_content IS NULL OR doc_content = '' THEN + RETURN NEW; + END IF; + + -- Chunk the document + chunks := pgedge_vectorizer.chunk_text(doc_content, strategy, chunk_sz, overlap); + + -- Insert chunks and queue for embedding + FOR i IN 1..array_length(chunks, 1) LOOP + chunk_text := chunks[i]; + + -- Insert chunk + EXECUTE format(' + INSERT INTO %I (source_id, chunk_index, content, token_count) + VALUES ($1::%s, $2, $3, $4) + RETURNING id', chunk_table, pk_type) + USING source_id_val, i, chunk_text, + pgedge_vectorizer.count_tokens(chunk_text) + INTO chunk_id; + + -- Queue for embedding + INSERT INTO pgedge_vectorizer.queue (chunk_id, chunk_table, content, max_attempts) + VALUES (chunk_id, chunk_table, chunk_text, + current_setting('pgedge_vectorizer.max_retries')::INT); + END LOOP; + + -- Notify workers (they will pick up work via polling and SKIP LOCKED) + PERFORM pg_notify('pgedge_vectorizer_queue', source_id_val::TEXT); + + RETURN NEW; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.vectorization_trigger IS +'Trigger function that chunks text and queues for vectorization'; + +--------------------------------------------------------------------------- +-- Views for monitoring +--------------------------------------------------------------------------- + +-- Queue status summary +CREATE VIEW pgedge_vectorizer.queue_status AS +SELECT + chunk_table, + status, + COUNT(*) as count, + MIN(created_at) as oldest, + MAX(created_at) as newest, + AVG(EXTRACT(EPOCH FROM (COALESCE(processed_at, NOW()) - created_at))) as avg_processing_time_secs +FROM pgedge_vectorizer.queue +GROUP BY chunk_table, status +ORDER BY chunk_table, status; + +COMMENT ON VIEW pgedge_vectorizer.queue_status IS +'Summary of queue items by table and status'; + +-- Failed items view +CREATE VIEW pgedge_vectorizer.failed_items AS +SELECT + id, + chunk_table, + chunk_id, + attempts, + max_attempts, + error_message, + created_at, + next_retry_at, + LEFT(content, 100) as content_preview, + rate_limit_deferrals +FROM pgedge_vectorizer.queue +WHERE status = 'failed' +ORDER BY created_at DESC; + +COMMENT ON VIEW pgedge_vectorizer.failed_items IS +'Failed queue items with error details'; + +-- Pending items count +CREATE VIEW pgedge_vectorizer.pending_count AS +SELECT + COUNT(*) as pending_items, + COUNT(DISTINCT chunk_table) as affected_tables +FROM pgedge_vectorizer.queue +WHERE status = 'pending'; + +COMMENT ON VIEW pgedge_vectorizer.pending_count IS +'Count of pending items waiting for processing'; + +--------------------------------------------------------------------------- +-- Utility functions +--------------------------------------------------------------------------- + +-- Retry failed items +CREATE FUNCTION pgedge_vectorizer.retry_failed( + max_age_hours INT DEFAULT 24 +) RETURNS INT AS $$ +DECLARE + rows_affected INT; +BEGIN + UPDATE pgedge_vectorizer.queue + SET status = 'pending', + attempts = 0, + error_message = NULL, + next_retry_at = NULL, + rate_limit_deferrals = 0 + WHERE status = 'failed' + AND attempts < max_attempts + AND created_at > NOW() - (max_age_hours || ' hours')::INTERVAL; + + GET DIAGNOSTICS rows_affected = ROW_COUNT; + RETURN rows_affected; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.retry_failed IS +'Reset failed items to pending for retry'; + +-- Clear completed items +CREATE FUNCTION pgedge_vectorizer.clear_completed( + older_than_hours INT DEFAULT 24 +) RETURNS INT AS $$ +DECLARE + rows_affected INT; +BEGIN + DELETE FROM pgedge_vectorizer.queue + WHERE status = 'completed' + AND processed_at < NOW() - (older_than_hours || ' hours')::INTERVAL; + + GET DIAGNOSTICS rows_affected = ROW_COUNT; + RETURN rows_affected; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.clear_completed IS +'Remove old completed items from the queue'; + +-- Reprocess chunks that don't have embeddings +CREATE FUNCTION pgedge_vectorizer.reprocess_chunks( + chunk_table_name TEXT +) RETURNS INT AS $$ +DECLARE + rows_affected INT := 0; + chunk_record RECORD; + hybrid_enabled BOOLEAN; +BEGIN + hybrid_enabled := COALESCE( + current_setting('pgedge_vectorizer.enable_hybrid', true), + 'false' + )::BOOLEAN; + + -- Queue chunks that need dense embeddings, and when hybrid is enabled, + -- also queue chunks missing sparse embeddings. + FOR chunk_record IN EXECUTE format( + 'SELECT id, content, (embedding IS NULL) AS needs_embedding, ' + ' (sparse_embedding IS NULL) AS needs_sparse ' + 'FROM %I ' + 'WHERE embedding IS NULL ' + ' OR (sparse_embedding IS NULL AND %L::boolean)', + chunk_table_name, + hybrid_enabled + ) + LOOP + -- Check if already queued + PERFORM 1 FROM pgedge_vectorizer.queue + WHERE chunk_id = chunk_record.id + AND chunk_table = chunk_table_name + AND status IN ('pending', 'processing'); + + -- Only queue if not already queued + IF NOT FOUND THEN + INSERT INTO pgedge_vectorizer.queue (chunk_id, chunk_table, content, metadata, max_attempts) + VALUES ( + chunk_record.id, + chunk_table_name, + chunk_record.content, + CASE + WHEN NOT chunk_record.needs_embedding AND chunk_record.needs_sparse + THEN jsonb_build_object('sparse_only', true) + ELSE NULL + END, + current_setting('pgedge_vectorizer.max_retries')::INT + ); + + rows_affected := rows_affected + 1; + END IF; + END LOOP; + + RAISE NOTICE 'Queued % chunks from % for processing', rows_affected, chunk_table_name; + RETURN rows_affected; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.reprocess_chunks IS +'Queue existing chunks without embeddings for processing'; + +-- Recreate all chunks from scratch +CREATE FUNCTION pgedge_vectorizer.recreate_chunks( + source_table_name REGCLASS, + source_column_name NAME +) RETURNS INT AS $$ +DECLARE + chunk_table_name TEXT; + rows_affected INT := 0; + trigger_name TEXT; + trigger_exists BOOLEAN; +BEGIN + -- Prefer authoritative mapping from vectorizers registry. + SELECT v.chunk_table + INTO chunk_table_name + FROM pgedge_vectorizer.vectorizers v + WHERE v.source_table = source_table_name::TEXT + AND v.source_column = source_column_name; + + -- Fallback to legacy default naming if no registry row exists. + IF chunk_table_name IS NULL THEN + chunk_table_name := source_table_name::TEXT || '_' || source_column_name || '_chunks'; + END IF; + + -- Verify chunk table exists + IF to_regclass(chunk_table_name) IS NULL THEN + RAISE EXCEPTION 'Chunk table % does not exist. Use enable_vectorization() first.', chunk_table_name; + END IF; + + -- Verify trigger exists + trigger_name := source_table_name::TEXT || '_' || source_column_name || '_vectorization_trigger'; + SELECT EXISTS ( + SELECT 1 FROM pg_trigger t + JOIN pg_class c ON t.tgrelid = c.oid + WHERE c.oid = source_table_name + AND t.tgname = trigger_name + ) INTO trigger_exists; + + IF NOT trigger_exists THEN + RAISE EXCEPTION 'Vectorization trigger % does not exist. Use enable_vectorization() first.', trigger_name; + END IF; + + RAISE NOTICE 'Recreating chunks for %.% -> %', source_table_name, source_column_name, chunk_table_name; + + -- Delete all existing chunks and reset IDF stats. + -- Truncating _idf_stats is safe here because recreate_chunks rebuilds + -- all chunks from scratch; the worker will repopulate IDF stats as it + -- processes the newly queued chunks. + EXECUTE format('DELETE FROM %I', chunk_table_name); + EXECUTE format('TRUNCATE TABLE %I', chunk_table_name || '_idf_stats'); + GET DIAGNOSTICS rows_affected = ROW_COUNT; + RAISE NOTICE 'Deleted % existing chunks', rows_affected; + + -- Delete all queue items for this chunk table (with retry logic) + BEGIN + -- Try to delete with a lock timeout + SET LOCAL lock_timeout = '5s'; + DELETE FROM pgedge_vectorizer.queue WHERE chunk_table = chunk_table_name; + RAISE NOTICE 'Cleared queue for %', chunk_table_name; + EXCEPTION WHEN lock_not_available OR deadlock_detected THEN + -- If we can't get the lock, just mark them for cleanup + RAISE WARNING 'Could not clear queue due to concurrent access, continuing anyway'; + END; + + -- Manually process each row to bypass trigger's unchanged-content optimization + DECLARE + row_record RECORD; + doc_content TEXT; + chunks TEXT[]; + chunk_text TEXT; + i INT; + chunk_id BIGINT; + rows_processed INT := 0; + actual_strategy TEXT; + actual_chunk_size INT; + actual_chunk_overlap INT; + pk_col TEXT; + pk_type TEXT; + BEGIN + -- Get chunking configuration from trigger arguments + -- In PostgreSQL 17+, tgargs is bytea and needs to be decoded + DECLARE + tgargs_array TEXT[]; + BEGIN + SELECT string_to_array(encode(t.tgargs, 'escape'), E'\\000') + INTO tgargs_array + FROM pg_trigger t + JOIN pg_class c ON t.tgrelid = c.oid + WHERE c.oid = source_table_name + AND t.tgname = trigger_name; + + -- Arguments: 1=content_col, 2=chunk_table, 3=strategy, 4=size, 5=overlap, 6=pk_col, 7=pk_type + actual_strategy := tgargs_array[3]; + actual_chunk_size := tgargs_array[4]::INT; + actual_chunk_overlap := tgargs_array[5]::INT; + pk_col := COALESCE(tgargs_array[6], 'id'); + pk_type := COALESCE(tgargs_array[7], 'bigint'); + END; + + RAISE NOTICE 'Re-chunking with strategy=%, size=%, overlap=%', + actual_strategy, actual_chunk_size, actual_chunk_overlap; + + -- pk_val is cast to text so that row_record.pk_val is always the same + -- type across calls in a session, whatever the source table's actual + -- primary key type. See the identical comment in enable_vectorization() + -- for why: PL/pgSQL fixes a RECORD field's parameter type the first + -- time a dynamic EXECUTE ... USING statement evaluates it, and this + -- statement's own "$1::%s" cast below does not protect it, because + -- that cast is applied after PostgreSQL has already bound the record + -- field's raw runtime type (issue #39). + FOR row_record IN EXECUTE format( + 'SELECT %I::text as pk_val, %I as content FROM %s WHERE %I IS NOT NULL AND %I != ''''', + pk_col, source_column_name, source_table_name, source_column_name, source_column_name + ) + LOOP + doc_content := row_record.content; + + -- Chunk the document + chunks := pgedge_vectorizer.chunk_text(doc_content, actual_strategy, actual_chunk_size, actual_chunk_overlap); + + -- Insert chunks and queue for embedding + FOR i IN 1..array_length(chunks, 1) LOOP + chunk_text := chunks[i]; + + -- Insert chunk + -- pk_type uses %s: value from format_type() is system-controlled (see enable_vectorization) + EXECUTE format(' + INSERT INTO %I (source_id, chunk_index, content, token_count) + VALUES ($1::%s, $2, $3, $4) + RETURNING id', chunk_table_name, pk_type) + USING row_record.pk_val, i, chunk_text, + pgedge_vectorizer.count_tokens(chunk_text) + INTO chunk_id; + + -- Queue for embedding + INSERT INTO pgedge_vectorizer.queue (chunk_id, chunk_table, content, max_attempts) + VALUES (chunk_id, chunk_table_name, chunk_text, + current_setting('pgedge_vectorizer.max_retries')::INT); + END LOOP; + + rows_processed := rows_processed + 1; + END LOOP; + + RAISE NOTICE 'Processed % rows', rows_processed; + RETURN rows_processed; + END; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.recreate_chunks IS +'Delete all chunks and recreate from source table (complete rebuild)'; + +-- Get configuration summary +CREATE FUNCTION pgedge_vectorizer.show_config() +RETURNS TABLE ( + setting TEXT, + value TEXT +) AS $$ +BEGIN + RETURN QUERY + SELECT + name::TEXT as setting, + current_setting(name)::TEXT as value + FROM pg_settings + WHERE name LIKE 'pgedge_vectorizer.%' + ORDER BY name; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.show_config IS +'Show all pgedge_vectorizer configuration settings'; + +--------------------------------------------------------------------------- +-- Grants (for non-superuser usage - optional) +--------------------------------------------------------------------------- + +-- Grant usage on schema to public (optional - comment out if not desired) +-- GRANT USAGE ON SCHEMA pgedge_vectorizer TO PUBLIC; +-- GRANT SELECT ON ALL TABLES IN SCHEMA pgedge_vectorizer TO PUBLIC; +-- GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA pgedge_vectorizer TO PUBLIC; +-- hybrid.sql +-- Hybrid BM25 + dense vector search using Reciprocal Rank Fusion (RRF). +-- +-- Requires: pgedge_vectorizer.enable_hybrid = true in postgresql.conf +-- and pgvector >= 0.7.0 for sparsevec support. + +--------------------------------------------------------------------------- +-- hybrid_search() — main user-facing function +--------------------------------------------------------------------------- + +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_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 INTO v_chunk_table + 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 + 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; + + -- Generate dense query vector via the existing C function + v_query_dense := pgedge_vectorizer.generate_embedding(p_query); + + -- 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.'; + +--------------------------------------------------------------------------- +-- hybrid_search_simple() — convenience wrapper +--------------------------------------------------------------------------- + +CREATE OR REPLACE FUNCTION pgedge_vectorizer.hybrid_search_simple( + p_source_table REGCLASS, + p_query TEXT, + p_limit INT DEFAULT 10, + p_source_column NAME DEFAULT NULL +) +RETURNS TABLE ( + source_id TEXT, + chunk TEXT, + rrf_score FLOAT8 +) +LANGUAGE sql AS $$ + SELECT source_id, chunk, rrf_score + FROM pgedge_vectorizer.hybrid_search( + p_source_table, p_query, p_limit, + p_source_column => p_source_column); +$$; + +COMMENT ON FUNCTION pgedge_vectorizer.hybrid_search_simple IS +'Convenience wrapper for hybrid_search() returning only source_id, chunk, and rrf_score'; diff --git a/src/embed.c b/src/embed.c index 384b89d..a6ddb98 100644 --- a/src/embed.c +++ b/src/embed.c @@ -21,6 +21,49 @@ * This function takes a text query and returns a vector embedding using * the configured provider (OpenAI, Voyage, or Ollama). */ +/* + * Resolve the provider for one call. + * + * A NULL or empty name means fall back to the GUC, which is exactly what a + * vectorizer with no override records, so the same rule serves both these + * SQL-callable functions and the worker. + */ +static EmbeddingProvider * +resolve_provider(const char *name) +{ + const char *use; + EmbeddingProvider *provider; + + use = (name != NULL && name[0] != '\0') ? name : pgedge_vectorizer_provider; + + if (use == NULL || use[0] == '\0') + elog(ERROR, "pgedge_vectorizer.provider is not set"); + + provider = get_embedding_provider(use); + if (provider == NULL) + elog(ERROR, "embedding provider \"%s\" is not available", use); + + return provider; +} + +/* As resolve_provider(), for the model name. */ +static const char * +resolve_model(const char *model) +{ + return (model != NULL && model[0] != '\0') + ? model : pgedge_vectorizer_model; +} + +/* Read an optional text argument as a cstring, or NULL if it was not given. */ +static char * +optional_text_arg(FunctionCallInfo fcinfo, int argno) +{ + if (PG_NARGS() <= argno || PG_ARGISNULL(argno)) + return NULL; + + return text_to_cstring(PG_GETARG_TEXT_PP(argno)); +} + PG_FUNCTION_INFO_V1(pgedge_vectorizer_generate_embedding); PG_FUNCTION_INFO_V1(pgedge_vectorizer_detect_embedding_dimension); @@ -37,6 +80,7 @@ pgedge_vectorizer_generate_embedding(PG_FUNCTION_ARGS) int ret; bool isnull; Datum result; + const char *model; /* Check for NULL input */ if (PG_ARGISNULL(0)) @@ -56,13 +100,12 @@ pgedge_vectorizer_generate_embedding(PG_FUNCTION_ARGS) PG_RETURN_NULL(); } - /* Get the current provider */ - provider = get_current_provider(); - if (provider == NULL) - { - elog(ERROR, "no embedding provider configured"); - PG_RETURN_NULL(); - } + /* + * Provider and model fall back to the GUCs when not given, so an + * existing one-argument call behaves exactly as it did. + */ + provider = resolve_provider(optional_text_arg(fcinfo, 1)); + model = resolve_model(optional_text_arg(fcinfo, 2)); /* Initialize the provider if needed */ if (provider->init != NULL) @@ -77,7 +120,7 @@ pgedge_vectorizer_generate_embedding(PG_FUNCTION_ARGS) } /* Generate embedding */ - embedding = provider->generate(query, &dim, &error_msg); + embedding = provider->generate(query, model, &dim, &error_msg); if (embedding == NULL) { elog(ERROR, "failed to generate embedding: %s", @@ -157,14 +200,16 @@ pgedge_vectorizer_detect_embedding_dimension(PG_FUNCTION_ARGS) float *embedding; int dim = 0; char *error_msg = NULL; + const char *model; - /* Get the current provider */ - provider = get_current_provider(); - if (provider == NULL) - { - elog(ERROR, "no embedding provider configured"); - PG_RETURN_NULL(); - } + /* + * The probe has to use the provider and model whose dimension is being + * asked about, which is not necessarily the configured one: a vectorizer + * created with an override needs the dimension of that model, not of + * whatever the GUCs happen to name. + */ + provider = resolve_provider(optional_text_arg(fcinfo, 0)); + model = resolve_model(optional_text_arg(fcinfo, 1)); /* Initialize the provider if needed */ if (provider->init != NULL) @@ -179,7 +224,7 @@ pgedge_vectorizer_detect_embedding_dimension(PG_FUNCTION_ARGS) } /* Generate a probe embedding to detect dimension */ - embedding = provider->generate("dimension probe", &dim, &error_msg); + embedding = provider->generate("dimension probe", model, &dim, &error_msg); if (embedding == NULL) { elog(ERROR, "failed to detect embedding dimension: %s", diff --git a/src/pgedge_vectorizer.h b/src/pgedge_vectorizer.h index 9b34ff9..f9bf33a 100644 --- a/src/pgedge_vectorizer.h +++ b/src/pgedge_vectorizer.h @@ -151,8 +151,10 @@ typedef struct EmbeddingProvider const char *name; bool (*init)(char **error_msg); void (*cleanup)(void); - float *(*generate)(const char *text, int *dim, char **error_msg); - float **(*generate_batch)(const char **texts, int count, int *dim, char **error_msg); + float *(*generate)(const char *text, const char *model, + int *dim, char **error_msg); + float **(*generate_batch)(const char **texts, int count, const char *model, + int *dim, char **error_msg); } EmbeddingProvider; /* @@ -181,6 +183,7 @@ extern EmbeddingProvider GeminiProvider; /* tokenizer.c */ int count_tokens(const char *text, const char *model); +Datum pgedge_vectorizer_count_tokens(PG_FUNCTION_ARGS); int *tokenize_text(const char *text, const char *model, int *token_count); char *detokenize_tokens(const int *tokens, int token_count, const char *model); int get_char_offset_for_tokens(const char *text, int target_tokens, const char *model); diff --git a/src/provider_common.c b/src/provider_common.c index 1f31d35..74af977 100644 --- a/src/provider_common.c +++ b/src/provider_common.c @@ -12,6 +12,7 @@ */ #include "provider_common.h" +#include #include #include #include @@ -442,11 +443,48 @@ provider_build_openai_request(const char **texts, int count, const char *model) appendStringInfo(&request_buf, "\"%s\"", escaped); pfree(escaped); } - appendStringInfo(&request_buf, "],\"model\":\"%s\"}", model); + { + /* + * The model is no longer only a GUC: it can come from a vectorizer's + * registry row or straight from a SQL argument, so it is escaped like + * any other value rather than trusted into the body. + */ + char *escaped_model = provider_escape_json_string(model); + + appendStringInfo(&request_buf, "],\"model\":\"%s\"}", escaped_model); + pfree(escaped_model); + } return request_buf.data; } +/* + * Percent-encode a string for use as one URL path segment. + * + * The model name reaches the Gemini URL, and it is user-supplied: a '/' would + * add a path segment and a '?' would start a query string, either of which + * sends the request somewhere other than intended. Everything outside the + * unreserved set of RFC 3986 is encoded. + */ +char * +provider_url_encode_segment(const char *str) +{ + StringInfoData buf; + const unsigned char *p; + + initStringInfo(&buf); + + for (p = (const unsigned char *) str; *p; p++) + { + if (isalnum(*p) || *p == '-' || *p == '.' || *p == '_' || *p == '~') + appendStringInfoChar(&buf, (char) *p); + else + appendStringInfo(&buf, "%%%02X", *p); + } + + return buf.data; +} + /* * Count dimensions in a JSON float array by counting commas. * Pointer should be positioned just after the opening '['. diff --git a/src/provider_common.h b/src/provider_common.h index 346d424..fb7227b 100644 --- a/src/provider_common.h +++ b/src/provider_common.h @@ -80,6 +80,9 @@ char *provider_expand_tilde(const char *path); /* Escape a string for safe inclusion in JSON */ char *provider_escape_json_string(const char *str); +/* Percent-encode a string for use as a single URL path segment */ +char *provider_url_encode_segment(const char *str); + /* Build OpenAI-format request body: {"input":[...], "model":"..."} */ char *provider_build_openai_request(const char **texts, int count, const char *model); diff --git a/src/provider_gemini.c b/src/provider_gemini.c index 0ff5c00..b345d5f 100644 --- a/src/provider_gemini.c +++ b/src/provider_gemini.c @@ -28,8 +28,10 @@ static bool provider_initialized = false; */ static bool gemini_init(char **error_msg); static void gemini_cleanup(void); -static float *gemini_generate(const char *text, int *dim, char **error_msg); -static float **gemini_generate_batch(const char **texts, int count, int *dim, +static float *gemini_generate(const char *text, const char *model, + int *dim, char **error_msg); +static float **gemini_generate_batch(const char **texts, int count, + const char *model, int *dim, char **error_msg); /* Gemini-specific response parser */ @@ -98,13 +100,14 @@ gemini_cleanup(void) * Generate a single embedding */ static float * -gemini_generate(const char *text, int *dim, char **error_msg) +gemini_generate(const char *text, const char *model, + int *dim, char **error_msg) { const char *texts[1] = {text}; float **embeddings; float *result; - embeddings = gemini_generate_batch(texts, 1, dim, error_msg); + embeddings = gemini_generate_batch(texts, 1, model, dim, error_msg); if (embeddings == NULL) return NULL; @@ -122,12 +125,15 @@ gemini_generate(const char *text, int *dim, char **error_msg) * Response: {"embeddings":[{"values":[0.1,0.2,...]}, ...]} */ static float ** -gemini_generate_batch(const char **texts, int count, int *dim, char **error_msg) +gemini_generate_batch(const char **texts, int count, const char *model, + int *dim, char **error_msg) { char *json_request; char *url; + char *url_model; const char *base_url; char *auth_header; + char *escaped_model; StringInfoData request_buf; ResponseBuffer response; float **embeddings; @@ -141,6 +147,7 @@ gemini_generate_batch(const char **texts, int count, int *dim, char **error_msg) /* Build JSON request - Gemini batch format */ initStringInfo(&request_buf); appendStringInfo(&request_buf, "{\"requests\":["); + escaped_model = provider_escape_json_string(model); for (int i = 0; i < count; i++) { char *escaped = provider_escape_json_string(texts[i]); @@ -149,9 +156,10 @@ gemini_generate_batch(const char **texts, int count, int *dim, char **error_msg) appendStringInfo(&request_buf, "{\"model\":\"models/%s\"," "\"content\":{\"parts\":[{\"text\":\"%s\"}]}}", - pgedge_vectorizer_model, escaped); + escaped_model, escaped); pfree(escaped); } + pfree(escaped_model); appendStringInfo(&request_buf, "]}"); json_request = request_buf.data; @@ -160,8 +168,14 @@ gemini_generate_batch(const char **texts, int count, int *dim, char **error_msg) pgedge_vectorizer_api_url[0] != '\0') ? pgedge_vectorizer_api_url : GEMINI_DEFAULT_BASE_URL; - url = psprintf("%s/models/%s:batchEmbedContents", base_url, - pgedge_vectorizer_model); + /* + * The model lands in the URL's path, so it is percent-encoded rather than + * pasted in: a '/' or '?' in a model name would otherwise change which + * endpoint the request reaches. + */ + url_model = provider_url_encode_segment(model); + url = psprintf("%s/models/%s:batchEmbedContents", base_url, url_model); + pfree(url_model); /* Build auth header */ auth_header = psprintf("x-goog-api-key: %s", api_key); diff --git a/src/provider_ollama.c b/src/provider_ollama.c index 4dc87d0..0a4c327 100644 --- a/src/provider_ollama.c +++ b/src/provider_ollama.c @@ -26,8 +26,10 @@ static bool provider_initialized = false; */ static bool ollama_init(char **error_msg); static void ollama_cleanup(void); -static float *ollama_generate(const char *text, int *dim, char **error_msg); -static float **ollama_generate_batch(const char **texts, int count, int *dim, +static float *ollama_generate(const char *text, const char *model, + int *dim, char **error_msg); +static float **ollama_generate_batch(const char **texts, int count, + const char *model, int *dim, char **error_msg); /* Ollama-specific response parser */ @@ -80,12 +82,14 @@ ollama_cleanup(void) * Generate a single embedding */ static float * -ollama_generate(const char *text, int *dim, char **error_msg) +ollama_generate(const char *text, const char *model, + int *dim, char **error_msg) { char *json_request; char *url; const char *base_url; char *escaped; + char *escaped_model; StringInfoData request_buf; ResponseBuffer response; float *embedding; @@ -99,8 +103,10 @@ ollama_generate(const char *text, int *dim, char **error_msg) /* Build JSON request - Ollama API format */ initStringInfo(&request_buf); escaped = provider_escape_json_string(text); + escaped_model = provider_escape_json_string(model); appendStringInfo(&request_buf, "{\"model\":\"%s\",\"prompt\":\"%s\"}", - pgedge_vectorizer_model, escaped); + escaped_model, escaped); + pfree(escaped_model); pfree(escaped); json_request = request_buf.data; @@ -138,7 +144,8 @@ ollama_generate(const char *text, int *dim, char **error_msg) * endpoint multiple times. */ static float ** -ollama_generate_batch(const char **texts, int count, int *dim, char **error_msg) +ollama_generate_batch(const char **texts, int count, const char *model, + int *dim, char **error_msg) { float **embeddings; int i; @@ -153,7 +160,7 @@ ollama_generate_batch(const char **texts, int count, int *dim, char **error_msg) for (i = 0; i < count; i++) { - embeddings[i] = ollama_generate(texts[i], dim, error_msg); + embeddings[i] = ollama_generate(texts[i], model, dim, error_msg); if (embeddings[i] == NULL) { for (int j = 0; j < i; j++) diff --git a/src/provider_openai.c b/src/provider_openai.c index 9626add..cb0e7bc 100644 --- a/src/provider_openai.c +++ b/src/provider_openai.c @@ -27,8 +27,10 @@ static bool provider_initialized = false; */ static bool openai_init(char **error_msg); static void openai_cleanup(void); -static float *openai_generate(const char *text, int *dim, char **error_msg); -static float **openai_generate_batch(const char **texts, int count, int *dim, +static float *openai_generate(const char *text, const char *model, + int *dim, char **error_msg); +static float **openai_generate_batch(const char **texts, int count, + const char *model, int *dim, char **error_msg); /* @@ -111,13 +113,14 @@ openai_cleanup(void) * Generate a single embedding */ static float * -openai_generate(const char *text, int *dim, char **error_msg) +openai_generate(const char *text, const char *model, + int *dim, char **error_msg) { const char *texts[1] = {text}; float **embeddings; float *result; - embeddings = openai_generate_batch(texts, 1, dim, error_msg); + embeddings = openai_generate_batch(texts, 1, model, dim, error_msg); if (embeddings == NULL) return NULL; @@ -130,7 +133,8 @@ openai_generate(const char *text, int *dim, char **error_msg) * Generate embeddings in batch */ static float ** -openai_generate_batch(const char **texts, int count, int *dim, char **error_msg) +openai_generate_batch(const char **texts, int count, const char *model, + int *dim, char **error_msg) { char *json_request; char *url; @@ -147,7 +151,7 @@ openai_generate_batch(const char **texts, int count, int *dim, char **error_msg) /* Build request body */ json_request = provider_build_openai_request(texts, count, - pgedge_vectorizer_model); + model); /* Build URL */ base_url = (pgedge_vectorizer_api_url != NULL && diff --git a/src/provider_voyage.c b/src/provider_voyage.c index 929da93..f7bab52 100644 --- a/src/provider_voyage.c +++ b/src/provider_voyage.c @@ -27,8 +27,10 @@ static bool provider_initialized = false; */ static bool voyage_init(char **error_msg); static void voyage_cleanup(void); -static float *voyage_generate(const char *text, int *dim, char **error_msg); -static float **voyage_generate_batch(const char **texts, int count, int *dim, +static float *voyage_generate(const char *text, const char *model, + int *dim, char **error_msg); +static float **voyage_generate_batch(const char **texts, int count, + const char *model, int *dim, char **error_msg); /* @@ -92,13 +94,14 @@ voyage_cleanup(void) * Generate a single embedding */ static float * -voyage_generate(const char *text, int *dim, char **error_msg) +voyage_generate(const char *text, const char *model, + int *dim, char **error_msg) { const char *texts[1] = {text}; float **embeddings; float *result; - embeddings = voyage_generate_batch(texts, 1, dim, error_msg); + embeddings = voyage_generate_batch(texts, 1, model, dim, error_msg); if (embeddings == NULL) return NULL; @@ -111,7 +114,8 @@ voyage_generate(const char *text, int *dim, char **error_msg) * Generate embeddings in batch */ static float ** -voyage_generate_batch(const char **texts, int count, int *dim, char **error_msg) +voyage_generate_batch(const char **texts, int count, const char *model, + int *dim, char **error_msg) { char *json_request; char *url; @@ -128,7 +132,7 @@ voyage_generate_batch(const char **texts, int count, int *dim, char **error_msg) /* Build request body */ json_request = provider_build_openai_request(texts, count, - pgedge_vectorizer_model); + model); /* Build URL */ base_url = (pgedge_vectorizer_api_url != NULL && diff --git a/src/tokenizer.c b/src/tokenizer.c index d08bc6a..7e4bd49 100644 --- a/src/tokenizer.c +++ b/src/tokenizer.c @@ -57,6 +57,24 @@ count_tokens(const char *text, const char *model) return token_estimate; } +/* + * SQL-callable wrapper around count_tokens() + * + * Exposed so that the plpgsql chunking paths compute token_count with exactly + * the same rule as the C chunker, rather than open-coding the approximation + * and disagreeing with it by a token. + */ +PG_FUNCTION_INFO_V1(pgedge_vectorizer_count_tokens); + +Datum +pgedge_vectorizer_count_tokens(PG_FUNCTION_ARGS) +{ + text *input = PG_GETARG_TEXT_PP(0); + const char *text_str = TextDatumGetCString(PointerGetDatum(input)); + + PG_RETURN_INT32(count_tokens(text_str, pgedge_vectorizer_model)); +} + /* * Tokenize text into token IDs * diff --git a/src/worker.c b/src/worker.c index f1f11d9..f2d681d 100644 --- a/src/worker.c +++ b/src/worker.c @@ -1492,6 +1492,73 @@ pgedge_vectorizer_worker_main(Datum main_arg) proc_exit(0); } +/* + * Swap two items of a fetched batch, moving every parallel array together. + */ +static void +swap_batch_items(int a, int b, int64 *queue_ids, int64 *chunk_ids, + char **chunk_tables, const char **contents, + int *content_lens, int *attempts, int *max_attempts, + bool *sparse_only, char **providers, char **models) +{ +#define SWAP(type, arr) do { type tmp_ = (arr)[a]; \ + (arr)[a] = (arr)[b]; \ + (arr)[b] = tmp_; } while (0) + SWAP(int64, queue_ids); + SWAP(int64, chunk_ids); + SWAP(char *, chunk_tables); + SWAP(const char *, contents); + SWAP(int, content_lens); + SWAP(int, attempts); + SWAP(int, max_attempts); + SWAP(bool, sparse_only); + SWAP(char *, providers); + SWAP(char *, models); +#undef SWAP +} + +/* + * Group a fetched batch by (provider, model). + * + * A batch is selected by age across every vectorizer at once, so items headed + * for different models interleave. One request carries one model, and merely + * breaking the run wherever the model changes would give requests of one item + * whenever two tables' work alternates in time. Sorting first keeps requests + * as full as they can be. + * + * An insertion sort is enough for batch_size items and is stable, so the age + * ordering survives within each group. The batch is still selected by + * created_at, so this changes only the order of requests within one batch, + * not which items are picked up. + */ +static void +sort_batch_by_model(int n_items, int64 *queue_ids, int64 *chunk_ids, + char **chunk_tables, const char **contents, + int *content_lens, int *attempts, int *max_attempts, + bool *sparse_only, char **providers, char **models) +{ + for (int i = 1; i < n_items; i++) + { + int j = i; + + while (j > 0) + { + int cmp = strcmp(providers[j - 1], providers[j]); + + if (cmp == 0) + cmp = strcmp(models[j - 1], models[j]); + + if (cmp <= 0) + break; + + swap_batch_items(j - 1, j, queue_ids, chunk_ids, chunk_tables, + contents, content_lens, attempts, max_attempts, + sparse_only, providers, models); + j--; + } + } +} + /* * How many of the items starting at `start` may be sent as one request. * @@ -1502,10 +1569,14 @@ pgedge_vectorizer_worker_main(Datum main_arg) * * Sparse-only items are grouped with their like because a request is skipped * only when every item in it already has its dense embedding. + * + * A request also carries exactly one provider and model, so the run breaks + * where either changes. sort_batch_by_model() has already grouped the batch, + * so this only marks the boundaries rather than fragmenting anything. */ static int batch_extent(int start, int n_items, const int *attempts, - const bool *sparse_only) + const bool *sparse_only, char **providers, char **models) { int count = 1; @@ -1514,7 +1585,9 @@ batch_extent(int start, int n_items, const int *attempts, while (start + count < n_items && attempts[start + count] == 0 && - sparse_only[start + count] == sparse_only[start]) + sparse_only[start + count] == sparse_only[start] && + strcmp(providers[start + count], providers[start]) == 0 && + strcmp(models[start + count], models[start]) == 0) count++; return count; @@ -1555,12 +1628,35 @@ process_queue_batch(const char *dbname) SPI_connect(); /* Fetch pending items using FOR UPDATE SKIP LOCKED */ + /* + * The left join resolves each item's provider and model, with the + * vectorizer's setting overriding the GUC and NULL meaning inherit. + * NULLIF puts an empty string on the same footing as NULL, which is the + * rule resolve_provider() and resolve_model() already apply in embed.c. + * Doing it here keeps the inheritance rule out of the C entirely, and + * an item whose vectorizer has since been disabled falls back to the + * GUCs through the same expression rather than needing a special case. + * + * 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. + */ ret = SPI_execute(psprintf( - "SELECT id, chunk_id, chunk_table, content, attempts, max_attempts, " - " COALESCE((metadata->>'sparse_only')::boolean, false) AS sparse_only " - "FROM pgedge_vectorizer.queue " - "WHERE status = 'pending' " - "AND (next_retry_at IS NULL OR next_retry_at <= NOW()) " + "SELECT q.id, q.chunk_id, q.chunk_table, q.content, q.attempts, " + " q.max_attempts, " + " COALESCE((q.metadata->>'sparse_only')::boolean, false) " + " AS sparse_only, " + " COALESCE(NULLIF(v.provider, ''), " + " current_setting('pgedge_vectorizer.provider')) " + " AS provider, " + " COALESCE(NULLIF(v.model, ''), " + " current_setting('pgedge_vectorizer.model')) " + " AS model " + "FROM pgedge_vectorizer.queue q " + "LEFT JOIN pgedge_vectorizer.vectorizers v " + " ON v.chunk_table = q.chunk_table " + "WHERE q.status = 'pending' " + "AND (q.next_retry_at IS NULL OR q.next_retry_at <= NOW()) " /* * Oldest first. Ordering by attempts DESC put the items that had * failed most at the head of every batch, so a provider outage left @@ -1568,9 +1664,9 @@ process_queue_batch(const char *dbname) * items exhausted max_attempts. next_retry_at already spaces retries * out; age is the only ordering the queue needs. */ - "ORDER BY created_at " + "ORDER BY q.created_at " "LIMIT %d " - "FOR UPDATE SKIP LOCKED", + "FOR UPDATE OF q SKIP LOCKED", batch_size), false, batch_size); @@ -1585,6 +1681,8 @@ process_queue_batch(const char *dbname) int *attempts = palloc(n_items * sizeof(int)); int *max_attempts = palloc(n_items * sizeof(int)); bool *sparse_only = palloc(n_items * sizeof(bool)); + char **providers = palloc(n_items * sizeof(char *)); + char **models = palloc(n_items * sizeof(char *)); float **embeddings = NULL; int dim = 0; int batch_count = 0; @@ -1622,6 +1720,12 @@ process_queue_batch(const char *dbname) val = SPI_getbinval(SPI_tuptable->vals[i], SPI_tuptable->tupdesc, 7, &isnull); sparse_only[i] = (!isnull && DatumGetBool(val)); + val = SPI_getbinval(SPI_tuptable->vals[i], SPI_tuptable->tupdesc, 8, &isnull); + providers[i] = TextDatumGetCString(val); + + val = SPI_getbinval(SPI_tuptable->vals[i], SPI_tuptable->tupdesc, 9, &isnull); + models[i] = TextDatumGetCString(val); + if (attempts[i] > 0) has_retries = true; } @@ -1658,6 +1762,15 @@ process_queue_batch(const char *dbname) has_sparse_only = true; } + /* + * Group the batch so that each request carries one provider and + * model. Safe here: every item is independent of its neighbours, and + * the arrays move together. + */ + sort_batch_by_model(n_items, queue_ids, chunk_ids, chunk_tables, + contents, content_lens, attempts, max_attempts, + sparse_only, providers, models); + /* * Stop charging the last probed item. What follows — marking the * batch, resolving the provider, generating embeddings — either @@ -1688,25 +1801,19 @@ process_queue_batch(const char *dbname) false, 0); } - /* Get the provider */ - provider = get_current_provider(); - if (provider == NULL) - { - elog(ERROR, "No provider configured"); - } - - /* Initialize provider if needed */ - if (!provider->init(&error_msg)) - { - elog(ERROR, "Failed to initialize provider: %s", - error_msg ? error_msg : "unknown error"); - } + /* + * The provider is resolved per request rather than once per batch, + * because a batch may hold items for vectorizers configured with + * different providers. Each provider caches its own initialisation + * in a file-static, so init() per request costs nothing after the + * first. + */ /* Process items in requests as large as batch_extent() allows */ for (int batch_start = 0; batch_start < n_items; batch_start += batch_count) { batch_count = batch_extent(batch_start, n_items, attempts, - sparse_only); + sparse_only, providers, models); /* Skip dense generation when every item in this batch is sparse-only. */ { @@ -1735,10 +1842,23 @@ process_queue_batch(const char *dbname) * provider that fails before reaching the network records * nothing of its own. */ + provider = get_embedding_provider(providers[batch_start]); + if (provider == NULL) + elog(ERROR, "embedding provider \"%s\" is not available", + providers[batch_start]); + + if (!provider->init(&error_msg)) + elog(ERROR, "failed to initialise provider \"%s\": %s", + providers[batch_start], + error_msg ? error_msg : "unknown error"); + provider_reset_rate_limit(); - /* Generate embeddings for this batch */ - embeddings = provider->generate_batch(&contents[batch_start], batch_count, &dim, &error_msg); + /* Generate embeddings for this request */ + embeddings = provider->generate_batch(&contents[batch_start], + batch_count, + models[batch_start], + &dim, &error_msg); } } @@ -1986,16 +2106,37 @@ process_queue_batch(const char *dbname) char *next_try = rate_limit_backoff_expr(ratelimit->retry_after); int cooldown; - int deferred = n_items - batch_start; + int deferred = 0; /* - * The rest of the pull is deferred too. It would meet the - * same limit, and it was marked 'processing' before the - * loop began: nothing reclaims an item left that way at - * commit. + * The rest of this provider's pull is deferred too. It + * would meet the same limit, and it was marked + * 'processing' before the loop began: nothing reclaims an + * item left that way at commit. + * + * Only this provider's items, though. A batch can now + * span providers, and charging another provider's work + * for this one's 429 would spend deferrals it never used + * and, once they ran out, fail it outright. Those items + * go straight back to pending, uncharged, for the next + * pull to take. */ for (int idx = batch_start; idx < n_items; idx++) { + if (strcmp(providers[idx], + providers[batch_start]) != 0) + { + SPI_execute(psprintf( + "UPDATE pgedge_vectorizer.queue " + "SET status = 'pending' " + "WHERE id = %ld", + queue_ids[idx]), + false, 0); + continue; + } + + deferred++; + SPI_execute(psprintf( "UPDATE pgedge_vectorizer.queue " "SET status = CASE WHEN rate_limit_deferrals + 1 >= %d " diff --git a/test/expected/count_tokens.out b/test/expected/count_tokens.out new file mode 100644 index 0000000..5a41c8c --- /dev/null +++ b/test/expected/count_tokens.out @@ -0,0 +1,157 @@ +-- count_tokens test +-- +-- count_tokens() exposes the estimate the C chunking code has always used, so +-- that the plpgsql paths which write the token_count column can call it rather +-- than open-coding the same rule and rounding it the other way. The point of +-- the tests below is therefore as much the agreement between the two as the +-- values themselves. +--------------------------------------------------------------------------- +-- The estimate itself: UTF-8 characters divided by four, rounded up +--------------------------------------------------------------------------- +-- 11 characters, so 3 tokens. +SELECT pgedge_vectorizer.count_tokens('hello world') AS eleven_chars; + eleven_chars +-------------- + 3 +(1 row) + +-- Exactly 4 characters is exactly 1 token; anything shorter still rounds up +-- to 1, which is the difference from the truncating arithmetic this replaces. +SELECT pgedge_vectorizer.count_tokens('test') AS four_chars, + pgedge_vectorizer.count_tokens('abc') AS three_chars, + pgedge_vectorizer.count_tokens('a') AS one_char; + four_chars | three_chars | one_char +------------+-------------+---------- + 1 | 1 | 1 +(1 row) + +-- Empty text is the one case that is genuinely zero. +SELECT pgedge_vectorizer.count_tokens('') AS empty; + empty +------- + 0 +(1 row) + +-- STRICT, so NULL in, NULL out. +SELECT pgedge_vectorizer.count_tokens(NULL) IS NULL AS null_is_null; + null_is_null +-------------- + t +(1 row) + +-- Characters, not bytes: four Han characters are twelve bytes but one token. +SELECT pgedge_vectorizer.count_tokens('你好世界') AS four_han_chars, + octet_length('你好世界') AS bytes; + four_han_chars | bytes +----------------+------- + 1 | 12 +(1 row) + +-- Declared STABLE rather than IMMUTABLE: the estimate is defined in terms of +-- pgedge_vectorizer.model, which will matter once the counter is model-aware, +-- and an index or cached plan built on an IMMUTABLE promise would then be +-- wrong. Pin the volatility so that cannot be relaxed by accident. +SELECT provolatile + FROM pg_proc + WHERE proname = 'count_tokens' + AND pronamespace = 'pgedge_vectorizer'::regnamespace; + provolatile +------------- + s +(1 row) + +--------------------------------------------------------------------------- +-- Agreement with what the chunking paths store +--------------------------------------------------------------------------- +CREATE TABLE count_tokens_docs ( + id BIGSERIAL PRIMARY KEY, + content TEXT +); +-- Chunks written by enable_vectorization() back-filling an existing table. +INSERT INTO count_tokens_docs (content) +VALUES ('abc'), + ('Short document.'), + (repeat('The quick brown fox jumps over the lazy dog. ', 20)); +SELECT pgedge_vectorizer.enable_vectorization( + 'count_tokens_docs'::regclass, + 'content', + 'token_based', + 100, + 10, + 1536 +); +NOTICE: Using primary key column: id (bigint) +NOTICE: column "sparse_embedding" of relation "count_tokens_docs_content_chunks" already exists, skipping +NOTICE: Vectorization enabled: count_tokens_docs -> count_tokens_docs_content_chunks +NOTICE: Strategy: token_based, chunk_size: 100, overlap: 10 +NOTICE: Processing existing rows... +NOTICE: Processed 3 existing rows + enable_vectorization +---------------------- + +(1 row) + +SELECT count(*) AS mismatched_on_backfill + FROM count_tokens_docs_content_chunks + WHERE token_count IS DISTINCT FROM pgedge_vectorizer.count_tokens(content); + mismatched_on_backfill +------------------------ + 0 +(1 row) + +-- The three-character row is the regression: length('abc') / 4 stored 0, which +-- the BM25 scoring path then had to clamp back up to 1. +SELECT token_count AS short_row_token_count + FROM count_tokens_docs_content_chunks c + JOIN count_tokens_docs d ON d.id = c.source_id + WHERE d.content = 'abc'; + short_row_token_count +----------------------- + 1 +(1 row) + +-- Chunks written by the insert trigger. +INSERT INTO count_tokens_docs (content) +VALUES ('xy'), + ('A document added after vectorization was enabled.'); +SELECT count(*) AS mismatched_on_insert + FROM count_tokens_docs_content_chunks + WHERE token_count IS DISTINCT FROM pgedge_vectorizer.count_tokens(content); + mismatched_on_insert +---------------------- + 0 +(1 row) + +-- Chunks written by recreate_chunks(). +SELECT pgedge_vectorizer.recreate_chunks('count_tokens_docs'::regclass, 'content'); +NOTICE: Recreating chunks for count_tokens_docs.content -> count_tokens_docs_content_chunks +NOTICE: Deleted 0 existing chunks +NOTICE: Cleared queue for count_tokens_docs_content_chunks +NOTICE: Re-chunking with strategy=token_based, size=100, overlap=10 +NOTICE: Processed 5 rows + recreate_chunks +----------------- + 5 +(1 row) + +SELECT count(*) AS mismatched_on_recreate + FROM count_tokens_docs_content_chunks + WHERE token_count IS DISTINCT FROM pgedge_vectorizer.count_tokens(content); + mismatched_on_recreate +------------------------ + 0 +(1 row) + +--------------------------------------------------------------------------- +-- Cleanup +--------------------------------------------------------------------------- +SELECT pgedge_vectorizer.disable_vectorization('count_tokens_docs'::regclass, + 'content', TRUE); +NOTICE: Vectorization disabled and chunk table dropped: count_tokens_docs_content_chunks + disable_vectorization +----------------------- + +(1 row) + +DROP TABLE count_tokens_docs; +DELETE FROM pgedge_vectorizer.queue; diff --git a/test/expected/embedding.out b/test/expected/embedding.out index 7e6cfe6..c2f70d7 100644 --- a/test/expected/embedding.out +++ b/test/expected/embedding.out @@ -2,7 +2,9 @@ -- Tests the generate_embedding() function with actual API calls -- Gracefully skips tests if API keys are not available -- Test 1: Verify function exists -SELECT pg_get_functiondef('pgedge_vectorizer.generate_embedding(text)'::regprocedure) IS NOT NULL AS function_exists; +SELECT pg_get_functiondef( + 'pgedge_vectorizer.generate_embedding(text, text, text)'::regprocedure + ) IS NOT NULL AS function_exists; function_exists ----------------- t @@ -192,7 +194,7 @@ EXCEPTION WHEN OTHERS THEN RAISE NOTICE 'NULL input correctly raises error: %', SQLERRM; END; $$; -NOTICE: WARNING: NULL input should have raised an error +NOTICE: NULL input correctly raises error: query text cannot be NULL -- Test empty string (should work if API key available) -- Configure for empty string test SET pgedge_vectorizer.provider = 'openai'; diff --git a/test/expected/embedding_1.out b/test/expected/embedding_1.out index 5b6ae69..94df9d2 100644 --- a/test/expected/embedding_1.out +++ b/test/expected/embedding_1.out @@ -2,7 +2,9 @@ -- Tests the generate_embedding() function with actual API calls -- Gracefully skips tests if API keys are not available -- Test 1: Verify function exists -SELECT pg_get_functiondef('pgedge_vectorizer.generate_embedding(text)'::regprocedure) IS NOT NULL AS function_exists; +SELECT pg_get_functiondef( + 'pgedge_vectorizer.generate_embedding(text, text, text)'::regprocedure + ) IS NOT NULL AS function_exists; function_exists ----------------- t @@ -196,7 +198,7 @@ EXCEPTION WHEN OTHERS THEN RAISE NOTICE 'NULL input correctly raises error: %', SQLERRM; END; $$; -NOTICE: WARNING: NULL input should have raised an error +NOTICE: NULL input correctly raises error: query text cannot be NULL -- Test empty string (should work if API key available) -- Configure for empty string test SET pgedge_vectorizer.provider = 'openai'; diff --git a/test/expected/embedding_2.out b/test/expected/embedding_2.out index b373c79..0d4e0a4 100644 --- a/test/expected/embedding_2.out +++ b/test/expected/embedding_2.out @@ -2,7 +2,9 @@ -- Tests the generate_embedding() function with actual API calls -- Gracefully skips tests if API keys are not available -- Test 1: Verify function exists -SELECT pg_get_functiondef('pgedge_vectorizer.generate_embedding(text)'::regprocedure) IS NOT NULL AS function_exists; +SELECT pg_get_functiondef( + 'pgedge_vectorizer.generate_embedding(text, text, text)'::regprocedure + ) IS NOT NULL AS function_exists; function_exists ----------------- t @@ -196,7 +198,7 @@ EXCEPTION WHEN OTHERS THEN RAISE NOTICE 'NULL input correctly raises error: %', SQLERRM; END; $$; -NOTICE: WARNING: NULL input should have raised an error +NOTICE: NULL input correctly raises error: query text cannot be NULL -- Test empty string (should work if API key available) -- Configure for empty string test SET pgedge_vectorizer.provider = 'openai'; diff --git a/test/expected/embedding_3.out b/test/expected/embedding_3.out index a4a8eef..66e5c63 100644 --- a/test/expected/embedding_3.out +++ b/test/expected/embedding_3.out @@ -2,7 +2,9 @@ -- Tests the generate_embedding() function with actual API calls -- Gracefully skips tests if API keys are not available -- Test 1: Verify function exists -SELECT pg_get_functiondef('pgedge_vectorizer.generate_embedding(text)'::regprocedure) IS NOT NULL AS function_exists; +SELECT pg_get_functiondef( + 'pgedge_vectorizer.generate_embedding(text, text, text)'::regprocedure + ) IS NOT NULL AS function_exists; function_exists ----------------- t @@ -196,7 +198,7 @@ EXCEPTION WHEN OTHERS THEN RAISE NOTICE 'NULL input correctly raises error: %', SQLERRM; END; $$; -NOTICE: WARNING: NULL input should have raised an error +NOTICE: NULL input correctly raises error: query text cannot be NULL -- Test empty string (should work if API key available) -- Configure for empty string test SET pgedge_vectorizer.provider = 'openai'; diff --git a/test/expected/embedding_4.out b/test/expected/embedding_4.out index da6732a..4bf6578 100644 --- a/test/expected/embedding_4.out +++ b/test/expected/embedding_4.out @@ -2,7 +2,9 @@ -- Tests the generate_embedding() function with actual API calls -- Gracefully skips tests if API keys are not available -- Test 1: Verify function exists -SELECT pg_get_functiondef('pgedge_vectorizer.generate_embedding(text)'::regprocedure) IS NOT NULL AS function_exists; +SELECT pg_get_functiondef( + 'pgedge_vectorizer.generate_embedding(text, text, text)'::regprocedure + ) IS NOT NULL AS function_exists; function_exists ----------------- t @@ -196,7 +198,7 @@ EXCEPTION WHEN OTHERS THEN RAISE NOTICE 'NULL input correctly raises error: %', SQLERRM; END; $$; -NOTICE: WARNING: NULL input should have raised an error +NOTICE: NULL input correctly raises error: query text cannot be NULL -- Test empty string (should work if API key available) -- Configure for empty string test SET pgedge_vectorizer.provider = 'openai'; diff --git a/test/expected/hybrid_test.out b/test/expected/hybrid_test.out index f2d59ce..2216166 100644 --- a/test/expected/hybrid_test.out +++ b/test/expected/hybrid_test.out @@ -611,8 +611,8 @@ NOTICE: Dropped trigger: hybrid_multi_test_title_1b7742e4_vectorization_truncat NOTICE: Dropped trigger: hybrid_multi_test_body_vectorization_trigger NOTICE: Dropped trigger: hybrid_multi_test_body_69c04c43_vectorization_delete_trigger NOTICE: Dropped trigger: hybrid_multi_test_body_69c04c43_vectorization_truncate_trigger -NOTICE: Vectorization disabled and chunk table dropped: hybrid_multi_test_title_chunks NOTICE: Vectorization disabled and chunk table dropped: hybrid_multi_test_body_chunks +NOTICE: Vectorization disabled and chunk table dropped: hybrid_multi_test_title_chunks disable_vectorization ----------------------- diff --git a/test/expected/per_table_model.out b/test/expected/per_table_model.out new file mode 100644 index 0000000..01116cc --- /dev/null +++ b/test/expected/per_table_model.out @@ -0,0 +1,355 @@ +-- per_table_model test +-- +-- A vectorizer may name its own embedding provider and model, with NULL in +-- either registry column meaning "inherit the GUC". Nothing here reaches a +-- provider: every call that would otherwise probe for a dimension passes one +-- explicitly, and the one provider-resolution case below is rejected on the +-- name before any request is built. +--------------------------------------------------------------------------- +-- Provider resolution happens by name, before any request +--------------------------------------------------------------------------- +-- A provider that does not exist is rejected as such, rather than failing +-- later as a connection error. +DO $$ +BEGIN + PERFORM pgedge_vectorizer.detect_embedding_dimension('nosuchprovider'); + RAISE EXCEPTION 'expected an error, got none'; +EXCEPTION WHEN OTHERS THEN + RAISE NOTICE '%', SQLERRM; +END; +$$; +WARNING: Embedding provider 'nosuchprovider' not found +NOTICE: embedding provider "nosuchprovider" is not available +-- The same for the embedding function, which takes the provider second. +DO $$ +BEGIN + PERFORM pgedge_vectorizer.generate_embedding('some text', 'nosuchprovider'); + RAISE EXCEPTION 'expected an error, got none'; +EXCEPTION WHEN OTHERS THEN + RAISE NOTICE '%', SQLERRM; +END; +$$; +WARNING: Embedding provider 'nosuchprovider' not found +NOTICE: embedding provider "nosuchprovider" is not available +-- Both functions kept their original arity through defaults, so an existing +-- call still resolves. Only the signatures are asserted here; calling either +-- one for real needs a provider. +SELECT p.proname, pg_get_function_arguments(p.oid) AS args + FROM pg_proc p + WHERE p.pronamespace = 'pgedge_vectorizer'::regnamespace + AND p.proname IN ('generate_embedding', 'detect_embedding_dimension') + ORDER BY p.proname; + proname | args +----------------------------+---------------------------------------------------------------------------------- + detect_embedding_dimension | provider text DEFAULT NULL::text, model text DEFAULT NULL::text + generate_embedding | query_text text, provider text DEFAULT NULL::text, model text DEFAULT NULL::text +(2 rows) + +-- Neither may be STRICT: NULL has to reach the C, where it means "fall back +-- to the GUC", and a STRICT function would return NULL before getting there. +SELECT p.proname, p.proisstrict + FROM pg_proc p + WHERE p.pronamespace = 'pgedge_vectorizer'::regnamespace + AND p.proname IN ('generate_embedding', 'detect_embedding_dimension') + ORDER BY p.proname; + proname | proisstrict +----------------------------+------------- + detect_embedding_dimension | f + generate_embedding | f +(2 rows) + +--------------------------------------------------------------------------- +-- What enable_vectorization() records +--------------------------------------------------------------------------- +CREATE TABLE ptm_inherits (id BIGSERIAL PRIMARY KEY, body TEXT); +INSERT INTO ptm_inherits (body) VALUES ('A document that inherits the GUCs.'); +-- No override, so both columns stay NULL and the vectorizer inherits. +SELECT pgedge_vectorizer.enable_vectorization( + 'ptm_inherits'::regclass, 'body', 'token_based', 100, 10, 1536); +NOTICE: Using primary key column: id (bigint) +NOTICE: column "sparse_embedding" of relation "ptm_inherits_body_chunks" already exists, skipping +NOTICE: Vectorization enabled: ptm_inherits -> ptm_inherits_body_chunks +NOTICE: Strategy: token_based, chunk_size: 100, overlap: 10 +NOTICE: Processing existing rows... +NOTICE: Processed 1 existing rows + enable_vectorization +---------------------- + +(1 row) + +CREATE TABLE ptm_pinned (id BIGSERIAL PRIMARY KEY, body TEXT); +INSERT INTO ptm_pinned (body) VALUES ('A document with a pinned model.'); +-- With an override, both are recorded exactly as passed. +SELECT pgedge_vectorizer.enable_vectorization( + 'ptm_pinned'::regclass, 'body', 'token_based', 100, 10, 1536, + NULL, NULL, 'ollama', 'nomic-embed-text'); +NOTICE: Using primary key column: id (bigint) +NOTICE: column "sparse_embedding" of relation "ptm_pinned_body_chunks" already exists, skipping +NOTICE: Vectorization enabled: ptm_pinned -> ptm_pinned_body_chunks +NOTICE: Strategy: token_based, chunk_size: 100, overlap: 10 +NOTICE: Processing existing rows... +NOTICE: Processed 1 existing rows + enable_vectorization +---------------------- + +(1 row) + +SELECT source_table, source_column, provider, model + FROM pgedge_vectorizer.vectorizers + WHERE source_table LIKE 'ptm_%' + ORDER BY source_table; + source_table | source_column | provider | model +--------------+---------------+----------+------------------ + ptm_inherits | body | | + ptm_pinned | body | ollama | nomic-embed-text +(2 rows) + +-- Named notation reads the same on both functions, and skips the +-- positional parameters nobody wants to spell out. +CREATE TABLE ptm_named (id BIGSERIAL PRIMARY KEY, body TEXT); +SELECT pgedge_vectorizer.enable_vectorization( + 'ptm_named'::regclass, 'body', + embedding_dimension => 1536, + model => 'text-embedding-3-large'); +NOTICE: Using primary key column: id (bigint) +NOTICE: column "sparse_embedding" of relation "ptm_named_body_chunks" already exists, skipping +NOTICE: Vectorization enabled: ptm_named -> ptm_named_body_chunks +NOTICE: Strategy: token_based, chunk_size: 400, overlap: 50 +NOTICE: Processing existing rows... +NOTICE: Processed 0 existing rows + enable_vectorization +---------------------- + +(1 row) + +SELECT provider, model + FROM pgedge_vectorizer.vectorizers WHERE source_table = 'ptm_named'; + provider | model +----------+------------------------ + | text-embedding-3-large +(1 row) + +--------------------------------------------------------------------------- +-- set_embedding_model() +--------------------------------------------------------------------------- +-- A table with no vectorizer is an error, not a silent no-op. +DO $$ +BEGIN + PERFORM pgedge_vectorizer.set_embedding_model( + 'ptm_named'::regclass, 'nosuchcolumn', 'some-model'); + RAISE EXCEPTION 'expected an error, got none'; +EXCEPTION WHEN OTHERS THEN + RAISE NOTICE '%', SQLERRM; +END; +$$; +NOTICE: no vectorizer registered for ptm_named.nosuchcolumn +SET pgedge_vectorizer.provider = 'openai'; +SET pgedge_vectorizer.model = 'text-embedding-3-small'; +-- ptm_named was pinned to text-embedding-3-large at creation, so this is a +-- real change. It goes through without complaint because the vectorizer has +-- no chunks yet: there is nothing embedded for it to invalidate. +SELECT pgedge_vectorizer.set_embedding_model( + 'ptm_named'::regclass, 'body', 'text-embedding-3-small', + embedding_dimension => 1536) AS requeued; + requeued +---------- + 0 +(1 row) + +SELECT provider, model + FROM pgedge_vectorizer.vectorizers WHERE source_table = 'ptm_named'; + provider | model +----------+------------------------ + | text-embedding-3-small +(1 row) + +-- An empty vectorizer still has its column rewidened. Skipping that because +-- there was nothing to re-embed would leave the table at its old width and +-- fail every embedding the worker later tried to write, which is exactly the +-- failure this function exists to prevent. +SELECT format_type(a.atttypid, a.atttypmod) AS before_width + FROM pg_attribute a + WHERE a.attrelid = 'ptm_named_body_chunks'::regclass + AND a.attname = 'embedding'; + before_width +-------------- + vector(1536) +(1 row) + +SELECT pgedge_vectorizer.set_embedding_model( + 'ptm_named'::regclass, 'body', 'nomic-embed-text', + provider => 'ollama', embedding_dimension => 768) AS requeued; +NOTICE: Embedding dimension changed from 1536 to 768 + requeued +---------- + 0 +(1 row) + +SELECT format_type(a.atttypid, a.atttypmod) AS after_width + FROM pg_attribute a + WHERE a.attrelid = 'ptm_named_body_chunks'::regclass + AND a.attname = 'embedding'; + after_width +------------- + vector(768) +(1 row) + +-- Put it back, pinning both this time, so that the reset below starts from a +-- vectorizer that has actually overridden something and can be seen to give +-- both up. The provider named here matches the GUC, so the effective values +-- do not move and nothing is re-embedded. +SELECT pgedge_vectorizer.set_embedding_model( + 'ptm_named'::regclass, 'body', 'text-embedding-3-small', + provider => 'openai', embedding_dimension => 1536) AS requeued; +NOTICE: Embedding dimension changed from 768 to 1536 + requeued +---------- + 0 +(1 row) + +SELECT provider, model + FROM pgedge_vectorizer.vectorizers WHERE source_table = 'ptm_named'; + provider | model +----------+------------------------ + openai | text-embedding-3-small +(1 row) + +-- Now the genuine no-op. Reverting to the GUC is a NULL model, and whilst the +-- GUC names what was pinned the effective model does not move, so nothing is +-- requeued even though the stored value changes. +SELECT pgedge_vectorizer.set_embedding_model( + 'ptm_named'::regclass, 'body', NULL) AS requeued; +NOTICE: Effective provider and model unchanged (openai/text-embedding-3-small) + requeued +---------- + 0 +(1 row) + +SELECT provider, model + FROM pgedge_vectorizer.vectorizers WHERE source_table = 'ptm_named'; + provider | model +----------+------- + | +(1 row) + +-- Populate a vectorizer and mark it embedded, as the worker would. +UPDATE ptm_inherits_body_chunks + SET embedding = array_fill(0.1::real, ARRAY[1536])::vector, + sparse_embedding = '{1:0.5}/65536'::sparsevec; +SELECT count(*) AS chunks, + count(embedding) AS embedded, + count(sparse_embedding) AS sparse, + count(token_count) AS counted + FROM ptm_inherits_body_chunks; + chunks | embedded | sparse | counted +--------+----------+--------+--------- + 1 | 1 | 1 | 1 +(1 row) + +-- Now the refusal. The message names both settings and the number of chunks. +DO $$ +BEGIN + PERFORM pgedge_vectorizer.set_embedding_model( + 'ptm_inherits'::regclass, 'body', 'nomic-embed-text', + provider => 'ollama', embedding_dimension => 768); + RAISE EXCEPTION 'expected an error, got none'; +EXCEPTION WHEN OTHERS THEN + RAISE NOTICE '%', SQLERRM; +END; +$$; +NOTICE: changing the embedding model for ptm_inherits.body would leave 1 chunks embedded with openai/text-embedding-3-small whilst everything after uses ollama/nomic-embed-text +-- Nothing was touched by the refusal. +SELECT count(embedding) AS still_embedded FROM ptm_inherits_body_chunks; + still_embedded +---------------- + 1 +(1 row) + +SELECT provider, model + FROM pgedge_vectorizer.vectorizers WHERE source_table = 'ptm_inherits'; + provider | model +----------+------- + | +(1 row) + +-- With force_reembed the change goes through: every embedding cleared, the +-- column rewidened, every chunk requeued, and the chunks themselves left +-- exactly as they were. +SELECT pgedge_vectorizer.set_embedding_model( + 'ptm_inherits'::regclass, 'body', 'nomic-embed-text', + provider => 'ollama', embedding_dimension => 768, + force_reembed => true) AS requeued; +NOTICE: Embedding dimension changed from 1536 to 768 +NOTICE: Requeued 1 chunks for re-embedding with ollama/nomic-embed-text + requeued +---------- + 1 +(1 row) + +SELECT count(*) AS chunks, + count(embedding) AS embedded, + count(sparse_embedding) AS sparse_kept, + count(token_count) AS token_counts_kept + FROM ptm_inherits_body_chunks; + chunks | embedded | sparse_kept | token_counts_kept +--------+----------+-------------+------------------- + 1 | 0 | 1 | 1 +(1 row) + +SELECT format_type(a.atttypid, a.atttypmod) AS embedding_type + FROM pg_attribute a + WHERE a.attrelid = 'ptm_inherits_body_chunks'::regclass + AND a.attname = 'embedding'; + embedding_type +---------------- + vector(768) +(1 row) + +SELECT count(*) AS queued + FROM pgedge_vectorizer.queue + WHERE chunk_table = 'ptm_inherits_body_chunks' AND status = 'pending'; + queued +-------- + 1 +(1 row) + +SELECT provider, model + FROM pgedge_vectorizer.vectorizers WHERE source_table = 'ptm_inherits'; + provider | model +----------+------------------ + ollama | nomic-embed-text +(1 row) + +--------------------------------------------------------------------------- +-- Cleanup +--------------------------------------------------------------------------- +SELECT pgedge_vectorizer.disable_vectorization('ptm_inherits'::regclass, + 'body', TRUE); +NOTICE: Vectorization disabled and chunk table dropped: ptm_inherits_body_chunks + disable_vectorization +----------------------- + +(1 row) + +SELECT pgedge_vectorizer.disable_vectorization('ptm_pinned'::regclass, + 'body', TRUE); +NOTICE: Vectorization disabled and chunk table dropped: ptm_pinned_body_chunks + disable_vectorization +----------------------- + +(1 row) + +SELECT pgedge_vectorizer.disable_vectorization('ptm_named'::regclass, + 'body', TRUE); +NOTICE: Vectorization disabled and chunk table dropped: ptm_named_body_chunks + disable_vectorization +----------------------- + +(1 row) + +DROP TABLE ptm_inherits; +DROP TABLE ptm_pinned; +DROP TABLE ptm_named; +DELETE FROM pgedge_vectorizer.queue; +RESET pgedge_vectorizer.provider; +RESET pgedge_vectorizer.model; diff --git a/test/expected/pk_types.out b/test/expected/pk_types.out index 152b612..9f52e7d 100644 --- a/test/expected/pk_types.out +++ b/test/expected/pk_types.out @@ -378,7 +378,7 @@ SELECT pgedge_vectorizer.enable_vectorization( 1536 ); ERROR: Table test_composite_pk has a composite primary key (2 columns), which is not supported by auto-detection. Use the source_pk parameter to specify a single column. -CONTEXT: PL/pgSQL function pgedge_vectorizer.enable_vectorization(regclass,name,text,integer,integer,integer,text,name) line 40 at RAISE +CONTEXT: PL/pgSQL function pgedge_vectorizer.enable_vectorization(regclass,name,text,integer,integer,integer,text,name,text,text) line 43 at RAISE -- Clean up (no vectorization to disable, just drop the table) DROP TABLE test_composite_pk; -- ============================================================================ @@ -466,7 +466,7 @@ SELECT pgedge_vectorizer.enable_vectorization( 1536 ); ERROR: Table test_no_pk has no primary key. Use the source_pk parameter to specify the column to use as document identifier. -CONTEXT: PL/pgSQL function pgedge_vectorizer.enable_vectorization(regclass,name,text,integer,integer,integer,text,name) line 35 at RAISE +CONTEXT: PL/pgSQL function pgedge_vectorizer.enable_vectorization(regclass,name,text,integer,integer,integer,text,name,text,text) line 38 at RAISE -- Clean up DROP TABLE test_no_pk; -- ============================================================================ diff --git a/test/sql/count_tokens.sql b/test/sql/count_tokens.sql new file mode 100644 index 0000000..06d3e95 --- /dev/null +++ b/test/sql/count_tokens.sql @@ -0,0 +1,99 @@ +-- count_tokens test +-- +-- count_tokens() exposes the estimate the C chunking code has always used, so +-- that the plpgsql paths which write the token_count column can call it rather +-- than open-coding the same rule and rounding it the other way. The point of +-- the tests below is therefore as much the agreement between the two as the +-- values themselves. + +--------------------------------------------------------------------------- +-- The estimate itself: UTF-8 characters divided by four, rounded up +--------------------------------------------------------------------------- + +-- 11 characters, so 3 tokens. +SELECT pgedge_vectorizer.count_tokens('hello world') AS eleven_chars; + +-- Exactly 4 characters is exactly 1 token; anything shorter still rounds up +-- to 1, which is the difference from the truncating arithmetic this replaces. +SELECT pgedge_vectorizer.count_tokens('test') AS four_chars, + pgedge_vectorizer.count_tokens('abc') AS three_chars, + pgedge_vectorizer.count_tokens('a') AS one_char; + +-- Empty text is the one case that is genuinely zero. +SELECT pgedge_vectorizer.count_tokens('') AS empty; + +-- STRICT, so NULL in, NULL out. +SELECT pgedge_vectorizer.count_tokens(NULL) IS NULL AS null_is_null; + +-- Characters, not bytes: four Han characters are twelve bytes but one token. +SELECT pgedge_vectorizer.count_tokens('你好世界') AS four_han_chars, + octet_length('你好世界') AS bytes; + +-- Declared STABLE rather than IMMUTABLE: the estimate is defined in terms of +-- pgedge_vectorizer.model, which will matter once the counter is model-aware, +-- and an index or cached plan built on an IMMUTABLE promise would then be +-- wrong. Pin the volatility so that cannot be relaxed by accident. +SELECT provolatile + FROM pg_proc + WHERE proname = 'count_tokens' + AND pronamespace = 'pgedge_vectorizer'::regnamespace; + +--------------------------------------------------------------------------- +-- Agreement with what the chunking paths store +--------------------------------------------------------------------------- + +CREATE TABLE count_tokens_docs ( + id BIGSERIAL PRIMARY KEY, + content TEXT +); + +-- Chunks written by enable_vectorization() back-filling an existing table. +INSERT INTO count_tokens_docs (content) +VALUES ('abc'), + ('Short document.'), + (repeat('The quick brown fox jumps over the lazy dog. ', 20)); + +SELECT pgedge_vectorizer.enable_vectorization( + 'count_tokens_docs'::regclass, + 'content', + 'token_based', + 100, + 10, + 1536 +); + +SELECT count(*) AS mismatched_on_backfill + FROM count_tokens_docs_content_chunks + WHERE token_count IS DISTINCT FROM pgedge_vectorizer.count_tokens(content); + +-- The three-character row is the regression: length('abc') / 4 stored 0, which +-- the BM25 scoring path then had to clamp back up to 1. +SELECT token_count AS short_row_token_count + FROM count_tokens_docs_content_chunks c + JOIN count_tokens_docs d ON d.id = c.source_id + WHERE d.content = 'abc'; + +-- Chunks written by the insert trigger. +INSERT INTO count_tokens_docs (content) +VALUES ('xy'), + ('A document added after vectorization was enabled.'); + +SELECT count(*) AS mismatched_on_insert + FROM count_tokens_docs_content_chunks + WHERE token_count IS DISTINCT FROM pgedge_vectorizer.count_tokens(content); + +-- Chunks written by recreate_chunks(). +SELECT pgedge_vectorizer.recreate_chunks('count_tokens_docs'::regclass, 'content'); + +SELECT count(*) AS mismatched_on_recreate + FROM count_tokens_docs_content_chunks + WHERE token_count IS DISTINCT FROM pgedge_vectorizer.count_tokens(content); + +--------------------------------------------------------------------------- +-- Cleanup +--------------------------------------------------------------------------- + +SELECT pgedge_vectorizer.disable_vectorization('count_tokens_docs'::regclass, + 'content', TRUE); +DROP TABLE count_tokens_docs; +DELETE FROM pgedge_vectorizer.queue; diff --git a/test/sql/embedding.sql b/test/sql/embedding.sql index 1bbee31..35cdded 100644 --- a/test/sql/embedding.sql +++ b/test/sql/embedding.sql @@ -3,7 +3,9 @@ -- Gracefully skips tests if API keys are not available -- Test 1: Verify function exists -SELECT pg_get_functiondef('pgedge_vectorizer.generate_embedding(text)'::regprocedure) IS NOT NULL AS function_exists; +SELECT pg_get_functiondef( + 'pgedge_vectorizer.generate_embedding(text, text, text)'::regprocedure + ) IS NOT NULL AS function_exists; -- Helper function to check if API key file exists and is readable CREATE OR REPLACE FUNCTION test_api_key_available() RETURNS boolean AS $$ diff --git a/test/sql/per_table_model.sql b/test/sql/per_table_model.sql new file mode 100644 index 0000000..770402e --- /dev/null +++ b/test/sql/per_table_model.sql @@ -0,0 +1,222 @@ +-- per_table_model test +-- +-- A vectorizer may name its own embedding provider and model, with NULL in +-- either registry column meaning "inherit the GUC". Nothing here reaches a +-- provider: every call that would otherwise probe for a dimension passes one +-- explicitly, and the one provider-resolution case below is rejected on the +-- name before any request is built. + +--------------------------------------------------------------------------- +-- Provider resolution happens by name, before any request +--------------------------------------------------------------------------- + +-- A provider that does not exist is rejected as such, rather than failing +-- later as a connection error. +DO $$ +BEGIN + PERFORM pgedge_vectorizer.detect_embedding_dimension('nosuchprovider'); + RAISE EXCEPTION 'expected an error, got none'; +EXCEPTION WHEN OTHERS THEN + RAISE NOTICE '%', SQLERRM; +END; +$$; + +-- The same for the embedding function, which takes the provider second. +DO $$ +BEGIN + PERFORM pgedge_vectorizer.generate_embedding('some text', 'nosuchprovider'); + RAISE EXCEPTION 'expected an error, got none'; +EXCEPTION WHEN OTHERS THEN + RAISE NOTICE '%', SQLERRM; +END; +$$; + +-- Both functions kept their original arity through defaults, so an existing +-- call still resolves. Only the signatures are asserted here; calling either +-- one for real needs a provider. +SELECT p.proname, pg_get_function_arguments(p.oid) AS args + FROM pg_proc p + WHERE p.pronamespace = 'pgedge_vectorizer'::regnamespace + AND p.proname IN ('generate_embedding', 'detect_embedding_dimension') + ORDER BY p.proname; + +-- Neither may be STRICT: NULL has to reach the C, where it means "fall back +-- to the GUC", and a STRICT function would return NULL before getting there. +SELECT p.proname, p.proisstrict + FROM pg_proc p + WHERE p.pronamespace = 'pgedge_vectorizer'::regnamespace + AND p.proname IN ('generate_embedding', 'detect_embedding_dimension') + ORDER BY p.proname; + +--------------------------------------------------------------------------- +-- What enable_vectorization() records +--------------------------------------------------------------------------- + +CREATE TABLE ptm_inherits (id BIGSERIAL PRIMARY KEY, body TEXT); +INSERT INTO ptm_inherits (body) VALUES ('A document that inherits the GUCs.'); + +-- No override, so both columns stay NULL and the vectorizer inherits. +SELECT pgedge_vectorizer.enable_vectorization( + 'ptm_inherits'::regclass, 'body', 'token_based', 100, 10, 1536); + +CREATE TABLE ptm_pinned (id BIGSERIAL PRIMARY KEY, body TEXT); +INSERT INTO ptm_pinned (body) VALUES ('A document with a pinned model.'); + +-- With an override, both are recorded exactly as passed. +SELECT pgedge_vectorizer.enable_vectorization( + 'ptm_pinned'::regclass, 'body', 'token_based', 100, 10, 1536, + NULL, NULL, 'ollama', 'nomic-embed-text'); + +SELECT source_table, source_column, provider, model + FROM pgedge_vectorizer.vectorizers + WHERE source_table LIKE 'ptm_%' + ORDER BY source_table; + +-- Named notation reads the same on both functions, and skips the +-- positional parameters nobody wants to spell out. +CREATE TABLE ptm_named (id BIGSERIAL PRIMARY KEY, body TEXT); + +SELECT pgedge_vectorizer.enable_vectorization( + 'ptm_named'::regclass, 'body', + embedding_dimension => 1536, + model => 'text-embedding-3-large'); + +SELECT provider, model + FROM pgedge_vectorizer.vectorizers WHERE source_table = 'ptm_named'; + +--------------------------------------------------------------------------- +-- set_embedding_model() +--------------------------------------------------------------------------- + +-- A table with no vectorizer is an error, not a silent no-op. +DO $$ +BEGIN + PERFORM pgedge_vectorizer.set_embedding_model( + 'ptm_named'::regclass, 'nosuchcolumn', 'some-model'); + RAISE EXCEPTION 'expected an error, got none'; +EXCEPTION WHEN OTHERS THEN + RAISE NOTICE '%', SQLERRM; +END; +$$; + +SET pgedge_vectorizer.provider = 'openai'; +SET pgedge_vectorizer.model = 'text-embedding-3-small'; + +-- ptm_named was pinned to text-embedding-3-large at creation, so this is a +-- real change. It goes through without complaint because the vectorizer has +-- no chunks yet: there is nothing embedded for it to invalidate. +SELECT pgedge_vectorizer.set_embedding_model( + 'ptm_named'::regclass, 'body', 'text-embedding-3-small', + embedding_dimension => 1536) AS requeued; + +SELECT provider, model + FROM pgedge_vectorizer.vectorizers WHERE source_table = 'ptm_named'; + +-- An empty vectorizer still has its column rewidened. Skipping that because +-- there was nothing to re-embed would leave the table at its old width and +-- fail every embedding the worker later tried to write, which is exactly the +-- failure this function exists to prevent. +SELECT format_type(a.atttypid, a.atttypmod) AS before_width + FROM pg_attribute a + WHERE a.attrelid = 'ptm_named_body_chunks'::regclass + AND a.attname = 'embedding'; + +SELECT pgedge_vectorizer.set_embedding_model( + 'ptm_named'::regclass, 'body', 'nomic-embed-text', + provider => 'ollama', embedding_dimension => 768) AS requeued; + +SELECT format_type(a.atttypid, a.atttypmod) AS after_width + FROM pg_attribute a + WHERE a.attrelid = 'ptm_named_body_chunks'::regclass + AND a.attname = 'embedding'; + +-- Put it back, pinning both this time, so that the reset below starts from a +-- vectorizer that has actually overridden something and can be seen to give +-- both up. The provider named here matches the GUC, so the effective values +-- do not move and nothing is re-embedded. +SELECT pgedge_vectorizer.set_embedding_model( + 'ptm_named'::regclass, 'body', 'text-embedding-3-small', + provider => 'openai', embedding_dimension => 1536) AS requeued; + +SELECT provider, model + FROM pgedge_vectorizer.vectorizers WHERE source_table = 'ptm_named'; + +-- Now the genuine no-op. Reverting to the GUC is a NULL model, and whilst the +-- GUC names what was pinned the effective model does not move, so nothing is +-- requeued even though the stored value changes. +SELECT pgedge_vectorizer.set_embedding_model( + 'ptm_named'::regclass, 'body', NULL) AS requeued; + +SELECT provider, model + FROM pgedge_vectorizer.vectorizers WHERE source_table = 'ptm_named'; + +-- Populate a vectorizer and mark it embedded, as the worker would. +UPDATE ptm_inherits_body_chunks + SET embedding = array_fill(0.1::real, ARRAY[1536])::vector, + sparse_embedding = '{1:0.5}/65536'::sparsevec; + +SELECT count(*) AS chunks, + count(embedding) AS embedded, + count(sparse_embedding) AS sparse, + count(token_count) AS counted + FROM ptm_inherits_body_chunks; + +-- Now the refusal. The message names both settings and the number of chunks. +DO $$ +BEGIN + PERFORM pgedge_vectorizer.set_embedding_model( + 'ptm_inherits'::regclass, 'body', 'nomic-embed-text', + provider => 'ollama', embedding_dimension => 768); + RAISE EXCEPTION 'expected an error, got none'; +EXCEPTION WHEN OTHERS THEN + RAISE NOTICE '%', SQLERRM; +END; +$$; + +-- Nothing was touched by the refusal. +SELECT count(embedding) AS still_embedded FROM ptm_inherits_body_chunks; +SELECT provider, model + FROM pgedge_vectorizer.vectorizers WHERE source_table = 'ptm_inherits'; + +-- With force_reembed the change goes through: every embedding cleared, the +-- column rewidened, every chunk requeued, and the chunks themselves left +-- exactly as they were. +SELECT pgedge_vectorizer.set_embedding_model( + 'ptm_inherits'::regclass, 'body', 'nomic-embed-text', + provider => 'ollama', embedding_dimension => 768, + force_reembed => true) AS requeued; + +SELECT count(*) AS chunks, + count(embedding) AS embedded, + count(sparse_embedding) AS sparse_kept, + count(token_count) AS token_counts_kept + FROM ptm_inherits_body_chunks; + +SELECT format_type(a.atttypid, a.atttypmod) AS embedding_type + FROM pg_attribute a + WHERE a.attrelid = 'ptm_inherits_body_chunks'::regclass + AND a.attname = 'embedding'; + +SELECT count(*) AS queued + FROM pgedge_vectorizer.queue + WHERE chunk_table = 'ptm_inherits_body_chunks' AND status = 'pending'; + +SELECT provider, model + FROM pgedge_vectorizer.vectorizers WHERE source_table = 'ptm_inherits'; + +--------------------------------------------------------------------------- +-- Cleanup +--------------------------------------------------------------------------- + +SELECT pgedge_vectorizer.disable_vectorization('ptm_inherits'::regclass, + 'body', TRUE); +SELECT pgedge_vectorizer.disable_vectorization('ptm_pinned'::regclass, + 'body', TRUE); +SELECT pgedge_vectorizer.disable_vectorization('ptm_named'::regclass, + 'body', TRUE); +DROP TABLE ptm_inherits; +DROP TABLE ptm_pinned; +DROP TABLE ptm_named; +DELETE FROM pgedge_vectorizer.queue; +RESET pgedge_vectorizer.provider; +RESET pgedge_vectorizer.model; diff --git a/test/t/011_per_table_model.pl b/test/t/011_per_table_model.pl new file mode 100644 index 0000000..69d637a --- /dev/null +++ b/test/t/011_per_table_model.pl @@ -0,0 +1,228 @@ +# Copyright (c) 2025 - 2026, pgEdge, Inc. +# +# Verify that each vectorizer's chunks are embedded with its own model, and +# that no single request ever carries two. +# +# The model used to come straight from pgedge_vectorizer.model inside each +# provider, so every table in a database was embedded with the same one. A +# vectorizer may now pin its own, with NULL in the registry meaning inherit, +# which puts two demands on the worker: the right model has to reach the +# request, and a batch that spans two vectorizers has to be split, because a +# request carries one model for every text in it. +# +# The second is the one worth a test with a real worker. A batch is selected +# by age across every vectorizer at once, so the two tables' items interleave +# in the queue; the worker groups them before issuing anything. Both tables are +# populated in one transaction so that a single poll is guaranteed to see all +# six items, which is the case that would otherwise send one model's text under +# the other's name. + +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 = 'per_table_model'; +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 $requestlog = "$tempdir/requests.log"; + +my $server_pid = fork(); +die "fork failed: $!" unless defined $server_pid; + +if ($server_pid == 0) +{ + fake_provider($listener, $requestlog); + 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_per_table_model'); +$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 = 'aaa-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)'); + +# One inherits the GUC, the other pins its own. The names sort in the order the +# worker groups them, so the request log below is deterministic. +$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 => 'zzz-pinned-model'))); + +my $registry = $node->safe_psql($dbname, + q(SELECT string_agg(source_table || '=' || COALESCE(model, 'NULL'), ' ' + ORDER BY source_table) + FROM pgedge_vectorizer.vectorizers)); + +is($registry, 'inherits=NULL pinned=zzz-pinned-model', + 'a vectorizer records its own model, and NULL where it inherits'); + +# 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; + +# One transaction, so a single poll is guaranteed to see both tables' items. +# That is the case being tested: a batch holding two models at once. +$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'); + +# Each line is one request: the model it named and how many texts it carried. +my @requests = split /\n/, slurp_file($requestlog); + +is(scalar(@requests), 2, + 'a batch spanning two models is split into one request per model'); + +is($requests[0], "aaa-inherited-model $rows", + 'the inheriting table is embedded with the GUC\'s model, all in one request'); +is($requests[1], "zzz-pinned-model $rows", + 'the pinned table is embedded with its own model, all in one request'); + +# Nothing failed, which is what a request carrying the wrong model would have +# risked once the dimensions differed. +my $failed = $node->safe_psql($dbname, + "SELECT count(*) FROM pgedge_vectorizer.queue WHERE status = 'failed'"); + +is($failed, '0', 'no item fails on the way'); + +$node->stop; + +kill 'TERM', $server_pid; +waitpid $server_pid, 0; + +done_testing(); + +# Answer every request, logging the model it named and the number of texts it +# carried. Those two together are what the assertions above read: a request +# naming one model but carrying another table's text would show up as a count +# that does not match, and a request mixing the two cannot be represented at +# all, because the provider API takes one model per request. +sub fake_provider +{ + my ($socket, $logfile) = @_; + + open my $log, '>', $logfile or die "could not write $logfile: $!"; + $log->autoflush(1); + + while (my $conn = $socket->accept()) + { + my $headers = ''; + my $body = ''; + my $length; + my $inputs = 0; + my $texts; + my $model; + 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; + + ($model) = $body =~ /"model"\s*:\s*"([^"]*)"/; + $model = '(none)' unless defined $model; + + $log->print("$model $inputs\n"); + + $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; + } +} diff --git a/test/t/012_upgrade_1_1_to_1_2.pl b/test/t/012_upgrade_1_1_to_1_2.pl new file mode 100644 index 0000000..4adc070 --- /dev/null +++ b/test/t/012_upgrade_1_1_to_1_2.pl @@ -0,0 +1,130 @@ +# Copyright (c) 2025 - 2026, pgEdge, Inc. +# +# Verify that an installation created at 1.1 upgrades to 1.2 and ends up with +# the same objects a fresh 1.2 install has. +# +# pg_regress cannot check this. Its database installs whatever the control +# file's default_version says, so every regression test only ever exercises a +# fresh install, and the upgrade script goes untested however wrong it is. That +# is not hypothetical: adding a defaulted parameter to enable_vectorization() +# with CREATE OR REPLACE defined a second function rather than replacing the +# old one, leaving two overloads behind, an eight-argument call reaching a body +# that knew nothing of the registry's new columns, and COMMENT ON FUNCTION +# failing outright as ambiguous. A fresh install was perfect throughout. +# +# The comparison below is deliberately structural rather than a list of names +# to keep in step: whatever a fresh 1.2 install has, an upgraded one must have +# too. + +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; + +my $node = PostgreSQL::Test::Cluster->new('vectorizer_upgrade'); +$node->init; +$node->append_conf( + 'postgresql.conf', qq( +shared_preload_libraries = 'pgedge_vectorizer' +max_worker_processes = 16 +)); +$node->start; + +# No databases are named, so no worker ever runs against either of these and +# nothing tries to reach a provider. +for my $db ('upgraded', 'fresh') +{ + $node->safe_psql('postgres', "CREATE DATABASE $db"); + $node->safe_psql($db, 'CREATE EXTENSION vector'); +} + +$node->safe_psql('upgraded', + "CREATE EXTENSION pgedge_vectorizer VERSION '1.1'"); + +is($node->safe_psql('upgraded', + "SELECT extversion FROM pg_extension WHERE extname = 'pgedge_vectorizer'"), + '1.1', 'the extension installs at 1.1'); + +# Real state before the upgrade, so the script runs against a populated +# registry and a chunk table rather than an empty schema. +$node->safe_psql('upgraded', q( +CREATE TABLE docs (id BIGSERIAL PRIMARY KEY, body TEXT); +INSERT INTO docs (body) VALUES ('Written while the extension was at 1.1.'); +)); +$node->safe_psql('upgraded', + q(SELECT pgedge_vectorizer.enable_vectorization('docs', 'body', + 'token_based', 100, 10, 1536))); + +$node->safe_psql('upgraded', + "ALTER EXTENSION pgedge_vectorizer UPDATE TO '1.2'"); + +is($node->safe_psql('upgraded', + "SELECT extversion FROM pg_extension WHERE extname = 'pgedge_vectorizer'"), + '1.2', 'the extension upgrades to 1.2'); + +$node->safe_psql('fresh', 'CREATE EXTENSION pgedge_vectorizer'); + +is($node->safe_psql('fresh', + "SELECT extversion FROM pg_extension WHERE extname = 'pgedge_vectorizer'"), + '1.2', 'a fresh install is 1.2, so the two are comparable'); + +# Every function, by name and argument list. An overload left behind by a +# CREATE OR REPLACE that should have been a DROP shows up here as an extra row. +my $signatures = q( + SELECT string_agg(p.proname || '(' || pg_get_function_arguments(p.oid) || ')', + E'\n' ORDER BY p.proname, pg_get_function_arguments(p.oid)) + FROM pg_proc p + WHERE p.pronamespace = 'pgedge_vectorizer'::regnamespace +); + +is($node->safe_psql('upgraded', $signatures), + $node->safe_psql('fresh', $signatures), + 'an upgraded install has exactly the functions a fresh one has'); + +# Columns of the extension's own tables, so a missed ALTER TABLE is caught. +my $columns = q( + SELECT string_agg(c.relname || '.' || a.attname || ' ' || + format_type(a.atttypid, a.atttypmod), + E'\n' ORDER BY c.relname, a.attname) + FROM pg_class c + JOIN pg_attribute a ON a.attrelid = c.oid + WHERE c.relnamespace = 'pgedge_vectorizer'::regnamespace + AND c.relkind = 'r' + AND a.attnum > 0 + AND NOT a.attisdropped +); + +is($node->safe_psql('upgraded', $columns), + $node->safe_psql('fresh', $columns), + 'an upgraded install has the same table columns as a fresh one'); + +# Views too, since those are replaced rather than altered. +my $views = q( + SELECT string_agg(c.relname, E'\n' ORDER BY c.relname) + FROM pg_class c + WHERE c.relnamespace = 'pgedge_vectorizer'::regnamespace + AND c.relkind = 'v' +); + +is($node->safe_psql('upgraded', $views), + $node->safe_psql('fresh', $views), + 'an upgraded install has the same views as a fresh 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)), + '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'); + +$node->stop; + +done_testing();