rails_console_ai 0.13.0

data/lib/rails_console_ai/providers/bedrock.rb

@@ -0,0 +1,171 @@
+module RailsConsoleAI
+  module Providers
+    class Bedrock < Base
+      def chat(messages, system_prompt: nil)
+        call_api(messages, system_prompt: system_prompt)
+      end
+
+      def chat_with_tools(messages, tools:, system_prompt: nil)
+        call_api(messages, system_prompt: system_prompt, tools: tools)
+      end
+
+      def format_assistant_message(result)
+        content = []
+        content << { text: result.text } if result.text && !result.text.empty?
+        (result.tool_calls || []).each do |tc|
+          content << {
+            tool_use: {
+              tool_use_id: tc[:id],
+              name: tc[:name],
+              input: tc[:arguments]
+            }
+          }
+        end
+        { role: 'assistant', content: content }
+      end
+
+      def format_tool_result(tool_call_id, result_string)
+        {
+          role: 'user',
+          content: [
+            {
+              tool_result: {
+                tool_use_id: tool_call_id,
+                content: [{ text: result_string.to_s }]
+              }
+            }
+          ]
+        }
+      end
+
+      private
+
+      def call_api(messages, system_prompt: nil, tools: nil)
+        params = {
+          model_id: config.resolved_model,
+          messages: format_messages(messages),
+          inference_config: {
+            max_tokens: config.resolved_max_tokens,
+            temperature: config.temperature
+          }
+        }
+        if system_prompt
+          sys_blocks = [{ text: system_prompt }]
+          sys_blocks << { cache_point: { type: 'default' } } if cache_supported?
+          params[:system] = sys_blocks
+        end
+        if tools
+          bedrock_tools = tools.to_bedrock_format
+          bedrock_tools << { cache_point: { type: 'default' } } if bedrock_tools.any? && cache_supported?
+          params[:tool_config] = { tools: bedrock_tools }
+        end
+
+        debug_bedrock_request(params)
+        response = client.converse(params)
+        debug_bedrock_response(response)
+
+        tool_calls = extract_tool_calls(response)
+        stop = response.stop_reason == 'tool_use' ? :tool_use : :end_turn
+
+        usage = response.usage
+        ChatResult.new(
+          text: extract_text(response),
+          input_tokens: usage&.input_tokens,
+          output_tokens: usage&.output_tokens,
+          cache_read_input_tokens: usage.respond_to?(:cache_read_input_token_count) ? usage.cache_read_input_token_count : nil,
+          cache_write_input_tokens: usage.respond_to?(:cache_write_input_token_count) ? usage.cache_write_input_token_count : nil,
+          tool_calls: tool_calls,
+          stop_reason: stop
+        )
+      rescue aws_error_class => e
+        raise ProviderError, "AWS Bedrock error: #{e.message}"
+      end
+
+      def client
+        @client ||= begin
+          unless defined?(Aws::BedrockRuntime::Client)
+            begin
+              require 'aws-sdk-bedrockruntime'
+            rescue LoadError
+              raise ProviderError,
+                "aws-sdk-bedrockruntime gem is required for the :bedrock provider. Add it to your Gemfile."
+            end
+          end
+          client_opts = {}
+          region = config.respond_to?(:bedrock_region) && config.bedrock_region
+          client_opts[:region] = region if region && !region.empty?
+          t = config.respond_to?(:resolved_timeout) ? config.resolved_timeout : config.timeout
+          client_opts[:http_read_timeout] = t
+          Aws::BedrockRuntime::Client.new(client_opts)
+        end
+      end
+
+      def cache_supported?
+        model = config.resolved_model
+        model.include?('anthropic')
+      end
+
+      def aws_error_class
+        if defined?(Aws::BedrockRuntime::Errors::ServiceError)
+          Aws::BedrockRuntime::Errors::ServiceError
+        else
+          # Fallback if the gem isn't loaded yet (shouldn't happen after client init)
+          StandardError
+        end
+      end
+
+      def format_messages(messages)
+        messages.map do |msg|
+          content = if msg[:content].is_a?(Array)
+                      msg[:content]
+                    else
+                      [{ text: msg[:content].to_s }]
+                    end
+          { role: msg[:role].to_s, content: content }
+        end
+      end
+
+      def extract_text(response)
+        content = response.output&.message&.content
+        return '' unless content.is_a?(Array)
+
+        content.select { |c| c.respond_to?(:text) && c.text }
+               .map(&:text)
+               .join("\n")
+      end
+
+      def extract_tool_calls(response)
+        content = response.output&.message&.content
+        return [] unless content.is_a?(Array)
+
+        content.select { |c| c.respond_to?(:tool_use) && c.tool_use }
+               .map do |c|
+          tu = c.tool_use
+          {
+            id: tu.tool_use_id,
+            name: tu.name,
+            arguments: tu.input || {}
+          }
+        end
+      end
+
+      def debug_bedrock_request(params)
+        return unless config.debug
+
+        msg_count = params[:messages]&.length || 0
+        sys_len = params.dig(:system, 0, :text).to_s.length
+        tool_count = params.dig(:tool_config, :tools)&.length || 0
+        $stderr.puts "\e[33m[debug] Bedrock converse | model: #{params[:model_id]} | #{msg_count} msgs | system: #{sys_len} chars | #{tool_count} tools\e[0m"
+      end
+
+      def debug_bedrock_response(response)
+        return unless config.debug
+
+        usage = response.usage
+        if usage
+          $stderr.puts "\e[36m[debug] response: #{response.stop_reason} | in: #{usage.input_tokens} out: #{usage.output_tokens}\e[0m"
+        end
+      end
+    end
+  end
+end

data/lib/rails_console_ai/providers/local.rb

@@ -0,0 +1,112 @@
+module RailsConsoleAI
+  module Providers
+    class Local < OpenAI
+      private
+
+      def call_api(messages, system_prompt: nil, tools: nil)
+        base_url = config.local_url
+
+        headers = { 'Content-Type' => 'application/json' }
+        api_key = config.local_api_key
+        if api_key && api_key != 'no-key' && !api_key.empty?
+          headers['Authorization'] = "Bearer #{api_key}"
+        end
+
+        conn = build_connection(base_url, headers)
+
+        formatted = []
+        formatted << { role: 'system', content: system_prompt } if system_prompt
+        formatted.concat(format_messages(messages))
+
+        body = {
+          model: config.resolved_model,
+          max_tokens: config.resolved_max_tokens,
+          temperature: config.temperature,
+          messages: formatted
+        }
+        body[:tools] = tools.to_openai_format if tools
+
+        estimated_input_tokens = estimate_tokens(formatted, system_prompt, tools)
+
+        json_body = JSON.generate(body)
+        debug_request("#{base_url}/v1/chat/completions", body)
+        response = conn.post('/v1/chat/completions', json_body)
+        debug_response(response.body)
+        data = parse_response(response)
+        usage = data['usage'] || {}
+
+        prompt_tokens = usage['prompt_tokens']
+        if prompt_tokens && estimated_input_tokens > 0 && prompt_tokens < estimated_input_tokens * 0.5
+          raise ProviderError,
+            "Context truncated by local server: sent ~#{estimated_input_tokens} estimated tokens " \
+            "but server only used #{prompt_tokens}. Increase the model's context window " \
+            "(e.g. num_ctx for Ollama) or reduce conversation length."
+        end
+
+        choice = (data['choices'] || []).first || {}
+        message = choice['message'] || {}
+        finish_reason = choice['finish_reason']
+
+        tool_calls = extract_tool_calls(message)
+
+        # Fallback: some local models (e.g. Ollama) emit tool calls as JSON
+        # in the content field instead of using the structured tool_calls format.
+        # Only match when the JSON "name" is a known tool name to avoid false positives.
+        if tool_calls.empty? && tools
+          tool_names = tools.to_openai_format.map { |t| t.dig('function', 'name') }.compact
+          text_calls = extract_tool_calls_from_text(message['content'], tool_names)
+          if text_calls.any?
+            tool_calls = text_calls
+            finish_reason = 'tool_calls'
+            message['content'] = ''
+          end
+        end
+
+        stop = finish_reason == 'tool_calls' ? :tool_use : :end_turn
+
+        ChatResult.new(
+          text: message['content'] || '',
+          input_tokens: usage['prompt_tokens'],
+          output_tokens: usage['completion_tokens'],
+          tool_calls: tool_calls,
+          stop_reason: stop
+        )
+      end
+
+      def estimate_tokens(messages, system_prompt, tools)
+        chars = system_prompt.to_s.length
+        messages.each { |m| chars += m[:content].to_s.length + (m[:tool_calls].to_s.length) }
+        chars += tools.to_openai_format.to_s.length if tools
+        chars / 4
+      end
+
+      # Parse tool calls emitted as JSON text in the content field.
+      # Only recognizes calls whose "name" matches a known tool name.
+      def extract_tool_calls_from_text(content, tool_names)
+        return [] if content.nil? || content.strip.empty?
+
+        text = content.strip
+        parsed = begin
+                   JSON.parse(text)
+                 rescue JSON::ParserError
+                   match = text.match(/```(?:json)?\s*(\{[\s\S]*?\}|\[[\s\S]*?\])\s*```/)
+                   match ? (JSON.parse(match[1]) rescue nil) : nil
+                 end
+
+        return [] unless parsed
+
+        calls = parsed.is_a?(Array) ? parsed : [parsed]
+        calls.filter_map do |call|
+          next unless call.is_a?(Hash) && tool_names.include?(call['name'])
+          {
+            id: "local_#{SecureRandom.hex(4)}",
+            name: call['name'],
+            arguments: call['arguments'] || {}
+          }
+        end
+      rescue
+        []
+      end
+    end
+  end
+end

data/lib/rails_console_ai/providers/openai.rb

@@ -0,0 +1,114 @@
+module RailsConsoleAI
+  module Providers
+    class OpenAI < Base
+      API_URL = 'https://api.openai.com'.freeze
+
+      def chat(messages, system_prompt: nil)
+        call_api(messages, system_prompt: system_prompt)
+      end
+
+      def chat_with_tools(messages, tools:, system_prompt: nil)
+        call_api(messages, system_prompt: system_prompt, tools: tools)
+      end
+
+      def format_assistant_message(result)
+        msg = { role: 'assistant' }
+        msg[:content] = result.text if result.text && !result.text.empty?
+        if result.tool_calls && !result.tool_calls.empty?
+          msg[:tool_calls] = result.tool_calls.map do |tc|
+            {
+              'id' => tc[:id],
+              'type' => 'function',
+              'function' => {
+                'name' => tc[:name],
+                'arguments' => JSON.generate(tc[:arguments] || {})
+              }
+            }
+          end
+        end
+        msg
+      end
+
+      def format_tool_result(tool_call_id, result_string)
+        {
+          role: 'tool',
+          tool_call_id: tool_call_id,
+          content: result_string.to_s
+        }
+      end
+
+      private
+
+      def call_api(messages, system_prompt: nil, tools: nil)
+        conn = build_connection(API_URL, {
+          'Authorization' => "Bearer #{config.resolved_api_key}"
+        })
+
+        formatted = []
+        formatted << { role: 'system', content: system_prompt } if system_prompt
+        formatted.concat(format_messages(messages))
+
+        body = {
+          model: config.resolved_model,
+          max_tokens: config.resolved_max_tokens,
+          temperature: config.temperature,
+          messages: formatted
+        }
+        body[:tools] = tools.to_openai_format if tools
+
+        json_body = JSON.generate(body)
+        debug_request("#{API_URL}/v1/chat/completions", body)
+        response = conn.post('/v1/chat/completions', json_body)
+        debug_response(response.body)
+        data = parse_response(response)
+        usage = data['usage'] || {}
+
+        choice = (data['choices'] || []).first || {}
+        message = choice['message'] || {}
+        finish_reason = choice['finish_reason']
+
+        tool_calls = extract_tool_calls(message)
+        stop = finish_reason == 'tool_calls' ? :tool_use : :end_turn
+
+        ChatResult.new(
+          text: message['content'] || '',
+          input_tokens: usage['prompt_tokens'],
+          output_tokens: usage['completion_tokens'],
+          tool_calls: tool_calls,
+          stop_reason: stop
+        )
+      end
+
+      def format_messages(messages)
+        messages.map do |msg|
+          base = { role: msg[:role].to_s }
+          if msg[:content]
+            base[:content] = msg[:content].is_a?(Array) ? JSON.generate(msg[:content]) : msg[:content].to_s
+          end
+          base[:tool_calls] = msg[:tool_calls] if msg[:tool_calls]
+          base[:tool_call_id] = msg[:tool_call_id] if msg[:tool_call_id]
+          base
+        end
+      end
+
+      def extract_tool_calls(message)
+        calls = message['tool_calls']
+        return [] unless calls.is_a?(Array)
+
+        calls.map do |tc|
+          func = tc['function'] || {}
+          args = begin
+                   JSON.parse(func['arguments'] || '{}')
+                 rescue
+                   {}
+                 end
+          {
+            id: tc['id'],
+            name: func['name'],
+            arguments: args
+          }
+        end
+      end
+    end
+  end
+end

data/lib/rails_console_ai/railtie.rb

@@ -0,0 +1,34 @@
+require 'rails_console_ai'
+
+module RailsConsoleAI
+  class Railtie < Rails::Railtie
+    rake_tasks do
+      load File.expand_path('../tasks/rails_console_ai.rake', __dir__)
+    end
+
+    console do
+      require 'rails_console_ai/console_methods'
+
+      # Inject into IRB if available
+      if defined?(IRB::ExtendCommandBundle)
+        IRB::ExtendCommandBundle.include(RailsConsoleAI::ConsoleMethods)
+      end
+
+      # Extend TOPLEVEL_BINDING's receiver as well
+      TOPLEVEL_BINDING.eval('self').extend(RailsConsoleAI::ConsoleMethods)
+
+      # Welcome message
+      if $stdout.respond_to?(:tty?) && $stdout.tty?
+        $stdout.puts "\e[36m[RailsConsoleAI v#{RailsConsoleAI::VERSION}] AI assistant loaded. Try: ai \"show me all tables\"\e[0m"
+      end
+
+      # Pre-build context in background
+      Thread.new do
+        require 'rails_console_ai/context_builder'
+        RailsConsoleAI::ContextBuilder.new.build
+      rescue => e
+        RailsConsoleAI.logger.debug("RailsConsoleAI: background context build failed: #{e.message}")
+      end
+    end
+  end
+end

data/lib/rails_console_ai/repl.rb

@@ -0,0 +1,65 @@
+require 'rails_console_ai/channel/console'
+require 'rails_console_ai/conversation_engine'
+
+module RailsConsoleAI
+  class Repl
+    def initialize(binding_context)
+      @binding_context = binding_context
+      @channel = Channel::Console.new
+      @engine = ConversationEngine.new(binding_context: binding_context, channel: @channel)
+    end
+
+    def one_shot(query)
+      @engine.one_shot(query)
+    end
+
+    def explain(query)
+      @engine.explain(query)
+    end
+
+    def init_guide
+      @engine.init_guide
+    end
+
+    def interactive
+      @channel.interactive_loop(@engine)
+    end
+
+    def resume(session)
+      @channel.resume_interactive(@engine, session)
+    end
+
+    # Expose engine internals for specs that inspect state
+    def instance_variable_get(name)
+      case name
+      when :@history
+        @engine.history
+      when :@executor
+        @engine.instance_variable_get(:@executor)
+      else
+        super
+      end
+    end
+
+    # Allow specs to set internal state
+    def instance_variable_set(name, value)
+      case name
+      when :@history
+        @engine.instance_variable_set(:@history, value)
+      else
+        super
+      end
+    end
+
+    private
+
+    # Expose send methods for spec compatibility
+    def send_query(query, conversation: nil)
+      @engine.send(:send_query, query, conversation: conversation)
+    end
+
+    def trim_old_outputs(messages)
+      @engine.send(:trim_old_outputs, messages)
+    end
+  end
+end

data/lib/rails_console_ai/safety_guards.rb

@@ -0,0 +1,207 @@
+module RailsConsoleAI
+  # Raised by safety guards to block dangerous operations.
+  # Host apps should raise this error in their custom guards.
+  # RailsConsoleAI will catch it and guide the user to use 'd' or /danger.
+  class SafetyError < StandardError
+    attr_reader :guard, :blocked_key
+
+    def initialize(message, guard: nil, blocked_key: nil)
+      super(message)
+      @guard = guard
+      @blocked_key = blocked_key
+    end
+  end
+
+  class SafetyGuards
+    attr_reader :guards
+
+    def initialize
+      @guards = {}
+      @enabled = true
+      @allowlist = {}  # { guard_name => [String or Regexp, ...] }
+    end
+
+    def add(name, &block)
+      @guards[name.to_sym] = block
+    end
+
+    def remove(name)
+      @guards.delete(name.to_sym)
+    end
+
+    def enabled?
+      @enabled
+    end
+
+    def enable!
+      @enabled = true
+    end
+
+    def disable!
+      @enabled = false
+    end
+
+    def empty?
+      @guards.empty?
+    end
+
+    def names
+      @guards.keys
+    end
+
+    def allow(guard_name, key)
+      guard_name = guard_name.to_sym
+      @allowlist[guard_name] ||= []
+      @allowlist[guard_name] << key unless @allowlist[guard_name].include?(key)
+    end
+
+    def allowed?(guard_name, key)
+      entries = @allowlist[guard_name.to_sym]
+      return false unless entries
+
+      entries.any? do |entry|
+        case entry
+        when Regexp then key.match?(entry)
+        else entry.to_s == key.to_s
+        end
+      end
+    end
+
+    def allowlist
+      @allowlist
+    end
+
+    # Compose all guards around a block of code.
+    # Each guard is an around-block: guard.call { inner }
+    # Result: guard_1 { guard_2 { guard_3 { yield } } }
+    def wrap(&block)
+      return yield unless @enabled && !@guards.empty?
+
+      @guards.values.reduce(block) { |inner, guard|
+        -> { guard.call(&inner) }
+      }.call
+    end
+  end
+
+  # Built-in guard: database write prevention
+  # Works on all Rails versions (5+) and all database adapters.
+  # Prepends a write-intercepting module once, controlled by a thread-local flag.
+  module BuiltinGuards
+    # Blocks INSERT, UPDATE, DELETE, DROP, CREATE, ALTER, TRUNCATE
+    module WriteBlocker
+      WRITE_PATTERN = /\A\s*(INSERT|UPDATE|DELETE|DROP|CREATE|ALTER|TRUNCATE)\b/i
+      TABLE_PATTERN = /\b(?:INTO|FROM|UPDATE|TABLE|TRUNCATE)\s+[`"]?(\w+)[`"]?/i
+
+      private
+
+      def rails_console_ai_check_write!(sql)
+        return unless Thread.current[:rails_console_ai_block_writes] && sql.match?(WRITE_PATTERN)
+
+        table = sql.match(TABLE_PATTERN)&.captures&.first
+        guards = RailsConsoleAI.configuration.safety_guards
+        return if table && guards.allowed?(:database_writes, table)
+
+        raise RailsConsoleAI::SafetyError.new(
+          "Database write blocked: #{sql.strip.split(/\s+/).first(3).join(' ')}...",
+          guard: :database_writes,
+          blocked_key: table
+        )
+      end
+
+      public
+
+      def execute(sql, *args, **kwargs)
+        rails_console_ai_check_write!(sql)
+        super
+      end
+
+      def exec_delete(sql, *args, **kwargs)
+        rails_console_ai_check_write!(sql)
+        super
+      end
+
+      def exec_update(sql, *args, **kwargs)
+        rails_console_ai_check_write!(sql)
+        super
+      end
+    end
+
+    def self.database_writes
+      ->(& block) {
+        ensure_write_blocker_installed!
+        Thread.current[:rails_console_ai_block_writes] = true
+        begin
+          block.call
+        ensure
+          Thread.current[:rails_console_ai_block_writes] = false
+        end
+      }
+    end
+
+    def self.ensure_write_blocker_installed!
+      return if @write_blocker_installed
+
+      connection = ActiveRecord::Base.connection
+      unless connection.class.ancestors.include?(WriteBlocker)
+        connection.class.prepend(WriteBlocker)
+      end
+      @write_blocker_installed = true
+    end
+
+    # Blocks non-safe HTTP requests (POST, PUT, PATCH, DELETE, etc.) via Net::HTTP.
+    # Since most Ruby HTTP libraries (HTTParty, RestClient, Faraday) use Net::HTTP
+    # under the hood, this covers them all.
+    module HttpBlocker
+      SAFE_METHODS = %w[GET HEAD OPTIONS TRACE].freeze
+
+      def request(req, *args, &block)
+        if Thread.current[:rails_console_ai_block_http] && !SAFE_METHODS.include?(req.method)
+          host = @address.to_s
+          guards = RailsConsoleAI.configuration.safety_guards
+          unless guards.allowed?(:http_mutations, host)
+            raise RailsConsoleAI::SafetyError.new(
+              "HTTP #{req.method} blocked (#{host}#{req.path})",
+              guard: :http_mutations,
+              blocked_key: host
+            )
+          end
+        end
+        super
+      end
+    end
+
+    def self.http_mutations
+      ->(&block) {
+        ensure_http_blocker_installed!
+        Thread.current[:rails_console_ai_block_http] = true
+        begin
+          block.call
+        ensure
+          Thread.current[:rails_console_ai_block_http] = false
+        end
+      }
+    end
+
+    def self.mailers
+      ->(&block) {
+        old_value = ActionMailer::Base.perform_deliveries
+        ActionMailer::Base.perform_deliveries = false
+        begin
+          block.call
+        ensure
+          ActionMailer::Base.perform_deliveries = old_value
+        end
+      }
+    end
+
+    def self.ensure_http_blocker_installed!
+      return if @http_blocker_installed
+
+      require 'net/http'
+      unless Net::HTTP.ancestors.include?(HttpBlocker)
+        Net::HTTP.prepend(HttpBlocker)
+      end
+      @http_blocker_installed = true
+    end
+  end
+end