swarm_sdk 2.7.14 → 3.0.0.alpha2

data/lib/swarm_sdk/v3/memory/retriever.rb

@@ -0,0 +1,289 @@
+# frozen_string_literal: true
+
+module SwarmSDK
+  module V3
+    module Memory
+      # Hybrid search retriever combining semantic and keyword search
+      #
+      # Uses reciprocal rank fusion (RRF) to combine results from:
+      # 1. Semantic search via FAISS vector similarity
+      # 2. Keyword search via text matching
+      # 3. Graph expansion via 1-hop edge traversal
+      #
+      # @example
+      #   retriever = Retriever.new(adapter: adapter, embedder: embedder)
+      #   results = retriever.search("JWT authentication", top_k: 10)
+      class Retriever
+        # @param adapter [Adapters::Base] Storage adapter
+        # @param embedder [Embedder] Text embedder
+        # @param semantic_weight [Float] Weight for semantic results (0.0-1.0)
+        # @param keyword_weight [Float] Weight for keyword results (0.0-1.0)
+        def initialize(adapter:, embedder:, semantic_weight: 0.5, keyword_weight: 0.5)
+          @adapter = adapter
+          @embedder = embedder
+          @semantic_weight = semantic_weight
+          @keyword_weight = keyword_weight
+          @config = Configuration.instance
+        end
+
+        # Search for relevant memory cards
+        #
+        # @param query [String] Search query
+        # @param top_k [Integer] Number of results to return
+        # @return [Array<Card>] Ranked cards
+        #
+        # @example
+        #   cards = retriever.search("How does auth work?", top_k: 10)
+        def search(query, top_k: 15)
+          DebugLog.log("retriever", "search: query=#{query[0..60].inspect}, top_k=#{top_k}")
+
+          semantic_results = DebugLog.time("retriever", "semantic_search(#{top_k * 2})") do
+            semantic_search(query, top_k: top_k * 2)
+          end
+
+          keyword_results = DebugLog.time("retriever", "keyword_search(#{top_k * 2})") do
+            keyword_search(query, top_k: top_k * 2)
+          end
+
+          DebugLog.log("retriever", "semantic=#{semantic_results.size} keyword=#{keyword_results.size}")
+
+          # Reciprocal rank fusion
+          fused = DebugLog.time("retriever", "rrf") do
+            reciprocal_rank_fusion(semantic_results, keyword_results)
+          end
+
+          # Graph expansion: add 1-hop neighbors of top results
+          expanded = DebugLog.time("retriever", "graph_expand") do
+            graph_expand(fused.take(top_k), top_k: top_k)
+          end
+
+          DebugLog.log("retriever", "fused=#{fused.size} expanded=#{expanded.size} returning=#{[expanded.size, top_k].min}")
+
+          # Return unique cards, limited to top_k
+          expanded.take(top_k)
+        end
+
+        private
+
+        # Semantic search via vector similarity
+        #
+        # @param query [String] Search query
+        # @param top_k [Integer] Number of results
+        # @return [Array<String>] Ranked card IDs
+        def semantic_search(query, top_k:)
+          embedding = @embedder.embed(query)
+          results = @adapter.vector_search(embedding, top_k: top_k, threshold: 0.1)
+          results.sort_by { |r| -r[:similarity] }.map { |r| r[:id] }
+        end
+
+        # BM25 keyword search
+        #
+        # Implements Okapi BM25 scoring for keyword-based retrieval.
+        # Catches exact names, IDs, error messages, and other literal matches
+        # that semantic search might miss.
+        #
+        # @param query [String] Search query
+        # @param top_k [Integer] Number of results
+        # @return [Array<String>] Ranked card IDs
+        def keyword_search(query, top_k:)
+          terms = tokenize(query)
+          return [] if terms.empty?
+
+          cards = @adapter.list_cards
+          return [] if cards.empty?
+
+          # Precompute corpus statistics for BM25
+          avg_dl = cards.sum { |c| tokenize(c.text).size }.to_f / cards.size
+          doc_freq = compute_document_frequencies(terms, cards)
+          n = cards.size
+
+          scored = cards.filter_map do |card|
+            card_terms = tokenize(card.text)
+            dl = card_terms.size
+
+            # BM25 score
+            bm25 = terms.sum do |term|
+              tf = card_terms.count(term)
+              next 0.0 if tf == 0
+
+              df = doc_freq[term] || 0
+              idf = Math.log((n - df + 0.5) / (df + 0.5) + 1.0)
+              numerator = tf * (@config.bm25_k1 + 1)
+              denominator = tf + @config.bm25_k1 * (1 - @config.bm25_b + @config.bm25_b * dl / avg_dl)
+              idf * numerator / denominator
+            end
+
+            # Boost for entity matches (exact term in entities list)
+            entity_boost = card.entities.count do |entity|
+              terms.any? { |t| entity.downcase.include?(t) }
+            end
+            bm25 += entity_boost * @config.bm25_entity_boost
+
+            { id: card.id, score: bm25 } if bm25 > 0.0
+          end
+
+          scored.sort_by { |s| -s[:score] }.take(top_k).map { |s| s[:id] }
+        end
+
+        # Combine ranked lists using reciprocal rank fusion
+        #
+        # @param semantic_ids [Array<String>] Semantic search results
+        # @param keyword_ids [Array<String>] Keyword search results
+        # @return [Array<Card>] Fused and ranked cards
+        def reciprocal_rank_fusion(semantic_ids, keyword_ids)
+          scores = Hash.new(0.0)
+
+          semantic_ids.each_with_index do |id, rank|
+            scores[id] += @semantic_weight * (1.0 / (@config.rrf_k + rank + 1))
+          end
+
+          keyword_ids.each_with_index do |id, rank|
+            scores[id] += @keyword_weight * (1.0 / (@config.rrf_k + rank + 1))
+          end
+
+          ranked_ids = scores.sort_by { |_, score| -score }.map(&:first)
+          ranked_ids.filter_map { |id| @adapter.read_card(id) }
+        end
+
+        # Expand results with 1-hop graph neighbors
+        #
+        # Limits to `config.max_neighbors_per_seed` neighbors per seed card
+        # to prevent any single highly-connected card from dominating
+        # the expanded results.
+        #
+        # @param cards [Array<Card>] Initial result cards
+        # @param top_k [Integer] Maximum total results
+        # @return [Array<Card>] Expanded results
+        def graph_expand(cards, top_k:)
+          seen_ids = Set.new(cards.map(&:id))
+          expanded = cards.dup
+
+          cards.each do |card|
+            break if expanded.size >= top_k
+
+            edges = @adapter.edges_for(card.id)
+            neighbors_added = 0
+
+            # Sort edges by weight descending to prefer strongest connections
+            edges.sort_by { |e| -e.weight }.each do |edge|
+              break if expanded.size >= top_k
+              break if neighbors_added >= @config.max_neighbors_per_seed
+
+              neighbor_id = edge.from_id == card.id ? edge.to_id : edge.from_id
+              next if seen_ids.include?(neighbor_id)
+
+              neighbor = @adapter.read_card(neighbor_id)
+              next unless neighbor
+
+              seen_ids.add(neighbor_id)
+              expanded << neighbor
+              neighbors_added += 1
+            end
+          end
+
+          expanded
+        end
+
+        # Tokenize text into lowercase terms, filtering stopwords and short tokens
+        #
+        # @param text [String] Text to tokenize
+        # @return [Array<String>] Lowercase tokens
+        def tokenize(text)
+          text.downcase
+            .gsub(/[^a-z0-9\s_-]/, " ")
+            .split(/\s+/)
+            .reject { |t| t.length < 2 || STOPWORDS.include?(t) }
+        end
+
+        # Compute document frequency for each query term
+        #
+        # @param terms [Array<String>] Query terms
+        # @param cards [Array<Card>] All cards in the corpus
+        # @return [Hash<String, Integer>] Term → document count
+        def compute_document_frequencies(terms, cards)
+          freq = Hash.new(0)
+          cards.each do |card|
+            card_terms = tokenize(card.text).to_set
+            terms.each { |term| freq[term] += 1 if card_terms.include?(term) }
+          end
+          freq
+        end
+
+        # Common English stopwords to filter from BM25 scoring
+        STOPWORDS = Set.new([
+          "the",
+          "a",
+          "an",
+          "is",
+          "are",
+          "was",
+          "were",
+          "be",
+          "been",
+          "being",
+          "have",
+          "has",
+          "had",
+          "do",
+          "does",
+          "did",
+          "will",
+          "would",
+          "shall",
+          "should",
+          "may",
+          "might",
+          "can",
+          "could",
+          "of",
+          "in",
+          "to",
+          "for",
+          "on",
+          "with",
+          "at",
+          "by",
+          "from",
+          "as",
+          "into",
+          "through",
+          "about",
+          "between",
+          "after",
+          "before",
+          "above",
+          "below",
+          "up",
+          "down",
+          "out",
+          "off",
+          "over",
+          "under",
+          "again",
+          "this",
+          "that",
+          "these",
+          "those",
+          "it",
+          "its",
+          "he",
+          "she",
+          "they",
+          "them",
+          "and",
+          "but",
+          "or",
+          "nor",
+          "not",
+          "no",
+          "so",
+          "if",
+          "then",
+          "than",
+          "too",
+          "very",
+        ]).freeze
+      end
+    end
+  end
+end