zvec-ruby 0.1.0

data/lib/zvec/active_record.rb

@@ -0,0 +1,113 @@
+require "zvec"
+require "active_support/concern"
+
+module Zvec
+  module ActiveRecord
+    # Rails concern that adds vector search capabilities to ActiveRecord models.
+    #
+    # Usage:
+    #   class Article < ApplicationRecord
+    #     include Zvec::ActiveRecord::Vectorize
+    #
+    #     vectorize :content,
+    #       dimensions: 1536,
+    #       prefix: "articles",
+    #       embed_with: ->(text) { OpenAI.embed(text) }
+    #   end
+    #
+    #   Article.vector_search([0.1, 0.2, ...], top_k: 5)
+    #   article.update_embedding!
+    #
+    module Vectorize
+      extend ActiveSupport::Concern
+
+      class_methods do
+        def vectorize(field, dimensions:, prefix: nil, embed_with: nil,
+                      metric: :cosine, zvec_path: nil)
+          prefix ||= table_name
+          zvec_path ||= Rails.root.join("tmp", "zvec", prefix).to_s if defined?(Rails)
+          zvec_path ||= File.join("tmp", "zvec", prefix)
+
+          class_attribute :zvec_config, instance_writer: false
+          self.zvec_config = {
+            field: field.to_s,
+            dimensions: dimensions,
+            prefix: prefix,
+            embed_with: embed_with,
+            metric: metric,
+            zvec_path: zvec_path
+          }
+
+          after_save :zvec_update_embedding!, if: -> { saved_change_to_attribute?(zvec_config[:field]) }
+          after_destroy :zvec_remove_embedding!
+
+          include InstanceMethods
+          extend SearchMethods
+        end
+      end
+
+      module InstanceMethods
+        def zvec_update_embedding!
+          cfg = self.class.zvec_config
+          text = send(cfg[:field])
+          return if text.blank?
+
+          embed_fn = cfg[:embed_with]
+          raise Zvec::Error, "No embed_with function configured" unless embed_fn
+
+          embedding = embed_fn.call(text)
+          store = self.class.zvec_store
+          store.add(id.to_s, embedding: embedding, content: text)
+          store.flush
+        end
+
+        def zvec_remove_embedding!
+          self.class.zvec_store.delete(id.to_s)
+        rescue => e
+          # Silently ignore if document doesn't exist
+        end
+
+        def zvec_embedding
+          result = self.class.zvec_store.fetch(id.to_s)
+          result[id.to_s]
+        end
+      end
+
+      module SearchMethods
+        def zvec_store
+          @zvec_store ||= begin
+            cfg = zvec_config
+            Zvec::RubyLLM::Store.new(
+              cfg[:zvec_path],
+              dimension: cfg[:dimensions],
+              metric: cfg[:metric]
+            )
+          end
+        end
+
+        def vector_search(query, top_k: 10, embed: true)
+          cfg = zvec_config
+
+          query_vector = if embed && query.is_a?(String) && cfg[:embed_with]
+                           cfg[:embed_with].call(query)
+                         elsif query.is_a?(Array)
+                           query
+                         else
+                           raise ArgumentError, "query must be a vector Array or a String with embed_with configured"
+                         end
+
+          results = zvec_store.search(query_vector, top_k: top_k)
+          ids = results.map { |r| r[:id] }
+          records = where(id: ids).index_by { |r| r.id.to_s }
+
+          results.filter_map do |r|
+            record = records[r[:id]]
+            next unless record
+            record.define_singleton_method(:zvec_score) { r[:score] }
+            record
+          end
+        end
+      end
+    end
+  end
+end

data/lib/zvec/collection.rb

@@ -0,0 +1,165 @@
+module Zvec
+  class Collection
+    attr_reader :schema
+
+    def initialize(ext_collection, schema: nil)
+      @ext = ext_collection
+      @schema = schema
+    end
+
+    # Create a new collection and open it.
+    def self.create_and_open(path, schema, read_only: false, enable_mmap: true)
+      opts = Ext::CollectionOptions.new
+      opts.read_only = read_only
+      opts.enable_mmap = enable_mmap
+      ext = Ext::Collection.create_and_open(path, schema.ext_schema, opts)
+      new(ext, schema: schema)
+    end
+
+    # Open an existing collection.
+    def self.open(path, read_only: false, enable_mmap: true)
+      opts = Ext::CollectionOptions.new
+      opts.read_only = read_only
+      opts.enable_mmap = enable_mmap
+      ext = Ext::Collection.open(path, opts)
+      new(ext)
+    end
+
+    def path
+      @ext.path
+    end
+
+    def stats
+      @ext.stats
+    end
+
+    def doc_count
+      @ext.stats.doc_count
+    end
+
+    # --- DDL ---
+
+    def create_index(field_name, index_params)
+      @ext.create_index(field_name.to_s, index_params)
+      self
+    end
+
+    def drop_index(field_name)
+      @ext.drop_index(field_name.to_s)
+      self
+    end
+
+    def optimize
+      @ext.optimize
+      self
+    end
+
+    def flush
+      @ext.flush
+      self
+    end
+
+    def destroy
+      @ext.destroy
+    end
+
+    # --- DML ---
+
+    def insert(docs)
+      docs = [docs] unless docs.is_a?(Array)
+      ext_docs = docs.map { |d| d.is_a?(Doc) ? d.ext_doc : d }
+      results = @ext.insert(ext_docs)
+      check_write_results!(results)
+    end
+
+    def upsert(docs)
+      docs = [docs] unless docs.is_a?(Array)
+      ext_docs = docs.map { |d| d.is_a?(Doc) ? d.ext_doc : d }
+      results = @ext.upsert(ext_docs)
+      check_write_results!(results)
+    end
+
+    def update(docs)
+      docs = [docs] unless docs.is_a?(Array)
+      ext_docs = docs.map { |d| d.is_a?(Doc) ? d.ext_doc : d }
+      results = @ext.update(ext_docs)
+      check_write_results!(results)
+    end
+
+    def delete(*pks)
+      pks = pks.flatten.map(&:to_s)
+      results = @ext.delete_pks(pks)
+      check_write_results!(results)
+    end
+
+    def delete_by_filter(filter)
+      @ext.delete_by_filter(filter)
+    end
+
+    # --- DQL ---
+
+    def query(field_name:, vector:, topk: 10, filter: nil,
+              include_vector: false, output_fields: nil, query_params: nil)
+      vq = VectorQuery.new(
+        field_name: field_name,
+        vector: vector,
+        topk: topk,
+        filter: filter,
+        include_vector: include_vector,
+        output_fields: output_fields,
+        query_params: query_params
+      )
+      raw_results = @ext.query(vq.ext_query)
+      raw_results.map do |h|
+        Doc.new(
+          pk: h["pk"],
+          fields: h.reject { |k, _| %w[pk score doc_id].include?(k) },
+          schema: @schema
+        ).tap { |d| d.instance_variable_set(:@score, h["score"]) }
+      end
+    end
+
+    def fetch(*pks)
+      pks = pks.flatten.map(&:to_s)
+      raw = @ext.fetch(pks)
+      raw.transform_values do |h|
+        Doc.new(pk: nil, fields: h, schema: @schema)
+      end
+    end
+
+    # Convenience: insert a hash directly
+    def add(pk:, **fields)
+      doc = Doc.new(pk: pk, fields: fields, schema: @schema)
+      insert(doc)
+    end
+
+    private
+
+    def check_write_results!(results)
+      results.each do |ok, msg|
+        raise Error, (msg.empty? ? "Write operation failed" : msg) unless ok
+      end
+      results
+    end
+
+    public
+
+    # Convenience: search with simpler API
+    def search(vector, field: nil, top_k: 10, filter: nil)
+      # Auto-detect vector field if not specified
+      fname = field&.to_s
+      unless fname
+        if @schema
+          vfield = @schema.ext_schema.vector_fields.first
+          raise Error, "No vector fields in schema" unless vfield
+          fname = vfield.name
+        else
+          vfields = @ext.schema.vector_fields
+          raise Error, "No vector fields in schema" if vfields.empty?
+          fname = vfields.first.name
+        end
+      end
+      query(field_name: fname, vector: vector, topk: top_k, filter: filter)
+    end
+  end
+end

data/lib/zvec/data_types.rb

@@ -0,0 +1,63 @@
+module Zvec
+  module DataTypes
+    # Re-export C++ enum values as Ruby-friendly constants
+    BINARY   = Ext::DataType::BINARY
+    STRING   = Ext::DataType::STRING
+    BOOL     = Ext::DataType::BOOL
+    INT32    = Ext::DataType::INT32
+    INT64    = Ext::DataType::INT64
+    UINT32   = Ext::DataType::UINT32
+    UINT64   = Ext::DataType::UINT64
+    FLOAT    = Ext::DataType::FLOAT
+    DOUBLE   = Ext::DataType::DOUBLE
+
+    VECTOR_FP32 = Ext::DataType::VECTOR_FP32
+    VECTOR_FP64 = Ext::DataType::VECTOR_FP64
+    VECTOR_FP16 = Ext::DataType::VECTOR_FP16
+    VECTOR_INT8 = Ext::DataType::VECTOR_INT8
+
+    SPARSE_VECTOR_FP32 = Ext::DataType::SPARSE_VECTOR_FP32
+    SPARSE_VECTOR_FP16 = Ext::DataType::SPARSE_VECTOR_FP16
+
+    ARRAY_STRING = Ext::DataType::ARRAY_STRING
+    ARRAY_INT32  = Ext::DataType::ARRAY_INT32
+    ARRAY_INT64  = Ext::DataType::ARRAY_INT64
+    ARRAY_FLOAT  = Ext::DataType::ARRAY_FLOAT
+    ARRAY_DOUBLE = Ext::DataType::ARRAY_DOUBLE
+    ARRAY_BOOL   = Ext::DataType::ARRAY_BOOL
+
+    # Metric types
+    L2     = Ext::MetricType::L2
+    IP     = Ext::MetricType::IP
+    COSINE = Ext::MetricType::COSINE
+
+    # Setter dispatch table: DataType -> Doc setter method name
+    SETTER_FOR = {
+      Ext::DataType::STRING => :set_string,
+      Ext::DataType::BOOL   => :set_bool,
+      Ext::DataType::INT32  => :set_int32,
+      Ext::DataType::INT64  => :set_int64,
+      Ext::DataType::UINT32 => :set_uint32,
+      Ext::DataType::UINT64 => :set_uint64,
+      Ext::DataType::FLOAT  => :set_float,
+      Ext::DataType::DOUBLE => :set_double,
+      Ext::DataType::VECTOR_FP32 => :set_float_vector,
+      Ext::DataType::VECTOR_FP64 => :set_double_vector,
+      Ext::DataType::ARRAY_STRING => :set_string_array,
+    }.freeze
+
+    GETTER_FOR = {
+      Ext::DataType::STRING => :get_string,
+      Ext::DataType::BOOL   => :get_bool,
+      Ext::DataType::INT32  => :get_int32,
+      Ext::DataType::INT64  => :get_int64,
+      Ext::DataType::UINT32 => :get_int32,
+      Ext::DataType::UINT64 => :get_int64,
+      Ext::DataType::FLOAT  => :get_float,
+      Ext::DataType::DOUBLE => :get_double,
+      Ext::DataType::VECTOR_FP32 => :get_float_vector,
+      Ext::DataType::VECTOR_FP64 => :get_double_vector,
+      Ext::DataType::ARRAY_STRING => :get_string_array,
+    }.freeze
+  end
+end

data/lib/zvec/doc.rb

@@ -0,0 +1,107 @@
+module Zvec
+  class Doc
+    attr_reader :ext_doc
+
+    def initialize(pk: nil, fields: {}, schema: nil)
+      @ext_doc = Ext::Doc.new
+      @ext_doc.pk = pk.to_s if pk
+      @schema = schema
+      fields.each { |k, v| set(k, v) } if schema
+    end
+
+    def pk
+      @ext_doc.pk
+    end
+
+    def pk=(value)
+      @ext_doc.pk = value.to_s
+    end
+
+    def score
+      @score || @ext_doc.score
+    end
+
+    def [](field_name)
+      get(field_name)
+    end
+
+    def []=(field_name, value)
+      set(field_name, value)
+    end
+
+    def set(field_name, value)
+      field_name = field_name.to_s
+      return @ext_doc.set_null(field_name) if value.nil?
+
+      if @schema
+        type = @schema.field_type(field_name)
+        if type
+          setter = DataTypes::SETTER_FOR[type]
+          return @ext_doc.send(setter, field_name, value) if setter
+        end
+      end
+
+      # Auto-detect type
+      case value
+      when String  then @ext_doc.set_string(field_name, value)
+      when Integer then @ext_doc.set_int64(field_name, value)
+      when Float   then @ext_doc.set_double(field_name, value)
+      when TrueClass, FalseClass then @ext_doc.set_bool(field_name, value)
+      when Array
+        if value.empty? || value.first.is_a?(Float) || value.first.is_a?(Integer)
+          @ext_doc.set_float_vector(field_name, value.map(&:to_f))
+        elsif value.first.is_a?(String)
+          @ext_doc.set_string_array(field_name, value)
+        end
+      end
+    end
+
+    def get(field_name)
+      field_name = field_name.to_s
+      return nil unless @ext_doc.has?(field_name)
+      return nil unless @ext_doc.has_value?(field_name)
+
+      if @schema
+        type = @schema.field_type(field_name)
+        if type
+          getter = DataTypes::GETTER_FOR[type]
+          return @ext_doc.send(getter, field_name) if getter
+        end
+      end
+
+      # Try common types in order
+      %i[get_string get_int64 get_float get_double get_bool
+         get_float_vector get_string_array].each do |m|
+        val = @ext_doc.send(m, field_name)
+        return val unless val.nil?
+      end
+      nil
+    end
+
+    def field_names
+      @ext_doc.field_names
+    end
+
+    def empty?
+      @ext_doc.empty?
+    end
+
+    def to_h
+      h = { "pk" => pk, "score" => score }
+      field_names.each { |f| h[f] = get(f) }
+      h
+    end
+
+    def to_s
+      @ext_doc.to_s
+    end
+
+    # Wrap a C++ Doc::Ptr into a Ruby Doc
+    def self.from_ext(ext_doc, schema: nil)
+      doc = allocate
+      doc.instance_variable_set(:@ext_doc, ext_doc)
+      doc.instance_variable_set(:@schema, schema)
+      doc
+    end
+  end
+end

data/lib/zvec/query.rb

@@ -0,0 +1,28 @@
+module Zvec
+  class VectorQuery
+    attr_reader :ext_query
+
+    def initialize(field_name:, vector:, topk: 10, filter: nil,
+                   include_vector: false, output_fields: nil, query_params: nil)
+      @ext_query = Ext::VectorQuery.new
+      @ext_query.field_name = field_name.to_s
+      @ext_query.topk = topk
+      @ext_query.set_query_vector(vector.map(&:to_f))
+      @ext_query.filter = filter if filter
+      @ext_query.include_vector = include_vector
+      @ext_query.set_output_fields(output_fields.map(&:to_s)) if output_fields
+      if query_params
+        case query_params
+        when Ext::HnswQueryParams
+          @ext_query.set_hnsw_query_params(query_params)
+        when Ext::IVFQueryParams
+          @ext_query.set_ivf_query_params(query_params)
+        when Ext::FlatQueryParams
+          @ext_query.set_flat_query_params(query_params)
+        else
+          raise ArgumentError, "Unknown query_params type: #{query_params.class}"
+        end
+      end
+    end
+  end
+end

data/lib/zvec/ruby_llm.rb

@@ -0,0 +1,108 @@
+require "zvec"
+
+module Zvec
+  module RubyLLM
+    # A vector store backend for the ruby_llm gem.
+    #
+    # Usage with ruby_llm:
+    #   store = Zvec::RubyLLM::Store.new("/path/to/db", dimension: 1536)
+    #   store.add("doc-1", embedding: [...], metadata: { title: "Hello" })
+    #   results = store.search([0.1, 0.2, ...], top_k: 5)
+    #
+    class Store
+      DEFAULT_VECTOR_FIELD = "embedding"
+      DEFAULT_CONTENT_FIELD = "content"
+
+      attr_reader :collection, :dimension
+
+      def initialize(path, dimension:, metric: :cosine, vector_field: DEFAULT_VECTOR_FIELD,
+                     content_field: DEFAULT_CONTENT_FIELD)
+        @vector_field = vector_field.to_s
+        @content_field = content_field.to_s
+        @dimension = dimension
+
+        metric_type = case metric.to_sym
+                      when :cosine then Zvec::DataTypes::COSINE
+                      when :l2     then Zvec::DataTypes::L2
+                      when :ip     then Zvec::DataTypes::IP
+                      else raise ArgumentError, "Unknown metric: #{metric}"
+                      end
+
+        cf = @content_field
+        vf = @vector_field
+        dim = dimension
+        schema = Zvec::Schema.new("ruby_llm_store") do
+          string cf, nullable: true
+          vector vf, dimension: dim,
+                 index: Zvec::Ext::HnswIndexParams.new(metric_type)
+        end
+
+        @schema = schema
+
+        if Dir.exist?(path)
+          @collection = Zvec::Collection.open(path)
+        else
+          @collection = Zvec::Collection.create_and_open(path, schema)
+        end
+      end
+
+      # Add a document with its embedding and optional metadata.
+      def add(id, embedding:, content: nil, metadata: {})
+        doc = Zvec::Doc.new(pk: id, schema: @schema)
+        doc[@vector_field] = embedding
+        doc[@content_field] = content if content
+        metadata.each { |k, v| doc[k] = v }
+        @collection.insert(doc)
+      end
+
+      # Batch-add documents.
+      # docs: array of { id:, embedding:, content:, metadata: {} }
+      def add_many(docs)
+        zvec_docs = docs.map do |d|
+          doc = Zvec::Doc.new(pk: d[:id], schema: @schema)
+          doc[@vector_field] = d[:embedding]
+          doc[@content_field] = d[:content] if d[:content]
+          (d[:metadata] || {}).each { |k, v| doc[k] = v }
+          doc
+        end
+        @collection.insert(zvec_docs)
+      end
+
+      # Search for similar vectors.
+      def search(query_vector, top_k: 10, filter: nil)
+        results = @collection.query(
+          field_name: @vector_field,
+          vector: query_vector,
+          topk: top_k,
+          filter: filter
+        )
+        results.map do |doc|
+          {
+            id: doc.pk,
+            score: doc.score,
+            content: doc[@content_field],
+            metadata: doc.to_h.reject { |k, _| ["pk", "score", @vector_field, @content_field].include?(k) }
+          }
+        end
+      end
+
+      # Delete documents by IDs.
+      def delete(*ids)
+        @collection.delete(*ids.flatten)
+      end
+
+      # Fetch documents by IDs.
+      def fetch(*ids)
+        @collection.fetch(*ids.flatten)
+      end
+
+      def flush
+        @collection.flush
+      end
+
+      def count
+        @collection.doc_count
+      end
+    end
+  end
+end

data/lib/zvec/schema.rb

@@ -0,0 +1,70 @@
+module Zvec
+  class Schema
+    attr_reader :ext_schema
+
+    def initialize(name, &block)
+      @ext_schema = Ext::CollectionSchema.new(name)
+      @field_types = {}
+      instance_eval(&block) if block
+    end
+
+    def field(name, type, dimension: nil, nullable: false, index: nil)
+      name = name.to_s
+      fs = Ext::FieldSchema.new(name, type)
+      fs.dimension = dimension if dimension
+      fs.nullable = nullable
+      fs.set_index_params(index) if index
+      @ext_schema.add_field(fs)
+      @field_types[name] = type
+      self
+    end
+
+    def vector(name, dimension:, type: DataTypes::VECTOR_FP32, index: nil)
+      field(name, type, dimension: dimension, index: index)
+    end
+
+    def string(name, **opts)
+      field(name, DataTypes::STRING, **opts)
+    end
+
+    def int32(name, **opts)
+      field(name, DataTypes::INT32, **opts)
+    end
+
+    def int64(name, **opts)
+      field(name, DataTypes::INT64, **opts)
+    end
+
+    def float(name, **opts)
+      field(name, DataTypes::FLOAT, **opts)
+    end
+
+    def double(name, **opts)
+      field(name, DataTypes::DOUBLE, **opts)
+    end
+
+    def bool(name, **opts)
+      field(name, DataTypes::BOOL, **opts)
+    end
+
+    def name
+      @ext_schema.name
+    end
+
+    def field_names
+      @ext_schema.field_names
+    end
+
+    def field_type(name)
+      @field_types[name.to_s]
+    end
+
+    def has_field?(name)
+      @ext_schema.has_field?(name.to_s)
+    end
+
+    def to_s
+      @ext_schema.to_s
+    end
+  end
+end

data/lib/zvec/version.rb

@@ -0,0 +1,3 @@
+module Zvec
+  VERSION = "0.1.0"
+end

data/lib/zvec.rb

@@ -0,0 +1,22 @@
+require_relative "zvec/version"
+
+# Try loading precompiled native extension first (platform gem),
+# then fall back to the compiled-from-source version.
+begin
+  ruby_version = RUBY_VERSION[/\d+\.\d+/]
+  require "zvec/#{ruby_version}/zvec_ext"
+rescue LoadError
+  require "zvec/zvec_ext"
+end
+
+require_relative "zvec/data_types"
+require_relative "zvec/schema"
+require_relative "zvec/doc"
+require_relative "zvec/query"
+require_relative "zvec/collection"
+
+module Zvec
+  class Error < StandardError; end
+
+  include DataTypes
+end