omni_agent 0.1.10 → 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/mcp/authenticator.rb

@@ -3,7 +3,8 @@ require "active_support/security_utils"
 module OmniAgent
   module MCP
     class Authenticator
-      Result = Struct.new(:authenticated, :principal) do
+      # error and access_token are only set by `authenticate :oauth`.
+      Result = Struct.new(:authenticated, :principal, :error, :access_token) do
         def authenticated?
           authenticated == true
         end
@@ -21,7 +22,7 @@ module OmniAgent
         if authentication.nil?
           raise OmniAgent::MCPError,
                 "#{@server_class.name} does not declare authentication. " \
-                "Add `authenticate :bearer, tokens: ...`, `authenticate { |request| ... }`, " \
+                "Add `authenticate :bearer, tokens: ...`, `authenticate :oauth, ...`, `authenticate { |request| ... }`, " \
                 "or explicitly opt out with `authenticate :none`."
         end
 
@@ -30,6 +31,8 @@ module OmniAgent
           success(nil)
         when :bearer
           authenticate_bearer(request, authentication[:tokens])
+        when :oauth
+          authentication[:resource_server].authenticate(request)
         when :custom
           principal = authentication[:block].call(request)
           principal ? success(principal) : failure

data/lib/omni_agent/mcp/oauth/access_token.rb

@@ -0,0 +1,60 @@
+module OmniAgent
+  module MCP
+    module OAuth
+      # A validated access token and its claims (JWT payload or introspection
+      # response). It is the principal unless `authenticate :oauth` has a block.
+      class AccessToken
+        attr_reader :token, :claims
+
+        def initialize(token, claims)
+          @token = token
+          @claims = claims.transform_keys(&:to_s).freeze
+        end
+
+        def [](key)
+          claims[key.to_s]
+        end
+
+        def subject
+          claims["sub"]
+        end
+
+        def issuer
+          claims["iss"]
+        end
+
+        def client_id
+          claims["client_id"] || claims["azp"]
+        end
+
+        def audience
+          Array(claims["aud"]).map(&:to_s)
+        end
+
+        # RFC 9068 / RFC 7662 use a space-separated "scope"; some providers
+        # (Okta, Entra ID) use "scp", as a string or an array.
+        def scopes
+          raw = claims.key?("scope") ? claims["scope"] : claims["scp"]
+          raw.is_a?(Array) ? raw.map(&:to_s) : raw.to_s.split
+        end
+
+        def scope?(*names)
+          (names.flatten.map(&:to_s) - scopes).empty?
+        end
+
+        def expires_at
+          Time.at(claims["exp"]) if claims["exp"].is_a?(Numeric)
+        end
+
+        def to_s
+          subject.to_s
+        end
+
+        # Never print the raw token.
+        def inspect
+          "#<#{self.class.name} sub=#{subject.inspect} client_id=#{client_id.inspect} scopes=#{scopes.inspect}>"
+        end
+      end
+    end
+  end
+end

data/lib/omni_agent/mcp/oauth/http_client.rb

@@ -0,0 +1,52 @@
+require "json"
+require "net/http"
+require "openssl"
+require "uri"
+
+module OmniAgent
+  module MCP
+    module OAuth
+      # Minimal JSON client for the authorization server (JWKS, introspection).
+      module HTTPClient
+        OPEN_TIMEOUT = 5
+        READ_TIMEOUT = 5
+
+        NETWORK_ERRORS = [
+          IOError, SystemCallError, SocketError, Timeout::Error, OpenSSL::SSL::SSLError, URI::InvalidURIError
+        ].freeze
+
+        module_function
+
+        def get_json(url)
+          uri = URI.parse(url.to_s)
+          request(uri, Net::HTTP::Get.new(uri))
+        end
+
+        def post_form_json(url, form, basic_auth: nil)
+          uri = URI.parse(url.to_s)
+          post = Net::HTTP::Post.new(uri)
+          post.set_form_data(form)
+          post.basic_auth(*basic_auth) if basic_auth
+          request(uri, post)
+        end
+
+        def request(uri, http_request)
+          http_request["accept"] = "application/json"
+
+          response = Net::HTTP.start(uri.host, uri.port, use_ssl: uri.scheme == "https",
+                                     open_timeout: OPEN_TIMEOUT, read_timeout: READ_TIMEOUT) do |http|
+            http.request(http_request)
+          end
+
+          unless response.is_a?(Net::HTTPSuccess)
+            raise VerificationError, "#{http_request.method} #{uri} returned HTTP #{response.code}"
+          end
+
+          JSON.parse(response.body.to_s)
+        rescue JSON::ParserError, *NETWORK_ERRORS => e
+          raise VerificationError, "#{http_request.method} #{uri} failed: #{e.class}: #{e.message}"
+        end
+      end
+    end
+  end
+end

data/lib/omni_agent/mcp/oauth/introspection_verifier.rb

@@ -0,0 +1,43 @@
+require "uri"
+
+module OmniAgent
+  module MCP
+    module OAuth
+      # Asks the authorization server whether a token is active (RFC 7662).
+      # Works with opaque tokens. Returns the introspection response, or nil
+      # when the token is not active.
+      class IntrospectionVerifier
+        def initialize(url:, client_id: nil, client_secret: nil)
+          raise ArgumentError, "introspection: requires url:" if url.nil?
+
+          @url = url
+          @client_id = client_id
+          @client_secret = client_secret
+        end
+
+        def call(token, _request = nil)
+          response = HTTPClient.post_form_json(
+            OAuth.resolve(@url),
+            { "token" => token, "token_type_hint" => "access_token" },
+            basic_auth: credentials
+          )
+          response.is_a?(Hash) && response["active"] == true ? response : nil
+        end
+
+        def inspect
+          "#<#{self.class.name} url=#{@url.inspect}>"
+        end
+
+        private
+
+        # client_secret_basic: RFC 6749 section 2.3.1 form-encodes both values.
+        def credentials
+          client_id = OAuth.resolve(@client_id)
+          return if client_id.nil?
+
+          [ client_id, OAuth.resolve(@client_secret) ].map { |value| URI.encode_www_form_component(value.to_s) }
+        end
+      end
+    end
+  end
+end