vectorstore 0.0.1

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: 38d0ad1caf146e2a67210910d242bed9032027e2cc7687840f3351db11cd3a6f
+  data.tar.gz: '06509c3c9626ba84274b434cb0c53bcf3cf48f683db07d7ce736dd7180d6bbdf'
+SHA512:
+  metadata.gz: 2694664e11a08dcc2df37939a5e785f71eb6ae62d1f4a54b03f80fd352f70ade8225038958a591d63d541b66affb558db1a1f29350ea7920349e7c1c15a14d60
+  data.tar.gz: a92e847ba349a1b0cd6e5065ce951b9a46794e50cb22bc88a4692f5ea08918cdb42e9e64407bbec1cf0f280f0b1756b1a428556c6234cb1d659b18a7a6e252b4

data/README.md

@@ -0,0 +1,101 @@
+# VectorStore
+
+A pure Ruby library for storing and querying vectors with optional 1 bit quantization. It provides an easy-to-use interface for adding vectors, computing cosine similarity, finding the closest vectors, and serializing to JSON. It also features quantized storage for **significantly** reduced memory requirements.
+
+## Features
+
+- **Vector Storage:** Easily add and retrieve vectors with unique keys.
+- **Closest Match:** Find the closest vectors to a given query vector using cosine similarity.
+- **Serialization:** Serialize the vector store to JSON for persistence.
+- **Quantization:** Optional 1-bit quantization to reduce memory footprint (significantly).
+- **Save/Load:** Persist vector store to disk and reload it.
+
+## Installation
+
+In a `Gemfile`:
+
+```ruby
+gem 'vectorstore'
+```
+
+Or directly:
+```bash
+gem install vectorstore
+```
+
+## Basic Usage
+
+```ruby
+require 'vectorstore'
+
+# Create a new VectorStore
+store = VectorStore.new
+
+# Add some vectors
+store.add("vector1", [1.0, 2.0, 3.0])
+store.add("vector2", [2.0, 3.0, 4.0])
+store.add("vector3", [3.0, 4.0, 5.0])
+store.add("vector4", [0.0, 0.0, 0.0])  # Zero vector edge case
+
+# Calculate cosine similarity between two vectors
+similarity = store.cosine_similarity([1.0, 2.0, 3.0], [2.0, 3.0, 4.0])
+puts "Cosine similarity: #{similarity}"
+
+# Find the closest vectors to a query vector
+closest = store.find_closest([2.0, 3.0, 4.0], 2)
+puts "Closest vectors: #{closest.inspect}"
+
+# Save the store to disk
+store.save("vector_store.json")
+
+# Load the store from disk
+loaded_store = VectorStore.new
+loaded_store.load("vector_store.json")
+```
+
+### Using OpenAI for vector embedding text
+
+VectorStore can integrate with OpenAI's API to generate embeddings for text inputs and queries. To use this feature with quantization (STRONGLY RECOMMENDED), initialize the store with quantized mode.
+
+> [!TIP]
+> Let me iterate again, using quantization with OpenAI embeddings is strongly recommended with VectorStore as the normal way we currently store the vectors is very space inefficient, particularly when serializing to disk.
+
+Example:
+
+```ruby
+store = VectorStore.new(quantized: true)
+
+store.add_with_openai("example", "Your sample text to embed")
+
+# You can query with text and have the text embedded automatically
+store.find_closest_with_openai("Your query text", 3)
+
+# You can also query by the key of the vector
+store.find_closest_with_key("example")
+```
+
+Supporting other embedding systems in a nice way would be good for the future, but I like OpenAI's embedding mechanism and it's cheap, so this is just step one. You can see example scripts in `examples/example_openai_*.rb` for a broader demo.
+
+> [!NOTE]  
+> For now, your API key is assumed to be in the `OPENAI_API_KEY` environment variable. The `text-embedding-3-small` model is also used by default but this can be overridden in calls by using the `embedding_model` keyword argument on `find_closest_with_openai` and `add_with_openai` calls.
+
+### Working with quantized vectors
+
+VectorStore supports 1 bit vector quantization so that vectors can be stored in a bitfield (using a ASCII-encoded string with 8 bits per character for portability) for a significant memory use reduction. The cost is accuracy, especially on low dimension vectors – high dimension vectors such as used for text embeddings from OpenAI's API (see above) will fare a LOT better. Initialize the store with the `quantized: true` option:
+
+```ruby
+store = VectorStore.new(quantized: true)
+store.add("vectorQ", [1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0])
+```
+
+## Running the Tests
+
+This project uses Minitest. To run the tests:
+
+```bash
+rake test
+```
+
+## License
+
+This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.

data/lib/vectorstore.rb

@@ -0,0 +1,181 @@
+require 'json'
+require 'base64'
+
+# If openai is available, offer it as an option for embedding text directly
+begin
+  require 'openai'
+rescue LoadError
+end
+
+
+class VectorStore
+  attr_reader :vectors
+  def initialize(quantized: false)
+    # Internal store mapping primary key to vector (array of numbers)
+    @vectors = {}
+    @quantized = quantized
+  end
+
+  # Add a vector with the given primary key. Overwrites any existing vector.
+  def add(key, vector)
+    if @quantized
+      @vectors[key] = quantize(vector)
+    else
+      @vectors[key] = vector
+    end
+  end
+
+  def add_with_openai(key, text, embedding_model: "text-embedding-3-small")
+    return false unless defined?(OpenAI)
+    add(key, get_openai_embedding(text, embedding_model: embedding_model))
+  end
+
+  # Remove a vector by its primary key.
+  def remove(key)
+    @vectors.delete(key)
+  end
+
+  # Retrieve a vector by its primary key.
+  def get(key)
+    @vectors[key]
+  end
+
+  # Compute the cosine similarity between two vectors.
+  def cosine_similarity(vec1, vec2)
+    if @quantized
+      vec1 = vec1.is_a?(String) ? vec1 : quantize(vec1)
+      vec2 = vec2.is_a?(String) ? vec2 : quantize(vec2)
+      return cosine_similarity_quantized(vec1, vec2)
+    end
+    if vec1.is_a?(String) && vec2.is_a?(String)
+      return cosine_similarity_quantized(vec1, vec2)
+    end
+
+    # Ensure vectors are of the same size
+    raise "Vector dimensions do not match" if vec1.size != vec2.size
+
+    dot_product = vec1.zip(vec2).map { |a, b| a * b }.sum
+    norm1 = Math.sqrt(vec1.map { |x| x * x }.sum)
+    norm2 = Math.sqrt(vec2.map { |x| x * x }.sum)
+    return 0.0 if norm1 == 0 || norm2 == 0
+
+    dot_product / (norm1 * norm2)
+  end
+
+  # Find the top k closest vectors to the query vector using cosine similarity.
+  # Returns an array of [key, similarity] pairs.
+  def find_closest(query_vector, k=1)
+    if @quantized
+      query_vector = quantize(query_vector)
+    end
+
+    similarities = @vectors.map do |key, vector|
+      similarity = cosine_similarity(query_vector, vector)
+      [key, similarity]
+    end
+    similarities.sort_by { |_, sim| -sim }.first(k)
+  end
+
+  def find_closest_with_key(key, k=1)
+    query_vector = @vectors[key]
+    find_closest(query_vector, k)
+  end
+
+  def get_openai_embedding(text, embedding_model: "text-embedding-3-small")
+    return false unless defined?(OpenAI)
+
+    client = OpenAI::Client.new(
+      access_token: ENV["OPENAI_API_KEY"],
+      log_errors: true
+    )
+    response = client.embeddings(
+      parameters: {
+        model: embedding_model,
+        input: text
+      }
+    )
+
+    response.dig("data", 0, "embedding")
+  end
+
+  def find_closest_with_openai(query_text, k=1, embedding_model: "text-embedding-3-small")
+    return false unless defined?(OpenAI)
+    query_vector = get_openai_embedding(query_text, embedding_model: embedding_model)
+    find_closest(query_vector, k)
+  end
+
+  # Compute cosine similarity for quantized vectors (bit strings).
+  def cosine_similarity_quantized(str1, str2)
+    dot = 0
+    total_ones_str1 = 0
+    total_ones_str2 = 0
+    str1.each_byte.with_index do |byte1, index|
+      byte2 = str2.getbyte(index)
+      dot += (byte1 & byte2).to_s(2).count("1")
+      total_ones_str1 += byte1.to_s(2).count("1")
+      total_ones_str2 += byte2.to_s(2).count("1")
+    end
+    return 0.0 if total_ones_str1 == 0 || total_ones_str2 == 0
+    sim = dot.to_f / (Math.sqrt(total_ones_str1) * Math.sqrt(total_ones_str2))
+    sim = 1.0 if (1.0 - sim).abs < 1e-6
+    sim
+  end
+
+  # Convert an array of floats to a 1-bit quantized bit string.
+  def quantize(vector)
+    # If it's already a string, it's already quantized
+    return vector if vector.is_a?(String)
+    
+    bits = vector.map { |x| x >= 0 ? 1 : 0 }
+    bytes = []
+    bits.each_slice(8) do |slice|
+      byte = slice.join.to_i(2)
+      bytes << byte.chr("ASCII-8BIT")
+    end
+    result = bytes.join
+    result.force_encoding("ASCII-8BIT")
+    result
+  end
+
+  # Serialize the internal vector store to a JSON string.
+  def serialize
+    if @quantized
+      encoded = {}
+      @vectors.each do |k, v|
+        encoded[k] = Base64.strict_encode64(v)
+      end
+      JSON.dump(encoded)
+    else
+      JSON.dump(@vectors)
+    end
+  end
+
+  # Deserialize a JSON string and update the internal store.
+  def deserialize(json_string)
+    data = JSON.parse(json_string)
+    # We need to detect if the data is quantized or not
+    # by seeing if the values are strings and not arrays
+    @quantized = data.values.first.is_a?(String)
+
+    if @quantized
+      decoded = {}
+      data.each do |k, v|
+        decoded[k] = Base64.decode64(v)
+      end
+      @vectors = decoded
+    else
+      @vectors = data
+    end
+  end
+
+  # Save the internal vector store to a file.
+  def save(filename)
+    File.write(filename, serialize)
+  end
+
+  # Load the internal vector store from a file.
+  def load(filename)
+    json_string = File.read(filename)
+    deserialize(json_string)
+  end
+end

metadata

@@ -0,0 +1,59 @@
+--- !ruby/object:Gem::Specification
+name: vectorstore
+version: !ruby/object:Gem::Version
+  version: 0.0.1
+platform: ruby
+authors:
+- Peter Cooper
+autorequire:
+bindir: bin
+cert_chain: []
+date: 2025-02-09 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: minitest
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '5.0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '5.0'
+description: A library for storing and handling vectors with optional quantization.
+email:
+- git@peterc.org
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- README.md
+- lib/vectorstore.rb
+homepage: http://github.com/peterc/vectorstore
+licenses:
+- MIT
+metadata: {}
+post_install_message:
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.5.23
+signing_key:
+specification_version: 4
+summary: A simple vector storage and querying library
+test_files: []