omni_agent 0.1.11 → 0.2.0

data/lib/omni_agent/chunkers/recursive.rb

@@ -0,0 +1,53 @@
+module OmniAgent
+  module Chunkers
+    # Splits on the coarsest boundary that fits: headings, then paragraphs,
+    # lines, sentences and words, cutting mid-word only as a last resort.
+    # Pieces are merged back up to `size` characters, and each chunk starts
+    # with up to `overlap` characters of whole pieces from the one before.
+    module Recursive
+      # Zero-width, so every separator stays attached to its text.
+      SEPARATORS = [
+        /(?=^\#{1,6}\s)/, # Markdown headings
+        /(?<=\n\n)/,      # paragraphs
+        /(?<=\n)/,        # lines
+        /(?<=[.!?]\s)/,   # sentences
+        /(?<=\s)/         # words
+      ].freeze
+
+      def self.call(text, size: 1000, overlap: 200)
+        merge(split(text.to_s, size, SEPARATORS), size, overlap)
+      end
+
+      def self.split(text, size, separators)
+        return [ text ] if text.length <= size
+
+        separator, *rest = separators
+        return text.scan(/.{1,#{size}}/m) if separator.nil?
+
+        text.split(separator).flat_map { |piece| split(piece, size, rest) }
+      end
+
+      def self.merge(pieces, size, overlap)
+        chunks = []
+        current = []
+
+        pieces.each do |piece|
+          if current.any? && length(current) + piece.length > size
+            chunks << current.join
+            current.shift while current.any? && (length(current) > overlap || length(current) + piece.length > size)
+          end
+          current << piece
+        end
+        chunks << current.join if current.any?
+
+        chunks.map(&:strip).reject(&:empty?)
+      end
+
+      def self.length(pieces)
+        pieces.sum(&:length)
+      end
+
+      private_class_method :split, :merge, :length
+    end
+  end
+end

data/lib/omni_agent/chunkers.rb

@@ -0,0 +1,31 @@
+module OmniAgent
+  # Strategies that split text into chunks for `has_chunks`. A chunker is
+  # anything that responds to `call(text, size:, overlap:)` and returns
+  # Strings, or Hashes with :content and optional :context and :metadata.
+  module Chunkers
+    def self.registry
+      {
+        recursive: Recursive,
+        markdown: Markdown,
+        hierarchical: Hierarchical
+      }
+    end
+
+    def self.fetch(chunker)
+      return chunker if chunker.respond_to?(:call)
+
+      registry.fetch(chunker.to_sym) do
+        raise ArgumentError, "Unknown chunker #{chunker.inspect}. Use #{registry.keys.join(', ')}, or an object that responds to call."
+      end
+    end
+
+    # Runs a chunker and normalizes its output to Hashes, dropping blank chunks.
+    def self.chunk(text, chunker:, size:, overlap:)
+      fetch(chunker).call(text, size: size, overlap: overlap).filter_map do |chunk|
+        chunk = chunk.is_a?(Hash) ? chunk.to_h.transform_keys(&:to_sym) : { content: chunk.to_s }
+        content = chunk[:content].to_s.strip
+        { content: content, context: chunk[:context].presence, metadata: chunk[:metadata] || {} } unless content.empty?
+      end
+    end
+  end
+end

data/lib/omni_agent/configuration.rb

@@ -2,7 +2,8 @@ module OmniAgent
   class Configuration
     attr_accessor :default_provider, :default_model, :max_retries, :retry_base_delay, :max_tool_iterations,
                   :max_delegation_depth, :eval_judge_provider, :eval_judge_model, :eval_cache_enabled, :eval_cache_path,
-                  :mcp_allowed_origins, :mcp_max_request_bytes
+                  :mcp_allowed_origins, :mcp_max_request_bytes,
+                  :embedding_provider, :embedding_model, :document_loaders
 
     def initialize
       @default_provider = :openai
@@ -17,6 +18,18 @@ module OmniAgent
       @eval_cache_path = "tmp/omni_agent_eval_cache.json"
       @mcp_allowed_origins = []
       @mcp_max_request_bytes = 1_048_576
+      @embedding_provider = nil
+      @embedding_model = nil
+      @embedding_job_options = {}
+      @document_loaders = {}
+    end
+
+    attr_reader :embedding_job_options
+
+    # ActiveJob `set` options for OmniAgent::EmbedJob (has_embedding later: true),
+    # e.g. { queue: :embeddings, priority: 10 }.
+    def embedding_job_options=(options)
+      @embedding_job_options = OmniAgent::Embeddable.validate_job_options!(options, "embedding_job_options")
     end
   end
 end

data/lib/omni_agent/docs.rb

@@ -0,0 +1,32 @@
+module OmniAgent
+  # The plain-Markdown documentation bundled in the gem's docs/ directory, so
+  # coding agents working in a host app can read it offline. Generated from
+  # omniagent-docs/ with `bundle exec rake docs:sync`.
+  module Docs
+    ROOT = File.expand_path("../../docs", __dir__)
+
+    class << self
+      def root
+        ROOT
+      end
+
+      # Page names relative to the docs root, without the .md extension
+      # (e.g. "tools", "mcp/authentication").
+      def pages
+        Dir.glob("**/*.md", base: root).map { |path| path.delete_suffix(".md") }.sort
+      end
+
+      def path(page = "index")
+        name = page.to_s.strip.delete_prefix("/").delete_suffix(".md")
+        return unless pages.include?(name)
+
+        File.join(root, "#{name}.md")
+      end
+
+      def read(page = "index")
+        file = path(page)
+        File.read(file) if file
+      end
+    end
+  end
+end

data/lib/omni_agent/embeddable.rb

@@ -0,0 +1,131 @@
+module OmniAgent
+  # ActiveRecord macro that keeps an embedding of some attributes up to date
+  # and adds `semantic_search`. Extended into every model by the engine.
+  #
+  #   class Document < ApplicationRecord
+  #     has_embedding :title, :content
+  #   end
+  #
+  #   Document.where(account_id: 1).semantic_search("refund policy", limit: 5)
+  #
+  # The default backend is the neighbor gem (pgvector, sqlite-vec, MariaDB).
+  # `later: true` embeds in OmniAgent::EmbedJob after commit instead; pass a
+  # Hash (`later: { queue: :embeddings }`) to override the global
+  # `embedding_job_options` for this model.
+  #
+  # `using: :ruby` stores vectors in a JSON column and ranks them in Ruby;
+  # it loads every row per search, so it only suits small tables.
+  module Embeddable
+    BACKENDS = %i[neighbor ruby].freeze
+    # The options ActiveJob's `set` takes, for `later:` and embedding_job_options.
+    JOB_OPTIONS = %i[queue priority wait wait_until].freeze
+
+    def self.validate_job_options!(options, setting)
+      unknown = options.keys.map(&:to_sym) - JOB_OPTIONS
+      return options.to_h.transform_keys(&:to_sym) if unknown.empty?
+
+      raise ArgumentError, "Unknown #{setting} option #{unknown.first.inspect}. Use #{JOB_OPTIONS.join(', ')}."
+    end
+
+    def has_embedding(*sources, column: :embedding, dimensions: nil, using: :neighbor, later: false)
+      raise ArgumentError, "has_embedding requires at least one attribute" if sources.empty?
+      raise ArgumentError, "has_embedding using: must be :neighbor or :ruby" unless BACKENDS.include?(using)
+      raise ArgumentError, "has_embedding later: must be true, false or a Hash of job options" unless [ true, false ].include?(later) || later.is_a?(Hash)
+      raise ArgumentError, "#{name} already uses has_chunks. A model uses has_embedding or has_chunks, not both." if respond_to?(:chunk_options)
+
+      class_attribute :embedding_options, instance_writer: false
+      self.embedding_options = {
+        sources: sources.map(&:to_sym),
+        column: column.to_sym,
+        using: using,
+        job_options: later.is_a?(Hash) ? Embeddable.validate_job_options!(later, "later:") : {}
+      }
+
+      if using == :neighbor
+        unless respond_to?(:has_neighbors)
+          raise OmniAgent::MissingDependencyError,
+                "has_embedding needs the neighbor gem: add `gem \"neighbor\"` to your Gemfile. " \
+                "Without a vector database, `using: :ruby` works for small tables only."
+        end
+
+        has_neighbors column, dimensions: dimensions
+      end
+
+      extend ClassMethods
+      include InstanceMethods
+
+      if later
+        after_commit :embed_later, on: %i[create update], if: -> { embedding_sources.any? { |source| saved_change_to_attribute?(source) } }
+      else
+        before_save :assign_embedding, if: -> { embedding_sources.any? { |source| will_save_change_to_attribute?(source) } }
+      end
+    end
+
+    module ClassMethods
+      def semantic_search(query, limit: 5)
+        vector = OmniAgent.embed(query)
+        column = embedding_options[:column]
+
+        if embedding_options[:using] == :ruby
+          # ponytail: loads every embedded row and ranks in Ruby, O(n) per search; use neighbor past ~10k rows.
+          all.where.not(column => nil).to_a.max_by(limit) { |record| cosine_similarity(record[column], vector) }
+        else
+          all.nearest_neighbors(column, vector, distance: "cosine").first(limit)
+        end
+      end
+
+      # Embeds records that have no vector yet, or every record with
+      # force: true (e.g. after switching embedding models).
+      def embed_all(force: false, batch_size: 100)
+        column = embedding_options[:column]
+        records = force ? all : all.where(column => nil)
+
+        records.find_in_batches(batch_size: batch_size) do |batch|
+          batch = batch.reject { |record| record.embedding_text.blank? }
+          next if batch.empty?
+
+          OmniAgent.embed(batch.map(&:embedding_text)).zip(batch) do |vector, record|
+            record.update_column(column, vector)
+          end
+        end
+      end
+
+      private
+
+      def cosine_similarity(a, b)
+        norms = Math.sqrt(a.sum { |x| x * x }) * Math.sqrt(b.sum { |x| x * x })
+        norms.zero? ? 0.0 : a.zip(b).sum { |x, y| x * y } / norms
+      end
+    end
+
+    module InstanceMethods
+      # The text that gets embedded, and what a `searches` tool shows the model.
+      def embedding_text
+        embedding_sources.filter_map { |source| public_send(source).presence&.to_s }.join("\n\n")
+      end
+
+      def embed!
+        update_column(embedding_options[:column], embedding_for(embedding_text))
+      end
+
+      private
+
+      def embedding_sources
+        embedding_options[:sources]
+      end
+
+      def embedding_for(text)
+        text.blank? ? nil : OmniAgent.embed(text)
+      end
+
+      def assign_embedding
+        self[embedding_options[:column]] = embedding_for(embedding_text)
+      end
+
+      def embed_later
+        job_options = OmniAgent.configuration.embedding_job_options.merge(embedding_options[:job_options])
+        OmniAgent::EmbedJob.set(**job_options).perform_later(self)
+      end
+    end
+  end
+end

data/lib/omni_agent/engine.rb

@@ -1,3 +1,7 @@
+# EmbedJob (has_embedding later:) needs ActiveJob. Loading its railtie here,
+# during Bundler.require, also works in apps that don't require it themselves.
+require "active_job/railtie"
+
 module OmniAgent
   class Engine < ::Rails::Engine
     isolate_namespace OmniAgent
@@ -5,5 +9,12 @@ module OmniAgent
     initializer "omni_agent.mcp_routing" do
       ActionDispatch::Routing::Mapper.include(OmniAgent::MCP::Routing)
     end
+
+    initializer "omni_agent.active_record" do
+      ActiveSupport.on_load(:active_record) do
+        extend OmniAgent::Embeddable
+        extend OmniAgent::Chunkable
+      end
+    end
   end
 end

data/lib/omni_agent/errors.rb

@@ -7,6 +7,18 @@ module OmniAgent
     class MaxDelegationDepthError < Error; end
     class EvalAssertionError < Error; end
     class MCPError < Error; end
+    # Carries the raw model output so callers can log, retry or repair it.
+    class StructuredOutputError < Error
+      attr_reader :content, :response
+
+      def initialize(message = nil, content: nil, response: nil)
+        super(message)
+        @content = content
+        @response = response
+      end
+    end
+    class UnknownStockAgentError < Error; end
+    class UnsupportedDocumentError < Error; end
   end
 
   Error = Errors::Error
@@ -16,4 +28,7 @@ module OmniAgent
   MaxDelegationDepthError = Errors::MaxDelegationDepthError
   EvalAssertionError = Errors::EvalAssertionError
   MCPError = Errors::MCPError
+  StructuredOutputError = Errors::StructuredOutputError
+  UnknownStockAgentError = Errors::UnknownStockAgentError
+  UnsupportedDocumentError = Errors::UnsupportedDocumentError
 end

data/lib/omni_agent/loaders/html.rb

@@ -0,0 +1,21 @@
+module OmniAgent
+  module Loaders
+    # Visible text only, with block elements on their own paragraphs. Uses
+    # Loofah, which Rails already depends on.
+    module HTML
+      EXTENSIONS = %w[.html .htm .xhtml].freeze
+      CONTENT_TYPES = %w[text/html application/xhtml+xml].freeze
+
+      def self.handles?(filename:, content_type:)
+        Loaders.matches?(filename, content_type, EXTENSIONS, CONTENT_TYPES)
+      end
+
+      def self.call(io, **)
+        require "loofah"
+        document = Loofah.document(io.read)
+        document.css("script, style, noscript, template").remove
+        document.to_text
+      end
+    end
+  end
+end

data/lib/omni_agent/loaders/markdown.rb

@@ -0,0 +1,18 @@
+module OmniAgent
+  module Loaders
+    # Reads the file as-is, so the Markdown and hierarchical chunkers can
+    # split on its headings.
+    module Markdown
+      EXTENSIONS = %w[.md .markdown .mdx].freeze
+      CONTENT_TYPES = %w[text/markdown text/x-markdown text/x-web-markdown].freeze
+
+      def self.handles?(filename:, content_type:)
+        Loaders.matches?(filename, content_type, EXTENSIONS, CONTENT_TYPES)
+      end
+
+      def self.call(io, **)
+        io.read
+      end
+    end
+  end
+end

data/lib/omni_agent/loaders/pdf.rb

@@ -0,0 +1,25 @@
+module OmniAgent
+  module Loaders
+    # Text-based PDFs, via the optional pdf-reader gem. Scanned PDFs have no
+    # text layer and come out empty; they need OCR, in a custom loader.
+    module PDF
+      EXTENSIONS = %w[.pdf].freeze
+      CONTENT_TYPES = %w[application/pdf].freeze
+
+      def self.handles?(filename:, content_type:)
+        Loaders.matches?(filename, content_type, EXTENSIONS, CONTENT_TYPES)
+      end
+
+      def self.call(io, **)
+        begin
+          require "pdf-reader"
+        rescue LoadError
+          raise OmniAgent::MissingDependencyError, "Reading PDFs needs the pdf-reader gem: add `gem \"pdf-reader\"` to your Gemfile."
+        end
+
+        # pdf-reader pads text with spaces to keep the page layout.
+        ::PDF::Reader.new(io).pages.map { |page| page.text.gsub(/[ \t]{2,}/, " ") }.join("\n\n")
+      end
+    end
+  end
+end

data/lib/omni_agent/loaders/text.rb

@@ -0,0 +1,15 @@
+module OmniAgent
+  module Loaders
+    module Text
+      EXTENSIONS = %w[.txt .text .csv .log].freeze
+
+      def self.handles?(filename:, content_type:)
+        content_type.to_s.start_with?("text/") || Loaders.matches?(filename, content_type, EXTENSIONS, [])
+      end
+
+      def self.call(io, **)
+        io.read
+      end
+    end
+  end
+end

data/lib/omni_agent/loaders.rb

@@ -0,0 +1,68 @@
+module OmniAgent
+  # Strategies that turn a document into text for `has_chunks`. A loader is
+  # anything that responds to `call(io, filename:, content_type:)` and returns
+  # a String. Loaders that also answer `handles?(filename:, content_type:)`
+  # take part in auto-detection; register your own in
+  # `config.document_loaders` (they're tried before the built-in ones).
+  module Loaders
+    # Content types that say nothing about the format, so detection skips them.
+    GENERIC_CONTENT_TYPES = [ "application/octet-stream", "binary/octet-stream" ].freeze
+    SNIFF_BYTES = 8192
+
+    def self.registry
+      built_in = { pdf: PDF, html: HTML, markdown: Markdown, text: Text }
+      OmniAgent.configuration.document_loaders.transform_keys(&:to_sym).merge(built_in) { |_name, custom, _built_in| custom }
+    end
+
+    def self.fetch(loader)
+      return loader if loader.respond_to?(:call)
+
+      registry.fetch(loader.to_sym) do
+        raise ArgumentError, "Unknown loader #{loader.inspect}. Use :auto, #{registry.keys.join(', ')}, or an object that responds to call."
+      end
+    end
+
+    # Reads a document into clean text, with the given loader or, for
+    # as: :auto, the first one that handles its filename or content type.
+    # Without either, it looks at the first bytes: PDF and HTML are
+    # recognized, anything else without NUL bytes is read as text.
+    def self.load(io, as: :auto, filename: nil, content_type: nil)
+      content_type = nil if GENERIC_CONTENT_TYPES.include?(content_type)
+      loader = as.to_s == "auto" ? detect(io, filename: filename, content_type: content_type) : fetch(as)
+
+      clean(loader.call(io, filename: filename, content_type: content_type))
+    end
+
+    def self.detect(io, filename:, content_type:)
+      loader = registry.values.find { |candidate| candidate.respond_to?(:handles?) && candidate.handles?(filename: filename, content_type: content_type) }
+      loader || sniff(io) || raise(
+        OmniAgent::UnsupportedDocumentError,
+        "Can't read #{filename || 'this document'} (#{content_type || 'unknown type'}). " \
+        "Pass `as:` a loader, or register one in config.document_loaders."
+      )
+    end
+
+    def self.sniff(io)
+      sample = io.read(SNIFF_BYTES).to_s
+      io.rewind
+
+      if sample.start_with?("%PDF-") then PDF
+      elsif sample.match?(/\A\s*<(!doctype html|html)/i) then HTML
+      elsif !sample.include?("\0") then Text
+      end
+    end
+
+    # Fixes encodings and strips what databases and chunkers trip on:
+    # invalid bytes, NULs (Postgres rejects them), CRLFs and runs of blank lines.
+    def self.clean(text)
+      text.to_s.dup.force_encoding(Encoding::UTF_8).scrub("").delete("\u0000")
+          .gsub(/\r\n?/, "\n").gsub(/[ \t]+$/, "").gsub(/\n{3,}/, "\n\n").strip
+    end
+
+    def self.matches?(filename, content_type, extensions, content_types)
+      content_types.include?(content_type.to_s.split(";").first) || extensions.include?(File.extname(filename.to_s).downcase)
+    end
+
+    private_class_method :detect, :sniff, :clean
+  end
+end

data/lib/omni_agent/providers/base.rb

@@ -13,6 +13,20 @@ module OmniAgent
         raise NotImplementedError, "Providers must implement #chat"
       end
 
+      # Returns one embedding vector (Array of Floats) per text in `texts`.
+      # `model` is the embedding model; it defaults per provider, never to the
+      # chat model.
+      def embed(texts, model: nil)
+        raise NotImplementedError, "#{self.class.name.split('::').last} does not support embeddings"
+      end
+
+      # Providers that return true receive an agent's output schema as the
+      # `output_schema:` keyword of #chat. Otherwise the schema is only
+      # described in the system prompt.
+      def supports_structured_output?
+        false
+      end
+
       def validate_messages!(messages, allowed_roles:)
         unless messages.is_a?(Array)
           raise OmniAgent::Error, "messages must be an array"

data/lib/omni_agent/providers/mock.rb

@@ -1,20 +1,25 @@
+require "json"
+require "digest"
+
 module OmniAgent
   module Providers
     class Mock < Base
       LOREM_IPSUM = "Lorem ipsum dolor sit amet, consectetur adipiscing elit."
 
-      def chat(messages:, tools: [], stream: nil, **_options)
+      def chat(messages:, tools: [], stream: nil, output_schema: nil, **_options)
         validate_messages!(messages, allowed_roles: %i[system user assistant tool])
 
+        content = output_schema ? JSON.generate(sample_output(output_schema)) : LOREM_IPSUM
+
         if stream
-          LOREM_IPSUM.split(" ").each_with_index do |word, index|
+          content.split(" ").each_with_index do |word, index|
             chunk = index.zero? ? word : " #{word}"
             stream.call(OmniAgent::Streaming::Event.text(chunk))
           end
         end
 
         OmniAgent::Providers::Response.new(
-          content: LOREM_IPSUM,
+          content: content,
           raw_request: {
             model: model,
             messages: messages,
@@ -24,7 +29,7 @@ module OmniAgent
             "choices" => [
               {
                 "message" => {
-                  "content" => LOREM_IPSUM
+                  "content" => content
                 }
               }
             ]
@@ -33,6 +38,16 @@ module OmniAgent
         )
       end
 
+      # Deterministic placeholder vectors derived from the text, so specs can
+      # index and search without calling a real embedding model.
+      def embed(texts, model: nil)
+        texts.map { |text| Digest::SHA256.digest(text.to_s).bytes.first(8).map { |byte| byte / 255.0 } }
+      end
+
+      def supports_structured_output?
+        true
+      end
+
       protected
 
       def default_api_key
@@ -42,6 +57,76 @@ module OmniAgent
       def default_model
         "mock"
       end
+
+      private
+
+      # Builds a placeholder value for every property so agents with an output
+      # schema get a parseable response without calling a real model. Raises
+      # when a field needs a value the mock can't invent (pattern, custom
+      # validate:, polymorphic resolve: true).
+      def sample_output(output_schema)
+        unresolvable = output_schema.polymorphics.find { |_name, group| group[:resolve] }
+        unsatisfiable!(unresolvable.first, "it is polymorphic with resolve: true") if unresolvable
+
+        sample = sample_object(output_schema.json_schema)
+
+        output_schema.validators.each do |name, validator|
+          passed = begin
+            validator.call(sample[name]) != false
+          rescue StandardError
+            false
+          end
+          unsatisfiable!(name, "its validate: proc rejects the placeholder #{sample[name].inspect}") unless passed
+        end
+
+        sample
+      end
+
+      def sample_object(schema, path = nil)
+        (schema[:properties] || {}).to_h do |name, property|
+          [ name, sample_value(property, [ path, name ].compact.join(".")) ]
+        end
+      end
+
+      def sample_value(property, path)
+        return property[:enum].first if property[:enum]
+
+        case property[:type]
+        when "string" then sample_string(property, path)
+        when "integer" then sample_number(property, integer: true)
+        when "number" then sample_number(property).to_f
+        when "boolean" then true
+        when "array" then Array.new([ property[:minItems].to_i, 1 ].max) { sample_value(property[:items] || {}, "#{path}[]") }
+        when "object" then sample_object(property, path)
+        end
+      end
+
+      def sample_number(property, integer: false)
+        minimum = property[:minimum]
+        maximum = property[:maximum]
+        return integer ? minimum.ceil : minimum if minimum
+        return integer ? maximum.floor : maximum if maximum && maximum < 0
+
+        0
+      end
+
+      def sample_string(property, path)
+        text = LOREM_IPSUM
+        text = text.ljust(property[:minLength], "x") if property[:minLength]
+        text = text[0, property[:maxLength]] if property[:maxLength]
+
+        if property[:pattern] && !text.match?(Regexp.new(property[:pattern]))
+          unsatisfiable!(path, "its value must match the pattern #{property[:pattern]}")
+        end
+
+        text
+      end
+
+      def unsatisfiable!(field, reason)
+        raise OmniAgent::Error,
+              "The mock provider can't generate structured output for `#{field}` because #{reason}. " \
+              "Stub the provider's #chat response in this test instead."
+      end
     end
   end
 end

data/lib/omni_agent/providers/ollama.rb

@@ -20,6 +20,10 @@ module OmniAgent
         ENV.fetch("OLLAMA_MODEL", "llama3.1")
       end
 
+      def default_embedding_model
+        ENV.fetch("OLLAMA_EMBEDDING_MODEL", "nomic-embed-text")
+      end
+
       private
 
       def base_url