maglev-rb 0.1.1 → 0.2.1

data/lib/maglev/search_result.rb

@@ -2,21 +2,25 @@
 
 module Maglev
   class SearchResult
-    attr_reader :owner, :content, :source, :distance, :chunk_index
+    attr_reader :owner, :content, :source, :distance, :chunk_index, :source_identity, :source_type, :score
 
-    def initialize(owner:, content:, source:, distance:, chunk_index: nil)
+    def initialize(owner:, content:, source:, distance:, chunk_index: nil, source_identity: nil, source_type: nil, score: nil)
       @owner = owner
       @content = content
       @source = source
       @distance = distance
       @chunk_index = chunk_index
+      @source_identity = source_identity || source
+      @source_type = (source_type || :snapshot).to_sym
+      @score = score
       freeze
     end
 
     def similarity
+      return score.clamp(0.0, 1.0) unless score.nil?
       return nil if distance.nil?
 
-      1.0 - distance.to_f
+      (1.0 - distance.to_f).clamp(0.0, 1.0)
     end
   end
 end

data/lib/maglev/snapshot.rb

@@ -2,13 +2,33 @@
 
 module Maglev
   class Snapshot
-    def initialize(lines)
+    attr_reader :metadata
+
+    def initialize(lines, metadata: {})
       @text = lines.join("\n").freeze
+      @metadata = deep_freeze(metadata)
       freeze
     end
 
     def to_s
       @text
     end
+
+    def truncated?
+      @metadata[:truncated] == true
+    end
+
+    private
+
+    def deep_freeze(value)
+      case value
+      when Hash
+        value.to_h { |key, item| [key, deep_freeze(item)] }.freeze
+      when Array
+        value.map { |item| deep_freeze(item) }.freeze
+      else
+        value.freeze
+      end
+    end
   end
 end

data/lib/maglev/snapshot_budget.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+module Maglev
+  class SnapshotBudget
+    LIMIT_METHODS = {
+      attribute: :snapshot_attribute_max_characters,
+      rich_text: :snapshot_attribute_max_characters,
+      attachment: :snapshot_attribute_max_characters,
+      related_record: :snapshot_related_record_max_characters,
+      whole_snapshot: :snapshot_max_characters
+    }.freeze
+
+    def initialize(configuration: Maglev.configuration)
+      @limits = {
+        attribute_characters: configuration.snapshot_attribute_max_characters,
+        related_record_characters: configuration.snapshot_related_record_max_characters,
+        snapshot_characters: configuration.snapshot_max_characters,
+        chunks: configuration.snapshot_max_chunks
+      }.freeze
+      @configuration = configuration
+      @sources = []
+    end
+
+    def truncate(text, kind:, path:)
+      value = text.to_s
+      limit = @configuration.public_send(LIMIT_METHODS.fetch(kind))
+      return value if value.length <= limit
+
+      @sources << {
+        kind: kind,
+        path: path,
+        original_characters: value.length,
+        retained_characters: limit
+      }.freeze
+      value[0, limit]
+    end
+
+    def record_chunk_truncation(original:, retained:, path: "snapshot.chunks")
+      return if original <= retained
+
+      @sources << {
+        kind: :chunks,
+        path: path,
+        original_chunks: original,
+        retained_chunks: retained
+      }.freeze
+    end
+
+    def metadata
+      {
+        truncated: @sources.any?,
+        limits: @limits,
+        sources: @sources.dup.freeze
+      }.freeze
+    end
+  end
+end

data/lib/maglev/snapshot_builder.rb

@@ -3,24 +3,30 @@
 require "cgi"
 
 require_relative "attachment_extractor"
+require_relative "relation_order"
 require_relative "snapshot"
+require_relative "snapshot_budget"
 
 module Maglev
   class SnapshotBuilder
-    def initialize(record, config, path: nil, visited: nil, attachment_extractor: nil, remaining_depth: Maglev.configuration.max_relation_depth)
+    def initialize(record, config, path: nil, visited: nil, attachment_extractor: nil, remaining_depth: Maglev.configuration.max_relation_depth, budget: nil)
       @record = record
       @config = config
       @path = path
       @visited = visited || {}
       @attachment_extractor = attachment_extractor || Maglev.configuration.attachment_extractor || AttachmentExtractor.new
       @remaining_depth = remaining_depth
+      @budget = budget || SnapshotBudget.new
     end
 
     def build
       return Snapshot.new([]) if visited?
 
       mark_visited
-      Snapshot.new(lines)
+      raw_lines = lines
+      text = raw_lines.join("\n")
+      text = @budget.truncate(text, kind: :whole_snapshot, path: "snapshot") unless @path
+      Snapshot.new(text.lines(chomp: true), metadata: @budget.metadata)
     end
 
     private
@@ -49,7 +55,10 @@ module Maglev
     def attribute_lines
       @config.exposed_attributes.filter_map do |attribute|
         value = @record.public_send(attribute)
-        "#{attribute_path(attribute)}: #{value}" unless value.nil?
+        next if value.nil?
+
+        text = @budget.truncate(value, kind: :attribute, path: attribute_path(attribute))
+        "#{attribute_path(attribute)}: #{text}"
       end
     end
 
@@ -67,14 +76,16 @@ module Maglev
         relation_records, collection = relation_records(relation)
         relation_records.each_with_index.flat_map do |related_record, index|
           relation_path = path_for(relation, index, collection: collection)
-          SnapshotBuilder.new(
+          child = SnapshotBuilder.new(
             related_record,
             related_record.class.maglev_config,
             path: relation_path,
             visited: @visited,
             attachment_extractor: @attachment_extractor,
-            remaining_depth: remaining_depth
-          ).build.to_s.split("\n")
+            remaining_depth: remaining_depth,
+            budget: @budget
+          ).build.to_s
+          @budget.truncate(child, kind: :related_record, path: relation_path).split("\n")
         end
       end
     end
@@ -84,7 +95,9 @@ module Maglev
         attachment_blobs(source.name).flat_map do |blob|
           document = extract_attachment(source.name, blob)
           if document.extracted?
-            ["#{document.source_identifier}.text: #{document.text}"]
+            path = document.source_identifier
+            text = @budget.truncate(document.text, kind: :attachment, path: path)
+            ["#{path}.text: #{text}"]
           else
             ["#{document.source_identifier}.skipped: #{document.metadata[:reason]}"]
           end
@@ -125,7 +138,9 @@ module Maglev
     def rich_text_lines
       @config.rich_text_sources.filter_map do |source|
         text = rich_text_value(source.name)
-        "#{rich_text_identifier(source.name)}.text: #{text}" unless text.nil? || text.empty?
+        path = rich_text_identifier(source.name)
+        text = @budget.truncate(text, kind: :rich_text, path: path) unless text.nil?
+        "#{path}.text: #{text}" unless text.nil? || text.empty?
       end
     end
 
@@ -169,10 +184,12 @@ module Maglev
       value = @record.public_send(relation.name)
       records, collection = if value.nil?
         [[], false]
+      elsif loaded_collection_proxy?(value)
+        [relation.order ? ordered_array(value.target, relation) : ordered_loaded_records(value), true]
+      elsif active_record_relation?(value)
+        [ordered_relation(value, relation).limit(relation.limit).to_a, true]
       elsif value.respond_to?(:to_ary)
-        [value.to_ary, true]
-      elsif value.respond_to?(:limit)
-        [value.limit(relation.limit).to_a, true]
+        [ordered_array(value.to_ary, relation), true]
       else
         [[value], false]
       end
@@ -180,6 +197,67 @@ module Maglev
       [records.first(relation.limit), collection]
     end
 
+    def active_record_relation?(value)
+      (defined?(ActiveRecord::Relation) && value.is_a?(ActiveRecord::Relation)) ||
+        (defined?(ActiveRecord::Associations::CollectionProxy) && value.is_a?(ActiveRecord::Associations::CollectionProxy))
+    end
+
+    def loaded_collection_proxy?(value)
+      defined?(ActiveRecord::Associations::CollectionProxy) &&
+        value.is_a?(ActiveRecord::Associations::CollectionProxy) && value.loaded?
+    end
+
+    def ordered_relation(value, relation)
+      if relation.order
+        order = effective_order(relation.order, value.klass)
+        return value.reorder(order)
+      end
+      return value unless value.order_values.empty?
+
+      primary_key = value.klass.primary_key
+      primary_key ? value.order(primary_key => :asc) : value
+    end
+
+    def effective_order(order, klass)
+      RelationOrder.with_primary_key(order, klass)
+    end
+
+    def ordered_array(records, relation)
+      return records unless relation.order
+
+      order = relation.order
+      records.sort do |left, right|
+        comparison = order.lazy.map do |attribute, direction|
+          value = compare_values(left.public_send(attribute), right.public_send(attribute))
+          (direction == :desc) ? -value : value
+        end.find { |value| !value.zero? } || 0
+        comparison.zero? ? compare_values(record_id_for(left), record_id_for(right)) : comparison
+      end
+    end
+
+    def compare_values(left, right)
+      return 0 if left == right
+      return -1 if left.nil?
+      return 1 if right.nil?
+
+      left <=> right
+    end
+
+    def record_id_for(record)
+      record.respond_to?(:id) ? record.id : record.object_id
+    end
+
+    def ordered_loaded_records(value)
+      records = value.to_ary
+      return records unless value.order_values.empty?
+
+      primary_key = value.klass.primary_key
+      return records unless primary_key
+
+      persisted, unsaved = records.partition(&:persisted?)
+      persisted.sort_by { |record| record.public_send(primary_key) } + unsaved
+    end
+
     def path_for(relation, index, collection:)
       segment = collection ? "#{relation.name}[#{index}]" : relation.name
       @path ? "#{@path}.#{segment}" : segment

data/lib/maglev/source_extractor.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+require_relative "source_fragment"
+
+module Maglev
+  class SourceExtractor
+    def call(snapshot)
+      related = {}
+      snapshot.to_s.lines(chomp: true).drop(1).filter_map do |line|
+        if (match = line.match(/\A(.+\[\d+\]) ([A-Za-z0-9_:]+)#([^ ]+)\z/))
+          related[match[1]] = [match[2], match[3]]
+          next
+        end
+        next unless line.include?(": ")
+
+        path = line.split(": ", 2).first
+        identity = related_identity(path, related) || path
+        next if identity.nil? || identity.empty?
+
+        SourceFragment.new(identity: identity, type: source_type(identity), content: line)
+      end.freeze
+    end
+
+    private
+
+    def related_identity(path, related)
+      prefix, record = related.find { |candidate, _identity| path.start_with?("#{candidate}.") }
+      return unless record
+
+      model, id = record
+      "related:#{model}:#{id}:#{path.delete_prefix("#{prefix}.")}"
+    end
+
+    def source_type(identity)
+      return :attachment if identity.include?("[blob:")
+      return :rich_text if identity.start_with?("rich_text.")
+      return :related_record if identity.start_with?("related:") || identity.include?("[") || identity.include?(".")
+      return :tag if identity == "tags"
+
+      :attribute
+    end
+  end
+end

data/lib/maglev/source_fragment.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+module Maglev
+  SourceFragment = Data.define(:identity, :type, :content) do
+    def initialize(identity:, type:, content:)
+      super(identity: identity.to_s.freeze, type: type.to_sym, content: content.to_s.freeze)
+    end
+  end
+end

data/lib/maglev/structured_answer_composer.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+require "json"
+require_relative "trace"
+
+module Maglev
+  class GroundingError < Error
+  end
+
+  class StructuredAnswerComposer
+    def initialize(generation_adapter: nil)
+      @generation_adapter = generation_adapter
+    end
+
+    def compose(result)
+      unless result.is_a?(StructuredResult)
+        raise GroundingError, "A structured result is required"
+      end
+
+      Trace.instrument(:composition, trace_id: result.trace_id,
+        resource: result.plan.resource, operation: result.plan.ir&.operation) do
+        (@generation_adapter && result.status == :succeeded) ? generate(result) : result.render
+      end
+    end
+
+    private
+
+    def generate(result)
+      evidence = {
+        "records" => result.evidence.records,
+        "scalar" => result.evidence.scalar,
+        "filters" => result.evidence.filters,
+        "date_ranges" => result.evidence.date_ranges,
+        "count" => result.evidence.count,
+        "truncated" => result.evidence.truncated
+      }.freeze
+      output = @generation_adapter.generate(JSON.generate(evidence))
+      unless output.is_a?(Hash) && output.keys == ["claim_paths"] && output["claim_paths"].is_a?(Array) &&
+          output["claim_paths"].any?
+        raise GroundingError, "Structured generation must return claim paths"
+      end
+
+      output.fetch("claim_paths").map { |path| render_claim(evidence, path) }.join(" ").freeze
+    end
+
+    def render_claim(evidence, path)
+      valid_segments = path.is_a?(Array) && path.all? { |segment| [String, Integer].any? { |type| segment.is_a?(type) } }
+      unless valid_segments && path.length.between?(1, 3)
+        raise GroundingError, "Generated claim is absent from structured evidence"
+      end
+
+      value = path.reduce(evidence) do |current, segment|
+        if current.is_a?(Hash) && segment.is_a?(String) && current.key?(segment)
+          current.fetch(segment)
+        elsif current.is_a?(Array) && segment.is_a?(Integer) && segment.between?(0, current.length - 1)
+          current.fetch(segment)
+        else
+          raise GroundingError, "Generated claim is absent from structured evidence"
+        end
+      end
+      raise GroundingError, "Generated claim must identify one scalar evidence value" if value.is_a?(Hash) || value.is_a?(Array)
+
+      label = path.last.to_s.tr("_", " ").capitalize
+      "#{label}: #{value}."
+    end
+  end
+end

data/lib/maglev/structured_evidence_builder.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+require "json"
+
+module Maglev
+  class StructuredEvidenceBuilder
+    def initialize(plan:, relation:, resource:, rows:, bytes:)
+      raise ArgumentError, "evidence rows must be positive" unless rows.is_a?(Integer) && rows.positive?
+      raise ArgumentError, "evidence bytes must be positive" unless bytes.is_a?(Integer) && bytes.positive?
+
+      @plan = plan
+      @relation = relation
+      @resource = resource
+      @rows = rows
+      @bytes = bytes
+    end
+
+    def build
+      @filters = @plan.ir.filters.map(&:to_h)
+      @date_ranges = @filters.select { |filter| filter["operator"] == "between" }
+      StructuredEvidence.new(filters: @filters,
+        date_ranges: @date_ranges,
+        loader: method(:load_records))
+    end
+
+    private
+
+    def load_records
+      Trace.instrument(:execution, trace_id: @plan.trace_id, resource: @plan.resource,
+        operation: @plan.ir.operation) do |payload|
+        query_limit = [@relation.limit_value || @rows + 1, @rows + 1].min
+        fields = @resource.fields.map(&:name)
+        projected = @relation.limit(query_limit).to_a.map do |record|
+          fields.to_h { |field| [field, record.public_send(field)] }
+        end
+        records = []
+        projected.first(@rows).each do |record|
+          candidate = [*records, record]
+          break if serialized_size(candidate, truncated: projected.length > candidate.length) > @bytes
+
+          records << record
+        end
+        truncated = projected.length > records.length
+        used_bytes = serialized_size(records, truncated: truncated)
+        payload[:row_count] = records.length
+        payload[:evidence_bytes] = used_bytes
+        [records, records.length, truncated]
+      end
+    end
+
+    def serialized_size(records, truncated:)
+      JSON.generate("records" => records, "scalar" => nil, "filters" => @filters,
+        "date_ranges" => @date_ranges, "count" => records.length, "truncated" => truncated).bytesize
+    end
+  end
+end

data/lib/maglev/structured_executor.rb

@@ -0,0 +1,157 @@
+# frozen_string_literal: true
+
+require_relative "query_compiler"
+require_relative "structured_result"
+require_relative "trace"
+require_relative "structured_evidence_builder"
+require "bigdecimal"
+
+module Maglev
+  class StructuredExecutionError < Error
+  end
+
+  class StructuredExecutor
+    DEFAULT_TIMEOUT = 5
+
+    def initialize(timeout: DEFAULT_TIMEOUT, role: nil, wrapper: nil)
+      @timeout = timeout
+      @role = role
+      @wrapper = wrapper
+      raise ArgumentError, "timeout must be positive" unless @timeout.respond_to?(:positive?) && @timeout.positive?
+    end
+
+    def execute(plan)
+      raise StructuredExecutionError, "A structured plan is required" unless plan.is_a?(StructuredPlan)
+      return protected_relation(plan.relation) if plan.records?
+
+      execute_with_policy(plan.relation) { execute_aggregate(plan) }
+    rescue StructuredExecutionError
+      raise
+    rescue ActiveRecord::ActiveRecordError => error
+      raise StructuredExecutionError, "Structured query execution failed: #{error.class.name}"
+    end
+
+    private
+
+    def protected_relation(relation)
+      executor = self
+      policy = Module.new do
+        define_method(:exec_queries) do |*arguments|
+          executor.send(:execute_with_policy, self) { super(*arguments) }
+        end
+        %i[calculate pluck pick exists? ids].each do |method_name|
+          define_method(method_name) do |*arguments|
+            executor.send(:execute_with_policy, self) { super(*arguments) }
+          end
+        end
+
+        define_method(:load_async) { load }
+      end
+      relation.extending(policy)
+    end
+
+    def execute_with_policy(relation, &query)
+      operation = proc { execute_read_only(relation, &query) }
+      if @wrapper
+        unwrapped_operation = operation
+        operation = proc { @wrapper.call(&unwrapped_operation) }
+      end
+      return ActiveRecord::Base.connected_to(role: @role, prevent_writes: true, &operation) if @role
+
+      ActiveRecord::Base.while_preventing_writes(&operation)
+    rescue StructuredExecutionError
+      raise
+    rescue ActiveRecord::ActiveRecordError => error
+      raise StructuredExecutionError, "Structured query execution failed: #{error.class.name}"
+    end
+
+    def execute_read_only(relation)
+      connection = relation.connection
+      connection.transaction(requires_new: true) do
+        apply_database_policy(connection)
+        yield
+      end
+    end
+
+    def apply_database_policy(connection)
+      unless connection.adapter_name.casecmp?("PostgreSQL")
+        raise StructuredExecutionError, "Structured execution requires an adapter with enforced statement timeouts"
+      end
+
+      milliseconds = (@timeout.to_f * 1000).ceil
+      connection.execute("SET LOCAL statement_timeout = #{milliseconds}")
+      connection.execute("SET TRANSACTION READ ONLY")
+    end
+
+    def execute_aggregate(plan)
+      aggregate = plan.aggregate
+      values = if aggregate.function == :count
+        plan.relation.pluck(plan.relation.klass.arel_table[plan.relation.klass.primary_key])
+      else
+        plan.relation.pluck(plan.aggregate_column)
+      end
+
+      case aggregate.function
+      when :count then values.length
+      when :sum then values.sum
+      when :average then values.empty? ? nil : BigDecimal(values.sum.to_s) / values.length
+      when :minimum then values.min
+      when :maximum then values.max
+      else raise StructuredExecutionError, "The aggregate is unavailable"
+      end
+    end
+  end
+
+  def self.structured_executor(configuration: Maglev.configuration)
+    StructuredExecutor.new(
+      timeout: configuration.structured_query_timeout,
+      role: configuration.structured_query_role,
+      wrapper: configuration.structured_query_executor_wrapper
+    )
+  end
+
+  def self.execute(plan, executor: structured_executor,
+    evidence_rows: Maglev.configuration.structured_evidence_max_rows,
+    evidence_bytes: Maglev.configuration.structured_evidence_max_bytes)
+    unless plan.is_a?(Planner::Plan)
+      raise StructuredExecutionError, "A structured planner result is required"
+    end
+
+    unless plan.ready?
+      status = (plan.status == :invalid) ? :failed : plan.status
+      return StructuredResult.new(status: status, kind: :none, warnings: plan.warnings,
+        plan: plan, trace_id: plan.trace_id)
+    end
+    unless plan.base_relation
+      raise StructuredExecutionError, "The structured plan has no authorized base relation"
+    end
+
+    compiled = Trace.instrument(:compilation, trace_id: plan.trace_id, resource: plan.resource,
+      operation: plan.ir.operation) do
+      QueryCompiler.new(snapshot: plan.validation.snapshot).compile(
+        validation: plan.validation, base_relation: plan.base_relation
+      )
+    end
+    resource = plan.validation.snapshot.resources.find { |candidate| candidate.identifier == plan.resource }
+    filters = plan.ir.filters.map(&:to_h)
+    date_ranges = filters.select { |filter| filter["operator"] == "between" }
+
+    if compiled.records?
+      value = executor.execute(compiled)
+      evidence = StructuredEvidenceBuilder.new(plan: plan, relation: value, resource: resource,
+        rows: evidence_rows, bytes: evidence_bytes).build
+      StructuredResult.new(status: :succeeded, kind: :relation, value: value, evidence: evidence,
+        interpretation: plan.explanation, warnings: plan.warnings, plan: plan, trace_id: plan.trace_id)
+    else
+      value = Trace.instrument(:execution, trace_id: plan.trace_id, resource: plan.resource,
+        operation: plan.ir.operation) { executor.execute(compiled) }
+      scalar_size = JSON.generate("records" => [], "scalar" => value, "filters" => filters,
+        "date_ranges" => date_ranges, "count" => 1, "truncated" => false).bytesize
+      bounded_scalar = scalar_size <= evidence_bytes
+      evidence = StructuredEvidence.new(scalar: bounded_scalar ? value : nil, filters: filters,
+        date_ranges: date_ranges, count: bounded_scalar ? 1 : 0, truncated: !bounded_scalar)
+      StructuredResult.new(status: :succeeded, kind: :aggregate, value: value, evidence: evidence,
+        interpretation: plan.explanation, warnings: plan.warnings, plan: plan, trace_id: plan.trace_id)
+    end
+  end
+end