tencentdb_ai.retrieve(
    embed_model    NAME,       -- Embedding model name
    question       TEXT,       -- Query question
    source_schema  TEXT,       -- schema of the knowledge base table
    source_table   TEXT,       -- Knowledge base table name
    chunk_col      TEXT,       -- Text chunk column name
    vector_col     TEXT,       -- Vector column name
    topn           INT  DEFAULT 10,    -- Return TopN results
    distance_type  TEXT DEFAULT 'L2'   -- Distance algorithm
) RETURNS SETOF tencentdb_ai.retrieve_result

retrieve_result AS (chunk TEXT, distance FLOAT8)

-- Euclidean distance search, returns Top 3
SELECT chunk, TRUNC(distance::numeric, 4) AS distance
FROM tencentdb_ai.retrieve(
    'kinfra-text-embedding-0.6b',
    Which data types are supported by PostgreSQL?
    'public', 'knowledge_base', 'chunk', 'embedding',
    3, 'L2'
)
ORDER BY distance, chunk;

-- Cosine distance search, returns Top 5
SELECT chunk, TRUNC(distance::numeric, 4) AS distance
FROM tencentdb_ai.retrieve(
    'kinfra-text-embedding-0.6b',
    How is vector search used?
    'public', 'knowledge_base', 'chunk', 'embedding',
    5, 'cosine'
)
ORDER BY distance, chunk;

-- Inner product search
SELECT chunk, TRUNC(distance::numeric, 4) AS distance
FROM tencentdb_ai.retrieve(
    'kinfra-text-embedding-0.6b',
    What is full-text search?
    'public', 'knowledge_base', 'chunk', 'embedding',
    10, 'inner product'
)
ORDER BY distance, chunk;

tencentdb_ai.rag(
    embed_model    NAME,       -- Embedding model name
    prompt_model   NAME,       -- Generation model name
    question       TEXT,       -- User question
    source_schema  TEXT,       -- schema of the knowledge base table
    source_table   TEXT,       -- Knowledge base table name
    chunk_col      TEXT,       -- Text chunk column name
    vector_col     TEXT,       -- Vector column name
    topn           INT  DEFAULT 10,    -- TopN recall
    distance_type  TEXT DEFAULT 'L2'   -- Distance algorithm
) RETURNS TEXT

User question → retrieve() recalls TopN document chunks
         → Concatenate all chunks to construct the context
         → Build the Prompt according to the template
         → Call chat_completions() to invoke the large language model for generation
         → Return the generated result

You are a rigorous Q&A assistant. Please answer the question based solely on the following reference materials.
If the available information is insufficient to answer, explicitly state that the question cannot be answered based on the available information.

[Reference Materials]
{All retrieved document chunks, separated by ---}

[Question]
{User question}

-- Install the extension (pgcrypto, vector, and pgmq are automatically installed as dependencies)
CREATE EXTENSION IF NOT EXISTS tencentdb_ai CASCADE;

-- Register the embedding model (json_path must be NULL, and the response is parsed by get_embedding itself)
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');

-- Register the generation model (json_path specifies the response parsing path)
SELECT tencentdb_ai.add_model('glm-5.2', NULL, NULL,
    '$.choices[0].message.content'::jsonpath, 'tokenhub');
SELECT tencentdb_ai.update_model_attr('glm-5.2', 'api_key', 'your_api_key');

-- Create the knowledge base table
CREATE TABLE kb_docs (
    id        bigserial PRIMARY KEY,
    title     text,
    chunk     text,
    embedding vector(1024)
);

-- Insert document chunks and vectorize them in real time
INSERT INTO kb_docs (title, chunk, embedding)
SELECT 
    'Introduction to PostgreSQL',
    'PostgreSQL is a powerful open-source object-relational database system with over 30 years of active development history, earning a strong reputation for reliability, feature robustness, and performance.',
    (SELECT e::vector(1024) FROM tencentdb_ai.get_embedding(
        'kinfra-text-embedding-0.6b', 
        ARRAY['PostgreSQL is a powerful open-source object-relational database system with over 30 years of active development history, earning a strong reputation for reliability, feature robustness, and performance.']
    ) AS e LIMIT 1);

INSERT INTO kb_docs (title, chunk, embedding)
SELECT 
    'PostgreSQL Features',
    'PostgreSQL supports a rich set of data types: numeric types (integer, numeric, real), character types (text, varchar), date/time types (timestamp, date), JSON/JSONB, arrays, range types, and more.'
    (SELECT e::vector(1024) FROM tencentdb_ai.get_embedding(
        'kinfra-text-embedding-0.6b',
        ARRAY['PostgreSQL supports a rich set of data types: numeric types (integer, numeric, real), character types (text, varchar), date/time types (timestamp, date), JSON/JSONB, arrays, range types, and more.']
    ) AS e LIMIT 1);

INSERT INTO kb_docs (title, chunk, embedding)
SELECT 
    'Introduction to pgvector',
    'pgvector is a vector extension for PostgreSQL that supports exact and approximate nearest neighbor search, as well as multiple similarity metrics such as L2 distance, inner product, and cosine distance.'
    (SELECT e::vector(1024) FROM tencentdb_ai.get_embedding(
        'kinfra-text-embedding-0.6b',
        ARRAY['pgvector is a vector extension for PostgreSQL that supports exact and approximate nearest neighbor search, as well as multiple similarity metrics such as L2 distance, inner product, and cosine distance.']
    ) AS e LIMIT 1);

-- Create a vector index (optional, to improve search performance)
CREATE INDEX ON kb_docs USING ivfflat (embedding vector_l2_ops) WITH (lists = 100);

SELECT chunk, TRUNC(distance::numeric, 6) AS distance
FROM tencentdb_ai.retrieve(
    'kinfra-text-embedding-0.6b',
    Which data types does PostgreSQL support?
    'public', 'kb_docs', 'chunk', 'embedding', 3, 'cosine'
)
ORDER BY distance, chunk;

-- One-stop RAG Q&A
SELECT tencentdb_ai.rag(
    'kinfra-text-embedding-0.6b',
    'glm-5.2',
    'Please introduce the main features of pgvector.',
    'public', 'kb_docs', 'chunk', 'embedding', 2, 'cosine'
) AS answer;

-- kb_empty is an empty table with the same structure as kb_docs
SELECT tencentdb_ai.rag(
    'kinfra-text-embedding-0.6b', 'glm-5.2', 'What is a database?',
    'public', 'kb_empty', 'chunk', 'embedding', 5, 'cosine'
) AS answer;

       answer
--------------------------
 Based on the available information, the question cannot be answered.
(1 row)