diff --git a/lib/music_library/records/record_embedding.ex b/lib/music_library/records/record_embedding.ex index a267d809..dd5a608a 100644 --- a/lib/music_library/records/record_embedding.ex +++ b/lib/music_library/records/record_embedding.ex @@ -10,7 +10,7 @@ defmodule MusicLibrary.Records.RecordEmbedding do schema "record_embeddings" do belongs_to :record, Record - field :embedding, MusicLibrary.Records.RecordEmbedding.EmbeddingType + field :embedding, SqliteVec.Ecto.Float32 field :text_representation, :string timestamps(type: :utc_datetime) @@ -20,28 +20,6 @@ defmodule MusicLibrary.Records.RecordEmbedding do record_embedding |> cast(attrs, [:record_id, :embedding, :text_representation]) |> validate_required([:record_id, :embedding, :text_representation]) - |> validate_embedding_dimensions() |> unique_constraint(:record_id) 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 diff --git a/lib/music_library/records/record_embedding/embedding_type.ex b/lib/music_library/records/record_embedding/embedding_type.ex deleted file mode 100644 index 813af778..00000000 --- a/lib/music_library/records/record_embedding/embedding_type.ex +++ /dev/null @@ -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 diff --git a/lib/music_library/records/similarity.ex b/lib/music_library/records/similarity.ex index b5007446..45049607 100644 --- a/lib/music_library/records/similarity.ex +++ b/lib/music_library/records/similarity.ex @@ -4,7 +4,9 @@ defmodule MusicLibrary.Records.Similarity do """ import Ecto.Query + import(SqliteVec.Ecto.Query) + alias MusicLibrary.Records alias MusicLibrary.Records.{Record, RecordEmbedding} alias MusicLibrary.Repo @@ -53,44 +55,30 @@ defmodule MusicLibrary.Records.Similarity do """ def find_similar(record_id, opts \\ []) do limit = Keyword.get(opts, :limit, 10) - min_similarity = Keyword.get(opts, :min_similarity, 0.0) scope = Keyword.get(opts, :scope) - with {:ok, source_embedding} <- get_embedding(record_id), - similar_records <- calculate_similarities(source_embedding, record_id, scope) do - 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 + record = Records.get_record!(record_id) + record_musicbrainz_id = record.musicbrainz_id - @doc """ - Calculates cosine similarity between two embedding vectors. + case get_embedding(record_id) do + {: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: - - 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 + query = apply_scope_filter(query, scope) - dot_product = - Enum.zip(vec_a, vec_b) - |> Enum.reduce(0.0, fn {a, b}, acc -> acc + a * b end) + query + |> Repo.all() - magnitude_a = calculate_magnitude(vec_a) - 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) + {:error, :not_found} -> + [] end end @@ -142,31 +130,6 @@ defmodule MusicLibrary.Records.Similarity do defp humanize_type(:other), do: "Other" 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 from [re, r] in query, where: not is_nil(r.purchased_at) end diff --git a/lib/music_library_web/components/record_components.ex b/lib/music_library_web/components/record_components.ex index bc019cfb..6cf09b04 100644 --- a/lib/music_library_web/components/record_components.ex +++ b/lib/music_library_web/components/record_components.ex @@ -422,7 +422,6 @@ defmodule MusicLibraryWeb.RecordComponents do - {Float.round(similarity * 100, 0)}% diff --git a/lib/sqlite_vec/ecto/float32.ex b/lib/sqlite_vec/ecto/float32.ex new file mode 100644 index 00000000..83462c63 --- /dev/null +++ b/lib/sqlite_vec/ecto/float32.ex @@ -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 diff --git a/lib/sqlite_vec/ecto/query.ex b/lib/sqlite_vec/ecto/query.ex new file mode 100644 index 00000000..503490e9 --- /dev/null +++ b/lib/sqlite_vec/ecto/query.ex @@ -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 diff --git a/lib/sqlite_vec/float32.ex b/lib/sqlite_vec/float32.ex new file mode 100644 index 00000000..f3a6cc76 --- /dev/null +++ b/lib/sqlite_vec/float32.ex @@ -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) + ...> <> = 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: <> + from_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(<>) + 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 + <> = vector.data + + for <>, 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 + <> = 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 diff --git a/priv/repo/migrations/20251011192421_create_record_embeddings.exs b/priv/repo/migrations/20251011192421_create_record_embeddings.exs index 0e1a2092..c36ae966 100644 --- a/priv/repo/migrations/20251011192421_create_record_embeddings.exs +++ b/priv/repo/migrations/20251011192421_create_record_embeddings.exs @@ -1,17 +1,23 @@ defmodule MusicLibrary.Repo.Migrations.CreateRecordEmbeddings do use Ecto.Migration - def change do - create table(:record_embeddings, primary_key: false) do - add :id, :binary_id, primary_key: true - add :record_id, references(:records, type: :binary_id, on_delete: :delete_all), null: false + def up do + execute(""" + CREATE TABLE record_embeddings ( + 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 - add :text_representation, :text, null: false + execute(""" + CREATE UNIQUE INDEX record_embeddings_record_id_index ON record_embeddings (record_id); + """) + end - timestamps(type: :utc_datetime) - end - - create unique_index(:record_embeddings, [:record_id]) + def down do + drop table(:record_embeddings) end end diff --git a/test/music_library/records/similarity_test.exs b/test/music_library/records/similarity_test.exs index c028fbd3..705ea2c3 100644 --- a/test/music_library/records/similarity_test.exs +++ b/test/music_library/records/similarity_test.exs @@ -107,57 +107,6 @@ defmodule MusicLibrary.Records.SimilarityTest do 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 test "stores and retrieves an embedding" do record = record() @@ -167,10 +116,7 @@ defmodule MusicLibrary.Records.SimilarityTest do assert {:ok, _} = Similarity.store_embedding(record.id, embedding, text_rep) assert {:ok, retrieved_embedding} = Similarity.get_embedding(record.id) - assert length(retrieved_embedding) == 1536 - # 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) + assert SqliteVec.Float32.new(embedding) == retrieved_embedding end 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, retrieved_embedding} = Similarity.get_embedding(record.id) - assert List.first(retrieved_embedding) == 0.7 + assert SqliteVec.Float32.new(embedding2) == retrieved_embedding end test "returns error for non-existent record" do @@ -227,13 +173,6 @@ defmodule MusicLibrary.Records.SimilarityTest do assert length(similar) == 1 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 record_without_embedding = record()