vectra-client 1.0.8 → 1.1.0

data/lib/vectra/client.rb

@@ -42,6 +42,37 @@ module Vectra
 
     attr_reader :config, :provider, :default_index, :default_namespace
 
+    class << self
+      # Get the global middleware stack
+      #
+      # @return [Array<Array>] Array of [middleware_class, options] pairs
+      def middleware
+        @middleware ||= []
+      end
+
+      # Add middleware to the global stack
+      #
+      # @param middleware_class [Class] Middleware class
+      # @param options [Hash] Options to pass to middleware constructor
+      #
+      # @example Add global logging middleware
+      #   Vectra::Client.use Vectra::Middleware::Logging
+      #
+      # @example Add middleware with options
+      #   Vectra::Client.use Vectra::Middleware::Retry, max_attempts: 5
+      #
+      def use(middleware_class, **options)
+        middleware << [middleware_class, options]
+      end
+
+      # Clear all global middleware
+      #
+      # @return [void]
+      def clear_middleware!
+        @middleware = []
+      end
+    end
+
     # Initialize a new Client
     #
     # @param provider [Symbol, nil] provider name (:pinecone, :qdrant, :weaviate)
@@ -51,12 +82,14 @@ module Vectra
     # @param options [Hash] additional options
     # @option options [String] :index default index name
     # @option options [String] :namespace default namespace
+    # @option options [Array<Class, Object>] :middleware instance-level middleware
     def initialize(provider: nil, api_key: nil, environment: nil, host: nil, **options)
       @config = build_config(provider, api_key, environment, host, options)
       @config.validate!
       @provider = build_provider
       @default_index = options[:index]
       @default_namespace = options[:namespace]
+      @middleware = build_middleware_stack(options[:middleware])
     end
 
     # Upsert vectors into an index
@@ -87,7 +120,7 @@ module Vectra
         index: index,
         metadata: { vector_count: vectors.size }
       ) do
-        provider.upsert(index: index, vectors: vectors, namespace: namespace)
+        @middleware.call(:upsert, index: index, vectors: vectors, namespace: namespace, provider: provider_name)
       end
     end
 
@@ -151,14 +184,16 @@ module Vectra
         index: index,
         metadata: { top_k: top_k }
       ) do
-        result = provider.query(
+        result = @middleware.call(
+          :query,
           index: index,
           vector: vector,
           top_k: top_k,
           namespace: namespace,
           filter: filter,
           include_values: include_values,
-          include_metadata: include_metadata
+          include_metadata: include_metadata,
+          provider: provider_name
         )
       end
 
@@ -188,7 +223,7 @@ module Vectra
         index: index,
         metadata: { id_count: ids.size }
       ) do
-        provider.fetch(index: index, ids: ids, namespace: namespace)
+        @middleware.call(:fetch, index: index, ids: ids, namespace: namespace, provider: provider_name)
       end
     end
 
@@ -222,12 +257,14 @@ module Vectra
         index: index,
         metadata: { has_metadata: !metadata.nil?, has_values: !values.nil? }
       ) do
-        provider.update(
+        @middleware.call(
+          :update,
           index: index,
           id: id,
           metadata: metadata,
           values: values,
-          namespace: namespace
+          namespace: namespace,
+          provider: provider_name
         )
       end
     end
@@ -265,12 +302,14 @@ module Vectra
         index: index,
         metadata: { id_count: ids&.size, delete_all: delete_all, has_filter: !filter.nil? }
       ) do
-        provider.delete(
+        @middleware.call(
+          :delete,
           index: index,
           ids: ids,
           namespace: namespace,
           filter: filter,
-          delete_all: delete_all
+          delete_all: delete_all,
+          provider: provider_name
         )
       end
     end
@@ -284,7 +323,7 @@ module Vectra
     #   indexes.each { |idx| puts idx[:name] }
     #
     def list_indexes
-      provider.list_indexes
+      @middleware.call(:list_indexes, provider: provider_name)
     end
 
     # Describe an index
@@ -299,7 +338,7 @@ module Vectra
     def describe_index(index: nil)
       index ||= default_index
       validate_index!(index)
-      provider.describe_index(index: index)
+      @middleware.call(:describe_index, index: index, provider: provider_name)
     end
 
     # Get index statistics
@@ -316,7 +355,7 @@ module Vectra
       index ||= default_index
       namespace ||= default_namespace
       validate_index!(index)
-      provider.stats(index: index, namespace: namespace)
+      @middleware.call(:stats, index: index, namespace: namespace, provider: provider_name)
     end
 
     # Create a new index
@@ -342,7 +381,7 @@ module Vectra
         index: name,
         metadata: { dimension: dimension, metric: metric }
       ) do
-        provider.create_index(name: name, dimension: dimension, metric: metric, **options)
+        @middleware.call(:create_index, name: name, dimension: dimension, metric: metric, provider: provider_name, **options)
       end
     end
 
@@ -365,7 +404,7 @@ module Vectra
         provider: provider_name,
         index: name
       ) do
-        provider.delete_index(name: name)
+        @middleware.call(:delete_index, name: name, provider: provider_name)
       end
     end
 
@@ -440,7 +479,8 @@ module Vectra
               "Hybrid search is not supported by #{provider_name} provider"
       end
 
-      provider.hybrid_search(
+      @middleware.call(
+        :hybrid_search,
         index: index,
         vector: vector,
         text: text,
@@ -449,7 +489,8 @@ module Vectra
         namespace: namespace,
         filter: filter,
         include_values: include_values,
-        include_metadata: include_metadata
+        include_metadata: include_metadata,
+        provider: provider_name
       )
     end
 
@@ -628,6 +669,21 @@ module Vectra
       end
     end
 
+    def build_middleware_stack(instance_middleware = nil)
+      # Combine class-level + instance-level middleware
+      all_middleware = self.class.middleware.map do |klass, opts|
+        klass.new(**opts)
+      end
+
+      if instance_middleware
+        all_middleware += Array(instance_middleware).map do |mw|
+          mw.is_a?(Class) ? mw.new : mw
+        end
+      end
+
+      Middleware::Stack.new(@provider, all_middleware)
+    end
+
     def validate_index!(index)
       raise ValidationError, "Index name cannot be nil" if index.nil?
       raise ValidationError, "Index name must be a string" unless index.is_a?(String)

data/lib/vectra/health_check.rb

@@ -29,7 +29,9 @@ module Vectra
     def health_check(index: nil, include_stats: false, timeout: 5)
       start_time = Time.now
 
-      indexes = with_timeout(timeout) { list_indexes }
+      # For health checks we bypass client middleware and call the provider
+      # directly to avoid interference from custom stacks.
+      indexes = with_timeout(timeout) { provider.list_indexes }
       index_name = index || indexes.first&.dig(:name)
 
       result = base_result(start_time, indexes)
@@ -70,7 +72,7 @@ module Vectra
     def add_index_stats(result, index_name, include_stats, timeout)
       return unless include_stats && index_name
 
-      stats = with_timeout(timeout) { stats(index: index_name) }
+      stats = with_timeout(timeout) { provider.stats(index: index_name) }
       result[:index] = index_name
       result[:stats] = {
         vector_count: stats[:total_vector_count],

data/lib/vectra/middleware/base.rb

@@ -0,0 +1,97 @@
+# frozen_string_literal: true
+
+module Vectra
+  module Middleware
+    # Base class for all middleware
+    #
+    # Middleware can hook into three lifecycle events:
+    # - before(request): Called before the next middleware/provider
+    # - after(request, response): Called after successful execution
+    # - on_error(request, error): Called when an error occurs
+    #
+    # @example Simple logging middleware
+    #   class LoggingMiddleware < Vectra::Middleware::Base
+    #     def before(request)
+    #       puts "Starting #{request.operation}"
+    #     end
+    #
+    #     def after(request, response)
+    #       puts "Completed #{request.operation}"
+    #     end
+    #   end
+    #
+    # @example Error handling middleware
+    #   class ErrorHandlerMiddleware < Vectra::Middleware::Base
+    #     def on_error(request, error)
+    #       ErrorTracker.notify(error, context: { operation: request.operation })
+    #     end
+    #   end
+    #
+    class Base
+      # Execute the middleware
+      #
+      # This is the main entry point called by the middleware stack.
+      # It handles the before/after/error lifecycle hooks.
+      #
+      # @param request [Request] The request object
+      # @param app [Proc] The next middleware in the chain
+      # @return [Response] The response object
+      def call(request, app)
+        # Before hook
+        before(request)
+
+        # Call next middleware
+        response = app.call(request)
+
+        # Check if response has an error
+        if response.error
+          on_error(request, response.error)
+        end
+
+        # After hook
+        after(request, response)
+
+        response
+      rescue StandardError => e
+        # Error handling hook (for exceptions raised directly)
+        on_error(request, e)
+        raise
+      end
+
+      protected
+
+      # Hook called before the next middleware
+      #
+      # Override this method to add logic before the operation executes.
+      #
+      # @param request [Request] The request object
+      # @return [void]
+      def before(request)
+        # Override in subclass
+      end
+
+      # Hook called after successful execution
+      #
+      # Override this method to add logic after the operation completes.
+      #
+      # @param request [Request] The request object
+      # @param response [Response] The response object
+      # @return [void]
+      def after(request, response)
+        # Override in subclass
+      end
+
+      # Hook called when an error occurs
+      #
+      # Override this method to add error handling logic.
+      # The error will be re-raised after this hook executes.
+      #
+      # @param request [Request] The request object
+      # @param error [Exception] The error that occurred
+      # @return [void]
+      def on_error(request, error)
+        # Override in subclass
+      end
+    end
+  end
+end

data/lib/vectra/middleware/cost_tracker.rb

@@ -0,0 +1,121 @@
+# frozen_string_literal: true
+
+module Vectra
+  module Middleware
+    # Cost tracking middleware for monitoring API usage costs
+    #
+    # Tracks estimated costs per operation based on provider pricing.
+    # Costs are stored in response metadata and can be aggregated.
+    #
+    # @example With default pricing
+    #   Vectra::Client.use Vectra::Middleware::CostTracker
+    #
+    # @example With custom pricing
+    #   custom_pricing = {
+    #     pinecone: { read: 0.0001, write: 0.0002 },
+    #     qdrant: { read: 0.00005, write: 0.0001 }
+    #   }
+    #   Vectra::Client.use Vectra::Middleware::CostTracker, pricing: custom_pricing
+    #
+    # @example With cost callback
+    #   Vectra::Client.use Vectra::Middleware::CostTracker, on_cost: ->(event) {
+    #     puts "Cost: $#{event[:cost_usd]} for #{event[:operation]}"
+    #   }
+    #
+    class CostTracker < Base
+      # Default pricing per operation (in USD)
+      # These are estimated values - check provider pricing for actual costs
+      DEFAULT_PRICING = {
+        pinecone: { read: 0.0001, write: 0.0002 },
+        qdrant: { read: 0.00005, write: 0.0001 },
+        weaviate: { read: 0.00008, write: 0.00015 },
+        pgvector: { read: 0.0, write: 0.0 }, # Self-hosted, no API costs
+        memory: { read: 0.0, write: 0.0 }    # In-memory, no costs
+      }.freeze
+
+      # @param pricing [Hash] Custom pricing structure
+      # @param on_cost [Proc] Callback to invoke with cost events
+      def initialize(pricing: DEFAULT_PRICING, on_cost: nil)
+        super()
+        @pricing = pricing
+        @on_cost = on_cost
+      end
+
+      def after(request, response)
+        return unless response.success?
+
+        provider = request.provider || :unknown
+        operation_type = write_operation?(request.operation) ? :write : :read
+
+        cost = calculate_cost(provider, operation_type, request)
+        response.metadata[:cost_usd] = cost
+
+        # Invoke callback if provided
+        return unless @on_cost
+
+        @on_cost.call(
+          operation: request.operation,
+          provider: provider,
+          index: request.index,
+          cost_usd: cost,
+          timestamp: Time.now
+        )
+      end
+
+      private
+
+      # Check if operation is a write operation
+      #
+      # @param operation [Symbol] The operation type
+      # @return [Boolean] true if write operation
+      def write_operation?(operation)
+        [:upsert, :delete, :update, :create_index, :delete_index].include?(operation)
+      end
+
+      # Calculate cost for operation
+      #
+      # @param provider [Symbol] Provider name
+      # @param operation_type [Symbol] :read or :write
+      # @param request [Request] The request object
+      # @return [Float] Cost in USD
+      def calculate_cost(provider, operation_type, request)
+        rate = @pricing.dig(provider, operation_type) || 0.0
+        multiplier = operation_multiplier(request)
+        rate * multiplier
+      end
+
+      # Calculate multiplier for operation based on batch size
+      #
+      # @param request [Request] The request object
+      # @return [Integer, Float] Multiplier for the base rate
+      def operation_multiplier(request)
+        return 100 if delete_all?(request)
+
+        case request.operation
+        when :upsert
+          collection_size(request.params[:vectors])
+        when :fetch, :delete
+          collection_size(request.params[:ids])
+        else
+          1 # Includes :query and all other operations
+        end
+      end
+
+      # Check if request is a delete_all operation
+      #
+      # @param request [Request] The request object
+      # @return [Boolean]
+      def delete_all?(request)
+        request.operation == :delete && request.params[:delete_all]
+      end
+
+      # Safely compute collection size with a default of 1
+      #
+      # @param collection [Enumerable, nil]
+      # @return [Integer]
+      def collection_size(collection)
+        collection&.size || 1
+      end
+    end
+  end
+end

data/lib/vectra/middleware/instrumentation.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+module Vectra
+  module Middleware
+    # Instrumentation middleware for metrics and monitoring
+    #
+    # Emits instrumentation events for all operations, compatible with
+    # Vectra's existing instrumentation system.
+    #
+    # @example Enable instrumentation middleware
+    #   Vectra::Client.use Vectra::Middleware::Instrumentation
+    #
+    # @example With custom event handler
+    #   Vectra.on_operation do |event|
+    #     puts "Operation: #{event[:operation]}, Duration: #{event[:duration_ms]}ms"
+    #   end
+    #
+    #   Vectra::Client.use Vectra::Middleware::Instrumentation
+    #
+    class Instrumentation < Base
+      def call(request, app)
+        start_time = Time.now
+
+        response = app.call(request)
+
+        duration_ms = ((Time.now - start_time) * 1000).round(2)
+
+        # Emit instrumentation event
+        Vectra::Instrumentation.instrument(
+          operation: request.operation,
+          provider: request.provider,
+          index: request.index,
+          namespace: request.namespace,
+          duration_ms: duration_ms,
+          success: response.success?,
+          error: response.error&.class&.name,
+          metadata: response.metadata
+        )
+
+        response
+      end
+    end
+  end
+end

data/lib/vectra/middleware/logging.rb

@@ -0,0 +1,62 @@
+# frozen_string_literal: true
+
+module Vectra
+  module Middleware
+    # Logging middleware for tracking operations
+    #
+    # Logs before and after each operation, including timing information.
+    #
+    # @example With default logger
+    #   Vectra::Client.use Vectra::Middleware::Logging
+    #
+    # @example With custom logger
+    #   logger = Logger.new($stdout)
+    #   Vectra::Client.use Vectra::Middleware::Logging, logger: logger
+    #
+    # @example Per-client logging
+    #   client = Vectra::Client.new(
+    #     provider: :qdrant,
+    #     middleware: [Vectra::Middleware::Logging]
+    #   )
+    #
+    class Logging < Base
+      def initialize(logger: nil)
+        super()
+        @logger = logger || Vectra.configuration.logger
+      end
+
+      def before(request)
+        return unless @logger
+
+        @start_time = Time.now
+        @logger.info(
+          "[Vectra] #{request.operation.upcase} " \
+          "index=#{request.index} " \
+          "namespace=#{request.namespace || 'default'}"
+        )
+      end
+
+      def after(request, response)
+        return unless @logger
+        return unless @start_time
+
+        duration_ms = ((Time.now - @start_time) * 1000).round(2)
+        response.metadata[:duration_ms] = duration_ms
+
+        if response.success?
+          @logger.info("[Vectra] ✅ #{request.operation} completed in #{duration_ms}ms")
+        else
+          @logger.error("[Vectra] ❌ #{request.operation} failed: #{response.error.message}")
+        end
+      end
+
+      def on_error(request, error)
+        return unless @logger
+
+        @logger.error(
+          "[Vectra] 💥 #{request.operation} exception: #{error.class} - #{error.message}"
+        )
+      end
+    end
+  end
+end

data/lib/vectra/middleware/pii_redaction.rb

@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+module Vectra
+  module Middleware
+    # PII Redaction middleware for protecting sensitive data
+    #
+    # Automatically redacts Personally Identifiable Information (PII) from
+    # metadata before upserting to vector databases.
+    #
+    # @example With default patterns (email, phone, SSN)
+    #   Vectra::Client.use Vectra::Middleware::PIIRedaction
+    #
+    # @example With custom patterns
+    #   custom_patterns = {
+    #     credit_card: /\b\d{4}[-\s]?\d{4}[-\s]?\d{4}[-\s]?\d{4}\b/,
+    #     api_key: /sk-[a-zA-Z0-9]{32}/
+    #   }
+    #   Vectra::Client.use Vectra::Middleware::PIIRedaction, patterns: custom_patterns
+    #
+    class PIIRedaction < Base
+      # Default PII patterns
+      DEFAULT_PATTERNS = {
+        email: /\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b/,
+        phone: /\b\d{3}[-.]?\d{3}[-.]?\d{4}\b/,
+        ssn: /\b\d{3}-\d{2}-\d{4}\b/,
+        credit_card: /\b\d{4}[-\s]?\d{4}[-\s]?\d{4}[-\s]?\d{4}\b/
+      }.freeze
+
+      def initialize(patterns: DEFAULT_PATTERNS)
+        super()
+        @patterns = patterns
+      end
+
+      def before(request)
+        return unless request.operation == :upsert
+        return unless request.params[:vectors]
+
+        # Redact PII from metadata in all vectors
+        request.params[:vectors].each do |vector|
+          next unless vector[:metadata]
+
+          vector[:metadata] = redact_metadata(vector[:metadata])
+        end
+      end
+
+      private
+
+      # Redact PII from metadata hash
+      #
+      # @param metadata [Hash] Metadata to redact
+      # @return [Hash] Redacted metadata
+      def redact_metadata(metadata)
+        metadata.transform_values do |value|
+          next value unless value.is_a?(String)
+
+          redacted = value.dup
+          @patterns.each do |type, pattern|
+            redacted.gsub!(pattern, "[REDACTED_#{type.upcase}]")
+          end
+          redacted
+        end
+      end
+    end
+  end
+end

data/lib/vectra/middleware/request.rb

@@ -0,0 +1,62 @@
+# frozen_string_literal: true
+
+module Vectra
+  module Middleware
+    # Request object passed through middleware chain
+    #
+    # @example Basic usage
+    #   request = Request.new(
+    #     operation: :upsert,
+    #     index: 'products',
+    #     namespace: 'prod',
+    #     vectors: [{ id: 'doc-1', values: [0.1, 0.2, 0.3] }]
+    #   )
+    #
+    #   request.operation # => :upsert
+    #   request.index     # => 'products'
+    #   request.namespace # => 'prod'
+    #   request.metadata[:custom_key] = 'custom_value'
+    #
+    class Request
+      attr_accessor :operation, :index, :namespace, :params, :metadata
+
+      # @param operation [Symbol] The operation type (:upsert, :query, :delete, etc.)
+      # @param params [Hash] All parameters for the operation
+      def initialize(operation:, **params)
+        @operation = operation
+        @index = params[:index]
+        @namespace = params[:namespace]
+        @params = params
+        @metadata = {}
+      end
+
+      # Convert request back to hash for provider call
+      #
+      # @return [Hash] Parameters hash
+      def to_h
+        params
+      end
+
+      # Get the provider from params
+      #
+      # @return [Symbol, nil] Provider name
+      def provider
+        params[:provider]
+      end
+
+      # Check if this is a write operation
+      #
+      # @return [Boolean]
+      def write_operation?
+        [:upsert, :delete, :update, :create_index, :delete_index].include?(operation)
+      end
+
+      # Check if this is a read operation
+      #
+      # @return [Boolean]
+      def read_operation?
+        [:query, :fetch, :list_indexes, :describe_index, :stats].include?(operation)
+      end
+    end
+  end
+end