ruby_llm 2.1.1 → 2.2.0

data/lib/ruby_llm/chunk.rb

@@ -5,12 +5,38 @@ module RubyLLM
   # yielded to the block passed to Chat#ask or Chat#complete when streaming.
   # A Chunk is a Message, so it responds to the same readers, but it holds
   # only the data received in that part of the stream, such as a piece of
-  # content, partial tool calls, or the finish reason on the final chunk.
+  # content, partial tool calls, or the finish reason.
   #
   #   chat.ask "Write a haiku about Ruby" do |chunk|
   #     print chunk.content
   #   end
   #
+  # Streamed chunks always have the +:assistant+ role, and most carry the
+  # model ID. #finish_reason, with predicates such as #max_tokens?,
+  # #content_filtered?, #tool_call_stop?, and #stopped?, arrives near the end
+  # of the stream, not necessarily on the last chunk. #tokens is often empty
+  # until then: +output+ counts the billable output tokens up to that chunk,
+  # including thinking tokens the provider bills as output, and providers
+  # differ in when they report it. Read the final reason and usage from the
+  # Message that Chat#ask returns.
+  #
+  # Pass +index:+ to build one with its position, such as in a test:
+  #
+  #   RubyLLM::Chunk.new(role: :assistant, content: "Hel", index: 0)
+  #
   class Chunk < Message
+    def initialize(options = {}) # :nodoc:
+      super
+      @index = options[:index]
+    end
+
+    # Returns the position of this chunk within the response it belongs
+    # to, starting at 0. Each response in a Chat#complete loop, such as
+    # the one after a tool call, starts again at 0. Counts every chunk,
+    # including those without content.
+    #
+    #   chat.ask("Count to three") { |chunk| puts "#{chunk.index}: #{chunk.content}" }
+    #
+    attr_accessor :index
   end
 end

data/lib/ruby_llm/citation.rb

@@ -41,7 +41,8 @@ module RubyLLM
     # content, or +nil+.
     attr_reader :end_index
 
-    # The provider's identifier for the cited source, or +nil+.
+    # The provider's identifier for the cited source, or +nil+. For a cited
+    # SearchResults result, this is its +:source_id+ when it has one.
     attr_reader :source_id
 
     # The 0-indexed position of the source document or search result,

data/lib/ruby_llm/code_execution/language.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  class CodeExecution
+    Language = Data.define(:name, :label, :filename, :program, :check) do # :nodoc:
+      def self.[](name)
+        self::ALL.fetch(name)
+      end
+
+      def command(path)
+        [program, path]
+      end
+
+      def check_command
+        [program, *check]
+      end
+    end
+
+    Language::ALL = {
+      ruby: Language.new(name: :ruby, label: 'Ruby', filename: 'main.rb', program: 'ruby', check: ['-e', 'exit']),
+      python: Language.new(name: :python, label: 'Python', filename: 'main.py', program: 'python3',
+                           check: ['-c', 'pass']),
+      javascript: Language.new(name: :javascript, label: 'JavaScript', filename: 'main.js', program: 'node',
+                               check: ['-e', '0']),
+      shell: Language.new(name: :shell, label: 'shell', filename: 'main.sh', program: 'sh', check: ['-c', 'exit 0'])
+    }.freeze
+  end
+end

data/lib/ruby_llm/code_execution/node_installation.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+require 'open3'
+
+module RubyLLM
+  class CodeExecution
+    # The Node.js the application's PATH finds, as code execution mounts
+    # it: the installation read-only, wherever a version manager put it.
+    # Asks +node+ itself for its executable, since a version manager's
+    # shim is a different program.
+    class NodeInstallation # :nodoc:
+      def self.current
+        @current ||= new
+      end
+
+      def initialize(node: 'node')
+        @node = node
+      end
+
+      def prefix
+        executable && File.dirname(executable, 2)
+      end
+
+      def bin_dir
+        executable && File.dirname(executable)
+      end
+
+      def version
+        details&.last
+      end
+
+      def readable
+        [prefix].compact
+      end
+
+      private
+
+      def executable
+        details&.first
+      end
+
+      def details
+        return @details if defined?(@details)
+
+        @details = begin
+          output, status = Open3.capture2(@node, '-p', 'process.execPath + "\\n" + process.version',
+                                          err: File::NULL)
+          lines = output.lines.map(&:chomp)
+          status.success? && lines.size == 2 ? [File.realpath(lines.first), lines.last] : nil
+        rescue SystemCallError
+          nil
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/code_execution/packages.rb

@@ -0,0 +1,241 @@
+# frozen_string_literal: true
+
+require 'digest'
+require 'fileutils'
+require 'rbconfig'
+require 'tmpdir'
+
+module RubyLLM
+  class CodeExecution
+    # The Ruby and Node.js installations, gems, Python packages, and npm
+    # packages code runs with. Declared packages are installed once per
+    # set, in the sandbox with only their registries reachable, into a
+    # directory under Configuration#package_cache_dir that every later run
+    # mounts read-only. A prebuilt GEM_HOME, virtualenv, or node_modules is
+    # mounted as it is. npm packages install without running their scripts.
+    class Packages # :nodoc:
+      SYSTEM_PATH = %w[/usr/local/bin /usr/bin /bin].freeze
+      RESOLVER_CONFIG = '/etc/resolv.conf'
+      INSTALL_TIMEOUT = 600
+      MAX_FAILURE_OUTPUT = 2000
+      REGISTRIES = {
+        ruby: %w[rubygems.org *.rubygems.org],
+        python: %w[pypi.org files.pythonhosted.org],
+        javascript: %w[registry.npmjs.org]
+      }.freeze
+
+      attr_reader :gems, :python, :npm
+
+      def initialize(sandbox:, cache_dir:, gems: [], python: [], npm: [], gem_home: nil, virtualenv: nil,
+                     node_modules: nil, node: nil, installation: RubyInstallation.current)
+        @sandbox = sandbox
+        @installation = installation
+        @node = node
+        @npm = Array(npm).map(&:to_s)
+        @node_modules = node_modules && File.expand_path(node_modules)
+        @cache_dir = cache_dir && File.expand_path(cache_dir)
+        @gems = Array(gems).map(&:to_s)
+        @python = Array(python).map(&:to_s)
+        @gem_home = gem_home && File.expand_path(gem_home)
+        @virtualenv = virtualenv && File.expand_path(virtualenv)
+        @lock = Mutex.new
+      end
+
+      def gem_names
+        return gems if gems.any? || @gem_home.nil?
+
+        Dir.glob(File.join(@gem_home, 'specifications', '*.gemspec'))
+           .map { |path| File.basename(path, '.gemspec').sub(/-\d[\w.]*(?:-[a-z][^.]*)?\z/, '') }.uniq.sort
+      end
+
+      def python_names
+        return python if python.any? || @virtualenv.nil?
+
+        Dir.glob(File.join(@virtualenv, 'lib', 'python*', 'site-packages', '*.dist-info'))
+           .map { |path| File.basename(path, '.dist-info').sub(/-[^-]+\z/, '') }.uniq.sort
+      end
+
+      def npm_names
+        return npm if npm.any? || @node_modules.nil?
+
+        manifest = File.join(File.dirname(@node_modules), 'package.json')
+        return [] unless File.file?(manifest)
+
+        JSON.parse(File.read(manifest)).fetch('dependencies', {}).keys.sort
+      rescue JSON::ParserError
+        []
+      end
+
+      def readable
+        home = gem_home
+        [*runtimes, home, (ruby_gems if home), virtualenv, node_modules].compact
+      end
+
+      def runtimes
+        [*@installation.readable, *@node&.readable]
+      end
+
+      # Returns +runner+ able to load the libraries Ruby links, when it can
+      # grant them.
+      def with_libraries(runner)
+        runner.respond_to?(:with_libraries_of) ? runner.with_libraries_of(*@installation.libraries) : runner
+      end
+
+      def denied
+        @installation.denied - [@gem_home]
+      end
+
+      def resolver_config(network)
+        return [] unless network && File.exist?(RESOLVER_CONFIG)
+
+        [RESOLVER_CONFIG, File.realpath(RESOLVER_CONFIG)].uniq
+      end
+
+      def env(home, tmp = home)
+        base_env(home, tmp).merge(
+          'PATH' => path,
+          'GEM_HOME' => gem_home,
+          'GEM_PATH' => gem_home && [gem_home, ruby_gems].join(File::PATH_SEPARATOR),
+          'VIRTUAL_ENV' => virtualenv
+        ).compact
+      end
+
+      def gem_home
+        return @gem_home if gems.empty?
+
+        @lock.synchronize { @installed_gem_home ||= install(:ruby, gems) }
+      end
+
+      def virtualenv
+        return @virtualenv if python.empty?
+
+        @lock.synchronize { @installed_virtualenv ||= install(:python, python) }
+      end
+
+      def node_modules
+        return @node_modules if npm.empty?
+
+        @lock.synchronize { @installed_node_modules ||= File.join(install(:javascript, npm), 'node_modules') }
+      end
+
+      def cache_dir
+        @cache_dir || raise(ConfigurationError, 'Set config.package_cache_dir to install code execution packages')
+      end
+
+      def base_env(home, tmp = home)
+        {
+          'PATH' => [RbConfig::CONFIG['bindir'], @node&.bin_dir, *SYSTEM_PATH].compact.join(File::PATH_SEPARATOR),
+          'HOME' => home,
+          'XDG_CACHE_HOME' => File.join(home, '.cache'),
+          'TMPDIR' => tmp,
+          'LANG' => 'C.UTF-8'
+        }
+      end
+
+      def ruby_gems
+        @installation.ruby_gems(@cache_dir)
+      end
+
+      def run(*command, home:, readable: [], writable: [], network: false, env: {})
+        sandbox = with_libraries(@sandbox.is_a?(Proc) ? @sandbox.call : @sandbox)
+        sandbox.run(*command, readable: [*runtimes, *resolver_config(network),
+                                         *readable], denied:,
+                              writable: [*writable, home], network:,
+                              env: base_env(home).merge(env), timeout: INSTALL_TIMEOUT)
+      end
+
+      private
+
+      def path
+        [virtualenv && File.join(virtualenv, 'bin'), gem_home && File.join(gem_home, 'bin'),
+         RbConfig::CONFIG['bindir'], @node&.bin_dir, *SYSTEM_PATH].compact.join(File::PATH_SEPARATOR)
+      end
+
+      def install(kind, packages)
+        directory = File.join(cache_dir, "#{kind}-#{cache_key(kind, packages)}")
+        return directory if File.exist?("#{directory}.ready")
+
+        FileUtils.mkdir_p(cache_dir)
+        File.open("#{directory}.lock", File::RDWR | File::CREAT, 0o644) do |lock|
+          lock.flock(File::LOCK_EX)
+          build(kind, directory, packages) unless File.exist?("#{directory}.ready")
+        end
+        directory
+      end
+
+      def build(kind, directory, packages)
+        FileUtils.rm_rf(directory)
+        FileUtils.mkdir_p(directory)
+        Dir.mktmpdir('ruby_llm-install-') do |home|
+          case kind
+          when :ruby then install_gems(directory, packages, home)
+          when :python then install_python(directory, packages, home)
+          else install_npm(directory, packages, home)
+          end
+        end
+        File.write("#{directory}.ready", "#{packages.sort.join("\n")}\n")
+      end
+
+      def install_gems(directory, gems, home)
+        result = run('ruby', '-S', 'gem', 'install', '--no-document', '--install-dir', directory,
+                     '--bindir', File.join(directory, 'bin'), *gems,
+                     home:, readable: [ruby_gems], writable: [directory], network: { allow: REGISTRIES[:ruby] },
+                     env: { 'GEM_HOME' => directory, 'GEM_PATH' => [directory, ruby_gems].join(File::PATH_SEPARATOR) })
+        verify(result, "gems #{gems.join(', ')}")
+      end
+
+      def install_python(directory, packages, home)
+        description = "Python packages #{packages.join(', ')}"
+        verify(run('python3', '-m', 'venv', directory, home:, writable: [directory]), description)
+        result = run(File.join(directory, 'bin', 'python'), '-m', 'pip', 'install', '--no-input',
+                     '--disable-pip-version-check', *packages,
+                     home:, writable: [directory], network: { allow: REGISTRIES[:python] })
+        verify(result, description)
+      end
+
+      def install_npm(directory, packages, home)
+        raise Error, 'Could not install npm packages: node is not on the PATH' unless @node&.bin_dir
+
+        result = run('npm', 'install', '--prefix', directory, '--ignore-scripts', '--no-audit', '--no-fund',
+                     '--no-update-notifier', *packages,
+                     home:, writable: [directory], network: { allow: REGISTRIES[:javascript] })
+        verify(result, "npm packages #{packages.join(', ')}")
+      end
+
+      def verify(result, description)
+        return if result.exit_code&.zero? && !result.timed_out?
+
+        raise Error, "Could not install #{description}: #{failure(result)}"
+      end
+
+      def failure(result)
+        return "timed out after #{INSTALL_TIMEOUT} seconds" if result.timed_out?
+
+        stderr = result.stderr.to_s.strip
+        blocked = result.respond_to?(:blocked) ? Array(result.blocked) : []
+        details = failure_details(stderr)
+        blocked.empty? ? details : "#{details} (blocked: #{blocked.join(', ')})"
+      end
+
+      def failure_details(stderr)
+        return stderr if stderr.length <= MAX_FAILURE_OUTPUT
+
+        half = MAX_FAILURE_OUTPUT / 2
+        "#{stderr[0, half]}\n[#{stderr.length - MAX_FAILURE_OUTPUT} characters omitted]\n#{stderr[-half..]}"
+      end
+
+      def cache_key(kind, packages)
+        Digest::SHA256.hexdigest(JSON.generate([kind, runtime_version(kind), packages.sort]))[0, 24]
+      end
+
+      def runtime_version(kind)
+        return [RUBY_VERSION, RUBY_PLATFORM, RbConfig.ruby] if kind == :ruby
+        return [@node&.version, RUBY_PLATFORM] if kind == :javascript
+
+        Dir.mktmpdir('ruby_llm-install-') do |home|
+          run('python3', '-c', 'import sys; print(sys.version)', home:).stdout.to_s.strip
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/code_execution/ruby_installation.rb

@@ -0,0 +1,170 @@
+# frozen_string_literal: true
+
+require 'digest'
+require 'fileutils'
+require 'rbconfig'
+require 'tmpdir'
+
+module RubyLLM
+  class CodeExecution
+    # The application's Ruby, as code execution mounts it: the installation
+    # read-only, with every gem in its gem directory hidden except Ruby's
+    # default and bundled gems, and every other gem directory the
+    # application uses hidden entirely, as are Bundler's git checkouts and
+    # RubyGems' cache and documents. RubyGems lists the plugins
+    # directory on every +gem+ command, so only the application's plugins
+    # are hidden, not the directory: macOS refuses to list a hidden one.
+    class RubyInstallation # :nodoc:
+      HIDDEN_DIRECTORIES = %w[bundler cache doc].freeze
+      BUILD_MARKER = 'gem.build_complete'
+
+      # Ruby's bundled gems that Gem::BUNDLED_GEMS::SINCE may leave out, by
+      # the Ruby version that first bundled them. SINCE lists only the gems
+      # that recently stopped being default gems, and an installed Ruby
+      # keeps no other list of them.
+      BUNDLED_GEMS = {
+        '3.2' => %w[minitest power_assert rake test-unit rexml rss net-ftp net-imap net-pop net-smtp matrix prime rbs
+                    typeprof debug],
+        '3.4' => %w[repl_type_completor],
+        '4.0' => %w[readline]
+      }.freeze
+
+      def self.current
+        @current ||= new
+      end
+
+      def initialize(prefix: RbConfig::CONFIG['prefix'], gem_dir: Gem.default_dir, gem_path: Gem.path)
+        @prefix = prefix
+        @gem_dir = gem_dir
+        @gem_path = gem_path
+      end
+
+      def readable
+        @gem_dir.start_with?(File.join(@prefix, '')) ? [@prefix] : [@prefix, @gem_dir]
+      end
+
+      # Returns Ruby, its shared library, and its extensions: the files
+      # whose own libraries, such as Homebrew's libyaml and OpenSSL on
+      # macOS, a sandbox lets Ruby load.
+      def libraries
+        @libraries ||= [RbConfig.ruby, *shared_library,
+                        *Dir.glob(File.join(RbConfig::CONFIG['archdir'], '**', "*.#{RbConfig::CONFIG['DLEXT']}"))]
+                       .select { |path| File.file?(path) }.freeze
+      end
+
+      def denied
+        @denied ||= [
+          *application_gems.flat_map { |stub| gem_paths(stub) },
+          *HIDDEN_DIRECTORIES.map { |name| File.join(@gem_dir, name) },
+          *(@gem_path - [@gem_dir])
+        ].select { |path| File.exist?(path) }.uniq.freeze
+      end
+
+      # Returns a directory, built once under +cache_dir+, or under a
+      # private temporary directory without one, that holds only Ruby's
+      # bundled gems: their gemspecs, and links to their files. A GEM_PATH
+      # with it instead of the gem directory finds no hidden application
+      # gem, so <tt>gem install</tt> installs every dependency and RubyGems
+      # never activates a gem whose files the sandbox hides. Default gems
+      # load from any GEM_PATH.
+      def ruby_gems(cache_dir)
+        cache_dir ||= (@private_dir ||= Dir.mktmpdir('ruby_llm-ruby-gems-'))
+        directory = File.join(cache_dir, "bundled-gems-#{ruby_gems_key}")
+        return directory if File.exist?("#{directory}.ready")
+
+        FileUtils.mkdir_p(cache_dir)
+        File.open("#{directory}.lock", File::RDWR | File::CREAT, 0o644) do |lock|
+          lock.flock(File::LOCK_EX)
+          link_ruby_gems(directory) unless File.exist?("#{directory}.ready")
+        end
+        directory
+      end
+
+      def ruby_gem?(name)
+        default_gems.include?(name) || bundled_gems.include?(name)
+      end
+
+      private
+
+      def shared_library
+        return unless RbConfig::CONFIG['ENABLE_SHARED'] == 'yes'
+
+        File.join(RbConfig::CONFIG['libdir'], RbConfig::CONFIG['LIBRUBY_SO'])
+      end
+
+      def application_gems
+        installed_stubs.reject { |stub| ruby_gem?(stub.name) }
+      end
+
+      def installed_stubs
+        Dir.glob(File.join(@gem_dir, 'specifications', '*.gemspec')).map do |path|
+          Gem::StubSpecification.gemspec_stub(path, @gem_dir, File.join(@gem_dir, 'gems'))
+        end
+      end
+
+      def bundled_gem_stubs
+        installed_stubs.select { |stub| bundled_gems.include?(stub.name) && !default_gems.include?(stub.name) }
+      end
+
+      def ruby_gems_key
+        Digest::SHA256.hexdigest([RbConfig.ruby, RUBY_VERSION, *bundled_gem_stubs.map(&:full_name)].join("\n"))[0, 24]
+      end
+
+      def link_ruby_gems(directory)
+        FileUtils.rm_rf(directory)
+        FileUtils.mkdir_p(File.join(directory, 'specifications'))
+        bundled_gem_stubs.each do |stub|
+          [File.join(@gem_dir, 'gems', stub.full_name), extensions_dir(stub)].each do |path|
+            next unless File.exist?(path)
+
+            link = File.join(directory, path.delete_prefix(File.join(@gem_dir, '')))
+            FileUtils.mkdir_p(File.dirname(link))
+            File.symlink(path, link)
+          end
+          FileUtils.cp(stub.loaded_from, File.join(directory, 'specifications'))
+        end
+        File.write("#{directory}.ready", '')
+      end
+
+      def extensions_dir(stub)
+        File.join(@gem_dir, 'extensions', Gem::Platform.local.to_s, Gem.extension_api_version, stub.full_name)
+      end
+
+      def gem_paths(stub)
+        extensions = extensions_dir(stub)
+        built = Dir.exist?(extensions) ? Dir.children(extensions) - [BUILD_MARKER] : []
+        [File.join(@gem_dir, 'gems', stub.full_name), *built.map { |name| File.join(extensions, name) },
+         File.join(@gem_dir, 'plugins', "#{stub.name}_plugin.rb")]
+      end
+
+      def default_gems
+        @default_gems ||= Dir.glob(File.join(@gem_dir, 'specifications', 'default', '*.gemspec')).map do |path|
+          Gem::StubSpecification.default_gemspec_stub(path, @gem_dir, File.join(@gem_dir, 'gems')).name
+        end
+      end
+
+      def bundled_gems
+        @bundled_gems ||= bundled_gems_file || (bundled_gems_table | bundled_gems_constant)
+      end
+
+      def bundled_gems_table
+        ruby = Gem::Version.new(RUBY_VERSION)
+        BUNDLED_GEMS.select { |version, _| ruby >= Gem::Version.new(version) }.values.flatten
+      end
+
+      def bundled_gems_file
+        path = File.join(RbConfig::CONFIG['srcdir'].to_s, 'gems', 'bundled_gems')
+        return unless File.file?(path)
+
+        File.readlines(path).filter_map { |line| line.split.first unless line.start_with?('#') || line.strip.empty? }
+      end
+
+      def bundled_gems_constant
+        require 'bundled_gems' unless defined?(Gem::BUNDLED_GEMS)
+        Gem::BUNDLED_GEMS::SINCE.keys
+      rescue LoadError, NameError
+        []
+      end
+    end
+  end
+end