sqlite_vec_vector_search_with_knn_query.py

1import sqlite3
2import sqlite_vec
3import struct
4from typing import List
5
6def serialize_float32(vector: List[float]) -> bytes:
7    """Serializes a list of floats into a compact byte format for sqlite-vec."""
8    return struct.pack(f"{len(vector)}f", *vector)
9
10# Connect to an in-memory database
11db = sqlite3.connect(":memory:")
12
13# Load the sqlite-vec extension
14db.enable_load_extension(True)
15sqlite_vec.load(db)
16db.enable_load_extension(False)
17
18# Create a virtual table for 3-dimensional float vectors
19db.execute("""
20  CREATE VIRTUAL TABLE vec_items USING vec0(
21    embedding float[3]
22  )
23""")
24
25# Insert sample vector data
26# Using a helper to convert the list of floats to bytes
27items = [
28    (serialize_float32([0.1, 0.1, 0.1]),),
29    (serialize_float32([0.2, 0.2, 0.2]),),
30    (serialize_float32([0.3, 0.3, 0.3]),),
31]
32
33with db:
34    db.executemany("INSERT INTO vec_items(embedding) VALUES (?)", items)
35
36# Query for the 3 nearest neighbors to a target vector
37query_vector = serialize_float32([0.11, 0.11, 0.11])
38rows = db.execute(
39    """
40    SELECT
41      rowid,
42      distance
43    FROM vec_items
44    WHERE embedding MATCH ?
45    ORDER BY distance
46    LIMIT 3
47    """,
48    (query_vector,),
49).fetchall()
50
51for row in rows:
52    print(f"ID: {row[0]}, Distance: {row[1]}")
53
54db.close()