ask-rag 0.2.3 → 0.4.0

data/lib/ask/rag/vector_store/base.rb

@@ -2,7 +2,8 @@
 
 module Ask
   module RAG
-    # Raised when embedding a query fails (e.g., no API key, network error).
+    # Raised when embedding a query or document fails (no API key, network
+    # error, a provider returning the wrong number of vectors).
     class EmbeddingError < StandardError; end
 
     # Abstract base class for vector stores.
@@ -11,15 +12,25 @@ module Ask
     # Subclasses implement the storage and retrieval logic for specific
     # backends (InMemory, PGVector, Chroma, etc.).
     #
+    # Two retrieval modes are part of the contract, not a backend extra:
+    # +similarity_search+ (dense) and +hybrid_search+ (dense + lexical,
+    # fused). Every store implements both, so switching backends cannot
+    # silently downgrade retrieval quality.
+    #
     # @example
     #   store = Ask::RAG::VectorStore::InMemory.new
     #   store.add(documents, model: "text-embedding-3-small")
     #   results = store.similarity_search("query text", limit: 5)
+    #   results = store.hybrid_search("query text", limit: 5)
     #
     class VectorStore
       # Add documents to the store. They are embedded using the given model
       # via ask-llm-providers or a custom embedding function.
       #
+      # Documents without an id are given a stable one derived from their
+      # source, chunk index, and content (see DocumentId), which makes
+      # re-adding the same pipeline output replace rather than duplicate.
+      #
       # @param documents [Array<Ask::Document>] documents to add
       # @param model [String] embedding model name (e.g. "text-embedding-3-small")
       # @param batch_size [Integer] number of documents to embed per API call
@@ -34,10 +45,12 @@ module Ask
       # @param limit [Integer] maximum number of results
       # @param filter [Hash, nil] metadata filter — only entries whose metadata
       #   matches all key/value pairs are considered
+      # @param min_score [Float, nil] drop results scoring below this
       # @param mmr [Boolean] apply Max Marginal Relevance for diversity
       # @param diversity_bonus [Float] MMR diversity factor
       # @return [Array<Ask::Document>] documents with +:score+ in metadata
-      def similarity_search(query, limit: 10, filter: nil, mmr: false, diversity_bonus: 0.3)
+      def similarity_search(query, limit: 10, filter: nil, min_score: nil,
+        mmr: false, diversity_bonus: Retrieval::MaxMarginalRelevance::DEFAULT_DIVERSITY_BONUS)
         raise NotImplementedError
       end
 
@@ -47,8 +60,29 @@ module Ask
       # @param vector [Array<Float>] the query vector
       # @param limit [Integer] maximum number of results
       # @param filter [Hash, nil] metadata filter
+      # @param min_score [Float, nil] drop results scoring below this
+      # @param mmr [Boolean] apply Max Marginal Relevance for diversity
+      # @param diversity_bonus [Float] MMR diversity factor
       # @return [Array<Ask::Document>] documents with +:score+ in metadata
-      def similarity_search_by_vector(vector, limit: 10, filter: nil)
+      def similarity_search_by_vector(vector, limit: 10, filter: nil, min_score: nil,
+        mmr: false, diversity_bonus: Retrieval::MaxMarginalRelevance::DEFAULT_DIVERSITY_BONUS)
+        raise NotImplementedError
+      end
+
+      # Search by dense and lexical retrieval at once, fused by Reciprocal
+      # Rank Fusion.
+      #
+      # Prefer this over +similarity_search+ for anything a person reads:
+      # dense retrieval alone misses exact terms (error codes, method names,
+      # identifiers), lexical alone misses paraphrase. A document both paths
+      # find outranks one only a single path found.
+      #
+      # @param query [String] the query text
+      # @param limit [Integer] maximum number of results
+      # @param filter [Hash, nil] metadata filter
+      # @param min_score [Float, nil] drop vector results scoring below this
+      # @return [Array<Ask::Document>] documents with +:rrf_score+ in metadata
+      def hybrid_search(query, limit: 10, filter: nil, min_score: nil)
         raise NotImplementedError
       end
 
@@ -58,6 +92,16 @@ module Ask
         raise NotImplementedError
       end
 
+      # Remove every document whose metadata matches all given pairs.
+      # This is how stale chunks are cleared after a source is edited or
+      # removed — the complement to `add`'s stable ids.
+      #
+      # @param filter [Hash] metadata key/value pairs to match
+      # @return [Integer] number of documents removed
+      def delete_by(filter)
+        raise NotImplementedError
+      end
+
       # Remove all documents from the store.
       def clear
         raise NotImplementedError

data/lib/ask/rag/vector_store/in_memory.rb

@@ -14,11 +14,12 @@ module Ask
       #   store.add(chunks, model: "text-embedding-3-small")
       #   results = store.similarity_search("query", limit: 5)
       #
-      # @example With metadata filtering
+      # @example With metadata filtering and a relevance floor
       #   results = store.similarity_search(
       #     "query",
       #     limit: 5,
-      #     filter: { source: "api_docs.md" }
+      #     filter: { source: "api_docs.md" },
+      #     min_score: 0.35
       #   )
       #
       # @example With MMR (diversified results)
@@ -29,9 +30,18 @@ module Ask
       #     diversity_bonus: 0.5
       #   )
       #
+      # @example Hybrid (dense + keyword, fused)
+      #   results = store.hybrid_search("query", limit: 5)
+      #
       class InMemory < VectorStore
+        include Embeddings
         Entry = Struct.new(:id, :document, :vector, keyword_init: true)
 
+        # How many dense results are pulled before MMR or fusion narrows
+        # them. Wider than any sane `limit` so MMR has candidates to choose
+        # between and RRF has ranks to merge.
+        CANDIDATE_LIMIT = 100
+
         def initialize
           @entries = {}
           @mutex = Mutex.new
@@ -41,30 +51,28 @@ module Ask
 
         # Add documents to the store.
         #
+        # An entry replaces any existing entry with the same id, so
+        # re-indexing the same chunks updates rather than duplicates them.
+        #
         # @param documents [Array<Ask::Document>] documents to add
-        # @param model [String] embedding model name
+        # @param model [String, nil] embedding model name (default:
+        #   configured embedding_model)
         # @param batch_size [Integer] texts per embed API call
         # @return [Array<String>] IDs of added documents
-        def add(documents, model:, batch_size: 20)
-          ids = documents.map { |d| d.id || SecureRandom.uuid }
-          provider = resolve_embedding_provider(model)
-
+        def add(documents, model: nil, batch_size: 20)
+          ids = documents.map { |document| DocumentId.for(document) }
+          model ||= default_embedding_model
           @embedding_model = model
 
-          documents.each_slice(batch_size).flat_map do |batch|
-            texts = batch.map(&:content)
-            raw_vectors = provider.embed(texts, model: model)
-            vectors = normalize_vectors(raw_vectors, texts)
+          documents.each_slice(batch_size).each_with_index.flat_map do |batch, batch_index|
+            vectors = embed_texts(batch.map(&:content), model: model)
+            offset = batch_index * batch_size
 
             @mutex.synchronize do
-              batch.each_with_index.map do |doc, idx|
-                entry_id = ids[documents.index(doc)]
-                @entries[entry_id] = Entry.new(
-                  id: entry_id,
-                  document: doc,
-                  vector: vectors[idx]
-                )
-                entry_id
+              batch.each_with_index.map do |document, index|
+                id = ids[offset + index]
+                @entries[id] = Entry.new(id: id, document: with_id(document, id), vector: vectors[index])
+                id
               end
             end
           end
@@ -72,51 +80,65 @@ module Ask
 
         # Search for documents similar to the query.
         #
+        # The query is embedded with the same model the store was populated
+        # with: a query embedded by a different model lands in a different
+        # vector space, where cosine similarity is meaningless.
+        #
         # @param query [String] the query text
         # @param limit [Integer] maximum results (default: 10)
         # @param filter [Hash, nil] metadata filter — only entries whose metadata
         #   matches all key/value pairs are considered
+        # @param min_score [Float, nil] drop results scoring below this
         # @param mmr [Boolean] apply Max Marginal Relevance for diversity
         # @param diversity_bonus [Float] MMR diversity factor (0 = pure relevance, 1 = pure diversity)
-        # @return [Array<Ask::Document>] documents with +:score+ and +:mmr_score+ in metadata
-        def similarity_search(query, limit: 10, filter: nil, mmr: false, diversity_bonus: 0.3)
-          query_vector = embed_query
-          search_by_vector(query_vector, limit: limit, filter: filter,
-                           mmr: mmr, diversity_bonus: diversity_bonus)
+        # @return [Array<Ask::Document>] documents with +:score+ in metadata
+        def similarity_search(query, limit: 10, filter: nil, min_score: nil,
+          mmr: false, diversity_bonus: Retrieval::MaxMarginalRelevance::DEFAULT_DIVERSITY_BONUS)
+          query_vector = embed_query(query)
+          search_by_vector(query_vector, limit: limit, filter: filter, min_score: min_score,
+            mmr: mmr, diversity_bonus: diversity_bonus)
         end
 
-        def similarity_search_by_vector(vector, limit: 10, filter: nil)
-          search_by_vector(vector, limit: limit, filter: filter)
+        def similarity_search_by_vector(vector, limit: 10, filter: nil, min_score: nil)
+          search_by_vector(vector, limit: limit, filter: filter, min_score: min_score)
         end
 
-        private
-
-        def search_by_vector(vector, limit:, filter: nil, mmr: false, diversity_bonus: 0.3)
-          candidates = @mutex.synchronize do
-            @entries.filter_map do |_id, entry|
-              next if filter && !matches_filter?(entry.document.metadata, filter)
-
-              score = cosine_similarity(vector, entry.vector)
-              [entry, score]
-            end
-          end
+        # Dense + keyword retrieval, fused by Reciprocal Rank Fusion.
+        #
+        # The keyword half is BM25 over the stored documents, which is what
+        # catches the exact terms an embedding blurs: identifiers, error
+        # codes, method names. A document both halves find outranks one that
+        # only a single half found.
+        #
+        # @param query [String] the query text
+        # @param limit [Integer] maximum results
+        # @param filter [Hash, nil] metadata filter
+        # @param min_score [Float, nil] drop vector results scoring below this
+        # @return [Array<Ask::Document>] documents with +:rrf_score+ in metadata
+        def hybrid_search(query, limit: 10, filter: nil, min_score: nil)
+          dense = similarity_search(query, limit: CANDIDATE_LIMIT, filter: filter, min_score: min_score)
+          # A query with no dense matches (all below min_score, or an empty
+          # store) can still have exact-term matches, so the keyword half
+          # runs regardless rather than short-circuiting on an empty list.
+          lexical = keyword_search(query, limit: CANDIDATE_LIMIT, filter: filter)
 
-          candidates.sort_by! { |_, score| -score }
+          Retrieval::ReciprocalRankFusion.fuse([dense, lexical], limit: limit)
+        end
 
-          if mmr
-            apply_mmr(candidates, vector, limit, diversity_bonus)
-          else
-            candidates.first(limit).map do |entry, score|
-              build_result(entry, score: score)
-            end
+        def delete(ids)
+          Array(ids).each do |id|
+            @mutex.synchronize { @entries.delete(id) }
           end
         end
 
-        public
-
-        def delete(ids)
+        # Remove every entry whose metadata matches all given pairs.
+        # @param filter [Hash] metadata key/value pairs to match
+        # @return [Integer] number of entries removed
+        def delete_by(filter)
           @mutex.synchronize do
-            ids.each { |id| @entries.delete(id) }
+            doomed = @entries.select { |_id, entry| matches_filter?(entry.document.metadata, filter) }.keys
+            doomed.each { |id| @entries.delete(id) }
+            doomed.length
           end
         end
 
@@ -130,66 +152,63 @@ module Ask
 
         private
 
-        # --- Embedding ---
+        # The stored document carries the id it was stored under, so
+        # `similarity_search` -> `delete` composes without bookkeeping.
+        def with_id(document, id)
+          return document if document.id == id
 
-        def resolve_embedding_provider(model)
-          info = Ask::ModelCatalog.find(model)
-          provider_class = Ask::Provider.resolve(info.provider)
-          config = build_provider_config(provider_class)
-          provider_class.new(config)
-        rescue Ask::ModelNotFound, Ask::UnknownProvider
-          Ask::Provider.resolve(:openai).new(
-            build_provider_config(Ask::Provider.resolve(:openai))
-          )
+          Ask::Document.new(content: document.content, metadata: document.metadata, id: id)
         end
 
-        def build_provider_config(provider_class)
-          config = Object.new
-          provider_class.configuration_requirements.each do |req|
-            val = ENV["#{req.to_s.upcase}"] || ENV["#{provider_class.slug.upcase}_#{req.to_s.upcase}"]
-            config.define_singleton_method(req) { val } if val
-          end
-          provider_class.configuration_options.each do |opt|
-            val = ENV["#{opt.to_s.upcase}"] || ENV["#{provider_class.slug.upcase}_#{opt.to_s.upcase}"]
-            config.define_singleton_method(opt) { val } if val
-          end
-          config
-        end
+        # --- Search ---
 
-        def embed_query
-          model = @embedding_model || "text-embedding-3-small"
-          provider = resolve_embedding_provider(model)
-          raw = provider.embed("query", model: model)
+        def search_by_vector(vector, limit:, filter: nil, min_score: nil, mmr: false,
+          diversity_bonus: Retrieval::MaxMarginalRelevance::DEFAULT_DIVERSITY_BONUS)
+          candidates = @mutex.synchronize do
+            @entries.each_value.filter_map do |entry|
+              next if filter && !matches_filter?(entry.document.metadata, filter)
 
-          if raw.is_a?(Ask::Result)
-            raw = raw.output
-          end
+              score = cosine_similarity(vector, entry.vector)
+              next if min_score && score < min_score
 
-          Array(raw).flatten.map(&:to_f)
-        rescue StandardError => e
-          raise EmbeddingError, "Failed to embed query with model #{model}: #{e.message}"
-        end
+              Retrieval::MaxMarginalRelevance::Candidate.new(
+                document: build_result(entry, score: score),
+                relevance: score,
+                vector: entry.vector
+              )
+            end
+          end
 
-        def normalize_vectors(raw_vectors, texts)
-          result = if raw_vectors.is_a?(Ask::Result)
-            raw_vectors.output
+          if mmr
+            Retrieval::MaxMarginalRelevance.select(candidates, limit: limit, diversity_bonus: diversity_bonus)
           else
-            raw_vectors
+            candidates.sort_by { |candidate| [-candidate.relevance, candidate.document.id.to_s] }
+              .first(limit)
+              .map(&:document)
           end
+        end
 
-          result = Array(result)
+        def keyword_search(query, limit:, filter: nil)
+          pairs = @mutex.synchronize do
+            @entries.each_value.filter_map do |entry|
+              next if filter && !matches_filter?(entry.document.metadata, filter)
 
-          vectors = if result.length == 1 && texts.length == 1
-            [Array(result.first).map(&:to_f)]
-          elsif result.first.is_a?(Array)
-            result.map { |v| Array(v).map(&:to_f) }
-          else
-            dims = @embedding_dimensions || result.length / texts.length
-            result.each_slice(dims).to_a
+              [entry.id, entry.document]
+            end
+          end
+          return [] if pairs.empty?
+
+          bm25 = Retrieval::BM25.new(pairs.map { |(id, document)| [id, document.content] })
+          by_id = pairs.to_h
+
+          bm25.search(query, limit: limit).map do |(id, score)|
+            document = by_id[id]
+            Ask::Document.new(
+              content: document.content,
+              metadata: document.metadata.merge(lexical_score: score),
+              id: id
+            )
           end
-
-          @embedding_dimensions ||= vectors.first&.length
-          vectors
         end
 
         # --- Filtering ---
@@ -202,86 +221,16 @@ module Ask
 
         # --- Result building ---
 
-        def build_result(entry, score:, mmr_score: nil)
-          meta = entry.document.metadata.merge(score: score)
-          meta[:mmr_score] = mmr_score if mmr_score
+        def build_result(entry, score:)
           Ask::Document.new(
             content: entry.document.content,
-            metadata: meta,
-            id: entry.document.id
+            metadata: entry.document.metadata.merge(score: score),
+            id: entry.id
           )
         end
 
-        # --- MMR (Max Marginal Relevance) ---
-
-        def apply_mmr(candidates, query_vector, limit, diversity_bonus)
-          selected = []
-          remaining = candidates.dup
-
-          limit = [limit, remaining.size].min
-
-          # Pick the first result by relevance
-          first = remaining.shift
-          selected << build_result(first[0], score: first[1], mmr_score: first[1])
-
-          while selected.size < limit && remaining.any?
-            best_idx = 0
-            best_mmr = -Float::INFINITY
-
-            remaining.each_with_index do |(entry, rel_score), idx|
-              # Find max similarity to any already-selected document
-              max_sim_to_selected = selected.map do |sel_doc|
-                sel_entry = @entries.values.find { |e| e.document.content == sel_doc.content }
-                sel_entry ? cosine_similarity(entry.vector, sel_entry.vector) : 0.0
-              end.max
-
-              mmr_score = rel_score - diversity_bonus * max_sim_to_selected
-
-              if mmr_score > best_mmr
-                best_mmr = mmr_score
-                best_idx = idx
-              end
-            end
-
-            entry, rel_score = remaining.delete_at(best_idx)
-            selected << build_result(entry, score: rel_score, mmr_score: best_mmr)
-          end
-
-          selected
-        end
-
-        # --- Cosine Similarity ---
-
-        def cosine_similarity(a, b)
-          return 0.0 if a.empty? || b.empty? || a.length != b.length
-
-          # Use Matrix for vectorized computation if available
-          if defined?(Matrix)
-            begin
-              va = Matrix.row_vector(a)
-              vb = Matrix.row_vector(b)
-              dot = (va * vb.transpose)[0, 0]
-              norm_a = Math.sqrt((va * va.transpose)[0, 0])
-              norm_b = Math.sqrt((vb * vb.transpose)[0, 0])
-              denom = norm_a * norm_b
-              return denom > 0 ? dot / denom : 0.0
-            rescue StandardError
-              # Fall through to pure Ruby
-            end
-          end
-
-          dot = 0.0
-          norm_a = 0.0
-          norm_b = 0.0
-
-          a.zip(b).each do |ai, bi|
-            dot += ai * bi
-            norm_a += ai * ai
-            norm_b += bi * bi
-          end
-
-          denom = Math.sqrt(norm_a) * Math.sqrt(norm_b)
-          denom > 0 ? dot / denom : 0.0
+        def cosine_similarity(left, right)
+          Retrieval::MaxMarginalRelevance.cosine_similarity(left, right)
         end
       end
     end