Use sqlite vector

This commit is contained in:
Claudio Ortolina
2025-10-11 22:48:27 +02:00
parent 14934b25d9
commit ee199272b1
9 changed files with 357 additions and 204 deletions
+1 -23
View File
@@ -10,7 +10,7 @@ defmodule MusicLibrary.Records.RecordEmbedding do
schema "record_embeddings" do schema "record_embeddings" do
belongs_to :record, Record belongs_to :record, Record
field :embedding, MusicLibrary.Records.RecordEmbedding.EmbeddingType field :embedding, SqliteVec.Ecto.Float32
field :text_representation, :string field :text_representation, :string
timestamps(type: :utc_datetime) timestamps(type: :utc_datetime)
@@ -20,28 +20,6 @@ defmodule MusicLibrary.Records.RecordEmbedding do
record_embedding record_embedding
|> cast(attrs, [:record_id, :embedding, :text_representation]) |> cast(attrs, [:record_id, :embedding, :text_representation])
|> validate_required([:record_id, :embedding, :text_representation]) |> validate_required([:record_id, :embedding, :text_representation])
|> validate_embedding_dimensions()
|> unique_constraint(:record_id) |> unique_constraint(:record_id)
end end
defp validate_embedding_dimensions(changeset) do
case get_change(changeset, :embedding) do
nil ->
changeset
embedding when is_list(embedding) ->
if length(embedding) == 1536 do
changeset
else
add_error(
changeset,
:embedding,
"must have exactly 1536 dimensions, got #{length(embedding)}"
)
end
_ ->
add_error(changeset, :embedding, "must be a list of floats")
end
end
end end
@@ -1,50 +0,0 @@
defmodule MusicLibrary.Records.RecordEmbedding.EmbeddingType do
@moduledoc """
Custom Ecto type for storing embedding vectors.
Embeddings are stored as JSON-encoded arrays of floats in the database,
but presented as Elixir lists in the application.
"""
use Ecto.Type
@impl true
def type, do: :string
@impl true
def cast(embedding) when is_list(embedding) do
if Enum.all?(embedding, &is_float/1) or Enum.all?(embedding, &is_number/1) do
# Convert all numbers to floats
{:ok, Enum.map(embedding, &to_float/1)}
else
:error
end
end
def cast(_), do: :error
@impl true
def load(json) when is_binary(json) do
case JSON.decode(json) do
{:ok, embedding} when is_list(embedding) ->
{:ok, Enum.map(embedding, &to_float/1)}
_ ->
:error
end
end
def load(_), do: :error
@impl true
def dump(embedding) when is_list(embedding) do
json = JSON.encode!(embedding)
{:ok, json}
rescue
_ -> :error
end
def dump(_), do: :error
defp to_float(n) when is_float(n), do: n
defp to_float(n) when is_integer(n), do: n * 1.0
end
+20 -57
View File
@@ -4,7 +4,9 @@ defmodule MusicLibrary.Records.Similarity do
""" """
import Ecto.Query import Ecto.Query
import(SqliteVec.Ecto.Query)
alias MusicLibrary.Records
alias MusicLibrary.Records.{Record, RecordEmbedding} alias MusicLibrary.Records.{Record, RecordEmbedding}
alias MusicLibrary.Repo alias MusicLibrary.Repo
@@ -53,44 +55,30 @@ defmodule MusicLibrary.Records.Similarity do
""" """
def find_similar(record_id, opts \\ []) do def find_similar(record_id, opts \\ []) do
limit = Keyword.get(opts, :limit, 10) limit = Keyword.get(opts, :limit, 10)
min_similarity = Keyword.get(opts, :min_similarity, 0.0)
scope = Keyword.get(opts, :scope) scope = Keyword.get(opts, :scope)
with {:ok, source_embedding} <- get_embedding(record_id), record = Records.get_record!(record_id)
similar_records <- calculate_similarities(source_embedding, record_id, scope) do record_musicbrainz_id = record.musicbrainz_id
similar_records
|> Enum.filter(fn {_record, similarity} -> similarity >= min_similarity end)
|> Enum.take(limit)
|> Enum.map(fn {record, similarity} -> {record, Float.round(similarity, 4)} end)
else
{:error, :not_found} -> []
end
end
@doc """ case get_embedding(record_id) do
Calculates cosine similarity between two embedding vectors. {:ok, source_embedding} ->
query =
from re in RecordEmbedding,
where: re.record_id != ^record_id,
join: r in Record,
on: r.id == re.record_id and r.musicbrainz_id != ^record_musicbrainz_id,
order_by: vec_distance_cosine(re.embedding, vec_f32(source_embedding)),
select: {r, re.embedding},
group_by: r.musicbrainz_id,
limit: ^limit
Returns a float between -1.0 and 1.0, where: query = apply_scope_filter(query, scope)
- 1.0 = identical vectors
- 0.0 = orthogonal vectors
- -1.0 = opposite vectors
"""
def cosine_similarity(vec_a, vec_b) when is_list(vec_a) and is_list(vec_b) do
if length(vec_a) != length(vec_b) do
raise ArgumentError, "Vectors must have the same length"
end
dot_product = query
Enum.zip(vec_a, vec_b) |> Repo.all()
|> Enum.reduce(0.0, fn {a, b}, acc -> acc + a * b end)
magnitude_a = calculate_magnitude(vec_a) {:error, :not_found} ->
magnitude_b = calculate_magnitude(vec_b) []
if magnitude_a == 0.0 or magnitude_b == 0.0 do
0.0
else
dot_product / (magnitude_a * magnitude_b)
end end
end end
@@ -142,31 +130,6 @@ defmodule MusicLibrary.Records.Similarity do
defp humanize_type(:other), do: "Other" defp humanize_type(:other), do: "Other"
defp humanize_type(_), do: "Unknown" defp humanize_type(_), do: "Unknown"
defp calculate_magnitude(vector) do
vector
|> Enum.reduce(0.0, fn x, acc -> acc + x * x end)
|> :math.sqrt()
end
defp calculate_similarities(source_embedding, source_record_id, scope) do
query =
from re in RecordEmbedding,
where: re.record_id != ^source_record_id,
join: r in Record,
on: r.id == re.record_id,
select: {r, re.embedding}
query = apply_scope_filter(query, scope)
query
|> Repo.all()
|> Enum.map(fn {record, embedding} ->
similarity = cosine_similarity(source_embedding, embedding)
{record, similarity}
end)
|> Enum.sort_by(fn {_record, similarity} -> similarity end, :desc)
end
defp apply_scope_filter(query, :collection) do defp apply_scope_filter(query, :collection) do
from [re, r] in query, where: not is_nil(r.purchased_at) from [re, r] in query, where: not is_nil(r.purchased_at)
end end
@@ -422,7 +422,6 @@ defmodule MusicLibraryWeb.RecordComponents do
</button> </button>
<span class="absolute top-2 right-2 rounded-full px-2 py-0.5 text-xs font-medium bg-zinc-900/75 text-white backdrop-blur-sm"> <span class="absolute top-2 right-2 rounded-full px-2 py-0.5 text-xs font-medium bg-zinc-900/75 text-white backdrop-blur-sm">
{Float.round(similarity * 100, 0)}%
</span> </span>
</div> </div>
+22
View File
@@ -0,0 +1,22 @@
defmodule SqliteVec.Ecto.Float32 do
@moduledoc """
`Ecto.Type` for `SqliteVec.Float32`
"""
use Ecto.Type
def type, do: :binary
def cast(value) do
{:ok, SqliteVec.Float32.new(value)}
end
def load(data) do
{:ok, SqliteVec.Float32.from_binary(data)}
end
def dump(%SqliteVec.Float32{} = vector) do
{:ok, SqliteVec.Float32.to_binary(vector)}
end
def dump(_), do: :error
end
+182
View File
@@ -0,0 +1,182 @@
defmodule SqliteVec.Ecto.Query do
@moduledoc """
Macros for Ecto
"""
@doc """
Creates a bit vector
"""
defmacro vec_bit(vector) do
quote do
fragment("vec_bit(?)", type(^unquote(vector).data, :binary))
end
end
@doc """
Creates an int8 vector
"""
defmacro vec_int8(vector) do
quote do
fragment("vec_int8(?)", type(^unquote(vector).data, :binary))
end
end
@doc """
Creates a float32 vector
"""
defmacro vec_f32(vector) do
quote do
fragment("vec_f32(?)", type(^unquote(vector).data, :binary))
end
end
@doc """
Calculates the L2 euclidian distance between vectors a and b. Only valid for float32 or int8 vectors.
Returns an error under the following conditions:
- a or b are invalid vectors
- a or b do not share the same vector element types (ex float32 or int8)
- a or b are bit vectors. Use vec_distance_hamming() for distance calculations between two bitvectors.
- a or b do not have the same length.
"""
# credo:disable-for-next-line Credo.Check.Readability.FunctionNames
defmacro vec_distance_L2(a, b) do
quote do
fragment("vec_distance_L2(?, ?)", unquote(a), unquote(b))
end
end
@doc """
Calculates the cosine distance between vectors a and b. Only valid for float32 or int8 vectors.
Returns an error under the following conditions:
- a or b are invalid vectors
- a or b do not share the same vector element types (ex float32 or int8)
- a or b are bit vectors. Use vec_distance_hamming() for distance calculations between two bitvectors.
- a or b do not have the same length
"""
defmacro vec_distance_cosine(a, b) do
quote do
fragment("vec_distance_cosine(?, ?)", unquote(a), unquote(b))
end
end
@doc """
Calculates the hamming distance between two bitvectors a and b. Only valid for bitvectors.
Returns an error under the following conditions:
- a or b are not bitvectors
- a and b do not share the same length
- Memory cannot be allocated
"""
defmacro vec_distance_hamming(a, b) do
quote do
fragment("vec_distance_hamming(?, ?)", unquote(a), unquote(b))
end
end
defmacro vec_match(a, b) do
quote do
fragment("? match ?", unquote(a), unquote(b))
end
end
@doc """
Returns the number of elements in the given vector
"""
defmacro vec_length(vector) do
quote do
fragment("vec_length(?)", unquote(vector))
end
end
@doc """
Returns the name of the type of `vector` as text
"""
defmacro vec_type(vector) do
quote do
fragment("vec_type(?)", unquote(vector))
end
end
@doc """
Adds every element in vector a with vector b, returning a new vector c.
Both vectors must be of the same type and same length.
Only float32 and int8 vectors are supported.
An error is raised if either a or b are invalid, or if they are not the same type or same length.
"""
defmacro vec_add(a, b) do
quote do
fragment("vec_add(?, ?)", unquote(a), unquote(b))
end
end
@doc """
Subtracts every element in vector a with vector b, returning a new vector c.
Both vectors must be of the same type and same length.
Only float32 and int8 vectors are supported.
An error is raised if either a or b are invalid, or if they are not the same type or same length.
"""
defmacro vec_sub(a, b) do
quote do
fragment("vec_sub(?, ?)", unquote(a), unquote(b))
end
end
@doc """
Performs L2 normalization on the given vector.
Only float32 vectors are currently supported.
Returns an error if the input is an invalid vector or not a float32 vector.
"""
defmacro vec_normalize(vector) do
quote do
fragment("vec_normalize(?)", unquote(vector))
end
end
@doc """
Extract a subset of vector from the start element (inclusive) to the end element (exclusive).
This is especially useful for Matryoshka embeddings, also known as "adaptive length" embeddings.
Use with vec_normalize() to get proper results.
Returns an error in the following conditions:
- If vector is not a valid vector
- If start is less than zero or greater than or equal to end
- If end is greater than the length of vector, or less than or equal to start.
- If vector is a bitvector, start and end must be divisible by 8.
"""
defmacro vec_slice(vector, start_index, end_index) do
quote do
fragment("vec_slice(?, ?, ?)", unquote(vector), unquote(start_index), unquote(end_index))
end
end
@doc """
Represents a vector as JSON text.
The input vector can be a vector BLOB or JSON text.
Returns an error if vector is an invalid vector, or when memory cannot be allocated.
"""
defmacro vec_to_json(vector) do
quote do
fragment("vec_to_json(?)", unquote(vector))
end
end
@doc """
Quantize a float32 or int8 vector into a bitvector.
For every element in the vector, a 1 is assigned to positive numbers and a 0 is assigned to negative numbers.
These values are then packed into a bit vector.
Returns an error if vector is invalid, or if vector is not a float32 or int8 vector.
"""
defmacro vec_quantize_binary(vector) do
quote do
fragment("vec_quantize_binary(?)", unquote(vector))
end
end
end
+114
View File
@@ -0,0 +1,114 @@
defmodule SqliteVec.Float32 do
@moduledoc """
A vector struct for float32 vectors.
Vectors are stored as binaries in the endianness of the system.
> ### Consider endianness {: .warning}
>
> `SqliteVec.Float32.Vector` holds data in system endianness.
> Therefore, the same vector data will be interpreted differently on another system with different endianness.
> Moreover, you must consider endianness when converting the binary data directly to a list of numbers.
iex> v = SqliteVec.Float32.new([-1.0, 2.0])
...> b = SqliteVec.Float32.to_binary(v)
...> <<f1::float-32, f2::float-32>> = b
...> [f1, f2]
case System.endianness() do
:big -> [-1.0, 2.0]
:little -> [4.618539608568165e-41, 8.96831017167883e-44]
end
"""
@type t :: %__MODULE__{data: binary()}
defstruct [:data]
@doc """
Creates a new vector from a vector, list, or tensor
The vector must be a `SqliteVec.Float32` vector.
The list may contain any number but the values will be converted to f32 format.
The tensor must have a rank of 1 and must be of type :f32.
## Examples
iex> SqliteVec.Float32.new([1.0, 2.0])
%SqliteVec.Float32{data: <<1.0::float-32-native, 2.0::float-32-native>>}
iex> v1 = SqliteVec.Float32.new([1, 2])
...> SqliteVec.Float32.new(v1)
%SqliteVec.Float32{data: <<1.0::float-32-native, 2.0::float-32-native>>}
iex> SqliteVec.Float32.new(Nx.tensor([1, 2], type: :f32))
%SqliteVec.Float32{data: <<1.0::float-32-native, 2.0::float-32-native>>}
"""
def new(vector_or_list_or_tensor)
def new(%SqliteVec.Float32{} = vector) do
vector
end
def new(list) when is_list(list) do
if list == [] do
raise ArgumentError, "list must not be empty"
end
bin = for v <- list, into: <<>>, do: <<v::float-32-native>>
from_binary(<<bin::binary>>)
end
if Code.ensure_loaded?(Nx) do
def new(tensor) when is_struct(tensor, Nx.Tensor) do
if Nx.rank(tensor) != 1 do
raise ArgumentError, "expected rank to be 1"
end
if Nx.type(tensor) != {:f, 32} do
raise ArgumentError, "expected type to be :f32"
end
bin = tensor |> Nx.to_binary()
from_binary(<<bin::binary>>)
end
end
@doc """
Creates a new vector from its binary representation
"""
def from_binary(binary) when is_binary(binary) do
%SqliteVec.Float32{data: binary}
end
@doc """
Converts the vector to its binary representation
"""
def to_binary(vector) when is_struct(vector, SqliteVec.Float32) do
vector.data
end
@doc """
Converts the vector to a list
"""
def to_list(vector) when is_struct(vector, SqliteVec.Float32) do
<<bin::binary>> = vector.data
for <<v::float-32-native <- bin>>, do: v
end
if Code.ensure_loaded?(Nx) do
@doc """
Converts the vector to a tensor
"""
def to_tensor(vector) when is_struct(vector, SqliteVec.Float32) do
<<bin::binary>> = vector.data
Nx.from_binary(bin, :f32)
end
end
end
defimpl Inspect, for: SqliteVec.Float32 do
import Inspect.Algebra
def inspect(vector, opts) do
concat(["vec_f32('", Inspect.List.inspect(SqliteVec.Float32.to_list(vector), opts), "')"])
end
end
@@ -1,17 +1,23 @@
defmodule MusicLibrary.Repo.Migrations.CreateRecordEmbeddings do defmodule MusicLibrary.Repo.Migrations.CreateRecordEmbeddings do
use Ecto.Migration use Ecto.Migration
def change do def up do
create table(:record_embeddings, primary_key: false) do execute("""
add :id, :binary_id, primary_key: true CREATE TABLE record_embeddings (
add :record_id, references(:records, type: :binary_id, on_delete: :delete_all), null: false id TEXT PRIMARY KEY,
record_id TEXT NOT NULL CONSTRAINT record_embeddings_record_id_fkey REFERENCES records(id) ON DELETE CASCADE,
embedding float[1536] NOT NULL,
text_representation TEXT NOT NULL,
inserted_at TEXT NOT NULL,
updated_at TEXT NOT NULL);
""")
add :embedding, :text, null: false execute("""
add :text_representation, :text, null: false CREATE UNIQUE INDEX record_embeddings_record_id_index ON record_embeddings (record_id);
""")
end
timestamps(type: :utc_datetime) def down do
end drop table(:record_embeddings)
create unique_index(:record_embeddings, [:record_id])
end end
end end
+2 -63
View File
@@ -107,57 +107,6 @@ defmodule MusicLibrary.Records.SimilarityTest do
end end
end end
describe "cosine_similarity/2" do
test "calculates similarity between identical vectors" do
vec = [1.0, 2.0, 3.0, 4.0]
similarity = Similarity.cosine_similarity(vec, vec)
assert_in_delta similarity, 1.0, 0.0001
end
test "calculates similarity between orthogonal vectors" do
vec_a = [1.0, 0.0, 0.0]
vec_b = [0.0, 1.0, 0.0]
similarity = Similarity.cosine_similarity(vec_a, vec_b)
assert_in_delta similarity, 0.0, 0.0001
end
test "calculates similarity between opposite vectors" do
vec_a = [1.0, 0.0, 0.0]
vec_b = [-1.0, 0.0, 0.0]
similarity = Similarity.cosine_similarity(vec_a, vec_b)
assert_in_delta similarity, -1.0, 0.0001
end
test "calculates similarity between similar vectors" do
vec_a = [1.0, 2.0, 3.0]
vec_b = [1.1, 2.1, 2.9]
similarity = Similarity.cosine_similarity(vec_a, vec_b)
# Should be close to 1.0 since vectors are similar
assert similarity > 0.99
end
test "raises error for vectors of different lengths" do
vec_a = [1.0, 2.0, 3.0]
vec_b = [1.0, 2.0]
assert_raise ArgumentError, fn ->
Similarity.cosine_similarity(vec_a, vec_b)
end
end
test "handles zero vectors" do
vec_a = [0.0, 0.0, 0.0]
vec_b = [1.0, 2.0, 3.0]
similarity = Similarity.cosine_similarity(vec_a, vec_b)
assert similarity == 0.0
end
end
describe "store_embedding/3 and get_embedding/1" do describe "store_embedding/3 and get_embedding/1" do
test "stores and retrieves an embedding" do test "stores and retrieves an embedding" do
record = record() record = record()
@@ -167,10 +116,7 @@ defmodule MusicLibrary.Records.SimilarityTest do
assert {:ok, _} = Similarity.store_embedding(record.id, embedding, text_rep) assert {:ok, _} = Similarity.store_embedding(record.id, embedding, text_rep)
assert {:ok, retrieved_embedding} = Similarity.get_embedding(record.id) assert {:ok, retrieved_embedding} = Similarity.get_embedding(record.id)
assert length(retrieved_embedding) == 1536 assert SqliteVec.Float32.new(embedding) == retrieved_embedding
# Check that embeddings are the same (within floating point precision)
Enum.zip(embedding, retrieved_embedding)
|> Enum.each(fn {a, b} -> assert_in_delta a, b, 0.0001 end)
end end
test "updates existing embedding on conflict" do test "updates existing embedding on conflict" do
@@ -182,7 +128,7 @@ defmodule MusicLibrary.Records.SimilarityTest do
assert {:ok, _} = Similarity.store_embedding(record.id, embedding2, "Text 2") assert {:ok, _} = Similarity.store_embedding(record.id, embedding2, "Text 2")
assert {:ok, retrieved_embedding} = Similarity.get_embedding(record.id) assert {:ok, retrieved_embedding} = Similarity.get_embedding(record.id)
assert List.first(retrieved_embedding) == 0.7 assert SqliteVec.Float32.new(embedding2) == retrieved_embedding
end end
test "returns error for non-existent record" do test "returns error for non-existent record" do
@@ -227,13 +173,6 @@ defmodule MusicLibrary.Records.SimilarityTest do
assert length(similar) == 1 assert length(similar) == 1
end end
test "respects min_similarity option", %{record1: record1} do
similar = Similarity.find_similar(record1.id, min_similarity: 0.99)
# Since we have slight variations, only very similar records pass
assert length(similar) <= 1
end
test "returns empty list for record without embedding" do test "returns empty list for record without embedding" do
record_without_embedding = record() record_without_embedding = record()