docker run \
  --name pgvecto-rs-demo \
  -e POSTGRES_PASSWORD=mysecretpassword \
  -p 5432:5432 \
  -d tensorchord/pgvecto-rs:pg16-v0.2.0

psql postgresql://postgres:mysecretpassword@localhost:5432/postgres

DROP EXTENSION IF EXISTS vectors;
CREATE EXTENSION vectors;

-- 创建包含向量列的表

CREATE TABLE items (
  id bigserial PRIMARY KEY,
  embedding vector(3) NOT NULL -- 3 维
);

-- 插入值

INSERT INTO items (embedding)
VALUES ('[1,2,3]'), ('[4,5,6]');

-- 或者将数组类型转换为向量后插入

INSERT INTO items (embedding)
VALUES (ARRAY[1, 2, 3]::real[]), (ARRAY[4, 5, 6]::real[]);

-- 通过操作符调用距离函数

-- 平方欧氏距离
SELECT '[1, 2, 3]'::vector <-> '[3, 2, 1]'::vector;
-- 负点积
SELECT '[1, 2, 3]'::vector <#> '[3, 2, 1]'::vector;
-- 余弦距离
SELECT '[1, 2, 3]'::vector <=> '[3, 2, 1]'::vector;

-- 查询相似的嵌入向量
SELECT * FROM items ORDER BY embedding <-> '[3,2,1]' LIMIT 5;

docker run \
  --name pgvecto-rs-demo \
  -e POSTGRES_PASSWORD=mysecretpassword \
  -p 5432:5432 \
  -d tensorchord/pgvecto-rs:pg16-v0.2.0

psql postgresql://postgres:mysecretpassword@localhost:5432/postgres

DROP EXTENSION IF EXISTS vectors;
CREATE EXTENSION vectors;

psql -U postgres -c 'ALTER SYSTEM SET shared_preload_libraries = "vectors.so"'
psql -U postgres -c 'ALTER SYSTEM SET search_path TO "$user", public, vectors'
# 需要重启 PostgreSQL 集群以使配置生效。
sudo systemctl restart postgresql.service   # 适用于通过 systemd 运行的 pgvecto.rs

DROP EXTENSION IF EXISTS vectors;
CREATE EXTENSION vectors;

# RHEL 9 dnf 示例
# 请根据您的包管理器进行检查
sudo dnf install -y unzip wget libpq-devel
sudo dnf module install -y postgresql:15/server
sudo postgresql-setup --initdb
sudo systemctl start postgresql.service
sudo systemctl enable postgresql.service

pg_config --pkglibdir
# 输出类似于：
# /usr/lib/postgresql/15/lib 或
# /usr/lib64/pgsql

pg_config --sharedir
# 输出类似于：
# /usr/share/postgresql/15 或
# /usr/share/pgsql

wget https://github.com/tensorchord/pgvecto.rs/releases/download/v0.2.1/vectors-pg15_x86_64-unknown-linux-gnu_0.2.1.zip -O vectors.zip
unzip vectors.zip -d vectors

# 将动态库复制到 `$pkglibdir`
sudo cp vectors/vectors.so $(pg_config --pkglibdir)/
# 将 schema 文件复制到 `$shardir`
sudo cp vectors/vectors--* $(pg_config --sharedir)/extension/
sudo cp vectors/vectors.control $(pg_config --sharedir)/extension/

psql -U postgres -c 'ALTER SYSTEM SET shared_preload_libraries = "vectors.so"'
psql -U postgres -c 'ALTER SYSTEM SET search_path TO "$user", public, vectors'
# 需要重启 PostgreSQL 集群以使配置生效。
sudo systemctl restart postgresql.service   # 适用于通过 systemd 运行的 pgvecto.rs

DROP EXTENSION IF EXISTS vectors;
CREATE EXTENSION vectors;

sed -e "s/@CARGO_VERSION@/0.0.0/g" ./vectors.control  # 这里的 0.0.0 是您所需的版本号
sudo cp -a ./sql/upgrade/. $(pg_config --sharedir)/extension
cargo pgrx install --sudo --release

psql -U postgres -c 'ALTER SYSTEM SET shared_preload_libraries = "vectors.so"'
psql -U postgres -c 'ALTER SYSTEM SET search_path TO "$user", public, vectors'
# 需要重启 PostgreSQL 集群以使配置生效。
sudo systemctl restart postgresql.service   # 适用于通过 systemd 运行的 pgvecto.rs
service postgresql restart  # 适用于在 envd 中运行的 pgvecto.rs

DROP EXTENSION IF EXISTS vectors;
CREATE EXTENSION vectors;

CREATE TABLE items (embedding vector(3));
INSERT INTO items (embedding) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000);

CREATE INDEX ON items USING vectors (embedding vector_l2_ops);

CREATE INDEX ON items USING vectors (embedding vector_l2_ops)
WITH (options = $$
[indexing.ivf]
$$);

SELECT * FROM items ORDER BY embedding <-> '[3,2,1]' LIMIT 5;

-- HNSW 算法，默认配置。

CREATE INDEX ON items USING vectors (embedding vector_l2_ops);

-- 使用 FP16 向量类型优化数据存储。

CREATE TABLE items (embedding vecf16(3));
CREATE INDEX ON items USING vectors (embedding vecf16_l2_ops)

--- 或者使用结合乘积量化（PQ）的暴力搜索（flat）。

CREATE INDEX ON items USING vectors (embedding vector_l2_ops)
WITH (options = $$
[indexing.flat]
quantization.product.ratio = "x16"
$$);

-- 或者使用结合标量量化（SQ）的索引。
CREATE INDEX ON items USING vectors (embedding vector_l2_ops)
WITH (options = "[indexing.hnsw.quantization.scalar]");

--- 或者使用 IVFPQ 算法。

CREATE INDEX ON items USING vectors (embedding vector_l2_ops)
WITH (options = $$
[indexing.ivf]
quantization.product.ratio = "x16"
$$);

-- 使用更多线程在后台构建索引。

CREATE INDEX ON items USING vectors (embedding vector_l2_ops)
WITH (options = $$
optimizing.optimizing_threads = 16
$$);

-- 倾向于使用更小的 HNSW 图分段。

CREATE INDEX ON items USING vectors (embedding vector_l2_ops)
WITH (options = $$
segment.max_growing_segment_size = 200000
$$);

SET vectors.hnsw_ef_search = 64;
SELECT * FROM items ORDER BY embedding <-> '[3,2,1]' LIMIT 5;

SELECT 1 FROM items WHERE category_id = 1 ORDER BY embedding <#> '[0.5,0.5,0.5]' limit 10

SET vectors.ivf_nprobe=1;

SET LOCAL vectors.hnsw_ef_search=40;

CREATE INDEX ON items USING vectors (embedding vector_l2_ops)
WITH (options = $$
[indexing.ivf]
quantization.product.ratio = "x16"
$$);

CREATE INDEX ON items USING vectors (embedding vector_l2_ops)
WITH (options = "[indexing.hnsw.quantization.scalar]");

DROP TABLE IF EXISTS t;
SET vectors.pgvector_compatibility=on;
SET hnsw.ef_search=40;
CREATE TABLE t (val vector(3));
INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000);
CREATE INDEX hnsw_cos_index ON t USING hnsw (val vector_cosine_ops);
SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 100) t2;
DROP INDEX hnsw_cos_index;

SET vectors.pgvector_compatibility=on;
SET hnsw.ef_search=40;
-- 使用默认选项的 [hnsw + vector_l2_ops] 索引
CREATE INDEX hnsw_l2_index ON t USING hnsw (val vector_l2_ops);
-- 指定单个 ef_construction 选项的 [hnsw + vector_cosine_ops] 索引
CREATE INDEX hnsw_cosine_index ON t USING hnsw (val vector_cosine_ops) WITH (ef_construction = 80);
-- 指定全部选项的匿名 [hnsw + vector_ip_ops] 索引
CREATE INDEX ON t USING hnsw (val vector_ip_ops) WITH (ef_construction = 80, m = 12);
SET ivfflat.probes=1;
-- 使用默认选项的 [ivfflat + vector_l2_ops] 索引
CREATE INDEX ivfflat_l2_index ON t USING ivfflat (val vector_l2_ops);
-- 指定全部选项的 [ivfflat + vector_ip_ops] 索引
CREATE INDEX ivfflat_ip_index ON t USING ivfflat (val vector_cosine_ops) WITH (nlist = 80);
-- 指定全部选项的匿名 [ivf + vector_ip_ops] 索引
CREATE INDEX ON t USING ivfflat (val vector_ip_ops) WITH (lists = 80)

SELECT 'a fat cat sat on a mat and ate a fat rat'::tsvector @@ 'cat & rat'::tsquery;
 ?column?
----------
 t

SELECT 'cat & rat'::tsquery @@ 'a fat cat sat on a mat and ate a fat rat'::tsvector;
 ?column?
----------
 t

CREATE TABLE items (
  id bigserial PRIMARY KEY,
  content text NOT NULL,
  embedding vector(3) NOT NULL -- 3 维
);

INSERT INTO items (content, embedding) VALUES
  ('a fat cat sat on a mat and ate a fat rat', '[1, 2, 3]'),
  ('a fat dog sat on a mat and ate a fat rat', '[4, 5, 6]'),
  ('a thin cat sat on a mat and ate a thin rat', '[7, 8, 9]'),
  ('a thin dog sat on a mat and ate a thin rat', '[10, 11, 12]');

SELECT content, embedding
FROM items
WHERE to_tsvector('english', content) @@ 'cat & rat'::tsquery;

SELECT content, embedding
FROM items
WHERE to_tsvector('english', content) @@ 'cat & rat'::tsquery
ORDER BY embedding <-> '[7, 8, 9]'::vector;

-- 查询从名为 "semantic_search" 的公用表表达式（CTE）开始。
WITH semantic_search AS (
    SELECT id, content, embedding, 
    RANK () OVER (ORDER BY embedding <-> '[7, 8, 9]') AS rank
    FROM items
    ORDER BY embedding <-> '[7, 8, 9]'::vector
    ), 
-- 定义另一个名为 "full_text_search" 的 CTE。
    full_text_search AS (
    SELECT id, content, embedding, 
    RANK () OVER (ORDER BY 
        ts_rank(to_tsvector('english', content), 
            'cat & rat'::tsquery) DESC) AS rank
    FROM items
    WHERE to_tsvector('english', content) @@ 'cat & rat'::tsquery
    ORDER BY ts_rank(
        to_tsvector('english', content), 'cat & rat'::tsquery) DESC
    )
-- 主查询从两个 CTE 中选择列并融合排名。
    SELECT 
        COALESCE(semantic_search.id, full_text_search.id) AS id,
        COALESCE(semantic_search.content, full_text_search.content) AS content,
        COALESCE(semantic_search.embedding, full_text_search.embedding) AS embedding,
        COALESCE(1.0 / (1 + semantic_search.rank), 0.0) + 
        COALESCE(1.0 / (1 + full_text_search.rank), 0.0) AS rank
    FROM semantic_search FULL OUTER JOIN full_text_search USING (id)
    ORDER BY rank DESC;

CREATE TABLE openai3072 (
  id bigserial PRIMARY KEY,
  text_embedding_3_large_3072_embedding vector(3072),
  text_embedding_3_large_3072_embedding bvector(3072)
);

-- 创建 256 维表达式索引
CREATE INDEX openai_vector_index_256 ON public.openai3072 USING vectors (((text_embedding_3_large_3072_embedding[0:256])::vector(256)) vector_l2_ops);

-- 创建 1024 维表达式索引
CREATE INDEX openai_vector_index_1024 on openai3072 using vectors((text_embedding_3_large_3072_embedding[0:1024]::vector(1024)) vector_l2_ops);

-- 创建 3072 维完整索引
CREATE INDEX openai_vector_index on openai3072 using vectors(text_embedding_3_large_3072_embedding vector_l2_ops);

CREATE INDEX openai_vector_index_bvector ON public.openai3072 USING vectors (text_embedding_3_large_3072_bvector bvector_l2_ops);

CREATE OR REPLACE FUNCTION match_documents_adaptive(
  query_embedding vector(3072),
  match_count int
)
RETURNS SETOF openai3072
LANGUAGE SQL
AS $$
WITH shortlist AS (
  SELECT *
  FROM openai3072
  ORDER BY (text_embedding_3_large_3072_embedding[0:256])::vector(256) <-> (query_embedding[0:256])::vector(256)
  LIMIT match_count * 2
)
SELECT *
FROM shortlist
ORDER BY text_embedding_3_large_3072_embedding <-> query_embedding
LIMIT match_count;
$$;

from transformers import CLIPProcessor, CLIPModel

processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")

def compute_image_embeddings(list_of_images):
    return model.get_image_features(**processor(images=list_of_images, return_tensors="pt", padding=True))

-- 创建表
CREATE TABLE images (
    id SERIAL PRIMARY KEY,
    author TEXT NOT NULL,
    vector vector(512) NOT NULL -- 512 维
);

INSERT INTO images (author, embedding) VALUES
  ('Allen', '[...]');

SELECT * FROM images ORDER BY embedding <-> '[...]' LIMIT 5;

SET vectors.hnsw_ef_search = 64;
SELECT * FROM images ORDER BY embedding <-> '[...]' WHERE author != 'Allen' LIMIT 64;

SELECT a.*, s.*, e.*
FROM asset_entity AS a
INNER JOIN smart_search AS s ON s.assetId = a.id
LEFT JOIN exif_info AS e ON e.assetId = a.id
WHERE a.ownerId IN (:...userIds)
  AND a.isArchived = false
  AND a.isVisible = true
  AND a.fileCreatedAt < NOW()
ORDER BY s.embedding <=> :embedding
LIMIT :numResults;

SELECT a.*, s.*, e.*

FROM asset_entity AS a
INNER JOIN smart_search AS s ON s.assetId = a.id
LEFT JOIN exif_info AS e ON e.assetId = a.id

WHERE a.ownerId IN (:...userIds)
  AND a.isArchived = false
  AND a.isVisible = true
  AND a.fileCreatedAt < NOW()

ORDER BY s.embedding <=> :embedding

LIMIT :numResults;

CREATE TABLE documents (
    id SERIAL PRIMARY KEY,
    user_id INTEGER NOT NULL,
    title TEXT,
    content TEXT,
    -- 其他元数据字段
);

CREATE TABLE chunks (
    id SERIAL PRIMARY KEY,
    document_id INTEGER REFERENCES documents(id),
    embedding vector(512) NOT NULL
);

CREATE INDEX idx_chunks_embedding ON chunks USING vectors (embedding vector_dot_ops);

SELECT chunks.id AS chunk_id, documents.id AS document_id, documents.title
FROM chunks
INNER JOIN documents ON chunks.document_id = documents.id
WHERE documents.user_id = 'alice' /* 替换为实际的 user_id */
ORDER BY chunks.embedding <-> '[3,2,1]' /* 替换为查询向量 */ LIMIT 5;

# pip install -U FlagEmbedding

from FlagEmbedding import BGEM3FlagModel

model = BGEM3FlagModel("BAAI/bge-m3", use_fp16=True)
sentences = ["What is BM25?", "Definition of BM25"]

output = model.encode(sentences, return_dense=True, return_sparse=True)

# 上述两段文本的稠密嵌入：
dense_embedding = output['dense_vecs']
# [[-0.02501617 -0.04525582 -0.01972856 ... -0.0099566   0.01677509
#   0.01292699]
# [-0.0106499  -0.04731942 -0.01477693 ... -0.0149255   0.01040097
#   0.00965083]

sparse_weight_id = output['lexical_weights']

# 您可以查看每个 token 对应的权重：
# print(model.convert_id_to_token(sparse_weight_id))
# 
# [{'What': 0.10430284, 'is': 0.10090457, 'BM': 0.2635918, '25': 0.3382988, '?': 0.052101523}, 
# {'Definition': 0.14298248, 'of': 0.076763615, 'BM': 0.2577639, '25': 0.33806682}]

# pip install -U numpy

import numpy as np

# 提取文本的稀疏向量下标与对应权重值
sparse_indices = [[int(token_id) for token_id in text] for text in sparse_weight_id]
# [[4865, 83, 90017, 2588, 32], [155455, 111, 90017, 2588]]
sparse_values = [[float(text[token_id]) for token_id in text] for text in sparse_weight_id]
# [[0.10430284, 0.10090457, 0.2635918, 0.3382988, 0.052101523], [0.14298248, 0.076763615, 0.2577639, 0.33806682]]

# 在 pgvecto.rs 中，每个稀疏向量的下标必须按升序排列
for i in range(len(sparse_weight_id)):
    sparse_values[i] = [x for _, x in sorted(zip(sparse_indices[i], sparse_values[i]))]
    sparse_indices[i] = sorted(sparse_indices[i])

docker run \
  --name pgvecto-rs-demo \
  -e POSTGRES_PASSWORD=mysecretpassword \
  -p 5432:5432 \
  -d tensorchord/pgvecto-rs:pg16-v0.3.0

print(len(model.tokenizer))
# 250002

# pip install -U psycopg

import psycopg
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

with psycopg.connect(URL) as conn:
    conn.execute("DROP TABLE IF EXISTS documents;")
    conn.execute(f"""
        CREATE TABLE documents (
            id SERIAL PRIMARY KEY, 
            text TEXT NOT NULL,
            dense vector(1024) NOT NULL,
            sparse svector({len(model.tokenizer)}) NOT NULL);
    """
    )

# pip install -U pgvecto_rs

import psycopg
from pgvecto_rs.psycopg import register_vector
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

with psycopg.connect(URL) as conn:
    register_vector(conn)
    for text, dense, sparse_ind, sparse_val in zip(sentences, dense_embedding, sparse_indices, sparse_values):
        conn.execute(
            "INSERT INTO documents (text, dense, sparse) VALUES (%s, %s, to_svector(%s, %s::real[]::int[], %s::real[]));",
            (text, dense, len(model.tokenizer), sparse_ind, sparse_val),
        )
    conn.execute("""
        CREATE INDEX ON documents 
        USING vectors (sparse svector_dot_ops) 
        WITH (options = \"[indexing.hnsw]\");
    """)
    conn.execute("""
        CREATE INDEX ON documents 
        USING vectors (dense vector_l2_ops) 
        WITH (options = \"[indexing.hnsw]\");
    """)

def embedding(sentence: str) -> tuple[np.ndarray, list[list[int]], list[list[float]]]:
    """
    从文本生成稠密嵌入和稀疏嵌入
    """
    output = model.encode(sentence, return_dense=True, return_sparse=True, return_colbert_vecs=False)
    dense_embedding = output['dense_vecs']
    sparse_weight_id = output['lexical_weights']
    sparse_indices = [int(token_id) for token_id in sparse_weight_id]
    sparse_values = [float(sparse_weight_id[token_id]) for token_id in sparse_weight_id]

    sparse_values = [x for _, x in sorted(zip(sparse_indices, sparse_values))]
    sparse_indices = sorted(sparse_indices)
    return dense_embedding, sparse_indices, sparse_values

dense_embedding, sparse_indices, sparse_values = embedding("The text you want to search...")

with psycopg.connect(URL) as conn:
    register_vector(conn)
    cur = conn.execute(
        "SELECT text FROM documents ORDER BY sparse <#> to_svector(%s, %s::real[]::int[], %s::real[]) LIMIT 50;",
        (len(model.tokenizer), sparse_indices, sparse_values),
    )
    sparse_result = cur.fetchall()
      
    cur = conn.execute(
              "SELECT text FROM documents ORDER BY dense <-> %s LIMIT 50;",
              (dense_embedding,),
          )
    dense_result = cur.fetchall()
    
# 合并稀疏向量与稠密向量的搜索结果
mix_text = set([r[0] for r in dense_result]).union([r[0] for r in sparse_result])

from FlagEmbedding import FlagReranker

reranker = FlagReranker('BAAI/bge-reranker-large', use_fp16=True)

# 获取相关性得分（得分越高表示相关性越强）
scores = reranker.compute_score([["The text you want to search...", candidate] for candidate in mix_text])
# 对结合了稠密嵌入与稀疏嵌入检索出的候选文本进行重排序
reranked_text = [t for _, t in sorted(zip(scores, mix_text))]

# pip install openai

import openai

openai.api_key = "YOUR_API_KEY_HERE"
def complete(prompt):
    result = openai.Completion.create(
        engine='gpt-3.5-turbo-instruct',
        prompt=prompt,
        temperature=0,
        max_tokens=400,
        top_p=1,
        frequency_penalty=0,
        presence_penalty=0,
        stop=None
    )
    return result['choices'][0]['text'].strip()

base_prompt = "What is the release year and overall plot of Skyfall?"
print(complete(base_prompt))

base_prompt_2 = "What is the release year and overall plot of Cat Napping?"
print(complete(base_prompt_2))

# pip install -U datasets
from datasets import load_dataset
dataset = load_dataset("vishnupriyavr/wiki-movie-plots-with-summaries-faiss-embeddings", split='train')

texts = dataset['train']['text']
# texts[0]
# 'Alice in Wonderland \n  starring May Clark \n  released in the year1903 \n  with the following plot  \n Alice follows a large white rabbit down a "Rabbit-hole". She finds a tiny door. When she finds a bottle labeled "Drink me", she does, and shrinks, but not enough to pass through the door. She then eats something labeled "Eat me" and grows larger. She finds a fan when enables her to shrink enough to get into the "Garden" and try to get a "Dog" to play with her. She enters the "White Rabbit\'s tiny House," but suddenly resumes her normal size. In order to get out, she has to use the "magic fan."\r\nShe enters a kitchen, in which there is a cook and a woman holding a baby. She persuades the woman to give her the child and takes the infant outside after the cook starts throwing things around. The baby then turns into a pig and squirms out of her grip. "The Duchess\'s Cheshire Cat" appears and disappears a couple of times to Alice and directs her to the Mad Hatter\'s "Mad Tea-Party." After a while, she leaves.\r\nThe Queen invites Alice to join the "ROYAL PROCESSION": a parade of marching playing cards and others headed by the White Rabbit. When Alice "unintentionally offends the Queen", the latter summons the "Executioner". Alice "boxes the ears", then flees when all the playing cards come for her. Then she wakes up and realizes it was all a dream.'

embeddings = dataset['train']['embeddings']

# pip install cohere
import getpass
import os
import cohere

if not os.getenv("COHERE_API_KEY"):
    os.environ["COHERE_API_KEY"] = getpass.getpass("Enter your Cohere API key: ")

co = cohere.Client()
# 嵌入模型可参考此文档：https://docs.cohere.com/docs/cohere-embed
embed = co.embed(
    texts=texts,
    input_type="search_query",
    model="embed-english-v3.0"
)
embeddings = embed.embeddings

docker run \
  --name pgvecto-rs-demo \
  -e POSTGRES_PASSWORD=mysecretpassword \
  -p 5432:5432 \
  -d tensorchord/pgvecto-rs:pg16-v0.3.0

#pip install "psycopg[binary,pool]"
import psycopg

URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"
# 启用扩展
with psycopg.connect(URL) as conn:
    conn.execute("DROP EXTENSION IF EXISTS vectors;")
    conn.execute("CREATE EXTENSION vectors;")

import psycopg
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

# 创建表
with psycopg.connect(URL) as conn:
    conn.execute("DROP TABLE IF EXISTS movies;")
    conn.execute(f"""
        CREATE TABLE movies (
            id SERIAL PRIMARY KEY, 
            text TEXT NOT NULL,
            embedding vector(768) NOT NULL);
    """
    )

# 可选安装，但推荐使用以显示进度条：
# pip install tqdm
import psycopg
from tqdm import tqdm
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

# 插入数据到表中
with psycopg.connect(URL) as conn:
    with conn.cursor() as cursor:    
        for text, embedding in tqdm(zip(texts, embeddings), total=len(texts)):
            cursor.execute(
                """INSERT INTO movies (text, embedding) 
                    VALUES (%s, %s::real[]);""",
                (text, embedding),
            )
        conn.commit()

import psycopg
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

# 为 embedding 列创建索引
with psycopg.connect(URL) as conn:
    cursor = conn.cursor()
    cursor.execute("""
        CREATE INDEX ON movies 
        USING vectors (embedding vector_l2_ops) 
        WITH (options = \"[indexing.hnsw]\");
    """)
    conn.commit()

# pip install sentence-transformers

from sentence_transformers import SentenceTransformer

# 可选模型：'paraphrase-mpnet-base-v2', 'all-mpnet-base-v2'
encoder = SentenceTransformer('all-mpnet-base-v2')

def encode_query(query): # query 为字符串
    # 预处理查询语句
    query = query.lower()
    query = query.split()
    query = [word for word in query]
    query = ' '.join(query)
    
    # 对查询进行编码
    vector = encoder.encode(query)
    return vector

import psycopg
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

def add_context(prompt):
    search_vector = encode_query(prompt).tolist()
    # 构造包含检索上下文的提示词
    prompt_start = (
        "Answer the question based on the following context:\n\n"
    )
    prompt_end = (
        f"\n\nQuestion: {prompt}\nAnswer:"
    )
    with psycopg.connect(URL) as conn:
        cursor = conn.cursor()
        # 根据希望提供给模型的电影文本数量调整 LIMIT 值
        cursor.execute("SELECT text FROM movies ORDER BY embedding <-> %s::real[] LIMIT 3;", 
                  (search_vector,),)
        result = cursor.fetchall()
    context = ""
    for r in result:
        context += r[0] + "\n\n"
    return prompt_start + context + prompt_end