-- Create the pgvector extension
CREATE EXTENSION IF NOT EXISTS vector;

-- Verify the installation
SELECT extversion FROM pg_extension WHERE extname = 'vector';

-- Create a table with a vector column (1536-dimensional, compatible with common Embedding models).
CREATE TABLE documents (
    id BIGSERIAL PRIMARY KEY,
    content TEXT NOT NULL,
    embedding vector(1536)
);

-- Insert vector data
INSERT INTO documents (content, embedding)
VALUES ('PostgreSQL is the most advanced open-source database', '[0.1, 0.2, 0.3, ...]'::vector(1536));

CREATE TABLE docs_halfvec (
    id BIGSERIAL PRIMARY KEY,
    embedding halfvec(1536)
);

CREATE TABLE docs_sparse (
    id BIGSERIAL PRIMARY KEY,
    embedding sparsevec(30000)
);

-- Use the {index:value}/dim format for sparse vectors.
INSERT INTO docs_sparse (embedding)
VALUES ('{1:0.5, 100:0.3, 999:0.8}/30000'::sparsevec);

-- Cosine similarity search (most commonly used)
SELECT id, content,
       1 - (embedding <=> query_vec) AS similarity
FROM documents
ORDER BY embedding <=> query_vec
LIMIT 10;

-- L2 distance search
SELECT id, content
FROM documents
ORDER BY embedding <-> query_vec
LIMIT 10;

-- Create an HNSW index (cosine distance)
CREATE INDEX ON documents
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 200);

-- Create an HNSW index (L2 distance)
CREATE INDEX ON documents
USING hnsw (embedding vector_l2_ops)
WITH (m = 16, ef_construction = 200);

-- Increase ef_search to improve recall (default: 40)
SET hnsw.ef_search = 100;

-- Data must exist before creation.
CREATE INDEX ON documents
USING ivfflat (embedding vector_cosine_ops)
WITH (lists = 100);  -- The recommended value for lists is the square root of the number of rows.

-- Set the probe count during queries.
SET ivfflat.probes = 10;

-- 1. Create a knowledge base table
CREATE TABLE knowledge_base (
    id BIGSERIAL PRIMARY KEY,
    title TEXT NOT NULL,
    content TEXT NOT NULL,
    embedding vector(1024),
    metadata JSONB DEFAULT '{}',
    created_at TIMESTAMPTZ DEFAULT NOW()
);

-- 2. Create an index
CREATE INDEX idx_kb_hnsw ON knowledge_base
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 200);

-- 3. Semantic search (obtain the query vector using tencentdb_ai)
SELECT title, content,
       1 - (embedding <=> query_embedding) AS similarity
FROM knowledge_base,
LATERAL (
    SELECT tencentdb_ai.get_embedding('bge-m3', 'What is PostgreSQL?')
    AS query_embedding
) q
WHERE embedding <=> query_embedding < 0.5
ORDER BY embedding <=> query_embedding
LIMIT 5;

-- Insert data first, then create the index (faster)
DROP INDEX IF EXISTS idx_kb_hnsw;
-- ... Batch COPY data ...
CREATE INDEX idx_kb_hnsw ON knowledge_base
USING hnsw (embedding vector_cosine_ops);