noiseless 0.7.1 → 0.8.0

data/lib/noiseless/adapters/execution_modules/opensearch_execution.rb

@@ -1,7 +1,6 @@
 # frozen_string_literal: true
 
 require "json"
-require_relative "es_compatible_execution"
 
 module Noiseless
   module Adapters
@@ -11,33 +10,82 @@ module Noiseless
 
         private
 
+        def ast_to_hash(ast_node)
+          result = super
+
+          if (knn = result.delete(:knn))
+            must = result.dig(:query, :bool, :must)
+            result[:query] = must ? { bool: { must: [*must, knn], filter: result.dig(:query, :bool, :filter) }.compact } : knn
+            result[:_source] = { excludes: [ast_node.vector.field.to_s] }
+          end
+
+          if ast_node.hybrid_search?
+            hybrid = ast_node.hybrid
+            result[:query][:hybrid][:pagination_depth] = [result[:from].to_i + result[:size].to_i, hybrid.vector.k].max
+            result[:search_pipeline] ||= rrf_pipeline(hybrid)
+            result[:_source] = { excludes: [hybrid.vector.field.to_s] }
+          end
+
+          result
+        end
+
+        def build_knn_query(vector_node, filters: nil)
+          options = { vector: vector_node.embedding, k: vector_node.k }
+          options[:filter] = { bool: { filter: filters } } if filters.present?
+          { knn: { vector_node.field.to_s => options } }
+        end
+
+        def build_hybrid_query(hybrid_node, query)
+          raise ArgumentError, "OpenSearch hybrid search takes filters, not must clauses" if query&.dig(:bool, :must).present?
+
+          text = { multi_match: { query: hybrid_node.text_query, fields: hybrid_node.fields.presence }.compact }
+          hybrid = { queries: [text, build_knn_query(hybrid_node.vector)] }
+          filters = query&.dig(:bool, :filter)
+          hybrid[:filter] = { bool: { filter: filters } } if filters.present?
+          { hybrid: hybrid }
+        end
+
+        def rrf_pipeline(hybrid_node)
+          weights = [hybrid_node.text_weight, hybrid_node.vector_weight]
+          ranker = { combination: { technique: "rrf", rank_constant: 60, parameters: { weights: weights } } }
+          { phase_results_processors: [{ "score-ranker-processor": ranker }] }
+        end
+
+        def encode_vector(values) = [values.pack("e*")].pack("m0")
+
         def execute_search(query_hash, indexes: [], **_opts)
           index_path = indexes.any? ? indexes.join(",") : "_all"
-          path = "/#{index_path}/_search"
-          body = JSON.generate(query_hash)
+          pipeline = query_hash[:search_pipeline]
+          body = pipeline.is_a?(Hash) ? query_hash : query_hash.except(:search_pipeline)
+          path = "/#{index_path}/_search#{query_string(search_pipeline: (pipeline unless pipeline.is_a?(Hash)))}"
 
-          response = post_request(path, body)
+          response = post_request(path, dump_json(body))
           parse_json_response!(response, error_class: Noiseless::SearchError, context: "search #{index_path}")
         ensure
           response&.close
         end
 
         def execute_create_index(index_name, mappings: nil, settings: nil, **opts)
+          settings = (settings || {}).deep_merge(index: { knn: true }) if knn_mapping?(mappings)
           body = opts.dup
           body[:mappings] = mappings if mappings
           body[:settings] = settings if settings
 
-          response = put_request("/#{index_name}", body.any? ? JSON.generate(body) : nil)
+          response = put_request("/#{index_name}", body.any? ? dump_json(body) : nil)
           parse_json_response!(response, context: "create index #{index_name}")
         ensure
           response&.close
         end
 
-        def execute_index_document(index, id, document, **_opts)
-          path = "/#{index}/_doc/#{id}"
-          body = JSON.generate(document)
+        def knn_mapping?(mappings)
+          properties = (mappings || {}).with_indifferent_access[:properties] || {}
+          properties.values.any? { it[:type].to_s == "knn_vector" }
+        end
+
+        def execute_index_document(index, id, document, refresh: nil, **_opts)
+          path = "/#{index}/_doc/#{id}#{query_string(refresh:)}"
 
-          response = put_request(path, body)
+          response = put_request(path, dump_json(encode_document(document)))
           parse_json_response!(response, context: "index document #{index}/#{id}")
         ensure
           response&.close
@@ -64,25 +112,11 @@ module Noiseless
           response&.close
         end
 
-        # OpenSearch-specific features
-        def execute_point_in_time_search(query_hash, pit_id:, **_opts)
-          # Point-in-time search for consistent pagination
-          enhanced_query = query_hash.merge(pit: { id: pit_id })
-          body = JSON.generate(enhanced_query)
-
-          response = post_request("/_search", body)
-          parse_json_response!(response, error_class: Noiseless::SearchError, context: "point-in-time search")
-        ensure
-          response&.close
-        end
+        def pit_path = "_search/point_in_time"
+        def pit_close_body(pit_id) = { pit_id: [pit_id] }
 
         def execute_search_template(template_id:, params: {}, **_opts)
-          # OpenSearch search templates
-          template_query = {
-            id: template_id,
-            params: params
-          }
-          body = JSON.generate(template_query)
+          body = dump_json(id: template_id, params: params)
 
           response = post_request("/_search/template", body)
           parse_json_response!(response, error_class: Noiseless::SearchError, context: "search template #{template_id}")
@@ -90,118 +124,72 @@ module Noiseless
           response&.close
         end
 
-        # ============================================
-        # Search Pipeline API (OpenSearch 3.x)
-        # ============================================
-
-        def execute_create_pipeline(name, request_processors:, response_processors:, description: nil)
-          body = {
-            description: description,
-            request_processors: request_processors,
-            response_processors: response_processors
-          }.compact
-
-          response = put_request("/_search/pipeline/#{name}", JSON.generate(body))
-          JSON.parse(response.read)
-        rescue StandardError => e
-          { acknowledged: false, error: { type: e.class.name, reason: e.message } }
+        def execute_json(verb, path, body = nil, context:)
+          response = case verb
+                     when :get then get_request(path)
+                     when :delete then delete_request(path)
+                     else send(:"#{verb}_request", path, body && dump_json(body))
+                     end
+          parse_json_response!(response, context:)
         ensure
           response&.close
         end
 
-        def execute_get_pipeline(name)
-          response = get_request("/_search/pipeline/#{name}")
-          JSON.parse(response.read)
-        rescue StandardError => e
-          { error: { type: e.class.name, reason: e.message } }
+        def execute_exists?(path, context:)
+          response = get_request(path)
+          exists_response?(response, context:)
         ensure
           response&.close
         end
 
-        def execute_list_pipelines
-          response = get_request("/_search/pipeline")
-          JSON.parse(response.read)
-        rescue StandardError => e
-          { error: { type: e.class.name, reason: e.message } }
-        ensure
-          response&.close
+        def execute_create_pipeline(name, request_processors: nil, response_processors: nil, phase_results_processors: nil,
+                                    description: nil)
+          body = { description:, request_processors:, response_processors:, phase_results_processors: }.compact
+          execute_json(:put, "/_search/pipeline/#{name}", body, context: "create pipeline #{name}")
         end
 
-        def execute_delete_pipeline(name)
-          response = delete_request("/_search/pipeline/#{name}")
-          JSON.parse(response.read)
-        rescue StandardError => e
-          { acknowledged: false, error: { type: e.class.name, reason: e.message } }
-        ensure
-          response&.close
-        end
+        def execute_get_pipeline(name) = execute_json(:get, "/_search/pipeline/#{name}", context: "get pipeline #{name}")
+        def execute_list_pipelines = execute_json(:get, "/_search/pipeline", context: "list pipelines")
+        def execute_delete_pipeline(name) = execute_json(:delete, "/_search/pipeline/#{name}", context: "delete pipeline #{name}")
+        def execute_pipeline_exists?(name) = execute_exists?("/_search/pipeline/#{name}", context: "pipeline exists #{name}")
 
-        def execute_pipeline_exists?(name)
-          response = head_request("/_search/pipeline/#{name}")
-          response.success?
-        rescue StandardError
-          false
-        ensure
-          response&.close
+        def execute_create_rule(feature_type, description:, value:, **attributes)
+          body = { description:, **attributes, feature_type => value }
+          execute_json(:put, "/_rules/#{feature_type}", body, context: "create #{feature_type} rule")
         end
 
-        # ============================================
-        # Query Rules API (OpenSearch 3.x)
-        # ============================================
+        def execute_update_rule(feature_type, id, **changes)
+          execute_json(:put, "/_rules/#{feature_type}/#{id}", changes, context: "update #{feature_type} rule #{id}")
+        end
 
-        def execute_create_rule(feature_type, rule_id, attributes:, feature_value:)
-          body = {
-            match_criteria: {
-              query: attributes
-            },
-            feature_value: feature_value
-          }
+        def execute_get_rule(feature_type, id)
+          execute_json(:get, "/_rules/#{feature_type}/#{id}", context: "get #{feature_type} rule #{id}").fetch("rules").first
+        end
 
-          response = put_request("/_rules/#{feature_type}/#{rule_id}", JSON.generate(body))
-          JSON.parse(response.read)
-        rescue StandardError => e
-          { acknowledged: false, error: { type: e.class.name, reason: e.message } }
-        ensure
-          response&.close
+        def execute_list_rules(feature_type, search_after: nil)
+          path = "/_rules/#{feature_type}#{query_string(search_after:)}"
+          execute_json(:get, path, context: "list #{feature_type} rules").fetch("rules")
         end
 
-        def execute_get_rule(feature_type, rule_id)
-          response = get_request("/_rules/#{feature_type}/#{rule_id}")
-          JSON.parse(response.read)
-        rescue StandardError => e
-          { error: { type: e.class.name, reason: e.message } }
-        ensure
-          response&.close
+        def execute_delete_rule(feature_type, id)
+          execute_json(:delete, "/_rules/#{feature_type}/#{id}", context: "delete #{feature_type} rule #{id}")
         end
 
-        def execute_list_rules(feature_type, search_after: nil)
-          path = "/_rules/#{feature_type}"
-          path += "?search_after=#{search_after}" if search_after
+        def execute_rule_exists?(feature_type, id)
+          execute_exists?("/_rules/#{feature_type}/#{id}", context: "#{feature_type} rule exists #{id}")
+        end
 
-          response = get_request(path)
-          JSON.parse(response.read)
-        rescue StandardError => e
-          { rules: [], error: { type: e.class.name, reason: e.message } }
-        ensure
-          response&.close
+        def execute_create_workload_group(name, resource_limits:, resiliency_mode: "soft")
+          body = { name:, resiliency_mode:, resource_limits: }
+          execute_json(:put, "/_wlm/workload_group", body, context: "create workload group #{name}")
         end
 
-        def execute_delete_rule(feature_type, rule_id)
-          response = delete_request("/_rules/#{feature_type}/#{rule_id}")
-          JSON.parse(response.read)
-        rescue StandardError => e
-          { acknowledged: false, error: { type: e.class.name, reason: e.message } }
-        ensure
-          response&.close
+        def execute_get_workload_group(name)
+          execute_json(:get, "/_wlm/workload_group/#{name}", context: "get workload group #{name}").fetch("workload_groups").first
         end
 
-        def execute_rule_exists?(feature_type, rule_id)
-          response = head_request("/_rules/#{feature_type}/#{rule_id}")
-          response.success?
-        rescue StandardError
-          false
-        ensure
-          response&.close
+        def execute_delete_workload_group(name)
+          execute_json(:delete, "/_wlm/workload_group/#{name}", context: "delete workload group #{name}")
         end
       end
     end

data/lib/noiseless/adapters/execution_modules/pgvector_support.rb

@@ -26,9 +26,9 @@ module Noiseless
         #
         def vector_search(scope, embedding, column: :embedding, limit: 20, distance_threshold: nil,
                           distance_metric: :cosine)
-          return scope unless pgvector_available?
+          return scope.none unless pgvector_available?
 
-          vector_string = "[#{embedding.join(',')}]"
+          vector_string = vector_literal(embedding)
           distance_op = distance_operator(distance_metric)
 
           # Build the query with distance calculation
@@ -63,30 +63,19 @@ module Noiseless
         #
         def hybrid_search(scope, text_query:, embedding:, text_fields:, vector_column: :embedding,
                           text_weight: 0.5, vector_weight: 0.5, limit: 20)
-          return scope unless pgvector_available?
-
-          vector_string = "[#{embedding.join(',')}]"
-          text_conditions = text_fields.map { |f| "similarity(#{quoted_column(f)}, ?)" }.join(" + ")
-          text_similarity_count = text_fields.size
-
-          # Normalized combined score
-          scope.select(
-            "#{scope.table_name}.*",
-            # Text similarity (0-1 per field, averaged)
-            Arel.sql(
-              "(#{text_conditions}) / #{text_similarity_count} * #{text_weight} AS text_score"
-            ),
-            # Vector similarity (convert distance to similarity: 1 - distance for cosine)
-            "(1 - (#{quoted_column(vector_column)} <=> '#{vector_string}')) * #{vector_weight} AS vector_score",
-            # Combined score
-            "(((#{text_conditions}) / #{text_similarity_count}) * #{text_weight} + " \
-            "(1 - (#{quoted_column(vector_column)} <=> '#{vector_string}')) * #{vector_weight}) AS combined_score"
-          ).where(
-            "#{text_conditions} > 0 OR #{quoted_column(vector_column)} IS NOT NULL",
-            *Array.new(text_similarity_count, text_query)
-          ).order(Arel.sql("combined_score DESC"))
+          return scope.none unless pgvector_available?
+          raise ArgumentError, "hybrid_search needs text_fields to run the text query against" if text_fields.empty?
+
+          column = quoted_column(vector_column)
+          text_score = "(#{text_fields.map { "similarity(#{quoted_column(it)}, :query)" }.join(' + ')}) / #{text_fields.size}"
+          vector_score = "(1 - (#{column} <=> '#{vector_literal(embedding)}'))"
+          combined = "#{text_score} * #{Float(text_weight)} + #{vector_score} * #{Float(vector_weight)}"
+          bind = ->(sql) { scope.sanitize_sql_array([sql, { query: text_query.to_s }]) }
+
+          scope.select("#{scope.table_name}.*", Arel.sql(bind.call("#{combined} AS combined_score")))
+               .where(bind.call("#{text_score} > 0 OR #{column} IS NOT NULL"))
+               .order(Arel.sql("combined_score DESC"))
                .limit(limit)
-               .tap { |s| s.bind_values.concat(Array.new(text_similarity_count, text_query)) }
         end
 
         # Execute a KNN (K-Nearest Neighbors) search
@@ -101,7 +90,7 @@ module Noiseless
         def knn_search(model, embedding, k: 10, column: :embedding, filters: {})
           return [] unless pgvector_available?
 
-          vector_string = "[#{embedding.join(',')}]"
+          vector_string = vector_literal(embedding)
 
           scope = model.all
           scope = scope.where(filters) if filters.any?
@@ -124,36 +113,36 @@ module Noiseless
         def store_embedding(record, embedding, column: :embedding)
           return false unless pgvector_available?
 
-          vector_string = "[#{embedding.join(',')}]"
-          record.update_column(column, vector_string)
+          record.update_column(column, vector_literal(embedding))
         end
 
         # Batch store embeddings
         #
         # @param model [Class] The ActiveRecord model
-        # @param embeddings [Hash<String, Array<Float>>] Map of ID -> embedding
+        # @param embeddings [Hash<Object, Array<Float>>] Map of primary key -> embedding
         # @param column [Symbol] The column to store embeddings
+        # @return [Integer] Number of rows updated
         #
         def batch_store_embeddings(model, embeddings, column: :embedding)
           return 0 unless pgvector_available?
-
-          # Use UPDATE FROM VALUES for efficient batch update
-          values = embeddings.map do |id, emb|
-            "(#{ActiveRecord::Base.connection.quote(id)}, '[#{emb.join(',')}]'::vector)"
-          end.join(",")
-
-          sql = <<~SQL.squish
-            UPDATE #{model.table_name}
-            SET #{column} = v.embedding
-            FROM (VALUES #{values}) AS v(id, embedding)
-            WHERE #{model.table_name}.id = v.id::uuid
-          SQL
-
-          ActiveRecord::Base.connection.execute(sql)
-          embeddings.size
-        rescue StandardError => e
-          Rails.logger.error("Failed to batch store embeddings: #{e.message}")
-          0
+          return 0 if embeddings.empty?
+
+          model.with_connection do |connection|
+            table = connection.quote_table_name(model.table_name)
+            primary_key = model.primary_key
+            pk_type = model.columns_hash.fetch(primary_key.to_s).sql_type
+
+            values = embeddings.map do |id, embedding|
+              "(#{connection.quote(id)}::#{pk_type}, '#{vector_literal(embedding)}'::vector)"
+            end
+
+            connection.exec_update(<<~SQL.squish, "Noiseless batch_store_embeddings")
+              UPDATE #{table}
+              SET #{connection.quote_column_name(column)} = v.embedding
+              FROM (VALUES #{values.join(', ')}) AS v(id, embedding)
+              WHERE #{table}.#{connection.quote_column_name(primary_key)} = v.id
+            SQL
+          end
         end
 
         # Find similar records to a given record
@@ -181,6 +170,12 @@ module Noiseless
 
         private
 
+        # pgvector operators take no bind parameter in SELECT or ORDER BY, so
+        # the literal is interpolated; coercion guarantees it holds only floats.
+        def vector_literal(embedding)
+          "[#{AST::Vector.coerce_embedding(embedding).join(',')}]"
+        end
+
         def distance_operator(metric)
           case metric
           when :l2, :euclidean