tencentdb_ai.add_incr_autoembedding_task(
    schema_name    NAME,        -- schema where the table resides
    table_name     NAME,        -- table name
    source_columns NAME[],      -- array of source text column names
    model_name     NAME,        -- embedding model name
    target_column  NAME  DEFAULT NULL,   -- target vector column name
    index_method   TEXT  DEFAULT 'hnsw', -- vector index method
    options        JSONB DEFAULT '{}'    -- extension options
) RETURNS BIGINT

-- Single-column embedding: Creates an embedding for the content column, with the target column automatically named content_embedding.
SELECT tencentdb_ai.add_incr_autoembedding_task(
    'public', 'articles',
    ARRAY['content'],
    'kinfra-text-embedding-0.6b'
);

-- Multi-column embedding: Merges title + body to generate embeddings, with the target column and index method explicitly specified.
SELECT tencentdb_ai.add_incr_autoembedding_task(
    'public', 'articles',
    ARRAY['title', 'body'],
    'kinfra-text-embedding-0.6b',
    target_column => 'doc_embedding',
    index_method => 'hnsw'
);

-- Override the vector dimension (when embedding_dim is not set in model_list).
SELECT tencentdb_ai.add_incr_autoembedding_task(
    'public', 'articles',
    ARRAY['content'],
    'kinfra-text-embedding-0.6b',
    options => '{"vector_dim": 1024}'
);

tencentdb_ai.add_backfill_autoembedding_task(
    schema_name    NAME,        -- schema where the table resides
    table_name     NAME,        -- table name
    source_columns NAME[],      -- array of source text column names
    model_name     NAME,        -- embedding model name
    target_column  NAME  DEFAULT NULL,   -- target vector column name
    start_now      BOOL  DEFAULT false,  -- whether to start backfill immediately
    options        JSONB DEFAULT '{}'    -- extension options
) RETURNS BIGINT

-- First create an incremental task (automatically creates the vector column + triggers)
SELECT tencentdb_ai.add_incr_autoembedding_task(
    'public', 'articles',
    ARRAY['content'],
    'kinfra-text-embedding-0.6b'
);

-- Then create a backfill task to immediately backfill historical data.
SELECT tencentdb_ai.add_backfill_autoembedding_task(
    'public', 'articles',
    ARRAY['content'],
    'kinfra-text-embedding-0.6b',
    start_now => true
);

-- Alternatively, you can register the task first (with start_now defaulting to false) and trigger it manually later.
SELECT tencentdb_ai.add_backfill_autoembedding_task(
    'public', 'articles',
    ARRAY['content'],
    'kinfra-text-embedding-0.6b'
);

-- Start manually later.
SELECT tencentdb_ai.run_backfill_autoembedding_task(1);

tencentdb_ai.drop_incr_autoembedding_task(task_id BIGINT) RETURNS VOID
tencentdb_ai.drop_backfill_autoembedding_task(task_id BIGINT) RETURNS VOID

-- Delete the incremental task.
SELECT tencentdb_ai.drop_incr_autoembedding_task(1);

-- Delete the backfill task.
SELECT tencentdb_ai.drop_backfill_autoembedding_task(2);

tencentdb_ai.run_backfill_autoembedding_task(
    task_id    BIGINT,           -- Backfill task ID
    batch_size INT DEFAULT 8192  -- Maximum number of rows a background worker scans and enqueues per round.
) RETURNS VOID

-- Manually start the backfill.
SELECT tencentdb_ai.run_backfill_autoembedding_task(1);

-- View the status of all tasks.
SELECT task_kind, task_id, schema_name, table_name, 
       target_column, status, backfill_state, 
       pending, backfilled_rows, failed_count
FROM tencentdb_ai.autoembedding_status
ORDER BY task_id;

-- View recent errors.
SELECT task_kind, task_id, row_id, error_message, created_at
FROM tencentdb_ai.autoembedding_error
ORDER BY created_at DESC
LIMIT 10;

-- Install extensions. pgmq / pgvector are installed automatically through CASCADE.
CREATE EXTENSION IF NOT EXISTS tencentdb_ai CASCADE;

-- Register an embedding model. The json_path parameter must be NULL.
SELECT tencentdb_ai.add_model('kinfra-text-embedding-0.6b', NULL, NULL, NULL, 'tokenhub');
SELECT tencentdb_ai.update_model_attr('kinfra-text-embedding-0.6b', 'api_key', 'your_api_key');

-- Set the embedding dimension. This parameter is required. If it is not set, specify it through options.vector_dim when creating a task.
UPDATE tencentdb_ai.model_list SET embedding_dim = 1024 WHERE model_name = 'kinfra-text-embedding-0.6b';

-- Create a knowledge base table.
CREATE TABLE kb_articles (
    id      bigserial PRIMARY KEY,
    title   text,
    content text
);

-- Insert test data.
INSERT INTO kb_articles (title, content) VALUES
    ('Introduction to PostgreSQL', 'PostgreSQL is a powerful open-source object-relational database system...'),
    ('Tencent Cloud VectorDB', 'Tencent Cloud VectorDB is a database system specifically designed for storing and searching high-dimensional vectors...'),
    ('RAG', 'Retrieval-Augmented Generation (RAG) is an AI technology that combines search and generation capabilities...'),

-- Create an incremental task: automatically add the content_embedding column and create a trigger.
SELECT tencentdb_ai.add_incr_autoembedding_task(
    'public', 'kb_articles',
    ARRAY['content'],
    'kinfra-text-embedding-0.6b'
) AS incr_task_id;

-- Create a backfill task: immediately backfill the existing 3 rows of historical data.
SELECT tencentdb_ai.add_backfill_autoembedding_task(
    'public', 'kb_articles',
    ARRAY['content'],
    'kinfra-text-embedding-0.6b',
    start_now => true
) AS backfill_task_id;

-- Wait for the background process to complete. By default, a poll is performed every second.
SELECT pg_sleep(3);

-- Check the backfill progress.
SELECT backfill_state, backfilled_rows
FROM tencentdb_ai.autoembedding_status
WHERE task_kind = 'backfill';

 backfill_state | backfilled_rows
----------------+-----------------
 done           |               3
(1 row)

-- Verify that the vector has been generated.
SELECT id, title, 
       content_embedding IS NOT NULL AS has_embedding,
       vector_dims(content_embedding) AS dims
FROM kb_articles
ORDER BY id;

 id |      title      | has_embedding | dims
----+-----------------+---------------+------
  1 | Introduction to PostgreSQL | t             | 1024
  2 | Tencent Cloud VectorDB      | t             | 1024
  3 | RAG        | t             | 1024
(3 rows)

-- View task status.
SELECT * FROM tencentdb_ai.autoembedding_status;

-- Insert new data. The trigger automatically enqueues it.
INSERT INTO kb_articles (title, content) VALUES
    ('Embedding model', 'An embedding model is a machine learning model that converts text into vector representations...');

-- Wait for the background process to complete.
SELECT pg_sleep(2);

-- Verify that the vector for the new data has been generated.
SELECT id, title, content_embedding IS NOT NULL AS has_embedding
FROM kb_articles
ORDER BY id;

 id |      title      | has_embedding
----+-----------------+---------------
  1 | Introduction to PostgreSQL | t
  2 | Tencent Cloud VectorDB      | t
  3 | RAG        | t
  4 | Embedding model        | t
(4 rows)

-- Check the queue backlog.
SELECT task_kind, task_id, pending
FROM tencentdb_ai.autoembedding_status
WHERE pending > 0;

-- View error details.
SELECT * FROM tencentdb_ai.autoembedding_error
ORDER BY created_at DESC LIMIT 20;