gemchat 0.1.0

data/lib/gemchat/store.rb

@@ -0,0 +1,555 @@
+# frozen_string_literal: true
+
+require "sqlite3"
+
+module Gemchat
+  # Local index store: SQLite with FTS5 for BM25. No server, no extension.
+  #
+  # Vectors are not stored here yet; the schema reserves `embedding` and
+  # `embed_model` so §7.1's store identity can refuse to open a store written by
+  # a different model rather than silently returning wrong neighbours.
+  class Store
+    SCHEMA_VERSION = "2"
+    BACKEND = "local"
+
+    # Fields that decide whether a store is readable at all (§7.1).
+    #
+    # `schema_version` is the hard gate: a newer or older layout cannot be
+    # queried correctly, and silently opening it produces wrong answers rather
+    # than an error. `gemchat_version` is recorded but tolerated — a point release
+    # that only adds a command must not strand a user's index. `embedding_model`
+    # and `embedding_dimensions` are placeholders for the embedding phase, where
+    # mixing vectors from different models *is* silently wrong.
+    IDENTITY_FIELDS = %w[schema_version gemchat_version embedding_model
+      embedding_dimensions backend].freeze
+
+    class IdentityError < Gemchat::Error; end
+
+    def self.default_path
+      File.join(Gemchat.root, "index.sqlite")
+    end
+
+    def initialize(path = self.class.default_path)
+      @path = path
+      @drift = []
+      FileUtils.mkdir_p(File.dirname(path))
+      @db = SQLite3::Database.new(path)
+      @db.results_as_hash = true
+      register_vector_functions
+      # Verified before migrating. migrate! rewrites schema_version on every
+      # open, so checking afterwards would compare the value against itself and
+      # always agree.
+      verify_identity!
+      migrate!
+      record_identity!
+    end
+
+    attr_reader :path
+
+    # Non-fatal differences worth surfacing in `status`.
+    attr_reader :drift
+
+    def verify_identity!
+      return unless table?("meta")
+
+      stored = meta("schema_version")
+      return if stored.nil? || stored == SCHEMA_VERSION
+
+      raise IdentityError, <<~MSG.strip
+        this index was written by a different gemchat (store schema #{stored}, this build expects #{SCHEMA_VERSION}).
+        Refusing to read it rather than returning wrong results.
+        Rebuild it with: gemchat reindex
+      MSG
+    end
+
+    def record_identity!
+      # Read before writing. Setting the field first and comparing after would
+      # compare the new value against itself and never report drift, which is
+      # the one thing this is here to notice.
+      previous_version = meta("gemchat_version")
+
+      set_meta("schema_version", SCHEMA_VERSION)
+      set_meta("gemchat_version", Gemchat::VERSION)
+      set_meta("backend", BACKEND)
+
+      @drift << "gemchat_version" if previous_version && previous_version != Gemchat::VERSION
+      drift
+    end
+
+    # SQLite has no dot product. This one computes it in C over the float32 blob,
+    # which matters: pulling every candidate vector into Ruby to multiply it
+    # would be a full scan of the corpus in the interpreter on every query.
+    #
+    # Dot product equals cosine similarity here because every stored vector is
+    # L2-normalised (Rllama::Model#embed normalises by default), so no magnitude
+    # term is needed and the result is directly comparable.
+    #
+    # The result is assigned to the FunctionProxy rather than returned. sqlite3
+    # 2.x wraps the block and reads `fp.result`; the block's own return value is
+    # discarded. Returning the float instead yields NULL for every row, and
+    # because the rank is negated on the way out, that surfaces as a silent 0.0
+    # for every result -- correct-looking output, no ordering at all.
+    def register_vector_functions
+      @db.create_function("dot", 2) do |fp, a, b|
+        left = a.to_s.unpack("e*")
+        right = b.to_s.unpack("e*")
+        fp.result = if left.empty? || left.size != right.size
+          nil
+        else
+          left.zip(right).sum { |x, y| x * y }
+        end
+      end
+    end
+
+    def table?(name)
+      !@db.get_first_value(
+        "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?", [name]
+      ).nil?
+    rescue SQLite3::SQLException
+      false
+    end
+
+    def migrate!
+      @db.execute_batch(<<~SQL)
+        CREATE TABLE IF NOT EXISTS meta (
+          key TEXT PRIMARY KEY,
+          value TEXT
+        );
+
+        CREATE TABLE IF NOT EXISTS gems (
+          name TEXT NOT NULL,
+          version TEXT NOT NULL,
+          source_type TEXT,
+          chunk_count INTEGER NOT NULL DEFAULT 0,
+          source_hash TEXT,
+          indexed_at TEXT,
+          PRIMARY KEY (name, version, source_type)
+        );
+
+        CREATE TABLE IF NOT EXISTS chunks (
+          id INTEGER PRIMARY KEY,
+          gem_name TEXT NOT NULL,
+          gem_version TEXT NOT NULL,
+          source_type TEXT NOT NULL,
+          title TEXT,
+          body TEXT,
+          class_name TEXT,
+          method_name TEXT,
+          method_type TEXT,
+          signature TEXT,
+          source_path TEXT,
+          source_line INTEGER,
+          content_digest TEXT NOT NULL,
+          embedding BLOB,
+          embed_model TEXT,
+          UNIQUE (gem_name, gem_version, source_type, class_name, method_name, title)
+        );
+
+        CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
+          title, body, signature, class_name, method_name,
+          content='chunks', content_rowid='id', tokenize='porter unicode61'
+        );
+      SQL
+    end
+
+    def set_meta(key, value)
+      @db.execute("INSERT INTO meta (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value", [key, value.to_s])
+    end
+
+    def meta(key)
+      row = @db.get_first_row("SELECT value FROM meta WHERE key = ?", [key])
+      row&.fetch("value")
+    end
+
+    def indexed?(name, version, source_type = "ri")
+      !@db.get_first_value(
+        "SELECT 1 FROM gems WHERE name = ? AND version = ? AND source_type = ?",
+        [name, version, source_type]
+      ).nil?
+    end
+
+    def gem_summary(name, version, source_type = "ri")
+      @db.get_first_row(
+        "SELECT * FROM gems WHERE name = ? AND version = ? AND source_type = ?",
+        [name, version, source_type]
+      )
+    end
+
+    # What a write actually did. "unchanged" is distinct from "indexed" with
+    # zero chunks: the first means the digest matched, the second would mean
+    # content was lost, and reporting them the same makes a broken run look
+    # like a successful no-op.
+    Outcome = Struct.new(:status, :chunks) do
+      def changed? = status != :unchanged
+    end
+
+    def replace_gem(name, version, source_type, chunks)
+      digest = Digest::SHA256.hexdigest(chunks.map { |c| c.content_digest }.join("|"))
+      existing = gem_summary(name, version, source_type)
+      return Outcome.new(:unchanged, 0) if existing && existing["source_hash"] == digest
+      status = existing ? :updated : :indexed
+
+      @db.transaction do
+        old_ids = @db.execute(
+          "SELECT id FROM chunks WHERE gem_name = ? AND gem_version = ? AND source_type = ?",
+          [name, version, source_type]
+        ).map { |r| r["id"] }
+        unless old_ids.empty?
+          @db.execute("DELETE FROM chunks_fts WHERE rowid IN (#{old_ids.map(&:to_i).join(",")})")
+          @db.execute("DELETE FROM chunks WHERE gem_name = ? AND gem_version = ? AND source_type = ?", [name, version, source_type])
+        end
+
+        now = Time.now.utc.iso8601
+        chunks.each do |chunk|
+          insert_sql = <<~SQL
+            INSERT OR IGNORE INTO chunks
+              (gem_name, gem_version, source_type, title, body, class_name, method_name,
+               method_type, signature, source_path, source_line, content_digest)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+          SQL
+          fts_sql = <<~SQL
+            INSERT INTO chunks_fts (rowid, title, body, signature, class_name, method_name)
+            VALUES (?, ?, ?, ?, ?, ?)
+          SQL
+          binds = [name, version, source_type, chunk.title, chunk.body,
+            chunk.class_name, chunk.method_name, chunk.method_type,
+            chunk.signature, chunk.source_path, chunk.source_line,
+            chunk.content_digest]
+          @db.execute(insert_sql, binds)
+          id = @db.last_insert_row_id
+          @db.execute(fts_sql, [id, chunk.title, chunk.body, chunk.signature, chunk.class_name, chunk.method_name])
+        end
+
+        @db.execute(<<~SQL, [name, version, source_type, chunks.size, digest, now])
+          INSERT INTO gems (name, version, source_type, chunk_count, source_hash, indexed_at)
+          VALUES (?, ?, ?, ?, ?, ?)
+          ON CONFLICT(name, version, source_type) DO UPDATE SET
+            chunk_count = excluded.chunk_count,
+            source_hash = excluded.source_hash,
+            indexed_at = excluded.indexed_at
+        SQL
+      end
+
+      Outcome.new(status, chunks.size)
+    end
+
+    # A gem has one row per source_type, so a plain COUNT(*) would report a gem
+    # twice as soon as it also had a README indexed.
+    def stats
+      {
+        gems: @db.get_first_value("SELECT COUNT(DISTINCT name || char(0) || version) FROM gems").to_i,
+        chunks: @db.get_first_value("SELECT COUNT(*) FROM chunks").to_i
+      }
+    end
+
+    def chunk_counts_by_source
+      @db.execute("SELECT source_type, COUNT(*) AS n FROM chunks GROUP BY source_type")
+        .to_h { |row| [row["source_type"], row["n"].to_i] }
+    end
+
+    # Explicit column list rather than a splat. `results_as_hash` makes every
+    # row carry integer keys too, so `{**row}` would leak 0, 1, 2... into the
+    # result and mix symbol and string key styles.
+    # `id` is first because fusion needs a stable key: BM25 and vector search
+    # return the same chunk as two separate rows, and the identity fields alone
+    # are not guaranteed unique. It is not in the CLI's RESULT_KEYS, so it does
+    # not leak into `--json` output.
+    COLUMNS = %w[id gem_name gem_version source_type title body class_name method_name signature source_path source_line].freeze
+
+    # Reciprocal rank fusion constant, matching gemchat_app's RRF_K and qmd.
+    RRF_K = 60
+
+    # FTS5 query syntax collides with Ruby naming: an unquoted `Puma::Server`
+    # parses as a column filter on `Puma` and fails with "no such column".
+    # Quoting each whitespace-separated term makes FTS5 treat it as a phrase and
+    # neutralises `:` , `*`, `-`, `"` and friends.
+    # The quoting is load-bearing: it is what stops FTS5 reading `Puma::Server` as
+    # a column filter, which fails with "no such column" and is a systematic
+    # collision with Ruby naming rather than a one-off.
+    #
+    # But joining the quoted terms with a *space* is FTS5's AND-of-a-phrase, and
+    # that was not a decision so much as a side effect of escaping. Measured: on
+    # the 236-chunk rake index, a phrase query returns **zero** rows for every
+    # natural-language question tried -- "how do I list all the available tasks",
+    # "syntax highlighting in the console" -- and a hit only when the query is a
+    # literal substring. So `search` was an exact-phrase finder wearing BM25's
+    # name, not a keyword search.
+    #
+    # :phrase  the quoted run. Exact substrings, high precision, near-zero recall.
+    #          The default for `search`, where a caller asking for
+    #          "Rake::Task#enhance" wants exactly that. Stopwords are never
+    #          dropped here: removing a word from a literal phrase changes what
+    #          that phrase matches, which is a different query, not a better one.
+    # :any     the same terms OR'd, with stopwords dropped. What a keyword search
+    #          means, and what the fusion's lexical arm needs. See §11.3 for what
+    #          it measures and §11.4 for the stopword pass and the eval gate.
+    def self.fts_query(raw, mode: :phrase)
+      terms = raw.to_s.split(/\s+/).reject(&:empty?)
+      return nil if terms.empty?
+
+      if mode == :any
+        kept = terms.uniq.reject { |term| Stopwords.only_stopwords?(term) }
+        # Every term was a stopword. Fall back to the full set rather than
+        # returning nil: the caller asked a real question, and an empty query
+        # would report "no results" as though the corpus were empty.
+        kept = terms.uniq if kept.empty?
+        kept.map { |term| %("#{term.gsub('"', '""')}") }.join(" OR ")
+      else
+        terms.uniq.map { |term| %("#{term.gsub('"', '""')}") }.join(" ")
+      end
+    end
+
+    # BM25's rank() is lower-is-better, so results are exposed as a positive
+    # `score` where higher is better. Every key is a String, including "score".
+    def search(query, limit: 10, min_score: 0.0, gem_names: nil, mode: :phrase)
+      match = self.class.fts_query(query, mode: mode)
+      return [] if match.nil?
+
+      binds = [match]
+      scope = ""
+      if gem_names && !gem_names.empty?
+        scope = "AND c.gem_name IN (#{(["?"] * gem_names.size).join(", ")})"
+        binds.concat(gem_names)
+      end
+      binds << limit
+
+      sql = <<~SQL
+        SELECT #{COLUMNS.map { |c| "c.#{c}" }.join(", ")},
+               bm25(chunks_fts) AS rank
+        FROM chunks_fts
+        JOIN chunks c ON c.id = chunks_fts.rowid
+        WHERE chunks_fts MATCH ? #{scope}
+        ORDER BY rank
+        LIMIT ?
+      SQL
+
+      @db.execute(sql, binds).map { |row|
+        COLUMNS.to_h { |c| [c, row[c]] }.merge("score" => -row["rank"].to_f)
+      }.select { |row| row["score"] >= min_score }
+    end
+
+    # --- vectors ------------------------------------------------------------
+
+    # Refuse to write vectors from a model this store has not seen. Two models in
+    # one table cannot be compared, and the failure would be silent: cosine
+    # between unrelated spaces still returns a number, just a meaningless one.
+    # Drops every vector and the metadata that identifies them. Needed because
+    # `unembedded_chunks` only returns rows where the vector IS NULL, so a model
+    # change cannot be reconciled by re-running `embed` -- without this the
+    # mismatch error would tell the user to run a command that does nothing.
+    def clear_vectors
+      @db.execute("UPDATE chunks SET embedding = NULL WHERE embedding IS NOT NULL")
+      dropped = @db.changes
+      @db.execute("DELETE FROM meta WHERE key IN ('embed_model', 'embed_dimensions', 'embed_at')")
+      dropped
+    end
+
+    def assert_embedding_model!(identity)
+      stored = meta("embed_model")
+      return true if stored.nil? || stored == identity
+
+      raise EmbeddingModelMismatch, <<~MSG.strip
+        this index already holds vectors from #{stored}, and the current model is #{identity}.
+        Vectors from different models cannot be compared, so the store will not mix them.
+        Discard the vectors with: gemchat embed --force
+      MSG
+    end
+
+    # float32 little-endian: 768 dimensions is 3072 bytes a chunk, so ~43MB for
+    # the 14k-chunk corpus the plan measured. Small enough to keep in the row.
+    def pack(vector)
+      vector.map { |v| v.to_f }.pack("e*")
+    end
+
+    def unpack(blob)
+      return [] if blob.nil? || blob.empty?
+
+      blob.unpack("e*")
+    end
+
+    def store_embeddings(rows, vectors, identity:, dimensions:)
+      assert_embedding_model!(identity)
+      return 0 if rows.empty?
+
+      unless vectors.size == rows.size
+        raise Error, "embedding count mismatch: #{vectors.size} vectors for #{rows.size} chunks"
+      end
+
+      bad = vectors.find { |v| v.size != dimensions }
+      if bad
+        raise Error, "expected #{dimensions}-dimension vectors, got #{bad.size}"
+      end
+
+      now = Time.now.utc.iso8601
+      @db.transaction do
+        rows.each_with_index do |id, i|
+          @db.execute("UPDATE chunks SET embedding = ?, embed_model = ? WHERE id = ?",
+            [pack(vectors[i]), identity, id])
+        end
+      end
+
+      set_meta("embed_model", identity)
+      set_meta("embedding_dimensions", dimensions)
+      set_meta("embedded_at", now)
+      vectors.size
+    end
+
+    # The chunks still needing a vector, with the text to embed. Ordered so a
+    # resumed run picks up where it stopped.
+    #
+    # The text is rebuilt through Chunker::Chunk#searchable_text rather than
+    # assembled in SQL. A query embedding and a document embedding must be
+    # produced by *identical* text, and a second implementation of that format in
+    # SQL would drift from the Ruby one the chunks were written with -- silently,
+    # because both still produce a valid-looking vector.
+    def unembedded_chunks(limit: nil)
+      sql = <<~SQL
+        SELECT id, gem_name, gem_version, source_type, title, body,
+               class_name, method_name, method_type, signature
+        FROM chunks
+        WHERE embedding IS NULL
+        ORDER BY id
+      SQL
+      sql += " LIMIT #{limit.to_i}" if limit
+
+      @db.execute(sql).map { |row|
+        chunk = Chunker::Chunk.new(
+          gem_name: row["gem_name"], gem_version: row["gem_version"],
+          source_type: row["source_type"], title: row["title"], body: row["body"],
+          class_name: row["class_name"], method_name: row["method_name"],
+          method_type: row["method_type"], signature: row["signature"]
+        )
+        {"id" => row["id"], "searchable" => chunk.searchable_text}
+      }
+    end
+
+    def embedded_stats
+      total = @db.get_first_value("SELECT COUNT(*) FROM chunks").to_i
+      embedded = @db.get_first_value("SELECT COUNT(*) FROM chunks WHERE embedding IS NOT NULL").to_i
+      {total:, embedded:, pending: total - embedded}
+    end
+
+    # Unit vectors, so a dot product is cosine similarity and the ordering falls
+    # out of a single pass. Scoped to the same model as the query was embedded
+    # with, so a mixed store cannot return a comparison across spaces.
+    def vector_search(query_vector, limit: 10, min_score: 0.0, gem_names: nil)
+      return [] if query_vector.nil? || query_vector.empty?
+
+      identity = meta("embed_model")
+      return [] if identity.nil?
+
+      # A mis-sized query makes dot() return NULL for every row, which sorts as
+      # 0.0 and would return an arbitrary slice of the corpus as if it were a
+      # ranked answer. Refusing up front turns a silent wrong answer into none.
+      stored_dimensions = meta("embedding_dimensions")&.to_i
+      if stored_dimensions && stored_dimensions != query_vector.size
+        return []
+      end
+
+      binds = [pack(query_vector), identity]
+      scope = ""
+      if gem_names && !gem_names.empty?
+        scope = "AND c.gem_name IN (#{(["?"] * gem_names.size).join(", ")})"
+        binds.concat(gem_names)
+      end
+      binds << limit
+
+      sql = <<~SQL
+        SELECT #{COLUMNS.map { |c| "c.#{c}" }.join(", ")},
+               -dot(embedding, ?) AS rank
+        FROM chunks c
+        WHERE c.embedding IS NOT NULL
+          AND c.embed_model = ? #{scope}
+        ORDER BY rank
+        LIMIT ?
+      SQL
+
+      @db.execute(sql, binds).map { |row|
+        COLUMNS.to_h { |c| [c, row[c]] }.merge("score" => -row["rank"].to_f)
+      }.select { |row| row["score"] >= min_score }
+    end
+
+    # --- hybrid -------------------------------------------------------------
+
+    # Runs whichever arms were asked for and fuses them by reciprocal rank.
+    #
+    # The UDF constraint from above does not bite here. It stopped `vector_search`
+    # from being narrowed by FTS inside a single statement, but each arm here
+    # returns a small ranked list on its own, so the fusion is a hash fold over a
+    # few hundred rows in Ruby. That is a different operation from the one that
+    # was infeasible, and worth distinguishing before assuming otherwise.
+    #
+    # RRF is used precisely because the two scores are not comparable: one is a
+    # negated BM25 rank on an unbounded scale, the other a cosine in [-1, 1].
+    # Fusing on *rank* sidesteps normalisation entirely, which is the reason to
+    # reach for it rather than for a weighted sum of the raw scores.
+    # `lex_mode` is a parameter rather than a constant because the obvious
+    # improvement to :phrase is wrong, and that is worth being able to measure
+    # rather than re-derive:
+    #
+    #   :phrase  the quoted run. Measured on 22,320 chunks from 34 gems, it
+    #            returns zero rows for every natural-language question tried, and
+    #            one exact hit for an identifier. High precision, no recall.
+    #   :any     the same terms OR'd. This is *worse*, not better: on five
+    #            known-answer queries -- queries whose answer definitely exists in
+    #            the corpus -- it located the right chunk in 1 of 5. Short ri
+    #            signatures containing one query word outrank the long prose chunk
+    #            that actually answers the question, because a rare term in a short
+    #            document scores enormously under BM25.
+    #
+    # So :any is not the default. Feeding it to a fusion is actively harmful,
+    # because RRF trusts rank position: a confidently wrong arm does not get
+    # ignored, it gets promoted. A real keyword arm needs term selection (drop
+    # the common words, require a rare term) before it belongs anywhere near
+    # :any. See §11.3 of the plan.
+    def hybrid_search(query, query_vector: nil, lex: true, lex_mode: :phrase,
+      limit: 10, pool: 50, k: RRF_K, gem_names: nil)
+      arms = []
+
+      if lex
+        found = search(query, limit: pool, gem_names: gem_names, mode: lex_mode)
+        arms << [:bm25, found] unless found.empty?
+      end
+
+      if query_vector
+        vec = vector_search(query_vector, limit: pool, gem_names: gem_names)
+        arms << [:vector, vec] unless vec.empty?
+      end
+
+      return [] if arms.empty?
+
+      rrf_fuse(arms, limit: limit, k: k)
+    end
+
+    # `arms` is [[name, ranked_rows], ...]. Each row carries "matched", naming
+    # the arms that found it, because two engines agreeing is signal in its own
+    # right and the fold would otherwise throw it away.
+    def rrf_fuse(arms, limit:, k: RRF_K)
+      scores = Hash.new(0.0)
+      matched = Hash.new { |h, key| h[key] = [] }
+      rows = {}
+
+      arms.each do |name, list|
+        list.each_with_index do |row, i|
+          id = row["id"]
+          scores[id] += 1.0 / (k + i + 1)
+          matched[id] << name
+          rows[id] ||= row
+        end
+      end
+
+      # Ties break on id so the order is deterministic across runs. Without this,
+      # two equally-scored chunks swap places between invocations, which makes a
+      # regression look like a change.
+      rows.keys.sort_by { |id| [-scores[id], id] }.first(limit).map { |id|
+        rows[id].merge("score" => scores[id], "matched" => matched[id])
+      }
+    end
+
+    def close
+      @db.close
+    end
+
+    class EmbeddingModelMismatch < Gemchat::Error; end
+  end
+end