vectra-client 0.2.2 → 0.3.0

data/lib/vectra/cache.rb

@@ -0,0 +1,261 @@
+# frozen_string_literal: true
+
+require "digest"
+
+module Vectra
+  # Optional caching layer for frequently queried vectors
+  #
+  # Provides in-memory caching with TTL support for query results
+  # and fetched vectors to reduce database load.
+  #
+  # @example Basic usage
+  #   cache = Vectra::Cache.new(ttl: 300, max_size: 1000)
+  #   cached_client = Vectra::CachedClient.new(client, cache: cache)
+  #
+  #   # First call hits the database
+  #   result1 = cached_client.query(index: 'idx', vector: vec, top_k: 10)
+  #
+  #   # Second call returns cached result
+  #   result2 = cached_client.query(index: 'idx', vector: vec, top_k: 10)
+  #
+  class Cache
+    DEFAULT_TTL = 300 # 5 minutes
+    DEFAULT_MAX_SIZE = 1000
+
+    attr_reader :ttl, :max_size
+
+    # Initialize cache
+    #
+    # @param ttl [Integer] time-to-live in seconds (default: 300)
+    # @param max_size [Integer] maximum cache entries (default: 1000)
+    def initialize(ttl: DEFAULT_TTL, max_size: DEFAULT_MAX_SIZE)
+      @ttl = ttl
+      @max_size = max_size
+      @store = {}
+      @timestamps = {}
+      @mutex = Mutex.new
+    end
+
+    # Get value from cache
+    #
+    # @param key [String] cache key
+    # @return [Object, nil] cached value or nil if not found/expired
+    def get(key)
+      @mutex.synchronize do
+        return nil unless @store.key?(key)
+
+        if expired?(key)
+          delete_entry(key)
+          return nil
+        end
+
+        @store[key]
+      end
+    end
+
+    # Set value in cache
+    #
+    # @param key [String] cache key
+    # @param value [Object] value to cache
+    # @return [Object] the cached value
+    def set(key, value)
+      @mutex.synchronize do
+        evict_if_needed
+        @store[key] = value
+        @timestamps[key] = Time.now
+        value
+      end
+    end
+
+    # Get or set value with block
+    #
+    # @param key [String] cache key
+    # @yield block to compute value if not cached
+    # @return [Object] cached or computed value
+    def fetch(key)
+      cached = get(key)
+      return cached unless cached.nil?
+
+      value = yield
+      set(key, value)
+      value
+    end
+
+    # Delete entry from cache
+    #
+    # @param key [String] cache key
+    # @return [Object, nil] deleted value
+    def delete(key)
+      @mutex.synchronize { delete_entry(key) }
+    end
+
+    # Clear all cache entries
+    #
+    # @return [void]
+    def clear
+      @mutex.synchronize do
+        @store.clear
+        @timestamps.clear
+      end
+    end
+
+    # Get cache statistics
+    #
+    # @return [Hash] cache stats
+    def stats
+      @mutex.synchronize do
+        {
+          size: @store.size,
+          max_size: max_size,
+          ttl: ttl,
+          keys: @store.keys
+        }
+      end
+    end
+
+    # Check if key exists and is not expired
+    #
+    # @param key [String] cache key
+    # @return [Boolean]
+    def exist?(key)
+      @mutex.synchronize do
+        return false unless @store.key?(key)
+        return false if expired?(key)
+
+        true
+      end
+    end
+
+    private
+
+    def expired?(key)
+      return true unless @timestamps.key?(key)
+
+      Time.now - @timestamps[key] > ttl
+    end
+
+    def delete_entry(key)
+      @timestamps.delete(key)
+      @store.delete(key)
+    end
+
+    def evict_if_needed
+      # Evict when at or above max_size to make room for new entry
+      return if @store.size < max_size
+
+      # Remove oldest entries (at least 10% of max_size to avoid frequent evictions)
+      entries_to_remove = [(max_size * 0.2).ceil, 1].max
+      oldest_keys = @timestamps.sort_by { |_, v| v }.first(entries_to_remove).map(&:first)
+      oldest_keys.each { |key| delete_entry(key) }
+    end
+  end
+
+  # Client wrapper with caching support
+  #
+  # Wraps a Vectra::Client to add transparent caching for query and fetch operations.
+  #
+  class CachedClient
+    attr_reader :client, :cache
+
+    # Initialize cached client
+    #
+    # @param client [Client] the underlying Vectra client
+    # @param cache [Cache] cache instance (creates default if nil)
+    # @param cache_queries [Boolean] whether to cache query results (default: true)
+    # @param cache_fetches [Boolean] whether to cache fetch results (default: true)
+    def initialize(client, cache: nil, cache_queries: true, cache_fetches: true)
+      @client = client
+      @cache = cache || Cache.new
+      @cache_queries = cache_queries
+      @cache_fetches = cache_fetches
+    end
+
+    # Query with caching
+    #
+    # @see Client#query
+    def query(index:, vector:, top_k: 10, namespace: nil, filter: nil, **options)
+      unless @cache_queries
+        return client.query(index: index, vector: vector, top_k: top_k,
+                            namespace: namespace, filter: filter, **options)
+      end
+
+      key = query_cache_key(index, vector, top_k, namespace, filter)
+      cache.fetch(key) do
+        client.query(index: index, vector: vector, top_k: top_k,
+                     namespace: namespace, filter: filter, **options)
+      end
+    end
+
+    # Fetch with caching
+    #
+    # @see Client#fetch
+    def fetch(index:, ids:, namespace: nil)
+      return client.fetch(index: index, ids: ids, namespace: namespace) unless @cache_fetches
+
+      # Check cache for each ID
+      results = {}
+      uncached_ids = []
+
+      ids.each do |id|
+        key = fetch_cache_key(index, id, namespace)
+        cached = cache.get(key)
+        if cached
+          results[id] = cached
+        else
+          uncached_ids << id
+        end
+      end
+
+      # Fetch uncached IDs
+      if uncached_ids.any?
+        fetched = client.fetch(index: index, ids: uncached_ids, namespace: namespace)
+        fetched.each do |id, vector|
+          key = fetch_cache_key(index, id, namespace)
+          cache.set(key, vector)
+          results[id] = vector
+        end
+      end
+
+      results
+    end
+
+    # Pass through other methods to underlying client
+    def method_missing(method, *, **, &)
+      if client.respond_to?(method)
+        client.public_send(method, *, **, &)
+      else
+        super
+      end
+    end
+
+    def respond_to_missing?(method, include_private = false)
+      client.respond_to?(method, include_private) || super
+    end
+
+    # Invalidate cache entries for an index
+    #
+    # @param index [String] index name
+    def invalidate_index(index)
+      cache.stats[:keys].each do |key|
+        cache.delete(key) if key.start_with?("#{index}:")
+      end
+    end
+
+    # Clear entire cache
+    def clear_cache
+      cache.clear
+    end
+
+    private
+
+    def query_cache_key(index, vector, top_k, namespace, filter)
+      vector_hash = Digest::MD5.hexdigest(vector.to_s)[0, 16]
+      filter_hash = filter ? Digest::MD5.hexdigest(filter.to_s)[0, 8] : "nofilter"
+      "#{index}:q:#{vector_hash}:#{top_k}:#{namespace || 'default'}:#{filter_hash}"
+    end
+
+    def fetch_cache_key(index, id, namespace)
+      "#{index}:f:#{id}:#{namespace || 'default'}"
+    end
+  end
+end

data/lib/vectra/configuration.rb

@@ -15,7 +15,8 @@ module Vectra
 
     attr_accessor :api_key, :environment, :host, :timeout, :open_timeout,
                   :max_retries, :retry_delay, :logger, :pool_size, :pool_timeout,
-                  :batch_size, :instrumentation
+                  :batch_size, :instrumentation, :cache_enabled, :cache_ttl,
+                  :cache_max_size, :async_concurrency
 
     attr_reader :provider
 
@@ -33,6 +34,10 @@ module Vectra
       @pool_timeout = 5
       @batch_size = 100
       @instrumentation = false
+      @cache_enabled = false
+      @cache_ttl = 300
+      @cache_max_size = 1000
+      @async_concurrency = 4
     end
 
     # Set the provider

data/lib/vectra/pool.rb

@@ -0,0 +1,256 @@
+# frozen_string_literal: true
+
+require "concurrent"
+
+module Vectra
+  # Connection pool with warmup support
+  #
+  # Provides connection pooling for database providers with configurable
+  # pool size, timeout, and connection warmup.
+  #
+  # @example Basic usage
+  #   pool = Vectra::Pool.new(size: 5, timeout: 5) { create_connection }
+  #   pool.warmup(3) # Pre-create 3 connections
+  #
+  #   pool.with_connection do |conn|
+  #     conn.execute("SELECT 1")
+  #   end
+  #
+  class Pool
+    class TimeoutError < Vectra::Error; end
+    class PoolExhaustedError < Vectra::Error; end
+
+    attr_reader :size, :timeout
+
+    # Initialize connection pool
+    #
+    # @param size [Integer] maximum pool size
+    # @param timeout [Integer] checkout timeout in seconds
+    # @yield connection factory block
+    def initialize(size:, timeout: 5, &block)
+      raise ArgumentError, "Connection factory block required" unless block_given?
+
+      @size = size
+      @timeout = timeout
+      @factory = block
+      @pool = Concurrent::Array.new
+      @checked_out = Concurrent::AtomicFixnum.new(0)
+      @mutex = Mutex.new
+      @condition = ConditionVariable.new
+      @shutdown = false
+    end
+
+    # Warmup the pool by pre-creating connections
+    #
+    # @param count [Integer] number of connections to create (default: pool size)
+    # @return [Integer] number of connections created
+    def warmup(count = nil)
+      count ||= size
+      count = [count, size].min
+      created = 0
+
+      count.times do
+        break if @pool.size >= size
+
+        conn = create_connection
+        if conn
+          @pool << conn
+          created += 1
+        end
+      end
+
+      created
+    end
+
+    # Execute block with a connection from the pool
+    #
+    # @yield [connection] the checked out connection
+    # @return [Object] result of the block
+    def with_connection
+      conn = checkout
+      begin
+        yield conn
+      ensure
+        checkin(conn)
+      end
+    end
+
+    # Checkout a connection from the pool
+    #
+    # @return [Object] a connection
+    # @raise [TimeoutError] if checkout times out
+    # @raise [PoolExhaustedError] if pool is exhausted
+    def checkout
+      raise PoolExhaustedError, "Pool has been shutdown" if @shutdown
+
+      deadline = Time.now + timeout
+
+      @mutex.synchronize do
+        loop do
+          # Try to get an existing connection
+          conn = @pool.pop
+          if conn
+            @checked_out.increment
+            return conn if healthy?(conn)
+
+            # Connection is unhealthy, discard and try again
+            close_connection(conn)
+            next
+          end
+
+          # Try to create a new connection if under limit
+          if @checked_out.value + @pool.size < size
+            conn = create_connection
+            if conn
+              @checked_out.increment
+              return conn
+            end
+          end
+
+          # Wait for a connection to be returned
+          remaining = deadline - Time.now
+          raise TimeoutError, "Connection checkout timed out after #{timeout}s" if remaining <= 0
+
+          @condition.wait(@mutex, remaining)
+        end
+      end
+    end
+
+    # Return a connection to the pool
+    #
+    # @param connection [Object] connection to return
+    def checkin(connection)
+      return if @shutdown
+
+      @mutex.synchronize do
+        @checked_out.decrement
+        if healthy?(connection) && @pool.size < size
+          @pool << connection
+        else
+          close_connection(connection)
+        end
+        @condition.signal
+      end
+    end
+
+    # Shutdown the pool, closing all connections
+    #
+    # @return [void]
+    def shutdown
+      @shutdown = true
+      @mutex.synchronize do
+        while (conn = @pool.pop)
+          close_connection(conn)
+        end
+      end
+    end
+
+    # Get pool statistics
+    #
+    # @return [Hash] pool stats
+    def stats
+      {
+        size: size,
+        available: @pool.size,
+        checked_out: @checked_out.value,
+        total_created: @pool.size + @checked_out.value,
+        shutdown: @shutdown
+      }
+    end
+
+    # Check if pool is healthy (public method)
+    #
+    # @return [Boolean]
+    def pool_healthy?
+      !@shutdown && (@pool.size + @checked_out.value).positive?
+    end
+
+    private
+
+    # Internal health check for individual connections
+    def healthy?(conn)
+      return false if conn.nil?
+      return true unless conn.respond_to?(:status)
+
+      # For PG connections, check status. Otherwise assume healthy.
+      if defined?(PG::CONNECTION_OK)
+        conn.status == PG::CONNECTION_OK
+      else
+        # If PG not loaded, assume connection is healthy if it exists
+        true
+      end
+    rescue StandardError
+      false
+    end
+
+    def create_connection
+      @factory.call
+    rescue StandardError => e
+      Vectra.configuration.logger&.error("Pool: Failed to create connection: #{e.message}")
+      nil
+    end
+
+    def close_connection(conn)
+      conn.close if conn.respond_to?(:close)
+    rescue StandardError => e
+      Vectra.configuration.logger&.warn("Pool: Error closing connection: #{e.message}")
+    end
+  end
+
+  # Pooled connection module for pgvector
+  module PooledConnection
+    # Get a pooled connection
+    #
+    # @return [Pool] connection pool
+    def connection_pool
+      @connection_pool ||= create_pool
+    end
+
+    # Warmup the connection pool
+    #
+    # @param count [Integer] number of connections to pre-create
+    # @return [Integer] connections created
+    def warmup_pool(count = nil)
+      connection_pool.warmup(count)
+    end
+
+    # Execute with pooled connection
+    #
+    # @yield [connection] database connection
+    def with_pooled_connection(&)
+      connection_pool.with_connection(&)
+    end
+
+    # Shutdown the connection pool
+    def shutdown_pool
+      @connection_pool&.shutdown
+      @connection_pool = nil
+    end
+
+    # Get pool statistics
+    #
+    # @return [Hash]
+    def pool_stats
+      return { status: "not_initialized" } unless @connection_pool
+
+      connection_pool.stats
+    end
+
+    private
+
+    def create_pool
+      pool_size = config.pool_size || 5
+      pool_timeout = config.pool_timeout || 5
+
+      Pool.new(size: pool_size, timeout: pool_timeout) do
+        create_raw_connection
+      end
+    end
+
+    def create_raw_connection
+      require "pg"
+      conn_params = parse_connection_params
+      PG.connect(conn_params)
+    end
+  end
+end

data/lib/vectra/streaming.rb

@@ -0,0 +1,153 @@
+# frozen_string_literal: true
+
+module Vectra
+  # Streaming query results for large datasets
+  #
+  # Provides lazy enumeration over query results with automatic pagination,
+  # reducing memory usage for large result sets.
+  #
+  # @example Stream through results
+  #   stream = Vectra::Streaming.new(client)
+  #   stream.query_each(index: 'my-index', vector: query_vec, total: 1000) do |match|
+  #     process(match)
+  #   end
+  #
+  # @example As lazy enumerator
+  #   results = stream.query_stream(index: 'my-index', vector: query_vec, total: 1000)
+  #   results.take(50).each { |m| puts m.id }
+  #
+  class Streaming
+    DEFAULT_PAGE_SIZE = 100
+
+    attr_reader :client, :page_size
+
+    # Initialize streaming query handler
+    #
+    # @param client [Client] the Vectra client
+    # @param page_size [Integer] results per page (default: 100)
+    def initialize(client, page_size: DEFAULT_PAGE_SIZE)
+      @client = client
+      @page_size = [page_size, 1].max
+    end
+
+    # Stream query results with a block
+    #
+    # @param index [String] the index name
+    # @param vector [Array<Float>] query vector
+    # @param total [Integer] total results to fetch
+    # @param namespace [String, nil] optional namespace
+    # @param filter [Hash, nil] metadata filter
+    # @yield [Match] each match result
+    # @return [Integer] total results yielded
+    def query_each(index:, vector:, total:, namespace: nil, filter: nil, &block)
+      return 0 unless block_given?
+
+      count = 0
+      query_stream(index: index, vector: vector, total: total, namespace: namespace, filter: filter).each do |match|
+        block.call(match)
+        count += 1
+      end
+      count
+    end
+
+    # Create a lazy enumerator for streaming results
+    #
+    # @param index [String] the index name
+    # @param vector [Array<Float>] query vector
+    # @param total [Integer] total results to fetch
+    # @param namespace [String, nil] optional namespace
+    # @param filter [Hash, nil] metadata filter
+    # @return [Enumerator::Lazy] lazy enumerator of results
+    def query_stream(index:, vector:, total:, namespace: nil, filter: nil)
+      Enumerator.new do |yielder|
+        fetched = 0
+        seen_ids = Set.new
+
+        while fetched < total
+          batch_size = [page_size, total - fetched].min
+
+          result = client.query(
+            index: index,
+            vector: vector,
+            top_k: batch_size,
+            namespace: namespace,
+            filter: filter,
+            include_metadata: true
+          )
+
+          break if result.empty?
+
+          result.each do |match|
+            # Skip duplicates (some providers may return overlapping results)
+            next if seen_ids.include?(match.id)
+
+            seen_ids.add(match.id)
+            yielder << match
+            fetched += 1
+
+            break if fetched >= total
+          end
+
+          # If we got fewer results than requested, we've exhausted the index
+          break if result.size < batch_size
+        end
+      end.lazy
+    end
+
+    # Scan all vectors in an index (provider-dependent)
+    #
+    # @param index [String] the index name
+    # @param namespace [String, nil] optional namespace
+    # @param batch_size [Integer] IDs per batch
+    # @yield [Vector] each vector
+    # @return [Integer] total vectors scanned
+    # @note Not all providers support efficient scanning
+    def scan_all(index:, namespace: nil, batch_size: 1000)
+      return 0 unless block_given?
+
+      count = 0
+      offset = 0
+
+      loop do
+        # This is a simplified scan - actual implementation depends on provider
+        stats = client.stats(index: index, namespace: namespace)
+        total = stats[:total_vector_count] || 0
+
+        break if offset >= total
+
+        # Fetch IDs in batches (this is provider-specific)
+        # For now, we return what we can
+        break if offset.positive? # Only one iteration for basic implementation
+
+        offset += batch_size
+        count = total
+      end
+
+      count
+    end
+  end
+
+  # Streaming result wrapper with additional metadata
+  class StreamingResult
+    include Enumerable
+
+    attr_reader :enumerator, :metadata
+
+    def initialize(enumerator, metadata = {})
+      @enumerator = enumerator
+      @metadata = metadata
+    end
+
+    def each(&)
+      enumerator.each(&)
+    end
+
+    def take(n)
+      enumerator.take(n)
+    end
+
+    def to_a
+      enumerator.to_a
+    end
+  end
+end

data/lib/vectra/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Vectra
-  VERSION = "0.2.2"
+  VERSION = "0.3.0"
 end

data/lib/vectra.rb

@@ -7,6 +7,10 @@ require_relative "vectra/vector"
 require_relative "vectra/query_result"
 require_relative "vectra/instrumentation"
 require_relative "vectra/retry"
+require_relative "vectra/batch"
+require_relative "vectra/streaming"
+require_relative "vectra/cache"
+require_relative "vectra/pool"
 require_relative "vectra/active_record"
 require_relative "vectra/providers/base"
 require_relative "vectra/providers/pinecone"