ask-rag 0.1.0 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 908700d8186c0433071a87f8b6f28ec3d60f857ef9bd0a5cf474a1804f888251
-  data.tar.gz: 9abbec0e77ea2ca596b065c67af63d92ea373da22dd33f418442119a0d32defd
+  metadata.gz: a55e3f456d3e444c8aa58ba2b4c0f909de86bffe87e3e08a0a50becdddd2232e
+  data.tar.gz: fd944ef839c03e0a5bf2b597e4867f05c835d62e6a3433529d8e5bfb4bc9ad54
 SHA512:
-  metadata.gz: 20d21031b52e4b0594f5485bbd8e709a71fa716115676ad493ebee998085e4ce8a378d4023a94b3c9587221837c4fcd18488c09c270c306714b4ff16ab9e7b1c
-  data.tar.gz: 2b2a31e86ae1d5adceb08782d186db4347b57aa1b15e899e8210c20dd65732c3ba7e32ff79b1109e8fe958b77ddd8e2a91361f0c389c14769a70f0bbb2bad323
+  metadata.gz: 698e26ae24b01752d2032ec7345b3671902e1ead593685902b308ebf81b3fdde203581184686e8a95deb02a981a42e46cf99cfb9be7b2a6bd3977f182228bd90
+  data.tar.gz: 722a7d7ac59603a15932adf20e4a99452f7f031e6f0623f63645701e36c1350730bc331e024666da953f185ac101bec115913fa53cc343f141e6a4eb0d220aba

data/CHANGELOG.md

@@ -1,3 +1,47 @@
+## [0.2.0] — 2026-07-26
+
+### Fixed
+
+- **Embedding model mismatch** — InMemory store now tracks the model used during `add` and reuses it for query embedding. Previously `embed_query` hardcoded `"text-embedding-3-small"`, which produced wrong-dimension vectors if documents were stored with a different model. Now raises `Ask::RAG::EmbeddingError` if query embedding fails instead of silently returning a zero vector.
+
+- **Batch embedding** — `add` now actually uses the `batch_size` parameter to split large document sets into multiple embed API calls instead of sending everything at once.
+
+### Added
+
+- **Metadata filtering** — `similarity_search` now accepts a `filter:` parameter. Only entries whose metadata matches all key/value pairs are considered for search results. Works for both InMemory and PGVector stores.
+
+  ```ruby
+  store.similarity_search("query", filter: { section: "api" })
+  ```
+
+- **MMR (Max Marginal Relevance)** — `similarity_search` now accepts `mmr: true` and `diversity_bonus:` parameters. MMR diversifies search results by penalizing documents too similar to already-selected ones. Results include both `:score` and `:mmr_score` in metadata.
+
+  ```ruby
+  store.similarity_search("query", limit: 5, mmr: true, diversity_bonus: 0.3)
+  ```
+
+- **Directory loader** — `Ask::RAG::Loader::Directory` auto-detects file types by extension and dispatches to the appropriate loader. Walks directories recursively, skips `.git`, `node_modules`, and other non-source directories. Supports optional glob pattern.
+
+  ```ruby
+  loader = Ask::RAG::Loader::Directory.new
+  docs = loader.load("docs/")
+
+  loader = Ask::RAG::Loader::Directory.new(pattern: "**/*.md")
+  ```
+
+- **Faster cosine similarity** — Uses Ruby's `matrix` gem for vectorized computation when available, with a pure Ruby fallback.
+
+- **`Ask::RAG::EmbeddingError`** — raised when embedding fails during query, with a clear error message.
+
+### Changed
+
+- `VectorStore#similarity_search` signature updated to accept `filter:`, `mmr:`, and `diversity_bonus:` keyword arguments.
+
+### Tested
+
+- 60 tests, 148 assertions, 0 failures
+- New tests for: metadata filtering, MMR scoring, directory loader, batch embedding, embedding model tracking, EmbeddingError
+
 ## [0.1.0] — 2026-07-26
 
 ### Added

data/README.md

@@ -52,6 +52,7 @@ Load files into `Ask::Document` objects (text + metadata).
 | `Ask::RAG::Loader::CSV` | CSV (1 doc/row) | None (`csv` gem) |
 | `Ask::RAG::Loader::HTML` | HTML (strips script/style/nav) | `nokogiri` |
 | `Ask::RAG::Loader::PDF` | PDF (1 doc/page) | `pdf-reader` |
+| `Ask::RAG::Loader::Directory` | Auto-detects file types by extension | None (uses other loaders internally) |
 
 Loaders without their dependency gem installed are skipped silently — install the gem when you need the format.
 
@@ -113,6 +114,31 @@ results.each do |doc|
 end
 ```
 
+### Metadata filtering
+
+Filter results by document metadata:
+
+```ruby
+results = store.similarity_search(
+  "authentication",
+  filter: { source: "api_docs.md", version: "2.0" }
+)
+```
+
+### MMR (diversified search)
+
+Avoid getting 5 near-identical chunks by applying Max Marginal Relevance:
+
+```ruby
+results = store.similarity_search(
+  "authentication",
+  limit: 5,
+  mmr: true,
+  diversity_bonus: 0.3
+)
+# Results include both :score and :mmr_score in metadata
+```
+
 ### High-Level RAG Query
 
 Retrieve + prompt + answer in one call:

data/lib/ask/rag/loader/directory.rb

@@ -0,0 +1,105 @@
+# frozen_string_literal: true
+
+module Ask
+  module RAG
+    module Loader
+      # Walks a directory tree and loads all supported file types.
+      #
+      # Automatically detects file types by extension and dispatches to the
+      # appropriate loader. Ignores common non-source directories and files.
+      #
+      # @example
+      #   loader = Ask::RAG::Loader::Directory.new
+      #   docs = loader.load("path/to/docs/")
+      #   # Loads all .md, .txt, .html, .csv, .pdf files recursively
+      #
+      # @example With glob pattern
+      #   loader = Ask::RAG::Loader::Directory.new(pattern: "**/*.md")
+      #   docs = loader.load("docs/")
+      #   # Only loads Markdown files
+      #
+      class Directory < Base
+        EXTENSION_MAP = {
+          ".txt" => :text,
+          ".md" => :markdown,
+          ".markdown" => :markdown,
+          ".html" => :html,
+          ".htm" => :html,
+          ".csv" => :csv,
+          ".pdf" => :pdf
+        }.freeze
+
+        # Directories to skip when walking (matches basename).
+        SKIP_DIRS = %w[
+          .git .svn .hg .bzr _darcs node_modules .bundle vendor
+          .yardoc coverage tmp log public/assets
+        ].freeze
+
+        # @param pattern [String, nil] optional glob pattern (e.g., "**/*.md")
+        #   When nil, walks all files with known extensions.
+        # @param skip_dirs [Array<String>] directory names to skip
+        def initialize(pattern: nil, skip_dirs: SKIP_DIRS)
+          @pattern = pattern
+          @skip_dirs = skip_dirs.map(&:freeze).freeze
+        end
+
+        # Load all supported files from a directory tree.
+        #
+        # @param path [String] root directory path
+        # @return [Array<Ask::Document>] all loaded documents
+        def load(path)
+          docs = []
+          files(path).each do |file|
+            loader = loader_for(file)
+            next unless loader
+
+            docs.concat(loader.load(file))
+          rescue StandardError => e
+            warn "[ask-rag] Skipping #{file}: #{e.class}: #{e.message}"
+          end
+          docs
+        end
+
+        private
+
+        def files(root)
+          if @pattern
+            Dir.glob(File.join(root, @pattern)).select { |f| File.file?(f) }
+          else
+            all_files(root)
+          end
+        end
+
+        def all_files(root)
+          result = []
+          Dir.entries(root).each do |entry|
+            next if entry.start_with?(".")
+            next if @skip_dirs.include?(entry)
+
+            full = File.join(root, entry)
+            if File.directory?(full)
+              result.concat(all_files(full))
+            elsif File.file?(full)
+              result << full
+            end
+          end
+          result
+        end
+
+        def loader_for(file)
+          ext = File.extname(file).downcase
+          type = EXTENSION_MAP[ext]
+          return nil unless type
+
+          case type
+          when :text then Loader::Text.new
+          when :markdown then Loader::Markdown.new
+          when :html then defined?(Loader::HTML) ? Loader::HTML.new : nil
+          when :csv then Loader::CSV.new
+          when :pdf then defined?(Loader::PDF) ? Loader::PDF.new : nil
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -2,6 +2,9 @@
 
 module Ask
   module RAG
+    # Raised when embedding a query fails (e.g., no API key, network error).
+    class EmbeddingError < StandardError; end
+
     # Abstract base class for vector stores.
     #
     # Vector stores store embedded documents and provide similarity search.
@@ -29,8 +32,12 @@ module Ask
       #
       # @param query [String] the query text
       # @param limit [Integer] maximum number of results
-      # @return [Array<Ask::Document>] documents with a +score+ attribute added
-      def similarity_search(query, limit: 10)
+      # @param filter [Hash, nil] metadata filter — only entries whose metadata
+      #   matches all key/value pairs are considered
+      # @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)
         raise NotImplementedError
       end
 

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

@@ -14,55 +14,91 @@ module Ask
       #   store.add(chunks, model: "text-embedding-3-small")
       #   results = store.similarity_search("query", limit: 5)
       #
+      # @example With metadata filtering
+      #   results = store.similarity_search(
+      #     "query",
+      #     limit: 5,
+      #     filter: { source: "api_docs.md" }
+      #   )
+      #
+      # @example With MMR (diversified results)
+      #   results = store.similarity_search(
+      #     "query",
+      #     limit: 5,
+      #     mmr: true,
+      #     diversity_bonus: 0.5
+      #   )
+      #
       class InMemory < VectorStore
         Entry = Struct.new(:id, :document, :vector, keyword_init: true)
 
         def initialize
           @entries = {}
           @mutex = Mutex.new
+          @embedding_model = nil
+          @embedding_dimensions = nil
         end
 
+        # Add documents to the store.
+        #
+        # @param documents [Array<Ask::Document>] documents to add
+        # @param model [String] embedding model name
+        # @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 }
-          texts = documents.map(&:content)
-
-          # Embed using ask-llm-providers
           provider = resolve_embedding_provider(model)
-          raw_vectors = provider.embed(texts, model: model)
 
-          vectors = normalize_vectors(raw_vectors, texts)
+          @embedding_model = model
 
-          @mutex.synchronize do
-            documents.each_with_index do |doc, idx|
-              @entries[ids[idx]] = Entry.new(
-                id: ids[idx],
-                document: doc,
-                vector: vectors[idx]
-              )
+          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)
+
+            @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
+              end
             end
           end
-
-          ids
         end
 
-        def similarity_search(query, limit: 10)
-          query_vector = embed_query(query)
+        # Search for documents similar to the query.
+        #
+        # @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 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
+
+          candidates = @mutex.synchronize do
+            @entries.filter_map do |_id, entry|
+              next if filter && !matches_filter?(entry.document.metadata, filter)
 
-          scored = @mutex.synchronize do
-            @entries.map do |_id, entry|
               score = cosine_similarity(query_vector, entry.vector)
               [entry, score]
             end
           end
 
-          scored.sort_by! { |_, score| -score }
-          scored.first(limit).map do |entry, score|
-            doc = entry.document
-            Ask::Document.new(
-              content: doc.content,
-              metadata: doc.metadata.merge(score: score),
-              id: doc.id
-            )
+          candidates.sort_by! { |_, score| -score }
+
+          if mmr
+            apply_mmr(candidates, query_vector, limit, diversity_bonus)
+          else
+            candidates.first(limit).map do |entry, score|
+              build_result(entry, score: score)
+            end
           end
         end
 
@@ -82,13 +118,14 @@ module Ask
 
         private
 
+        # --- Embedding ---
+
         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
-          # Fallback: try to use OpenAI-compatible directly
           Ask::Provider.resolve(:openai).new(
             build_provider_config(Ask::Provider.resolve(:openai))
           )
@@ -107,18 +144,18 @@ module Ask
           config
         end
 
-        def embed_query(query)
-          provider = resolve_embedding_provider("text-embedding-3-small")
-          raw = provider.embed(query, model: "text-embedding-3-small")
+        def embed_query
+          model = @embedding_model || "text-embedding-3-small"
+          provider = resolve_embedding_provider(model)
+          raw = provider.embed("query", model: model)
 
           if raw.is_a?(Ask::Result)
             raw = raw.output
           end
 
           Array(raw).flatten.map(&:to_f)
-        rescue StandardError
-          # Return a zero vector as fallback (will produce no meaningful matches)
-          [0.0] * 256
+        rescue StandardError => e
+          raise EmbeddingError, "Failed to embed query with model #{model}: #{e.message}"
         end
 
         def normalize_vectors(raw_vectors, texts)
@@ -130,29 +167,97 @@ module Ask
 
           result = Array(result)
 
-          # If the provider returned a flat vector for a single document
-          if result.length == 1 && texts.length == 1
+          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
-            # Flattened — try to split by dimensions
-            dims = infer_dimensions(result)
+            dims = @embedding_dimensions || result.length / texts.length
             result.each_slice(dims).to_a
           end
+
+          @embedding_dimensions ||= vectors.first&.length
+          vectors
         end
 
-        def infer_dimensions(flat)
-          # Common embedding dimensions
-          [3072, 1536, 1024, 768, 512, 384, 256].each do |d|
-            return d if flat.length % d == 0 && flat.length / d > 0
+        # --- Filtering ---
+
+        def matches_filter?(metadata, filter)
+          filter.all? do |key, value|
+            metadata.key?(key) && metadata[key] == value
           end
-          flat.length
         end
 
+        # --- 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
+          Ask::Document.new(
+            content: entry.document.content,
+            metadata: meta,
+            id: entry.document.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

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

@@ -61,14 +61,22 @@ module Ask
           ids
         end
 
-        def similarity_search(query, limit: 10)
+        def similarity_search(query, limit: 10, filter: nil, mmr: false, diversity_bonus: 0.3)
           query_vector = embed_query(query)
           column = "#{@table_name}.#{@embedding_column}"
           model_class = @model_class
 
-          records = model_class
+          scope = model_class
             .select("#{@table_name}.*, 1 - (#{column} <=> '#{query_vector.to_s}') AS score")
             .where("#{column} IS NOT NULL")
+
+          if filter
+            filter.each do |key, value|
+              scope = scope.where("#{@metadata_column} @> ?", { key.to_s => value }.to_json)
+            end
+          end
+
+          records = scope
             .order(Arel.sql("#{column} <=> '#{query_vector.to_s}'"))
             .limit(limit)
 

data/lib/ask/rag/version.rb

@@ -2,6 +2,6 @@
 
 module Ask
   module RAG
-    VERSION = "0.1.0"
+    VERSION = "0.2.0"
   end
 end

data/lib/ask/rag.rb

@@ -3,11 +3,18 @@
 require_relative "rag/version"
 require_relative "rag/document"
 
+# Optional: Matrix for faster cosine similarity
+begin
+  require "matrix"
+rescue LoadError # rubocop:disable Lint/SuppressedException
+end
+
 # Loaders
 require_relative "rag/loader/base"
 require_relative "rag/loader/text"
 require_relative "rag/loader/markdown"
 require_relative "rag/loader/csv"
+require_relative "rag/loader/directory"
 
 # Optional loaders (loaded only when their dependency gem is available)
 begin

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: ask-rag
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 0.2.0
 platform: ruby
 authors:
 - Kaka Ruto
@@ -109,6 +109,7 @@ files:
 - lib/ask/rag/document.rb
 - lib/ask/rag/loader/base.rb
 - lib/ask/rag/loader/csv.rb
+- lib/ask/rag/loader/directory.rb
 - lib/ask/rag/loader/html.rb
 - lib/ask/rag/loader/markdown.rb
 - lib/ask/rag/loader/pdf.rb