vectra-client 0.1.3 → 0.2.0

data/lib/generators/vectra/templates/enable_pgvector_extension.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+class EnablePgvectorExtension < ActiveRecord::Migration<%= migration_version %>
+  def up
+    enable_extension 'vector'
+  end
+
+  def down
+    disable_extension 'vector'
+  end
+end

data/lib/generators/vectra/templates/vectra.rb

@@ -0,0 +1,79 @@
+# frozen_string_literal: true
+
+# Vectra configuration
+#
+# For more information see: https://github.com/stokry/vectra
+
+Vectra.configure do |config|
+  # Provider configuration
+  config.provider = :<%= options[:provider] %>
+
+  <%- if options[:provider] == 'pinecone' -%>
+  # Pinecone credentials
+  config.api_key = Rails.application.credentials.dig(:pinecone, :api_key)
+  config.environment = Rails.application.credentials.dig(:pinecone, :environment) || 'us-east-1'
+  # Or use direct host:
+  # config.host = 'your-index-host.pinecone.io'
+
+  <%- elsif options[:provider] == 'pgvector' -%>
+  # PostgreSQL with pgvector extension
+  <%- if options[:database_url] -%>
+  config.host = '<%= options[:database_url] %>'
+  <%- else -%>
+  config.host = ENV['DATABASE_URL'] || Rails.configuration.database_configuration[Rails.env]['url']
+  <%- end -%>
+  config.api_key = nil  # pgvector uses connection URL for auth
+
+  # Connection pooling (recommended for production)
+  config.pool_size = ENV.fetch('VECTRA_POOL_SIZE', 10).to_i
+  config.pool_timeout = 5
+
+  # Batch operations
+  config.batch_size = ENV.fetch('VECTRA_BATCH_SIZE', 100).to_i
+
+  <%- elsif options[:provider] == 'qdrant' -%>
+  # Qdrant credentials
+  config.api_key = Rails.application.credentials.dig(:qdrant, :api_key)
+  config.host = Rails.application.credentials.dig(:qdrant, :host)
+
+  <%- elsif options[:provider] == 'weaviate' -%>
+  # Weaviate credentials
+  config.api_key = Rails.application.credentials.dig(:weaviate, :api_key)
+  config.host = Rails.application.credentials.dig(:weaviate, :host)
+
+  <%- end -%>
+  # Timeouts
+  config.timeout = 30
+  config.open_timeout = 10
+
+  # Retry configuration
+  config.max_retries = 3
+  config.retry_delay = 1
+
+  # Logging
+  config.logger = Rails.logger
+
+  <%- if options[:instrumentation] -%>
+  # Instrumentation (metrics and monitoring)
+  config.instrumentation = true
+
+  # Uncomment for New Relic:
+  # require 'vectra/instrumentation/new_relic'
+  # Vectra::Instrumentation::NewRelic.setup!
+
+  # Uncomment for Datadog:
+  # require 'vectra/instrumentation/datadog'
+  # Vectra::Instrumentation::Datadog.setup!(
+  #   host: ENV['DD_AGENT_HOST'] || 'localhost',
+  #   port: ENV['DD_DOGSTATSD_PORT']&.to_i || 8125
+  # )
+
+  # Custom instrumentation:
+  # Vectra.on_operation do |event|
+  #   Rails.logger.info "Vectra: #{event.operation} on #{event.provider} took #{event.duration}ms"
+  #   if event.failure?
+  #     Rails.logger.error "Vectra error: #{event.error.message}"
+  #   end
+  # end
+  <%- end -%>
+end

data/lib/vectra/active_record.rb

@@ -0,0 +1,195 @@
+# frozen_string_literal: true
+
+require "active_support/concern"
+
+module Vectra
+  # ActiveRecord integration for vector embeddings
+  #
+  # Provides ActiveRecord models with vector search capabilities.
+  #
+  # @example Basic usage
+  #   class Document < ApplicationRecord
+  #     include Vectra::ActiveRecord
+  #
+  #     has_vector :embedding,
+  #                dimension: 384,
+  #                provider: :pgvector,
+  #                index: 'documents'
+  #   end
+  #
+  #   # Auto-index on create/update
+  #   doc = Document.create!(title: 'Hello', embedding: [0.1, 0.2, ...])
+  #
+  #   # Search similar documents
+  #   results = Document.vector_search([0.1, 0.2, ...], limit: 10)
+  #
+  module ActiveRecord
+    extend ActiveSupport::Concern
+
+    included do
+      class_attribute :_vectra_config, default: {}
+      class_attribute :_vectra_client
+    end
+
+    class_methods do
+      # Define a vector attribute
+      #
+      # @param attribute [Symbol] The attribute name (e.g., :embedding)
+      # @param dimension [Integer] Vector dimension
+      # @param provider [Symbol] Provider name (:pinecone, :pgvector, etc.)
+      # @param index [String] Index/collection name
+      # @param auto_index [Boolean] Automatically index on save
+      # @param metadata_fields [Array<Symbol>] Fields to include in metadata
+      #
+      # @example
+      #   has_vector :embedding,
+      #              dimension: 384,
+      #              provider: :pgvector,
+      #              index: 'documents',
+      #              auto_index: true,
+      #              metadata_fields: [:title, :category, :status]
+      #
+      def has_vector(attribute, dimension:, provider: nil, index: nil, auto_index: true, metadata_fields: [])
+        self._vectra_config = {
+          attribute: attribute,
+          dimension: dimension,
+          provider: provider || Vectra.configuration.provider,
+          index: index || table_name,
+          auto_index: auto_index,
+          metadata_fields: metadata_fields
+        }
+
+        # Initialize client lazily
+        define_singleton_method(:vectra_client) do
+          @_vectra_client ||= Vectra::Client.new(provider: _vectra_config[:provider])
+        end
+
+        # Callbacks for auto-indexing
+        if auto_index
+          after_save :_vectra_index_vector
+          after_destroy :_vectra_delete_vector
+        end
+
+        # Class methods for search
+        define_singleton_method(:vector_search) do |query_vector, limit: 10, **options|
+          _vectra_search(query_vector, limit: limit, **options)
+        end
+
+        define_singleton_method(:similar_to) do |record, limit: 10, **options|
+          vector = record.send(_vectra_config[:attribute])
+          raise ArgumentError, "Record has no vector" if vector.nil?
+
+          _vectra_search(vector, limit: limit, **options)
+        end
+      end
+
+      # Search vectors
+      #
+      # @api private
+      def _vectra_search(query_vector, limit: 10, filter: {}, score_threshold: nil, load_records: true)
+        config = _vectra_config
+        results = vectra_client.query(
+          index: config[:index],
+          vector: query_vector,
+          top_k: limit,
+          filter: filter
+        )
+
+        # Filter by score if threshold provided
+        results = results.above_score(score_threshold) if score_threshold
+
+        return results unless load_records
+
+        # Load ActiveRecord objects
+        ids = results.map { |match| match.id.gsub("#{config[:index]}_", "").to_i }
+        records = where(id: ids).index_by(&:id)
+
+        results.map do |match|
+          id = match.id.gsub("#{config[:index]}_", "").to_i
+          record = records[id]
+          next unless record
+
+          record.instance_variable_set(:@_vectra_score, match.score)
+          record.define_singleton_method(:vector_score) { @_vectra_score }
+          record
+        end.compact
+      end
+    end
+
+    # Instance methods
+
+    # Index this record's vector
+    #
+    # @return [void]
+    def index_vector!
+      config = self.class._vectra_config
+      vector_data = send(config[:attribute])
+
+      raise ArgumentError, "#{config[:attribute]} is nil" if vector_data.nil?
+
+      metadata = config[:metadata_fields].each_with_object({}) do |field, hash|
+        hash[field.to_s] = send(field) if respond_to?(field)
+      end
+
+      self.class.vectra_client.upsert(
+        index: config[:index],
+        vectors: [{
+          id: _vectra_vector_id,
+          values: vector_data,
+          metadata: metadata
+        }]
+      )
+    end
+
+    # Delete this record's vector from index
+    #
+    # @return [void]
+    def delete_vector!
+      config = self.class._vectra_config
+
+      self.class.vectra_client.delete(
+        index: config[:index],
+        ids: [_vectra_vector_id]
+      )
+    end
+
+    # Find similar records
+    #
+    # @param limit [Integer] Number of results
+    # @param filter [Hash] Metadata filter
+    # @return [Array<ActiveRecord::Base>]
+    def similar(limit: 10, filter: {})
+      config = self.class._vectra_config
+      vector_data = send(config[:attribute])
+
+      raise ArgumentError, "#{config[:attribute]} is nil" if vector_data.nil?
+
+      self.class._vectra_search(vector_data, limit: limit + 1, filter: filter)
+        .reject { |record| record.id == id } # Exclude self
+        .first(limit)
+    end
+
+    private
+
+    # Auto-index callback
+    def _vectra_index_vector
+      return unless saved_change_to_attribute?(self.class._vectra_config[:attribute])
+
+      index_vector!
+    rescue StandardError => e
+      Rails.logger.error("Vectra auto-index failed: #{e.message}") if defined?(Rails)
+    end
+
+    # Auto-delete callback
+    def _vectra_delete_vector
+      delete_vector!
+    rescue StandardError => e
+      Rails.logger.error("Vectra auto-delete failed: #{e.message}") if defined?(Rails)
+    end
+
+    # Generate vector ID
+    def _vectra_vector_id
+      "#{self.class._vectra_config[:index]}_#{id}"
+    end
+  end
+end

data/lib/vectra/client.rb

@@ -60,7 +60,14 @@ module Vectra
       validate_index!(index)
       validate_vectors!(vectors)
 
-      provider.upsert(index: index, vectors: vectors, namespace: namespace)
+      Instrumentation.instrument(
+        operation: :upsert,
+        provider: provider_name,
+        index: index,
+        metadata: { vector_count: vectors.size }
+      ) do
+        provider.upsert(index: index, vectors: vectors, namespace: namespace)
+      end
     end
 
     # Query vectors by similarity
@@ -94,15 +101,25 @@ module Vectra
       validate_index!(index)
       validate_query_vector!(vector)
 
-      provider.query(
+      result = nil
+      Instrumentation.instrument(
+        operation: :query,
+        provider: provider_name,
         index: index,
-        vector: vector,
-        top_k: top_k,
-        namespace: namespace,
-        filter: filter,
-        include_values: include_values,
-        include_metadata: include_metadata
-      )
+        metadata: { top_k: top_k }
+      ) do
+        result = provider.query(
+          index: index,
+          vector: vector,
+          top_k: top_k,
+          namespace: namespace,
+          filter: filter,
+          include_values: include_values,
+          include_metadata: include_metadata
+        )
+      end
+
+      result
     end
 
     # Fetch vectors by IDs
@@ -120,7 +137,14 @@ module Vectra
       validate_index!(index)
       validate_ids!(ids)
 
-      provider.fetch(index: index, ids: ids, namespace: namespace)
+      Instrumentation.instrument(
+        operation: :fetch,
+        provider: provider_name,
+        index: index,
+        metadata: { id_count: ids.size }
+      ) do
+        provider.fetch(index: index, ids: ids, namespace: namespace)
+      end
     end
 
     # Update a vector's metadata or values
@@ -145,13 +169,20 @@ module Vectra
 
       raise ValidationError, "Must provide metadata or values to update" if metadata.nil? && values.nil?
 
-      provider.update(
+      Instrumentation.instrument(
+        operation: :update,
+        provider: provider_name,
         index: index,
-        id: id,
-        metadata: metadata,
-        values: values,
-        namespace: namespace
-      )
+        metadata: { has_metadata: !metadata.nil?, has_values: !values.nil? }
+      ) do
+        provider.update(
+          index: index,
+          id: id,
+          metadata: metadata,
+          values: values,
+          namespace: namespace
+        )
+      end
     end
 
     # Delete vectors
@@ -179,13 +210,20 @@ module Vectra
         raise ValidationError, "Must provide ids, filter, or delete_all"
       end
 
-      provider.delete(
+      Instrumentation.instrument(
+        operation: :delete,
+        provider: provider_name,
         index: index,
-        ids: ids,
-        namespace: namespace,
-        filter: filter,
-        delete_all: delete_all
-      )
+        metadata: { id_count: ids&.size, delete_all: delete_all, has_filter: !filter.nil? }
+      ) do
+        provider.delete(
+          index: index,
+          ids: ids,
+          namespace: namespace,
+          filter: filter,
+          delete_all: delete_all
+        )
+      end
     end
 
     # List all indexes

data/lib/vectra/configuration.rb

@@ -14,7 +14,8 @@ module Vectra
     SUPPORTED_PROVIDERS = %i[pinecone qdrant weaviate pgvector].freeze
 
     attr_accessor :api_key, :environment, :host, :timeout, :open_timeout,
-                  :max_retries, :retry_delay, :logger
+                  :max_retries, :retry_delay, :logger, :pool_size, :pool_timeout,
+                  :batch_size, :instrumentation
 
     attr_reader :provider
 
@@ -28,6 +29,10 @@ module Vectra
       @max_retries = 3
       @retry_delay = 1
       @logger = nil
+      @pool_size = 5
+      @pool_timeout = 5
+      @batch_size = 100
+      @instrumentation = false
     end
 
     # Set the provider

data/lib/vectra/instrumentation/datadog.rb

@@ -0,0 +1,82 @@
+# frozen_string_literal: true
+
+module Vectra
+  module Instrumentation
+    # Datadog instrumentation adapter
+    #
+    # Automatically reports Vectra metrics to Datadog using DogStatsD.
+    #
+    # @example Enable Datadog instrumentation
+    #   # config/initializers/vectra.rb
+    #   require 'vectra/instrumentation/datadog'
+    #
+    #   Vectra.configure do |config|
+    #     config.instrumentation = true
+    #   end
+    #
+    #   Vectra::Instrumentation::Datadog.setup!(
+    #     host: ENV['DD_AGENT_HOST'] || 'localhost',
+    #     port: ENV['DD_DOGSTATSD_PORT']&.to_i || 8125
+    #   )
+    #
+    module Datadog
+      class << self
+        attr_reader :statsd
+
+        # Setup Datadog instrumentation
+        #
+        # @param host [String] DogStatsD host
+        # @param port [Integer] DogStatsD port
+        # @param namespace [String] Metric namespace
+        # @return [void]
+        def setup!(host: "localhost", port: 8125, namespace: "vectra")
+          require "datadog/statsd"
+
+          @statsd = ::Datadog::Statsd.new(host, port, namespace: namespace)
+
+          Vectra::Instrumentation.on_operation do |event|
+            record_metrics(event)
+          end
+        rescue LoadError
+          warn "Datadog StatsD gem not found. Install with: gem 'dogstatsd-ruby'"
+        end
+
+        private
+
+        # Record metrics to Datadog
+        def record_metrics(event)
+          return unless statsd
+
+          tags = [
+            "provider:#{event.provider}",
+            "operation:#{event.operation}",
+            "index:#{event.index}",
+            "status:#{event.success? ? 'success' : 'error'}"
+          ]
+
+          # Record timing
+          statsd.timing("operation.duration", event.duration, tags: tags)
+
+          # Record count
+          statsd.increment("operation.count", tags: tags)
+
+          # Record result count if available
+          if event.metadata[:result_count]
+            statsd.gauge("operation.results", event.metadata[:result_count], tags: tags)
+          end
+
+          # Record vector count if available
+          if event.metadata[:vector_count]
+            statsd.gauge("operation.vectors", event.metadata[:vector_count], tags: tags)
+          end
+
+          # Record errors
+          return unless event.failure?
+
+          error_tags = tags + ["error_type:#{event.error.class.name}"]
+          statsd.increment("operation.error", tags: error_tags)
+        end
+      end
+    end
+  end
+end

data/lib/vectra/instrumentation/new_relic.rb

@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+module Vectra
+  module Instrumentation
+    # New Relic instrumentation adapter
+    #
+    # Automatically reports Vectra metrics to New Relic APM.
+    #
+    # @example Enable New Relic instrumentation
+    #   # config/initializers/vectra.rb
+    #   require 'vectra/instrumentation/new_relic'
+    #
+    #   Vectra.configure do |config|
+    #     config.instrumentation = true
+    #   end
+    #
+    #   Vectra::Instrumentation::NewRelic.setup!
+    #
+    module NewRelic
+      class << self
+        # Setup New Relic instrumentation
+        #
+        # @return [void]
+        def setup!
+          return unless defined?(::NewRelic::Agent)
+
+          Vectra::Instrumentation.on_operation do |event|
+            record_metrics(event)
+            record_transaction(event)
+          end
+        end
+
+        private
+
+        # Record custom metrics
+        def record_metrics(event)
+          prefix = "Custom/Vectra/#{event.provider}/#{event.operation}"
+
+          ::NewRelic::Agent.record_metric("#{prefix}/duration", event.duration)
+          ::NewRelic::Agent.record_metric("#{prefix}/calls", 1)
+
+          if event.success?
+            ::NewRelic::Agent.record_metric("#{prefix}/success", 1)
+
+            # Record result count if available
+            if event.metadata[:result_count]
+              ::NewRelic::Agent.record_metric("#{prefix}/results", event.metadata[:result_count])
+            end
+          else
+            ::NewRelic::Agent.record_metric("#{prefix}/error", 1)
+          end
+        end
+
+        # Add to transaction trace
+        def record_transaction(event)
+          ::NewRelic::Agent.add_custom_attributes(
+            vectra_operation: event.operation,
+            vectra_provider: event.provider,
+            vectra_index: event.index,
+            vectra_duration: event.duration
+          )
+
+          return unless event.failure?
+
+          ::NewRelic::Agent.notice_error(event.error)
+        end
+      end
+    end
+  end
+end