Using sqlite-vec with embeddings in sqlite-utils and Datasette#

Installing sqlite-vec#

A neat thing about sqlite-vec (and many of Alex’s other extensions) is that it’s packaged as both a raw SQLite extension and as packages for various different platforms.

Alex makes Python plugins available for both Datasette and sqlite-utils which bundle that extension and register it with those tools such that the functions become available to SQL queries.

For Datasette:

1
datasette install datasette-sqlite-vec

For sqlite-utils:

1
sqlite-utils install sqlite-utils-sqlite-vec

Both of these commands will make the various sqlite-vec functions available within those tools. Test it like this:

1
sqlite-utils memory 'select vec_version()'

For me that outputs:

1
[{"vec_version()": "v0.1.1"}]

For Datasette you can run that query using the Datasette web interface, or from the command-line like this:

1
datasette --get '/_memory.json?sql=select+vec_version()'

Or in Datasette 1.0a14 or higher:

1
datasette --get '/_memory/-/query.json?sql=select+vec_version()'

Returning:

1
{"ok": true, "rows": [{"vec_version()": "v0.1.1"}], "truncated": false}

Some example queries#

My TIL website has an embeddings table that stores embeddings for each of the TILs. It has two columns: id is the text ID for the TIL, and embedding is the binary LLM embedding for that text.

Here’s how to use the sqlite-vec vec_distance_cosine() function to find similar documents based on their embeddings:

1
with document_embedding as (
2
  select embedding as first_embedding from embeddings where id = :id
3
)
4
select
5
  id,
6
  vec_distance_cosine(embedding, first_embedding) as distance
7
from
8
  embeddings, document_embedding
9
order by distance limit 10

This accepts the id of a TIL and returns the 10 most similar TILs based on their embeddings. Try it out here.

Here’s a more fun query that also explores the vec_to_json() function - which turns that binary format into a readable JSON array of floats - the vec_slice() function for returning a shorter slice of that array and the vec_quantize_binary() function for quantizing a vector to binary - returning a 1 for values >0 and a -1 for <0.

1
with document_embedding as (
2
  select embedding as first_embedding from embeddings where id = :id
3
)
4
select
5
  id,
6
  vec_distance_cosine(embedding, first_embedding) as distance,
7
  vec_to_json(vec_slice(embedding, 0, 3)) as first_3,
8
  vec_to_json(vec_quantize_binary(vec_slice(embedding, 0, 8))) as binary_8
9
from
10
  embeddings, document_embedding
11
order by distance limit 5

Run that here. I get back these results:

Creating an index#

sqlite-vec also includes the ability to create an index for a collection of vectors.

Here’s how I created an index for my TILs. First, I created a virtual table using the vec0 mechanism provided by sqlite-vec - I told it to store an embedding column that was an array of 1536 floats (the size of the OpenAI embeddings I’ve been using for my TILs):

1
create virtual table vec_tils using vec0(
2
  embedding float[1536]
3
);

Then I populated it like this:

1
insert into vec_tils(rowid, embedding)
2
  select rowid, embedding from embeddings;

vec0 tables require an integer ID, so I used the rowid of the embeddings table. If I had my own numeric ID on that table I would use that instead.

Now I can run queries against this index like so:

1
with document_embedding as (
2
  select embedding as first_embedding from embeddings where id = :id
3
)
4
select
5
  (select id from embeddings where embeddings.rowid = vec_tils.rowid) as id,
6
  distance
7
from vec_tils, document_embedding
8
where embedding match first_embedding
9
and k = 5
10
order by distance;

Try it here. I get back:

The where embedding match first_embedding and k = 5 clause hooks into the magic of the underlying virtual table to run an efficient k-nearest-neighbors query against the index.

I’m using that (select id from embeddings where embeddings.rowid = vec_tils.rowid) as id trick to convert the numeric rowid into a human-readable id value by running a subquery against the embeddings table.

The vec_tils table is created for my TIL site by this step in my GitHub Actions workflow that deploys the application.