ask-rag 0.1.0 → 0.2.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 908700d8186c0433071a87f8b6f28ec3d60f857ef9bd0a5cf474a1804f888251
-  data.tar.gz: 9abbec0e77ea2ca596b065c67af63d92ea373da22dd33f418442119a0d32defd
+  metadata.gz: b79936f65b8b4b8f7cbfc1667703b645df6f411e84e7ee543071e4fabe9e6015
+  data.tar.gz: 474bac5cbc7f1d3f26b3ec0115591947577a3957aff4733a815f3f96b6f63dd6
 SHA512:
-  metadata.gz: 20d21031b52e4b0594f5485bbd8e709a71fa716115676ad493ebee998085e4ce8a378d4023a94b3c9587221837c4fcd18488c09c270c306714b4ff16ab9e7b1c
-  data.tar.gz: 2b2a31e86ae1d5adceb08782d186db4347b57aa1b15e899e8210c20dd65732c3ba7e32ff79b1109e8fe958b77ddd8e2a91361f0c389c14769a70f0bbb2bad323
+  metadata.gz: 3b7042d0c1f27aee6ceb4a21122f3c84691afb3b635a8b2f25f94a8c8cac3d633e09507f5a5b11347c241885016bdca1261be9322a0823034c4d97e11d8bce9e
+  data.tar.gz: 48b984042f7138262c13cd1458fa9191ed914909b35690b524cf3917e8837009644d71078908c47e1d1332b8895c4cee3ec5e453315442ef13091ff77420c6da

data/CHANGELOG.md

@@ -1,3 +1,78 @@
+## [0.2.1] — 2026-07-26
+
+### Fixed
+
+- **PGVector model tracking** — same fix as InMemory: now tracks the embedding model used during `add` and reuses it for query embedding. Also raises `EmbeddingError` instead of silently returning a zero vector on failure.
+
+- **PGVector document mutability** — no longer calls `define_singleton_method(:score)` on frozen `Ask::Document`. Score is now returned in `metadata[:score]` consistently with InMemory.
+
+- **PGVector embed_query** — now accepts the query text as an argument (was hardcoded to embed the literal string "query").
+
+### Added
+
+- **`similarity_search_by_vector`** — search using a raw vector instead of text. Available on both InMemory and PGVector stores. Accepts `limit:` and `filter:` params.
+
+  ```ruby
+  store.similarity_search_by_vector([0.1, 0.2, ...], limit: 5)
+  ```
+
+- **`Ask::RAG::Loader::JSON`** — loads JSON files. Supports arrays of objects (one document per object) and single objects. Configurable `content_key` and `metadata_keys`.
+
+  ```ruby
+  loader = Ask::RAG::Loader::JSON.new(content_key: "body")
+  docs = loader.load("articles.json")
+  ```
+
+- **Directory loader** now also recognizes `.json` files.
+
+### Tested
+
+- 65 tests, 163 assertions, 0 failures
+
+## [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,107 @@
+# 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,
+          ".json" => :json,
+          ".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 :json then Loader::JSON.new
+          when :pdf then defined?(Loader::PDF) ? Loader::PDF.new : nil
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+module Ask
+  module RAG
+    module Loader
+      # Loads a JSON file as documents.
+      #
+      # Supports two modes:
+      # 1. Array of objects — each object becomes one document.
+      #    Specify +content_key+ and +metadata_keys+ to map fields.
+      # 2. Single object — becomes one document.
+      #    Specify +content_key+ for which field holds the text.
+      #
+      # @example Array of objects
+      #   # [{ "title": "...", "body": "...", "tags": [...] }, ...]
+      #   loader = Ask::RAG::Loader::JSON.new(content_key: "body")
+      #   docs = loader.load("articles.json")
+      #
+      # @example With metadata mapping
+      #   loader = Ask::RAG::Loader::JSON.new(
+      #     content_key: "content",
+      #     metadata_keys: ["source", "date"]
+      #   )
+      #
+      class JSON < Base
+        # @param content_key [String] key whose value becomes document content
+        # @param metadata_keys [Array<String>, nil] keys to include in metadata
+        #   (nil = all keys except content_key)
+        def initialize(content_key: "content", metadata_keys: nil)
+          @content_key = content_key
+          @metadata_keys = metadata_keys
+        end
+
+        def load(path)
+          raw = ::JSON.parse(File.read(path))
+          items = raw.is_a?(Array) ? raw : [raw]
+
+          items.map { |item| build_document(item, path) }
+        end
+
+        private
+
+        def build_document(item, path)
+          content = item[@content_key].to_s
+          metadata = { source: path }
+
+          if @metadata_keys
+            @metadata_keys.each do |key|
+              metadata[key] = item[key] if item.key?(key)
+            end
+          else
+            item.each do |key, value|
+              next if key == @content_key
+              metadata[key] = value if value.is_a?(String) || value.is_a?(Numeric) || value == true || value == false
+            end
+          end
+
+          Ask::Document.new(content: content, metadata: metadata)
+        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,23 @@ 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
+
+      # Search by a raw vector instead of a text query.
+      # Useful when you've already embedded the query externally.
+      #
+      # @param vector [Array<Float>] the query vector
+      # @param limit [Integer] maximum number of results
+      # @param filter [Hash, nil] metadata filter
+      # @return [Array<Ask::Document>] documents with +:score+ in metadata
+      def similarity_search_by_vector(vector, limit: 10, filter: nil)
         raise NotImplementedError
       end
 

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

@@ -14,58 +14,106 @@ 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
+        end
+
+        # 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
+          search_by_vector(query_vector, limit: limit, filter: filter,
+                           mmr: mmr, diversity_bonus: diversity_bonus)
+        end
 
-          ids
+        def similarity_search_by_vector(vector, limit: 10, filter: nil)
+          search_by_vector(vector, limit: limit, filter: filter)
         end
 
-        def similarity_search(query, limit: 10)
-          query_vector = embed_query(query)
+        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)
 
-          scored = @mutex.synchronize do
-            @entries.map do |_id, entry|
-              score = cosine_similarity(query_vector, entry.vector)
+              score = cosine_similarity(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, vector, limit, diversity_bonus)
+          else
+            candidates.first(limit).map do |entry, score|
+              build_result(entry, score: score)
+            end
           end
         end
 
+        public
+
         def delete(ids)
           @mutex.synchronize do
             ids.each { |id| @entries.delete(id) }
@@ -82,13 +130,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 +156,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 +179,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

@@ -19,7 +19,7 @@ module Ask
       #
       # @example Usage
       #   store = Ask::RAG::VectorStore::PGVector.new(
-      #     table_name: :documents,
+      #     table_name: :embeddings,
       #     embedding_column: :embedding
       #   )
       #   store.add(chunks, model: "text-embedding-3-small")
@@ -38,18 +38,19 @@ module Ask
           @metadata_column = metadata_column.to_s
           @embedding_column = embedding_column.to_s
           @model_class = model_class || infer_model_class
+          @embedding_model = nil
         end
 
         def add(documents, model:, batch_size: 20)
           ids = []
-          model_class = @model_class
+          @embedding_model = model
 
           documents.each_slice(batch_size) do |batch|
             texts = batch.map(&:content)
             vectors = embed_texts(texts, model)
 
             batch.each_with_index do |doc, idx|
-              record = model_class.create!(
+              record = @model_class.create!(
                 @content_column => doc.content,
                 @metadata_column => doc.metadata,
                 @embedding_column => vectors[idx]
@@ -61,24 +62,39 @@ 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)
+          similarity_search_by_vector(query_vector, limit: limit, filter: filter)
+        end
+
+        def similarity_search_by_vector(vector, limit: 10, filter: nil)
+          vector_str = vector.is_a?(Array) ? "[#{vector.join(',')}]" : vector.to_s
           column = "#{@table_name}.#{@embedding_column}"
           model_class = @model_class
 
-          records = model_class
-            .select("#{@table_name}.*, 1 - (#{column} <=> '#{query_vector.to_s}') AS score")
+          scope = model_class
+            .select("#{@table_name}.*, 1 - (#{column} <=> '#{vector_str}') AS score")
             .where("#{column} IS NOT NULL")
-            .order(Arel.sql("#{column} <=> '#{query_vector.to_s}'"))
+
+          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} <=> '#{vector_str}'"))
             .limit(limit)
 
           records.map do |record|
-            doc = Ask::Document.new(
+            Ask::Document.new(
               content: record.send(@content_column),
-              metadata: (record.send(@metadata_column) || {}).merge(db_id: record.id)
+              metadata: (record.send(@metadata_column) || {}).merge(
+                score: record.score,
+                db_id: record.id
+              ),
+              id: record.id.to_s
             )
-            doc.define_singleton_method(:score) { record.score }
-            doc
           end
         end
 
@@ -109,13 +125,14 @@ module Ask
           Array(raw).map { |v| Array(v).map(&:to_f) }
         end
 
-        def embed_query(query)
-          provider = resolve_provider("text-embedding-3-small")
-          raw = provider.embed(query, model: "text-embedding-3-small")
+        def embed_query(query = "query")
+          model = @embedding_model || "text-embedding-3-small"
+          provider = resolve_provider(model)
+          raw = provider.embed(query, model: model)
           raw = raw.output if raw.is_a?(Ask::Result)
           Array(raw).flatten.map(&:to_f)
-        rescue StandardError
-          [0.0] * 1536
+        rescue StandardError => e
+          raise EmbeddingError, "Failed to embed query with model #{model}: #{e.message}"
         end
 
         def resolve_provider(model)

data/lib/ask/rag/version.rb

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

data/lib/ask/rag.rb

@@ -3,11 +3,19 @@
 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/json"
+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.1
 platform: ruby
 authors:
 - Kaka Ruto
@@ -109,7 +109,9 @@ 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/json.rb
 - lib/ask/rag/loader/markdown.rb
 - lib/ask/rag/loader/pdf.rb
 - lib/ask/rag/loader/text.rb