htm 0.0.30 → 0.0.32

data/lib/htm/long_term_memory/hybrid_search.rb

@@ -45,7 +45,7 @@ class HTM
       #
       def search_hybrid(timeframe:, query:, limit:, embedding_service:, prefilter_limit: 100, metadata: {})
         # Enforce limits to prevent DoS
-        safe_limit = [[limit.to_i, 1].max, MAX_HYBRID_LIMIT].min
+        safe_limit = limit.to_i.clamp(1, MAX_HYBRID_LIMIT)
         safe_prefilter = [prefilter_limit.to_i, 1].max
 
         start_time = Process.clock_gettime(Process::CLOCK_MONOTONIC)
@@ -66,11 +66,34 @@ class HTM
 
       private
 
-      # Hybrid search using Reciprocal Rank Fusion
+      def init_rrf_entry(result, rank)
+        {
+          'id'              => result['id'],
+          'content'         => result['content'],
+          'access_count'    => result['access_count'],
+          'created_at'      => result['created_at'],
+          'token_count'     => result['token_count'],
+          'similarity'      => 0.0,
+          'text_rank'       => 0.0,
+          'tag_depth_score' => 0.0,
+          'matched_tags'    => [],
+          'rrf_score'       => 1.0 / (RRF_K + rank),
+          'vector_rank'     => nil,
+          'fulltext_rank'   => nil,
+          'tag_rank'        => nil,
+          'sources'         => []
+        }
+      end
+
+      # Hybrid search using Reciprocal Rank Fusion with retrieve-then-rerank
       #
-      # Runs vector, fulltext, and tag searches independently, then merges
-      # results using RRF scoring. Nodes appearing in multiple searches
-      # get contributions from each, naturally boosting them.
+      # Uses a single SQL CTE query instead of three round-trips:
+      # 1. Fulltext + tag candidates are retrieved first (cheap)
+      # 2. Vector similarity is computed only for those candidates (expensive but scoped)
+      # 3. RRF scoring merges all three rankings in SQL
+      #
+      # Trade-off: Queries with high semantic relevance but zero keyword/tag
+      # overlap will be missed. Use :vector strategy for pure semantic search.
       #
       # @param timeframe [nil, Range, Array<Range>] Time range(s) to search
       # @param query [String] Search query
@@ -81,41 +104,241 @@ class HTM
       # @return [Array<Hash>] Merged results with RRF scores
       #
       def search_hybrid_rrf(timeframe:, query:, limit:, embedding_service:, candidate_limit:, metadata: {})
-        # Run all three searches independently
-        vector_results = fetch_vector_candidates(
-          query: query,
-          embedding_service: embedding_service,
-          timeframe: timeframe,
-          metadata: metadata,
-          limit: candidate_limit
+        tag_info   = extract_rrf_tag_info(query)
+        literals   = build_rrf_literals(query, tag_info[:search_tags], embedding_service)
+        filter_sql = build_rrf_filter_sql(timeframe, metadata)
+
+        sql = build_hybrid_cte_sql(
+          query_literal:    literals[:query_literal],
+          embedding_literal: literals[:embedding_literal],
+          tag_literals:     literals[:tag_literals],
+          additional_sql:   filter_sql[:additional_sql],
+          additional_sql_n: filter_sql[:additional_sql_n],
+          candidate_limit:  candidate_limit.to_i,
+          limit:            limit.to_i
         )
 
-        fulltext_results = fetch_fulltext_candidates(
-          query: query,
-          timeframe: timeframe,
-          metadata: metadata,
-          limit: candidate_limit
-        )
+        results     = HTM.db.fetch(sql).all
+        top_results = post_process_rrf_results(results, tag_info[:tag_depth_map])
+        track_access(top_results.map { |r| r['id'] })
+        top_results
+      end
 
-        # Extract tags from query and find matching nodes
-        tag_results = fetch_tag_candidates(
-          query: query,
-          timeframe: timeframe,
-          metadata: metadata,
-          limit: candidate_limit
-        )
+      def extract_rrf_tag_info(query)
+        extraction      = find_query_matching_tags(query, include_extracted: true)
+        extracted_tags  = extraction[:extracted] || []
+        matched_db_tags = extraction[:matched]   || []
+        {
+          search_tags:   matched_db_tags.any? ? matched_db_tags : extracted_tags,
+          tag_depth_map: build_tag_depth_map(extracted_tags)
+        }
+      end
+
+      def build_rrf_literals(query, search_tags, embedding_service)
+        query_literal   = HTM.db.literal(query)
+        query_embedding = embedding_service.embed(query)
+        has_embedding   = query_embedding.is_a?(Array) && query_embedding.any?
+        embedding_literal = if has_embedding
+                              padded = HTM::SqlBuilder.pad_embedding(query_embedding)
+                              HTM.db.literal(HTM::SqlBuilder.sanitize_embedding(padded))
+                            end
+        tag_literals = search_tags.any? ? search_tags.map { |t| HTM.db.literal(t) }.join(', ') : nil
+        { query_literal: query_literal, embedding_literal: embedding_literal, tag_literals: tag_literals }
+      end
 
-        # Merge using RRF
-        merged = merge_with_rrf(vector_results, fulltext_results, tag_results)
+      def build_rrf_filter_sql(timeframe, metadata)
+        tc = HTM::SqlBuilder.timeframe_condition(timeframe)
+        mc = HTM::SqlBuilder.metadata_condition(metadata)
+        tn = HTM::SqlBuilder.timeframe_condition(timeframe, table_alias: 'n')
+        mn = HTM::SqlBuilder.metadata_condition(metadata, table_alias: 'n')
+        {
+          additional_sql:   [tc, mc].compact.then { |p| p.any? ? "AND #{p.join(' AND ')}" : "" },
+          additional_sql_n: [tn, mn].compact.then { |p| p.any? ? "AND #{p.join(' AND ')}" : "" }
+        }
+      end
 
-        # Take top results
-        top_results = merged.first(limit)
+      def post_process_rrf_results(results, tag_depth_map)
+        results.map do |row|
+          r = row.transform_keys(&:to_s)
+          r['matched_tags'] = r.key?('matched_tags') ? parse_pg_array(r['matched_tags']) : []
+          r['tag_depth_score'] = r['matched_tags'].any? && tag_depth_map.any? ?
+                                   calculate_tag_depth_score(r['matched_tags'], tag_depth_map) : 0.0
+          sources = []
+          sources << 'vector'   if r['vector_rank']
+          sources << 'fulltext' if r['fulltext_rank']
+          sources << 'tags'     if r['tag_rank']
+          r['sources']      = sources
+          r['similarity']   = (r['similarity']  || 0.0).to_f
+          r['text_rank']    = (r['text_rank']   || 0.0).to_f
+          r['rrf_score']    = r['rrf_score'].to_f
+          r['vector_rank'] ||= nil
+          r['fulltext_rank'] ||= nil
+          r['tag_rank']     ||= nil
+
+          r
+        end
+      end
 
-        # Track access for retrieved nodes
-        node_ids = top_results.map { |r| r['id'] }
-        track_access(node_ids)
+      # Build the single-CTE SQL for hybrid search
+      #
+      # Conditionally includes/excludes CTEs based on available components:
+      # - Always: fulltext (tsvector + trigram) candidates
+      # - If tag_literals: tag candidates CTE
+      # - If embedding_literal: vector rerank CTE (only on candidate IDs)
+      #
+      # @param query_literal [String] SQL-escaped query string
+      # @param embedding_literal [String, nil] SQL-escaped embedding vector (nil to skip vector)
+      # @param tag_literals [String, nil] Comma-separated SQL-escaped tag names (nil to skip tags)
+      # @param additional_sql [String] Extra WHERE conditions (no table alias)
+      # @param additional_sql_n [String] Extra WHERE conditions (with 'n.' alias)
+      # @param candidate_limit [Integer] Max candidates per source
+      # @param limit [Integer] Final result limit
+      # @return [String] Complete SQL query
+      #
+      def build_hybrid_cte_sql(query_literal:, embedding_literal:, tag_literals:, additional_sql:, additional_sql_n:, candidate_limit:,
+                               limit:)
+        has_embedding = !embedding_literal.nil?
+        has_tags      = !tag_literals.nil?
+
+        ctes = fulltext_ctes_sql(query_literal, additional_sql, candidate_limit)
+        ctes << tag_candidates_cte_sql(tag_literals, additional_sql_n, candidate_limit) if has_tags
+        ctes << candidate_ids_cte_sql(has_tags)
+        ctes << vector_rerank_cte_sql(embedding_literal) if has_embedding
+        ctes.concat ranked_ctes_sql(has_tags, has_embedding)
+        ctes << rrf_scores_cte_sql(has_tags, has_embedding)
+
+        final_fields = build_final_select_fields(has_tags, has_embedding)
+
+        <<~SQL
+          WITH #{ctes.join(",\n")}
+          SELECT #{final_fields.join(",\n         ")}
+          FROM rrf_scores rrf
+          JOIN nodes n ON n.id = rrf.id
+          ORDER BY rrf.rrf_score DESC
+          LIMIT #{limit}
+        SQL
+      end
 
-        top_results
+      def fulltext_ctes_sql(query_literal, additional_sql, candidate_limit)
+        [<<~SQL.chomp, <<~SQL.chomp, <<~SQL.chomp]
+          tsvector_matches AS (
+            SELECT id,
+                   (1.0 + ts_rank(to_tsvector('english', content), plainto_tsquery('english', #{query_literal}))) as text_rank
+            FROM nodes
+            WHERE deleted_at IS NULL
+            AND to_tsvector('english', content) @@ plainto_tsquery('english', #{query_literal})
+            #{additional_sql}
+          )
+        SQL
+          trigram_matches AS (
+            SELECT id,
+                   similarity(content, #{query_literal}) as text_rank
+            FROM nodes
+            WHERE deleted_at IS NULL
+            AND similarity(content, #{query_literal}) >= 0.1
+            AND id NOT IN (SELECT id FROM tsvector_matches)
+            #{additional_sql}
+          )
+        SQL
+          fulltext_candidates AS (
+            SELECT * FROM tsvector_matches
+            UNION ALL
+            SELECT * FROM trigram_matches
+            ORDER BY text_rank DESC
+            LIMIT #{candidate_limit}
+          )
+        SQL
+      end
+
+      def tag_candidates_cte_sql(tag_literals, additional_sql_n, candidate_limit)
+        <<~SQL.chomp
+          tag_candidates AS (
+            SELECT n.id, array_agg(t.name) as matched_tags, count(t.name) as tag_match_count
+            FROM nodes n
+            JOIN node_tags nt ON nt.node_id = n.id
+            JOIN tags t ON t.id = nt.tag_id
+            WHERE n.deleted_at IS NULL
+            AND t.name IN (#{tag_literals})
+            #{additional_sql_n}
+            GROUP BY n.id
+            LIMIT #{candidate_limit}
+          )
+        SQL
+      end
+
+      def candidate_ids_cte_sql(has_tags)
+        id_sources = ["SELECT id FROM fulltext_candidates"]
+        id_sources << "SELECT id FROM tag_candidates" if has_tags
+        <<~SQL.chomp
+          all_candidate_ids AS (
+            #{id_sources.join("\n    UNION\n    ")}
+          )
+        SQL
+      end
+
+      def vector_rerank_cte_sql(embedding_literal)
+        <<~SQL.chomp
+          vector_rerank AS (
+            SELECT id,
+                   1 - (embedding <=> #{embedding_literal}::vector) as similarity
+            FROM nodes
+            WHERE id IN (SELECT id FROM all_candidate_ids)
+            AND embedding IS NOT NULL
+          )
+        SQL
+      end
+
+      def ranked_ctes_sql(has_tags, has_embedding)
+        ctes = [<<~SQL.chomp]
+          fulltext_ranked AS (
+            SELECT id, text_rank, ROW_NUMBER() OVER (ORDER BY text_rank DESC) as rank
+            FROM fulltext_candidates
+          )
+        SQL
+        ctes << <<~SQL.chomp if has_tags
+          tag_ranked AS (
+            SELECT id, matched_tags, ROW_NUMBER() OVER (ORDER BY tag_match_count DESC) as rank
+            FROM tag_candidates
+          )
+        SQL
+        ctes << <<~SQL.chomp if has_embedding
+          vector_ranked AS (
+            SELECT id, similarity, ROW_NUMBER() OVER (ORDER BY similarity DESC) as rank
+            FROM vector_rerank
+          )
+        SQL
+        ctes
+      end
+
+      def rrf_scores_cte_sql(has_tags, has_embedding)
+        rrf_id         = ["fr.id"]
+        rrf_id         << "tr.id" if has_tags
+        rrf_id         << "vr.id" if has_embedding
+        score_terms    = ["COALESCE(1.0/(#{RRF_K} + fr.rank), 0)"]
+        score_terms    << "COALESCE(1.0/(#{RRF_K} + tr.rank), 0)" if has_tags
+        score_terms    << "COALESCE(1.0/(#{RRF_K} + vr.rank), 0)" if has_embedding
+        fields         = ["COALESCE(#{rrf_id.join(', ')}) as id", "#{score_terms.join(' + ')} as rrf_score",
+                          "fr.rank as fulltext_rank", "fr.text_rank"]
+        fields         << "tr.rank as tag_rank" << "tr.matched_tags" if has_tags
+        fields         << "vr.rank as vector_rank" << "vr.similarity" if has_embedding
+        from           = "fulltext_ranked fr"
+        from += "\n      FULL OUTER JOIN tag_ranked tr ON fr.id = tr.id" if has_tags
+        coalesce_id = has_tags ? "COALESCE(fr.id, tr.id)" : "fr.id"
+        from += "\n      FULL OUTER JOIN vector_ranked vr ON #{coalesce_id} = vr.id" if has_embedding
+        <<~SQL.chomp
+          rrf_scores AS (
+            SELECT #{fields.join(",\n           ")}
+            FROM #{from}
+          )
+        SQL
+      end
+
+      def build_final_select_fields(has_tags, has_embedding)
+        fields = ["rrf.id", "n.content", "n.access_count", "n.created_at", "n.token_count",
+                  "rrf.rrf_score", "rrf.fulltext_rank", "COALESCE(rrf.text_rank, 0.0) as text_rank"]
+        fields << "rrf.tag_rank" << "rrf.matched_tags" if has_tags
+        fields << "rrf.vector_rank" << "COALESCE(rrf.similarity, 0.0) as similarity" if has_embedding
+        fields
       end
 
       # Fetch candidates using vector similarity search
@@ -149,7 +372,7 @@ class HTM
 
         where_clause = "WHERE #{conditions.join(' AND ')}"
 
-        # Note: Using Sequel.lit for the vector comparison since it needs special handling
+        # NOTE: Using Sequel.lit for the vector comparison since it needs special handling
         embedding_literal = HTM.db.literal(embedding_str)
         sql = <<~SQL
           SELECT id, content, access_count, created_at, token_count,
@@ -282,7 +505,8 @@ class HTM
           depth_score = calculate_tag_depth_score(matched_tags, tag_depth_map)
 
           result.transform_keys(&:to_s).merge('tag_depth_score' => depth_score, 'matched_tags' => matched_tags)
-        end.sort_by { |r| -r['tag_depth_score'] }
+        end
+        results.sort_by { |r| -r['tag_depth_score'] }
       end
 
       # Build a map of tag prefixes to their depth information
@@ -336,13 +560,13 @@ class HTM
             # Score is depth / max_depth
             # e.g., "database:postgresql" matching query "database:postgresql:extensions"
             # gives 2/3 = 0.67
-            score = info[:depth].to_f / info[:max_depth].to_f
+            score = info[:depth].to_f / info[:max_depth]
             best_score = [best_score, score].max
           else
             # Check if this tag is a parent of any extracted tag
             tag_depth_map.each do |prefix, info|
               if prefix.start_with?(tag + ':') || prefix == tag
-                score = tag.split(':').size.to_f / info[:max_depth].to_f
+                score = tag.split(':').size.to_f / info[:max_depth]
                 best_score = [best_score, score].max
               end
             end
@@ -388,103 +612,51 @@ class HTM
       # @return [Array<Hash>] Merged results sorted by RRF score
       #
       def merge_with_rrf(vector_results, fulltext_results, tag_results = [])
-        # Build RRF scores
-        # Key: node_id, Value: { node_data:, rrf_score:, sources: }
         merged = {}
+        vector_results.each_with_index   { |r, i| merge_vector_rrf_entry(merged, r, i + 1) }
+        fulltext_results.each_with_index  { |r, i| merge_fulltext_rrf_entry(merged, r, i + 1) }
+        tag_results.each_with_index       { |r, i| merge_tag_rrf_entry(merged, r, i + 1) }
+        merged.values.sort_by { |r| -r['rrf_score'] }
+      end
 
-        # Process vector results
-        vector_results.each_with_index do |result, index|
-          id = result['id']
-          rank = index + 1  # 1-based rank
-          rrf_contribution = 1.0 / (RRF_K + rank)
-
-          merged[id] = {
-            'id' => result['id'],
-            'content' => result['content'],
-            'access_count' => result['access_count'],
-            'created_at' => result['created_at'],
-            'token_count' => result['token_count'],
-            'similarity' => result['similarity'],
-            'text_rank' => 0.0,
-            'tag_depth_score' => 0.0,
-            'matched_tags' => [],
-            'rrf_score' => rrf_contribution,
-            'vector_rank' => rank,
-            'fulltext_rank' => nil,
-            'tag_rank' => nil,
-            'sources' => ['vector']
-          }
-        end
+      def merge_vector_rrf_entry(merged, result, rank)
+        merged[result['id']] = init_rrf_entry(result, rank).merge(
+          'similarity' => result['similarity'], 'vector_rank' => rank, 'sources' => ['vector']
+        )
+      end
 
-        # Process fulltext results
-        fulltext_results.each_with_index do |result, index|
-          id = result['id']
-          rank = index + 1  # 1-based rank
-          rrf_contribution = 1.0 / (RRF_K + rank)
-
-          if merged.key?(id)
-            # Node appears in both - add RRF contribution (this is the boost!)
-            merged[id]['rrf_score'] += rrf_contribution
-            merged[id]['text_rank'] = result['text_rank']
-            merged[id]['fulltext_rank'] = rank
-            merged[id]['sources'] << 'fulltext'
-          else
-            # Node only in fulltext
-            merged[id] = {
-              'id' => result['id'],
-              'content' => result['content'],
-              'access_count' => result['access_count'],
-              'created_at' => result['created_at'],
-              'token_count' => result['token_count'],
-              'similarity' => 0.0,
-              'text_rank' => result['text_rank'],
-              'tag_depth_score' => 0.0,
-              'matched_tags' => [],
-              'rrf_score' => rrf_contribution,
-              'vector_rank' => nil,
-              'fulltext_rank' => rank,
-              'tag_rank' => nil,
-              'sources' => ['fulltext']
-            }
-          end
+      def merge_fulltext_rrf_entry(merged, result, rank)
+        rrf = 1.0 / (RRF_K + rank)
+        id  = result['id']
+        if merged.key?(id)
+          merged[id]['rrf_score']    += rrf
+          merged[id]['text_rank']     = result['text_rank']
+          merged[id]['fulltext_rank'] = rank
+          merged[id]['sources'] << 'fulltext'
+        else
+          merged[id] = init_rrf_entry(result, rank).merge(
+            'text_rank' => result['text_rank'], 'fulltext_rank' => rank, 'sources' => ['fulltext']
+          )
         end
+      end
 
-        # Process tag results
-        tag_results.each_with_index do |result, index|
-          id = result['id']
-          rank = index + 1  # 1-based rank
-          rrf_contribution = 1.0 / (RRF_K + rank)
-
-          if merged.key?(id)
-            # Node already found - add RRF contribution (boost!)
-            merged[id]['rrf_score'] += rrf_contribution
-            merged[id]['tag_depth_score'] = result['tag_depth_score']
-            merged[id]['matched_tags'] = result['matched_tags']
-            merged[id]['tag_rank'] = rank
-            merged[id]['sources'] << 'tags'
-          else
-            # Node only found via tags
-            merged[id] = {
-              'id' => result['id'],
-              'content' => result['content'],
-              'access_count' => result['access_count'],
-              'created_at' => result['created_at'],
-              'token_count' => result['token_count'],
-              'similarity' => 0.0,
-              'text_rank' => 0.0,
-              'tag_depth_score' => result['tag_depth_score'],
-              'matched_tags' => result['matched_tags'],
-              'rrf_score' => rrf_contribution,
-              'vector_rank' => nil,
-              'fulltext_rank' => nil,
-              'tag_rank' => rank,
-              'sources' => ['tags']
-            }
-          end
+      def merge_tag_rrf_entry(merged, result, rank)
+        rrf = 1.0 / (RRF_K + rank)
+        id  = result['id']
+        if merged.key?(id)
+          merged[id]['rrf_score']       += rrf
+          merged[id]['tag_depth_score']  = result['tag_depth_score']
+          merged[id]['matched_tags']     = result['matched_tags']
+          merged[id]['tag_rank']         = rank
+          merged[id]['sources'] << 'tags'
+        else
+          merged[id] = init_rrf_entry(result, rank).merge(
+            'tag_depth_score' => result['tag_depth_score'],
+            'matched_tags'    => result['matched_tags'],
+            'tag_rank'        => rank,
+            'sources'         => ['tags']
+          )
         end
-
-        # Sort by RRF score descending
-        merged.values.sort_by { |r| -r['rrf_score'] }
       end
     end
   end

data/lib/htm/long_term_memory/node_operations.rb

@@ -24,7 +24,7 @@ class HTM
       # @return [Hash] { node_id:, is_new:, robot_node: }
       # @raise [ArgumentError] If metadata is not a Hash
       #
-      def add(content:, token_count: 0, robot_id:, embedding: nil, metadata: {})
+      def add(content:, robot_id:, token_count: 0, embedding: nil, metadata: {})
         # Validate metadata parameter
         unless metadata.is_a?(Hash)
           raise ArgumentError, "metadata must be a Hash, got #{metadata.class}"
@@ -169,7 +169,7 @@ class HTM
       # @return [Boolean] True if node exists
       #
       def exists?(node_id)
-        HTM::Models::Node.where(id: node_id).count > 0
+        HTM::Models::Node.where(id: node_id).any?
       end
 
       # Mark nodes as evicted from working memory