pgvector_psycopg_store_and_query_vectors_quickstart.py

1import psycopg
2from pgvector.psycopg import register_vector
3
4# Connect to an existing database
5conn = psycopg.connect(dbname='pgvector_example', autocommit=True)
6
7# Install extension
8conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
9
10# Register the vector type with psycopg
11register_vector(conn)
12
13# Create a table
14conn.execute('DROP TABLE IF EXISTS items')
15conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))')
16
17# Insert vectors
18embeddings = [
19    [1, 1, 1],
20    [2, 2, 2],
21    [1, 1, 2]
22]
23for embedding in embeddings:
24    conn.execute('INSERT INTO items (embedding) VALUES (%s)', (embedding,))
25
26# Get the nearest neighbors by L2 distance
27embedding = [1, 1, 1]
28results = conn.execute('SELECT * FROM items ORDER BY embedding <-> %s LIMIT 5', (embedding,)).fetchall()
29for row in results:
30    print(row)
31
32# Close the connection
33conn.close()