pgvector_psycopg3_vector_insert_and_nearest_neighbor_search.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# Register the vector type with psycopg
8register_vector(conn)
9
10# Create a table with a vector column (3 dimensions)
11conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
12conn.execute('DROP TABLE IF EXISTS items')
13conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))')
14
15# Insert vectors
16embedding1 = [1, 1, 1]
17embedding2 = [2, 2, 2]
18embedding3 = [1, 1, 2]
19conn.execute('INSERT INTO items (embedding) VALUES (%s), (%s), (%s)', (embedding1, embedding2, embedding3))
20
21# Get the nearest neighbors using L2 distance
22neighbor_embedding = [1, 1, 1]
23results = conn.execute('SELECT * FROM items ORDER BY embedding <-> %s LIMIT 5', (neighbor_embedding,)).fetchall()
24
25for row in results:
26    print(row)
27
28# Close the connection
29conn.close()