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_ajz72eb cascade;
drop table if exists test_retriever_ajz72eb_vector cascade;

-- Create source test table-- Create source test table
CREATE TABLE test_source_table_ajz72eb
(
    id               INT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
    content          TEXT NOT NULL,
    unrelated_column TEXT
);
INSERT INTO test_source_table_ajz72eb
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.create_model('simple_model_ajz72eb', 'bert_local');

SELECT aidb.create_retriever_for_table(
    name => 'test_retriever_ajz72eb',
    model_name => 'simple_model_ajz72eb',
    source_table => 'test_source_table_ajz72eb',
    source_data_column => 'content',
    source_data_type => 'Text'
);


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

SELECT * FROM aidb.retrievers;

SELECT aidb.bulk_embedding('test_retriever_ajz72eb');


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

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


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


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


-- check embeddings
SELECT id FROM test_retriever_ajz72eb_vector;

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

-- check embeddings
SELECT id FROM test_retriever_ajz72eb_vector;

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

select aidb.delete_retriever('test_retriever_ajz72eb');

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