Vector Search in Ruby

Storing embeddings

Databases most frequently used with Ruby on Rails applications aren't well-suited for storing vectors by default. Thankfully, there are custom extensions which can be enabled.

At the moment of writing this article (Nov, 2024), the most mature extension is available for PostgreSQL and is called pgvector. There are options for other relational databases, but they are in an experimental state - squlite-vec for SQLite, Vector for MariaDB and HeatWave, required for searching with MySQL.

Alternatively, you can use a NoSQL database, which is also well established in the Ruby on Rails ecosystem. You may use Redis with neighbour-redis gem.

Embeddings with pgvector and Ruby on Rails

In order to enable pgvector, you have to first install it on your machine and add an appropriate gem to the application:

$ bundle add neighbor

Then add and run a simple migration:

class EnablePgVector < ActiveRecord::Migration[8.0]
  def change
    enable_extension "vector"
  end
end

Add a column of vector type to your ActiveRecord model:

class AddEmbeddingToItems < ActiveRecord::Migration[8.0]
  def change
    add_column :items, :embedding, :vector, limit: 3072
  end
end

Generate an embedding and fill the new column of your model with data:

client = OpenAI::Client.new(access_token: OPENAI_API_KEY)
response = client.embeddings(
  parameters: {
    model: "text-embedding-3-large",
    input: "Ruby is great for writing AI software"
  }
)

item = Item.create(
  embedding: response.dig("data", 0, "embedding")
)

With data in the database, we can query for similar vectors using operators provided by the pgvector extension, for example:

SELECT items.* FROM items ORDER BY items.embedding <=> '[0.1, 0.4, 0.6]'

The above SQL code orders the found items by cosine similarity to the given embedding. The available operators are:

• <=> for cosine
• <#> for inner product
• <-> for euclidean
• <+> for taxicab

However, the goal is not to use SQL directly, but to embrace ActiveRecord and the Rails way.

Find the most similar item with Ruby on Rails

Thankfully, the same gem can be used for finding similar items using ActiveRecord. Its usage is very simple and follows the well-known syntax of Rails. It can be used with Postgres, SQLite, MariaDB or MySQL.

You've already added it to the Gemfile:

gem "neighbor"

Follow it up by declaring it in your ActiveRecord model definition:

class Item < ApplicationRecord
  has_neighbors :embedding
end

Finally, use the nearest_neighbors method to find the most similar items:

item = Item.first
item.nearest_neighbors(:embedding, distance: "cosine")

The result is an ActiveRecord_Relation, ordered by the distance to the item. You may use .first(3) to get the most similar ones.

Vector Search using Ruby on Rails

The neighbor gem also allows to query a collection of records, using a custom embedding.

embedding = [0.1, 0.7, -0.2]
most_similar_items = Item.nearest_neighbors(
  :embedding, embedding, distance: "cosine"
)

And that is it! A few lines of code allow us to grasp all the power of Vector Search!

Note: An alternative solution may be achieved using a more complex library - Langchain.rb. However, its usage is beyond the scope of this article.

Find the lost image

With the knowledge gathered in the article, let's create a Ruby on Rails application to find photos by their descriptions. The only additional tool you'll need is the image-to-text transcript as we need to create a text description of each image. Fortunately, this can also be done using an LLM.

Convert image to text using Ruby and ChatGPT

Using such a powerful tool as an LLM it's actually very easy. Let's take a cute dog image and see how to do it:

encoded_image = Base64.strict_encode64("/path/to/dog.jpg")

messages = [
  { "type": "text", "text": "What’s in this image?" },
  { "type": "image_url",
    "image_url": {
      "url": "data:image/png;base64,#{encoded_image}"
    }
  }
]

response = client.chat(
  parameters: {
    model: "gpt-4o-mini",
    messages: [ { role: "user", content: messages } ]
  }
)

puts response.dig("choices", 0, "message", "content")

# => "The image depicts a small dog with curly fur that appears to be playfully biting on a shoe. The dog has large, expressive eyes and is engaging with the shoe, which is tan with a light blue interior. The background is removed, making the focus solely on the dog and the shoe."

Vector Search application with Ruby on Rails

The rest of the application's design is straightforward and follows the Rails conventions. The best part is - you don't even need to code it yourself! Check out a ready application at https://github.com/pstrzalk/image-finder-vector-search

The main element of the solution is the Image model, which processes the attachment after creation. It gathers the image’s description and embeds it using ChatGPT.

# https://github.com/pstrzalk/image-finder-vector-search/blob/main/app/models/image.rb

class Image < ApplicationRecord
  has_one_attached :file

  has_neighbors :embedding
  after_create_commit :recalculate_embedding

  def recalculate_embedding
    return unless file.attachment.present?

    client = OpenAI::Client.new(
      access_token: Rails.configuration.x.openai_api_key
    )

    base64_encoded_image = Base64.strict_encode64(file.download)
    messages = [
      { "type": "text", "text": "What’s in this image?" },
      { "type": "image_url",
        "image_url": {
          "url": "data:image/png;base64,#{base64_encoded_image}"
        }
      }
    ]
    response = client.chat(
      parameters: {
        model: "gpt-4o-mini",
        messages: [ { role: "user", content: messages } ]
      }
    )
    self.description = response.dig("choices", 0, "message", "content")

    response = client.embeddings(
      parameters: {
        model: "text-embedding-3-large",
        input: description
      }
    )
    self.embedding = response.dig("data", 0, "embedding")
    self.save
  end
end

The second key component is the ImageController , which by default shows all the available images.

# https://github.com/pstrzalk/image-finder-vector-search/blob/main/app/controllers/images_controller.rb

class ImagesController < ApplicationController
    # ...

  def index
    if params[:query].present?
      client = OpenAI::Client.new(access_token: Rails.configuration.x.openai_api_key)

      embedding = client
        .embeddings(parameters: { model: "text-embedding-3-large", input: params[:query] })
        .dig("data", 0, "embedding")

      @images = Image.nearest_neighbors(:embedding, embedding, distance: "euclidean").first(1)
    else
      @images = Image.all
    end
  end

  # ...
end

When a query param is present, it:

• Embeds the query using the specified model
• Finds the image with the closest vector embedding to the query

Infinite possibilities

By combining the power of vector search with Ruby on Rails and tools like ChatGPT, you can build innovative and intuitive applications that go beyond traditional search capabilities. The only limit is your imagination. Check out the image-finder-vector-search repository to explore the example implementation and see how you can enhance your own projects with these cutting-edge techniques.