ask-rag 0.2.2 → 0.4.0

data/lib/ask/rag/embeddings.rb

@@ -0,0 +1,138 @@
+# frozen_string_literal: true
+
+module Ask
+  module RAG
+    # Embedding for the vector stores: resolving a model name to a provider,
+    # calling it, and validating what comes back.
+    #
+    # Every store needs this and none of it is store-specific — the same
+    # request goes to the same provider whether the vectors land in a Hash or
+    # a Postgres table. It lives here so a store adds storage, not an
+    # embeddings client.
+    #
+    # Including classes provide `@embedding_model` (the model the corpus was
+    # embedded with, reused for queries) and may provide
+    # `@embedding_dimensions` to help split a flat response into vectors.
+    module Embeddings
+      # Embed a batch of documents.
+      #
+      # @param texts [Array<String>] texts to embed
+      # @param model [String, nil] embedding model name (default: configured
+      #   embedding_model)
+      # @return [Array<Array<Float>>] one vector per text
+      def embed_texts(texts, model: nil)
+        model ||= default_embedding_model
+        raw = embedding_provider(model).embed(texts, model: model)
+        normalize_vectors(unwrap_result(raw, "#{texts.length} document(s)", model), texts, model)
+      rescue EmbeddingError
+        raise
+      rescue StandardError => e
+        raise EmbeddingError, "Failed to embed #{texts.length} document(s) with model #{model}: #{e.message}"
+      end
+
+      # Embed one query with the model the corpus was embedded with.
+      #
+      # The model matters: a query embedded by a different model lands in a
+      # different vector space, where the similarity scores are meaningless
+      # rather than merely worse.
+      #
+      # @param text [String] the query
+      # @param model [String, nil] override the corpus model
+      # @return [Array<Float>]
+      def embed_query(text, model: nil)
+        model ||= @embedding_model || default_embedding_model
+        raw = embedding_provider(model).embed(text.to_s, model: model)
+
+        Array(unwrap_result(raw, "query", model)).flatten.map(&:to_f)
+      rescue EmbeddingError
+        raise
+      rescue StandardError => e
+        raise EmbeddingError, "Failed to embed query with model #{model}: #{e.message}"
+      end
+
+      DEFAULT_EMBEDDING_MODEL = "text-embedding-3-small"
+
+      private
+
+      # The configured embedding model, or the historical default when the
+      # gem is used without configuration. This indirection keeps
+      # `model:` optional on every store method while preserving the
+      # previous default for plain Ruby users.
+      def default_embedding_model
+        Ask::RAG.configuration.embedding_model || DEFAULT_EMBEDDING_MODEL
+      end
+
+      # Providers return either a bare payload or an Ask::Result wrapping one.
+      # A failed Result must raise — coercing it to floats would silently make
+      # every score meaningless.
+      def unwrap_result(raw, subject, model)
+        return raw unless raw.is_a?(Ask::Result)
+
+        raise EmbeddingError, "Failed to embed #{subject} with model #{model}: #{raw.error}" unless raw.ok?
+
+        raw.output
+      end
+
+      # Providers vary in what they return for a batch: an array of vectors,
+      # or one flat array of floats that has to be sliced. The count is
+      # checked either way, because a mismatch would attach the wrong vector
+      # to a document and no later stage could detect it.
+      def normalize_vectors(result, texts, model)
+        vectors = split_vectors(Array(result), texts.length)
+        verify_vector_count(vectors, texts.length, model)
+
+        @embedding_dimensions ||= vectors.first&.length
+        vectors
+      end
+
+      def split_vectors(result, text_count)
+        return result.map { |vector| Array(vector).map(&:to_f) } if result.first.is_a?(Array)
+
+        flat = Array(result).flatten.map(&:to_f)
+        # A single text embeds to one flat vector. Otherwise the response is
+        # a batch flattened into one array and has to be sliced back into
+        # vectors, using the dimension seen on an earlier call when there is
+        # one — the provider does not label the boundaries.
+        return [flat] if text_count <= 1
+
+        dimensions = @embedding_dimensions || (flat.length / text_count)
+        return [flat] if dimensions <= 0
+
+        flat.each_slice(dimensions).to_a
+      end
+
+      def verify_vector_count(vectors, text_count, model)
+        return if vectors.length == text_count
+
+        raise EmbeddingError,
+          "model #{model} returned #{vectors.length} vector(s) for #{text_count} text(s)"
+      end
+
+      def embedding_provider(model)
+        info = Ask::ModelCatalog.find(model)
+        klass = Ask::Provider.resolve(info.provider)
+        klass.new(provider_config(klass))
+      rescue Ask::ModelNotFound, Ask::UnknownProvider
+        klass = Ask::Provider.resolve(:openai)
+        klass.new(provider_config(klass))
+      end
+
+      # Providers declare what they need as configuration; each
+      # requirement resolves from Rails credentials first
+      # (`ask.openai_api_key`), then the environment variable named after
+      # it, namespaced by provider first (`OPENAI_API_KEY`) then bare
+      # (`API_KEY`). Non-Rails apps skip the credentials lookup entirely.
+      def provider_config(klass)
+        config = Object.new
+        slug = klass.slug
+
+        (klass.configuration_requirements + klass.configuration_options).each do |option|
+          value = Credentials.fetch(slug, option)
+          config.define_singleton_method(option) { value } if value
+        end
+
+        config
+      end
+    end
+  end
+end

data/lib/ask/rag/query.rb

@@ -8,13 +8,22 @@ module Ask
     # answer a question based on those documents.
     #
     # @example
-    #   answer = Ask::RAG.query(
+    #   answer = Ask::RAG::Query.query(
     #     store: my_vector_store,
     #     question: "What is the default timeout?",
     #     model: "gpt-4o"
     #   )
     #   puts answer.content
     #
+    # @example Hybrid retrieval with a relevance floor
+    #   answer = Ask::RAG::Query.query(
+    #     store: store,
+    #     question: "What is the default timeout?",
+    #     model: "gpt-4o",
+    #     hybrid: true,
+    #     min_score: 0.35
+    #   )
+    #
     module Query
       module_function
 
@@ -23,34 +32,58 @@ module Ask
       #
       # @param store [Ask::RAG::VectorStore] the vector store to search
       # @param question [String] the user's question
-      # @param model [String] LLM model to use for answering
+      # @param model [String, nil] LLM model to use for answering (default:
+      #   configured chat_model)
       # @param limit [Integer] number of documents to retrieve (default: 5)
+      # @param min_score [Float, nil] drop retrieved documents scoring below
+      #   this — the main defense against answering from irrelevant context
+      # @param hybrid [Boolean] retrieve with hybrid search (dense + lexical)
+      #   instead of dense alone; recommended for anything a person reads
       # @param system_prompt [String, nil] custom system prompt override
       # @param provider [Symbol, nil] optional provider override
       # @return [Ask::Document, nil] the answer with metadata including sources
-      def query(store:, question:, model: nil, limit: 5, system_prompt: nil, provider: nil)
-        documents = store.similarity_search(question, limit: limit)
+      def query(store:, question:, model: nil, limit: 5, min_score: nil, hybrid: false,
+        system_prompt: nil, provider: nil)
+        documents = retrieve(store, question, limit: limit, min_score: min_score, hybrid: hybrid)
 
+        # No context is a legitimate answer to return rather than a nil the
+        # caller has to distinguish from "the model failed": the question may
+        # simply be outside this corpus, and answering it from nothing is how
+        # a RAG system invents an answer.
         return nil if documents.empty?
 
-        context = documents.map.with_index(1) do |doc, i|
-          "[#{i}] #{doc.content}\n"
-        end.join("\n")
-
-        prompt = system_prompt || build_default_prompt
-        full_prompt = "#{prompt}\n\nContext:\n#{context}\n\nQuestion: #{question}"
-
-        answer = ask_llm(full_prompt, model: model, provider: provider)
+        model ||= Ask::RAG.configuration.chat_model
+        answer = ask_llm(build_prompt(documents, question, system_prompt), model: model, provider: provider)
 
         Ask::Document.new(
           content: answer,
           metadata: {
-            sources: documents.map { |d| d.metadata[:source] }.compact.uniq,
+            sources: documents.map { |document| document.metadata[:source] }.compact.uniq,
             model: model
           }
         )
       end
 
+      # Retrieve the context documents a question will be answered from.
+      # Exposed so an agent can inspect retrieval without calling an LLM.
+      #
+      # @return [Array<Ask::Document>]
+      def retrieve(store, question, limit: 5, min_score: nil, hybrid: false)
+        if hybrid
+          store.hybrid_search(question, limit: limit, min_score: min_score)
+        else
+          store.similarity_search(question, limit: limit, min_score: min_score)
+        end
+      end
+
+      def build_prompt(documents, question, system_prompt = nil)
+        context = documents.map.with_index(1) do |document, index|
+          "[#{index}] #{document.content}\n"
+        end.join("\n")
+
+        "#{system_prompt || build_default_prompt}\n\nContext:\n#{context}\n\nQuestion: #{question}"
+      end
+
       # Generate an answer using the LLM.
       def ask_llm(prompt, model: nil, provider: nil) # rubocop:disable Metrics/MethodLength
         require "ask/agent" unless defined?(Ask::Agent)
@@ -73,13 +106,13 @@ module Ask
       end
 
       def fallback_ask(prompt, model)
-        model ||= "gpt-4o"
+        model ||= Ask::RAG.configuration.chat_model || "gpt-4o"
         info = Ask::ModelCatalog.find(model)
         klass = Ask::Provider.resolve(info.provider)
 
         config = Object.new
         klass.configuration_requirements.each do |req|
-          val = ENV["#{req.to_s.upcase}"]
+          val = Credentials.fetch(klass.slug, req)
           config.define_singleton_method(req) { val } if val
         end
 

data/lib/ask/rag/railtie.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+require "rails/railtie"
+
+module Ask
+  module RAG
+    # Rails integration for ask-rag, shipped inside this gem — there is no
+    # separate wrapper gem.
+    #
+    # This file is required from `lib/ask/rag.rb` only when `rails/railtie`
+    # is loadable, so plain Ruby apps never pay for it and never need the
+    # `railties` gem installed. (It stays a development dependency, the
+    # same arrangement Kamal uses.)
+    #
+    # @example config/application.rb (or an environment file)
+    #   config.ask_rag.embedding_model = "text-embedding-3-small"
+    #   config.ask_rag.chat_model = "gpt-4o"
+    #   config.ask_rag.table_name = :embeddings
+    class Railtie < ::Rails::Railtie
+      config.ask_rag = ActiveSupport::OrderedOptions.new
+
+      initializer "ask_rag.configure" do |app|
+        Ask::RAG.configuration.apply_ordered_options(app.config.ask_rag)
+      end
+    end
+  end
+end

data/lib/ask/rag/retrieval/bm25.rb

@@ -0,0 +1,137 @@
+# frozen_string_literal: true
+
+module Ask
+  module RAG
+    module Retrieval
+      # Okapi BM25 over an in-memory document set.
+      #
+      # This is the lexical half of hybrid retrieval: the exact terms —
+      # error codes, method names, identifiers, product numbers — that a
+      # dense embedding blurs away. It plays the same role for the
+      # InMemory store that full-text search plays for PGVector, and both
+      # exist for the same reason: neither half is reliable alone.
+      #
+      # Documents are tokenized and scored per call. The InMemory store
+      # holds every vector in a Ruby Hash by definition, so maintaining an
+      # inverted index in step with add/delete would add state without
+      # buying anything at that scale.
+      #
+      # @example
+      #   bm25 = Ask::RAG::Retrieval::BM25.new([["a", "Ruby on Rails"], ["b", "Python"]])
+      #   bm25.search("rails") # => [["a", 0.69...]]
+      #
+      class BM25
+        K1 = 1.2
+        B = 0.75
+
+        # Words carrying no retrieval signal: a question matched only on
+        # these has matched nothing, and returning a document for it is worse
+        # than returning none — the caller cannot tell the difference between
+        # "this page is about your question" and "this page also contains the
+        # word 'how'".
+        #
+        # IDF alone is supposed to handle this and does not at the corpus
+        # sizes this gem serves: in a folder of twenty documents a stopword
+        # appears in enough of them to score low but still positive, so it
+        # survives into the results. The standard remedy is to drop them
+        # before scoring, which is what this list does.
+        STOPWORDS = %w[
+          a about above after again against all am an and any are as at
+          be because been before being below between both but by can cannot
+          could did do does doing down during each few for from further had
+          has have having he her here hers herself him himself his how i if
+          in into is it its itself just me more most my myself no nor not of
+          off on once only or other ought our ours ourselves out over own
+          same she should so some such than that the their theirs them
+          themselves then there these they this those through to too under
+          until up very was we were what when where which while who whom why
+          will with would you your yours yourself yourselves
+        ].freeze
+
+        # @param documents [Array<Array(Object, String)>] id/content pairs
+        # @param k1 [Float] term-frequency saturation
+        # @param b [Float] length-normalization strength (0 = none, 1 = full)
+        # @param stopwords [Array<String>] terms to drop before scoring;
+        #   pass [] to keep every token (the raw Okapi BM25 behaviour)
+        def initialize(documents, k1: K1, b: B, stopwords: STOPWORDS) # rubocop:disable Naming/MethodParameterName -- k1 and b are the algorithm's own names
+          @documents = documents
+          @k1 = k1
+          @b = b
+          @stopwords = stopwords.to_h { |word| [word, true] }
+          # Documents are tokenized the same way queries are, so a query term
+          # can only match a token the index actually holds. Keeping stopwords
+          # in the documents would also inflate every length, weakening the
+          # length normalization that `b` controls.
+          @tokens = documents.map { |(_id, content)| significant_terms(content) }
+          @lengths = @tokens.map(&:length)
+          @average_length = [@lengths.sum.to_f / [@lengths.length, 1].max, 1.0].max
+          @document_frequency = build_document_frequency
+        end
+
+        # @param query [String] the search query
+        # @param limit [Integer, nil] maximum results (nil = all scoring above zero)
+        # @return [Array<Array(Object, Float)>] id/score pairs, best first
+        def search(query, limit: nil)
+          terms = significant_terms(query).uniq
+          return [] if terms.empty?
+
+          scored = @documents.each_index.filter_map do |index|
+            score = score_document(index, terms)
+            [@documents[index].first, score] if score.positive?
+          end
+
+          scored.sort_by! { |(id, score)| [-score, id.to_s] }
+          limit ? scored.first(limit) : scored
+        end
+
+        private
+
+        def score_document(index, terms)
+          counts = @tokens[index].tally
+          length_ratio = @lengths[index] / @average_length
+
+          terms.sum do |term|
+            frequency = counts[term]
+            next 0 unless frequency
+
+            # The Okapi BM25 term weight: saturation on term frequency
+            # (`k1`) and normalization by document length (`b`).
+            saturation = frequency * (@k1 + 1)
+            normalization = frequency + (@k1 * (1 - @b + (@b * length_ratio)))
+
+            idf(term) * (saturation / normalization)
+          end
+        end
+
+        # The (1 + ...) form keeps IDF positive even for terms that appear in
+        # most documents — the raw log form goes negative there, which would
+        # rank a common word below no match at all.
+        def idf(term)
+          total = @documents.length
+          documents_with_term = @document_frequency[term]
+
+          Math.log(1 + ((total - documents_with_term + 0.5) / (documents_with_term + 0.5)))
+        end
+
+        def build_document_frequency
+          frequency = Hash.new(0)
+          @tokens.each { |tokens| tokens.uniq.each { |term| frequency[term] += 1 } }
+          frequency
+        end
+
+        # Words, lowercased and stripped of punctuation. Digits are kept —
+        # an identifier like `error_4021` is exactly the kind of token the
+        # lexical half exists to match — and underscores separate, which is
+        # how Postgres tokenizes them too. Without that, the same query
+        # would match in PGVector and not in this store.
+        def tokenize(text)
+          text.to_s.downcase.scan(/[a-z0-9]+/)
+        end
+
+        def significant_terms(text)
+          tokenize(text).reject { |term| @stopwords.key?(term) }
+        end
+      end
+    end
+  end
+end

data/lib/ask/rag/retrieval/max_marginal_relevance.rb

@@ -0,0 +1,114 @@
+# frozen_string_literal: true
+
+module Ask
+  module RAG
+    module Retrieval
+      # Max Marginal Relevance: picks results that are relevant *and* unlike
+      # each other.
+      #
+      # Plain top-k returns five paraphrases of the same paragraph. The
+      # failure is worst exactly when retrieval works well, because
+      # near-duplicate chunks all score highly and crowd out the one passage
+      # that answers a different part of the question. MMR trades a little
+      # relevance for coverage: each pick maximizes
+      #
+      #   relevance - diversity_bonus * (similarity to the closest pick so far)
+      #
+      # diversity_bonus = 0 is plain relevance ranking, 1 is novelty above
+      # all. Both stores share this implementation so their results for the
+      # same parameters are the same shape — a documented option that
+      # silently does nothing in one backend is worse than no option.
+      module MaxMarginalRelevance
+        DEFAULT_DIVERSITY_BONUS = 0.3
+
+        # One retrieval hit on its way into selection.
+        # @!attribute document
+        #   @return [Ask::Document] the hit, already carrying its :score
+        # @!attribute relevance
+        #   @return [Float] the retrieval score that ranked it
+        # @!attribute vector
+        #   @return [Array<Float>, nil] the stored embedding, for pairwise similarity
+        # @!attribute mmr_score
+        #   @return [Float, nil] the adjusted score it was selected with
+        Candidate = Struct.new(:document, :relevance, :vector, :mmr_score, keyword_init: true)
+
+        # @param candidates [Array<Candidate>] hits to choose from, any order
+        # @param limit [Integer] how many to select
+        # @param diversity_bonus [Float] 0 = pure relevance, 1 = pure diversity
+        # @param similarity [#call, nil] pairwise vector similarity
+        # @return [Array<Ask::Document>] selected, best first, with :mmr_score
+        def self.select(candidates, limit:, diversity_bonus: DEFAULT_DIVERSITY_BONUS, similarity: nil)
+          # Ties broken by id, because two equally relevant documents must
+          # select in the same order on every run — an unstable sort makes
+          # retrieval results irreproducible.
+          remaining = candidates.sort_by { |candidate| [-candidate.relevance, candidate.document.id.to_s] }
+          return [] if remaining.empty? || limit <= 0
+
+          similarity ||= method(:cosine_similarity)
+          selected = [remaining.shift.tap { |first| first.mmr_score = first.relevance }]
+
+          while selected.length < limit && remaining.any?
+            pick = take_best(remaining, selected, diversity_bonus, similarity)
+            selected << pick
+          end
+
+          selected.map { |candidate| with_mmr_score(candidate) }
+        end
+
+        # Removes and returns the candidate maximizing relevance minus
+        # redundancy against what is already selected.
+        def self.take_best(remaining, selected, diversity_bonus, similarity)
+          scored = remaining.each_with_index.map do |candidate, index|
+            penalty = diversity_bonus * closest_similarity(candidate, selected, similarity)
+            [index, candidate.relevance - penalty]
+          end
+
+          index, score = scored.max_by { |(position, value)| [value, -position] }
+          remaining.delete_at(index).tap { |pick| pick.mmr_score = score }
+        end
+
+        # The similarity to the single most similar pick, not the average:
+        # a document needs to be different from *everything* already chosen,
+        # and averaging lets one close neighbour hide behind several distant
+        # ones.
+        def self.closest_similarity(candidate, selected, similarity)
+          selected.map { |pick| similarity.call(candidate.vector, pick.vector) }.max || 0.0
+        end
+        private_class_method :closest_similarity
+
+        def self.with_mmr_score(candidate)
+          Ask::Document.new(
+            content: candidate.document.content,
+            metadata: candidate.document.metadata.merge(mmr_score: candidate.mmr_score),
+            id: candidate.document.id
+          )
+        end
+        private_class_method :with_mmr_score
+
+        # Cosine similarity between two vectors, with every degenerate input
+        # (nil, empty, mismatched length, zero norm) scoring 0. A query
+        # embedded by a different model than the corpus produces a length
+        # mismatch, which must read as "no similarity" rather than raising
+        # on a search path.
+        def self.cosine_similarity(left, right)
+          return 0.0 if left.nil? || right.nil? || left.empty? || right.empty? || left.length != right.length
+
+          dot = 0.0
+          left_norm = 0.0
+          right_norm = 0.0
+
+          left.each_index do |index|
+            a = left[index]
+            b = right[index]
+            dot += a * b
+            left_norm += a * a
+            right_norm += b * b
+          end
+
+          denominator = Math.sqrt(left_norm) * Math.sqrt(right_norm)
+          denominator.positive? ? dot / denominator : 0.0
+        end
+      end
+    end
+  end
+end

data/lib/ask/rag/retrieval/reciprocal_rank_fusion.rb

@@ -0,0 +1,72 @@
+# frozen_string_literal: true
+
+module Ask
+  module RAG
+    module Retrieval
+      # Reciprocal Rank Fusion: merges several ranked result lists into one
+      # by summing 1/(k + rank) per document across the lists that found it.
+      #
+      # RRF is used here rather than normalized score addition because it
+      # assumes nothing about the scores being comparable — a cosine
+      # similarity and a ts_rank are not on the same scale, and any
+      # normalization between them is a guess. It only needs each list to be
+      # ordered correctly, which both retrieval paths guarantee.
+      #
+      # k damps the top ranks so that no single list's first result can
+      # dominate the fused order. 60 is the constant from the original TREC
+      # paper and is not worth tuning before the retrieval itself is.
+      #
+      # A document found by both a keyword and a vector search outranks a
+      # document found by only one, which is the whole point: agreement
+      # between two independent signals is the strongest relevance evidence
+      # available without a reranker.
+      module ReciprocalRankFusion
+        K = 60
+
+        # @param result_sets [Array<Array<Ask::Document>>] ranked lists, best first
+        # @param limit [Integer] maximum fused results
+        # @param damping [Integer] rank-damping constant
+        # @return [Array<Ask::Document>] fused, best first, with :rrf_score in metadata
+        def self.fuse(result_sets, limit:, damping: K)
+          scores = Hash.new(0.0)
+          documents = {}
+
+          result_sets.each do |results|
+            Array(results).each_with_index do |document, rank|
+              key = key_for(document)
+              next unless key
+
+              scores[key] += 1.0 / (damping + rank + 1)
+              documents[key] ||= document
+            end
+          end
+
+          scores
+            .sort_by { |key, score| [-score, key.to_s] }
+            .first(limit)
+            .map { |key, score| with_fusion_score(documents[key], score) }
+        end
+
+        # Identity for fusion. Stores return documents carrying the id they
+        # were added under; content is the fallback for a document that has
+        # none, so a hit found by both paths still merges.
+        def self.key_for(document)
+          document.id || (document.content.to_s.empty? ? nil : document.content)
+        end
+        private_class_method :key_for
+
+        # The source list's own score is kept alongside :rrf_score — the
+        # fused number orders the result, the original says how strongly
+        # each half matched.
+        def self.with_fusion_score(document, score)
+          Ask::Document.new(
+            content: document.content,
+            metadata: document.metadata.merge(rrf_score: score),
+            id: document.id
+          )
+        end
+        private_class_method :with_fusion_score
+      end
+    end
+  end
+end