engram 0.4.0 → 0.6.0

data/lib/engram/adapters/pgvector_store.rb

@@ -28,7 +28,8 @@ module Engram
         to_record(row)
       end
 
-      def search(embedding:, scope:, limit:, kinds: nil)
+      def search(embedding:, scope:, limit:, kinds: nil, embedding_metadata: nil)
+        Engram::EmbeddingMetadata.validate_query!(embedding, embedding_metadata)
         query = model.where(scope: scope)
         normalized_kinds = normalize_kinds(kinds)
         query = query.where(kind: normalized_kinds) if normalized_kinds
@@ -37,30 +38,56 @@ module Engram
           .nearest_neighbors(:embedding, embedding, distance: "cosine")
           .limit(limit)
           .map { |row| to_record(row) }
+          .tap { |records| validate_query_records!(records, embedding, embedding_metadata) }
       end
 
-      def all(scope:)
-        model.where(scope: scope).map { |row| to_record(row) }
+      def all(scope:, limit: nil, offset: 0, after_id: nil)
+        query = model.where(scope: scope).order(:id)
+        query = query.where("id > ?", after_id) if after_id
+        query = query.limit(limit) if limit
+        query = query.offset(offset) if offset > 0
+        query.map { |row| to_record(row) }
       end
 
-      def update(id:, record:)
-        row = model.find(id)
-        row.update!(
-          content: record.content,
-          kind: record.kind.to_s,
-          importance: record.importance,
-          metadata: record.metadata,
-          embedding: record.embedding
-        )
-        to_record(row)
+      def existing_ids(scope:, ids:)
+        persisted_ids = model.where(scope: scope, id: ids).pluck(:id)
+        id_type = model.type_for_attribute("id")
+        persisted_lookup = persisted_ids.each_with_object({}) { |id, lookup| lookup[id] = true }
+        matched_ids = {}
+
+        ids.select do |id|
+          cast_id = id_type.cast(id)
+          next false unless persisted_lookup[cast_id]
+          next false if matched_ids[cast_id]
+
+          matched_ids[cast_id] = true
+        end
+      end
+
+      def update(scope:, id:, record:)
+        raise Engram::Error, "cannot move memory across scopes" unless record.scope == scope
+
+        model.transaction do
+          row = model.lock.find_by!(id: id, scope: scope)
+          row.update!(
+            content: record.content,
+            kind: record.kind.to_s,
+            importance: record.importance,
+            metadata: record.metadata,
+            embedding: record.embedding
+          )
+          to_record(row)
+        end
+      rescue ActiveRecord::RecordNotFound
+        raise Engram::Error, "no memory with id #{id.inspect} in scope #{scope.inspect}"
       end
 
-      def delete(id:)
-        model.where(id: id).delete_all
+      def delete(scope:, id:)
+        model.where(id: id, scope: scope).delete_all
       end
 
-      def touch(id:, at: Time.now)
-        model.where(id: id).update_all(last_accessed_at: at)
+      def touch(scope:, id:, at: Time.now)
+        model.where(id: id, scope: scope).update_all(last_accessed_at: at)
       end
 
       private
@@ -69,6 +96,12 @@ module Engram
         raise Engram::Error, "memory scope cannot be nil" if scope.nil?
       end
 
+      def validate_query_records!(records, embedding, embedding_metadata)
+        records.each do |record|
+          Engram::EmbeddingMetadata.validate_record!(record, embedding, embedding_metadata)
+        end
+      end
+
       def model
         @model ||= resolve_default_model
       end

data/lib/engram/adapters/ruby_llm_embedder.rb

@@ -10,16 +10,32 @@ module Engram
       DEFAULT_MODEL = "text-embedding-3-small"
       DEFAULT_DIMENSIONS = 1536
 
-      def initialize(model: DEFAULT_MODEL, dimensions: DEFAULT_DIMENSIONS)
+      def initialize(model: DEFAULT_MODEL, dimensions: nil)
+        if dimensions && (!dimensions.is_a?(Integer) || !dimensions.positive?)
+          raise ArgumentError, "dimensions must be a positive integer"
+        end
+
         @model = model
-        @dimensions = dimensions
+        @requested_dimensions = dimensions
+        @dimensions = dimensions || DEFAULT_DIMENSIONS
       end
 
-      attr_reader :dimensions
+      attr_reader :dimensions, :model
+
+      def embedding_metadata
+        Engram::EmbeddingMetadata.build(
+          adapter: self.class.name,
+          provider: "ruby_llm",
+          model: model,
+          dimensions: dimensions
+        )
+      end
 
       def embed(text)
         ensure_ruby_llm!
-        RubyLLM.embed(text, model: @model).vectors
+        vectors = request_embedding(text)
+        validate_dimensions!(vectors)
+        vectors
       end
 
       private
@@ -30,6 +46,31 @@ module Engram
         raise Engram::Error,
           "RubyLLMEmbedder requires the `ruby_llm` gem. Add it to your Gemfile and configure it."
       end
+
+      # Older RubyLLM versions do not accept the dimensions keyword.
+      def request_embedding(text)
+        if @requested_dimensions && embed_accepts_dimensions?
+          RubyLLM.embed(text, model: @model, dimensions: @requested_dimensions).vectors
+        else
+          RubyLLM.embed(text, model: @model).vectors
+        end
+      end
+
+      def embed_accepts_dimensions?
+        RubyLLM.method(:embed).parameters.any? do |type, name|
+          type == :keyrest || ([:key, :keyreq].include?(type) && name == :dimensions)
+        end
+      end
+
+      def validate_dimensions!(vectors)
+        length = vectors.respond_to?(:length) ? vectors.length : nil
+        return if length == @dimensions
+
+        raise Engram::Error,
+          "embedding model #{@model.inspect} returned #{length ? "#{length}-dimension" : "non-vector"} output " \
+            "but the embedder is configured with dimensions: #{@dimensions}; align the dimensions option " \
+            "(and your vector column) with the model's output"
+      end
     end
   end
 end

data/lib/engram/consolidators/heuristic_consolidator.rb

@@ -15,7 +15,13 @@ module Engram
 
       def reconcile_all(candidates:, scope:)
         Array(candidates).map do |candidate|
-          nearest = @store.search(embedding: candidate.embedding, scope: scope, limit: 1).first
+          nearest = Engram::EmbeddingMetadata.search(
+            @store,
+            embedding: candidate.embedding,
+            embedding_metadata: Engram::EmbeddingMetadata.extract(candidate.metadata),
+            scope: scope,
+            limit: 1
+          ).first
           similarity = nearest ? Engram::Math.cosine_similarity(candidate.embedding, nearest.embedding) : 0.0
 
           if nearest && similarity >= @similarity_threshold

data/lib/engram/consolidators/llm_consolidator.rb

@@ -56,23 +56,35 @@ module Engram
         candidates = Array(candidates)
         return [] if candidates.empty?
 
+        neighbors = neighbor_map(candidates, scope)
         result = @completion.complete(
           system: SYSTEM,
-          user: JSON.generate(payload(candidates, scope)),
+          user: JSON.generate(payload(candidates, neighbors)),
           schema: SCHEMA
         )
-        map_decisions(decisions(result), candidates)
+        map_decisions(decisions(result), candidates, neighbors)
       end
 
       private
 
-      def payload(candidates, scope)
+      def neighbor_map(candidates, scope)
+        candidates.map do |candidate|
+          Engram::EmbeddingMetadata.search(
+            @store,
+            embedding: candidate.embedding,
+            embedding_metadata: Engram::EmbeddingMetadata.extract(candidate.metadata),
+            scope: scope,
+            limit: @neighbors
+          )
+        end
+      end
+
+      def payload(candidates, neighbors)
         items = candidates.each_with_index.map do |candidate, index|
-          existing = @store.search(embedding: candidate.embedding, scope: scope, limit: @neighbors)
           {
             index: index,
             candidate: candidate.content,
-            existing: existing.map { |r| {id: r.id, content: r.content} }
+            existing: neighbors[index].map { |r| {id: r.id, content: r.content} }
           }
         end
         {candidates: items}
@@ -81,19 +93,34 @@ module Engram
       def decisions(result)
         return [] unless result.is_a?(Hash)
 
-        result["decisions"] || result[:decisions] || []
+        decisions = result["decisions"] || result[:decisions] || []
+        decisions.is_a?(Array) ? decisions : []
       end
 
-      def map_decisions(raw, candidates)
+      def map_decisions(raw, candidates, neighbors)
+        seen = Set.new
         raw.filter_map do |decision|
+          next unless decision.is_a?(Hash)
+
           decision = decision.transform_keys(&:to_s)
           index = decision["index"]
-          next unless index && candidates[index]
+          next unless index.is_a?(Integer) && index >= 0 && candidates[index]
+          next if seen.include?(index)
+
+          action = (decision["action"] || "noop").to_s
+          next unless Engram::Decision::ACTIONS.include?(action.to_sym)
+
+          target_id = decision["target_id"]
+          if %w[update forget].include?(action)
+            next if target_id.nil?
+            next unless neighbors[index].any? { |neighbor| neighbor.id == target_id }
+          end
 
+          seen.add(index)
           Engram::Decision.new(
-            action: (decision["action"] || "noop").to_sym,
+            action: action.to_sym,
             candidate: candidates[index],
-            target_id: decision["target_id"],
+            target_id: target_id,
             reason: decision["reason"]
           )
         end

data/lib/engram/embedding_metadata.rb

@@ -0,0 +1,135 @@
+# frozen_string_literal: true
+
+require "digest"
+
+module Engram
+  # Helpers for storing and validating embedding provenance under Engram's
+  # reserved metadata namespace.
+  module EmbeddingMetadata
+    RESERVED_KEY = "_engram"
+    EMBEDDING_KEY = "embedding"
+
+    module_function
+
+    def build(adapter:, dimensions:, model: nil, provider: nil)
+      data = {
+        "adapter" => adapter.to_s,
+        "dimensions" => dimensions
+      }
+      data["provider"] = provider.to_s if provider
+      data["model"] = model.to_s if model
+      data["fingerprint"] = fingerprint(data)
+      data
+    end
+
+    def for_embedder(embedder, embedding: nil)
+      return nil unless embedder.respond_to?(:embedding_metadata)
+
+      metadata = stringify_keys(embedder.embedding_metadata || {})
+      return nil if metadata.empty?
+
+      dimensions = embedding.respond_to?(:length) ? embedding.length : metadata["dimensions"]
+      metadata = metadata.merge("dimensions" => dimensions) if dimensions
+      metadata["fingerprint"] = fingerprint(metadata)
+      metadata
+    end
+
+    def attach(record, embedder: nil, embedding_metadata: nil)
+      metadata = stringify_keys(embedding_metadata || for_embedder(embedder, embedding: record.embedding)) || {}
+      return record if metadata.empty?
+
+      record.with(metadata: merge(record.metadata, metadata))
+    end
+
+    def extract(metadata)
+      metadata = stringify_keys(metadata || {})
+      reserved = metadata[RESERVED_KEY]
+      return {} unless reserved.is_a?(Hash)
+
+      stringify_keys(reserved[EMBEDDING_KEY] || {})
+    end
+
+    def merge(metadata, embedding_metadata)
+      Engram::ReservedMetadata.attach(
+        metadata, EMBEDDING_KEY, embedding_metadata || {}, namespace_description: "embedding metadata"
+      )
+    end
+
+    def search(store, embedding:, embedding_metadata:, scope:, limit:, kinds: nil)
+      store.search(
+        embedding: embedding,
+        embedding_metadata: embedding_metadata,
+        scope: scope,
+        limit: limit,
+        kinds: kinds
+      )
+    rescue ArgumentError => error
+      raise unless unknown_embedding_metadata_keyword?(error)
+
+      store.search(embedding: embedding, scope: scope, limit: limit, kinds: kinds)
+    end
+
+    def validate_query!(embedding, embedding_metadata)
+      metadata = stringify_keys(embedding_metadata || {})
+      return if metadata.empty?
+
+      expected = metadata["dimensions"]
+      return unless expected && embedding.respond_to?(:length) && embedding.length != expected.to_i
+
+      raise Engram::Error,
+        "embedding dimension mismatch: query vector has #{embedding.length} dimensions, metadata declares #{expected}"
+    end
+
+    def validate_record!(record, query_embedding, query_metadata)
+      if query_embedding.respond_to?(:length) && record.embedding.respond_to?(:length) &&
+          query_embedding.length != record.embedding.length
+        raise Engram::Error,
+          "embedding dimension mismatch: query vector has #{query_embedding.length} dimensions, " \
+          "record #{record.id.inspect} has #{record.embedding.length}"
+      end
+
+      stored = extract(record.metadata)
+      return if stored.empty?
+
+      stored_dimensions = stored["dimensions"]
+      if stored_dimensions && record.embedding.respond_to?(:length) && record.embedding.length != stored_dimensions.to_i
+        raise Engram::Error,
+          "embedding metadata mismatch: record #{record.id.inspect} has #{record.embedding.length} dimensions, " \
+          "metadata declares #{stored_dimensions}"
+      end
+
+      query_metadata = stringify_keys(query_metadata || {})
+      return if query_metadata.empty?
+
+      conflicting_key = %w[adapter provider model dimensions fingerprint].find do |key|
+        stored.key?(key) && query_metadata.key?(key) && stored[key].to_s != query_metadata[key].to_s
+      end
+      return unless conflicting_key
+
+      raise Engram::Error,
+        "embedding metadata mismatch for record #{record.id.inspect}: #{conflicting_key} " \
+        "#{stored[conflicting_key].inspect} does not match query #{query_metadata[conflicting_key].inspect}"
+    end
+
+    def fingerprint(metadata)
+      relevant = stringify_keys(metadata).values_at("adapter", "provider", "model", "dimensions")
+      Digest::SHA256.hexdigest(relevant.map { |value| value.nil? ? "" : value.to_s }.join("\0"))
+    end
+
+    def stringify_keys(value)
+      case value
+      when Hash
+        value.each_with_object({}) do |(key, nested), out|
+          out[key.to_s] = stringify_keys(nested)
+        end
+      else
+        value
+      end
+    end
+
+    def unknown_embedding_metadata_keyword?(error)
+      error.message.include?("unknown keyword: :embedding_metadata") ||
+        error.message.include?("unknown keyword: \"embedding_metadata\"")
+    end
+  end
+end

data/lib/engram/extraction.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+module Engram
+  # An immutable extractor result that carries a candidate Record and its source provenance.
+  # Converting it creates a new record and leaves the caller-owned record and metadata untouched.
+  class Extraction
+    attr_reader :record, :provenance
+
+    def initialize(record:, provenance:)
+      is_record = Object.instance_method(:is_a?).bind_call(record, Engram::Record)
+      raise ArgumentError, "record must be an Engram::Record" unless is_record
+
+      is_provenance = Object.instance_method(:is_a?).bind_call(provenance, Engram::Provenance)
+      raise ArgumentError, "provenance must be an Engram::Provenance" unless is_provenance
+
+      @record = record
+      @provenance = provenance
+      freeze
+    end
+
+    def to_record
+      source = Engram::Internal::CandidateIntegrity.new.detach(record)
+      attributes = Engram::Record::STATE_READERS.to_h do |attribute|
+        [attribute, Engram::Record.instance_method(attribute).bind_call(source)]
+      end
+      attributes[:metadata] = Engram::Provenance.attach(attributes.fetch(:metadata), provenance)
+      Engram::Record.new(**attributes)
+    end
+  end
+end

data/lib/engram/extractors/llm_extractor.rb

@@ -58,14 +58,15 @@ module Engram
           next if content.empty?
           next if (fact["confidence"] || 1.0).to_f < @min_confidence
 
-          Engram::Record.new(
+          embedding = @embedder.embed(content)
+          Engram::EmbeddingMetadata.attach(Engram::Record.new(
             content: content,
             scope: scope,
             kind: fact["kind"] || "fact",
             importance: (fact["importance"] || 1.0).to_f,
             metadata: {confidence: (fact["confidence"] || 1.0).to_f},
-            embedding: @embedder.embed(content)
-          )
+            embedding: embedding
+          ), embedder: @embedder)
         end
       end