Retrievers end to end example usage

This is a full end-to-end example of using retrievers in EDB Postgres AI - AI Accelerator.

-- Create extension
DROP EXTENSION aidb CASCADE;
CREATE EXTENSION aidb CASCADE;

drop table if exists test_source_table cascade;
drop table if exists test_retriever_vector cascade;

-- Create source test table
CREATE TABLE test_source_table
(
    id               INT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
    content          TEXT NOT NULL,
    unrelated_column TEXT
);
INSERT INTO test_source_table
VALUES (43941, 'Catwalk Women Brown Heels'),
       (55018, 'Lakme 3 in 1 Orchid  Aqua Shine Lip Color'),
       (19337, 'United Colors of Benetton Men Stripes Black Jacket');


-- Register model
SELECT aidb.register_model('simple_model', 'bert_local');

SELECT aidb.register_retriever_for_table(
    p_name => 'test_retriever',
    p_model_name => 'simple_model',
    p_source_table_name => 'test_source_table',
    p_source_table_data_column => 'content',
    p_source_table_data_column_type => 'Text',
    p_source_table_key_column => 'id',                  -- Default
    p_vector_table_name => 'test_source_table_vector',  -- Defaults to `source_table_name + '_vector'`
    p_vector_table_vector_column => 'embeddings',       -- Default
    p_vector_table_key_column => 'id',                  -- Default
    p_topk => 1,                                        -- Default
    p_distance_operator => 'L2',                        -- Default
    p_options => '{}'::JSONB                            -- Default
);


-- expect "Table"
SELECT aidb.get_retriever_data_source('test_retriever');

SELECT * FROM aidb.retrievers;

SELECT aidb.bulk_embedding('test_retriever');


-- Perform retrieval similarity search for the closest `key`
SELECT * FROM aidb.retrieve_key('test_retriever', 'orchid');
SELECT * FROM aidb.retrieve_key('test_retriever', 'orchid', 2);  -- Limit to top 2 results

SELECT * FROM aidb.retrieve_text('test_retriever', 'orchid');
SELECT * FROM aidb.retrieve_text('test_retriever', 'orchid', 2);  -- Limit to top 2 results


-- enable the auto embedding
SELECT aidb.enable_auto_embedding_for_table('test_retriever');


-- add additional data to test auto-embedding
INSERT INTO test_source_table
VALUES (11211, 'Bicycle'),
       (11311, 'What is this?'),
       (11411, 'Elephants');


-- check embeddings
SELECT id FROM test_retriever_vector;

-- delete one of the source rows
DELETE FROM test_source_table WHERE id = 11211;

-- check embeddings
SELECT id FROM test_retriever_vector;

-- enable the auto embedding
SELECT aidb.disable_auto_embedding_for_table('test_retriever');
INSERT INTO test_source_table VALUES (212121, 'new value');


select aidb.delete_retriever('test_retriever');

Could this page be better? Report a problem or suggest an addition!