ai_guardrails 1.2.0

data/lib/ai_guardrails/background_job.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+module AiGuardrails
+  # Provides helper methods to run AiGuardrails in background jobs or CLI
+  module BackgroundJob
+    class << self
+      # Executes a task safely in background or CLI
+      #
+      # Example usage:
+      #   AiGuardrails::BackgroundJob.perform do
+      #     AiGuardrails::DSL.run(prompt: "...", schema: {...})
+      #   end
+      #
+      # Optional parameters:
+      #   logger: custom logger instance
+      #   debug: true/false for debug mode
+      def perform(logger: nil, debug: false, &block)
+        with_temp_logger(logger, debug, &block)
+      rescue StandardError => e
+        Logger.logger&.error("Background job failed: #{e.class} - #{e.message}")
+        raise e
+      end
+
+      def with_temp_logger(temp_logger, temp_debug, &block)
+        prev_logger = Logger.logger
+        prev_debug = Logger.debug_mode
+
+        Logger.logger = temp_logger if temp_logger
+        Logger.debug_mode = temp_debug
+
+        perform_with_error_logging(&block)
+      ensure
+        Logger.logger = prev_logger
+        Logger.debug_mode = prev_debug
+      end
+
+      private
+
+      def perform_with_error_logging(&block)
+        block.call
+      rescue StandardError => e
+        Logger.logger&.error("Background job failed: #{e.class} - #{e.message}")
+        raise e
+      end
+    end
+  end
+end

data/lib/ai_guardrails/cache.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+require "digest"
+
+module AiGuardrails
+  # Simple caching layer for AI responses
+  module Cache
+    class << self
+      attr_accessor :enabled, :store, :expires_in
+
+      # Setup cache
+      # store: any object responding to #read/#write (e.g., Rails.cache, ActiveSupport::Cache)
+      # expires_in: seconds
+      def configure(enabled: true, store: nil, expires_in: 300)
+        @enabled = enabled
+        @store = store || NullStore.new
+        @expires_in = expires_in
+      end
+
+      # Accept default or block, works with caching disabled
+      def fetch(key, default = nil)
+        return (block_given? ? yield : default) unless enabled
+
+        cached = store.read(key, expires_in: expires_in)
+        return cached if cached
+
+        result = block_given? ? yield : default
+        store.write(key, result, expires_in: expires_in)
+        result
+      end
+
+      # Generate a cache key from prompt + schema
+      def key(prompt, schema)
+        digest_input = "#{prompt}-#{schema}"
+        Digest::SHA256.hexdigest(digest_input)
+      end
+
+      # Null object if no cache store is provided
+      class NullStore
+        def read(_key, **_options)
+          nil
+        end
+
+        def write(_key, value, **_options)
+          value
+        end
+      end
+    end
+  end
+end

data/lib/ai_guardrails/cli.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module AiGuardrails
+  # Provides a CLI-friendly interface for running AiGuardrails safely
+  module CLI
+    # Runs AiGuardrails safely in CLI scripts
+    #
+    # Example:
+    #   AiGuardrails::CLI.run do
+    #     result = AiGuardrails::DSL.run(prompt: "...", schema: {...})
+    #     puts result
+    #   end
+    def self.run(debug: false, &block)
+      BackgroundJob.perform(logger: Logger.logger, debug: debug, &block)
+    end
+  end
+end

data/lib/ai_guardrails/config.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+module AiGuardrails
+  # Holds configuration options for AiGuardrails.
+  class Config
+    attr_accessor :logger, :debug
+
+    def initialize
+      @logger = nil
+      @debug = false
+    end
+  end
+end

data/lib/ai_guardrails/dsl.rb

@@ -0,0 +1,101 @@
+# frozen_string_literal: true
+
+module AiGuardrails
+  # Provides a simple developer-friendly interface
+  module DSL
+    class << self
+      # Main entry point used by developers.
+      # Run AI model with validation, retries, and safety checks.
+      def run(prompt:, schema:, schema_hint: nil, **options)
+        Cache.fetch(Cache.key(prompt, schema)) do
+          result = fetch_with_retries_and_correction(prompt, schema, schema_hint, options)
+
+          puts "result in DSL: #{result}"
+
+          # Apply JSON + schema auto-fix when hooks are given.
+          hooks = options.fetch(:auto_fix_hooks, [])
+          fix_schema = schema_hint || schema
+          result = apply_auto_fix(result, fix_schema, hooks) unless hooks.empty?
+
+          check_safety(result, options.fetch(:blocklist, []))
+          result
+        end
+      end
+
+      private
+
+      # Extracted to reduce run method length
+      def fetch_with_retries_and_correction(prompt, schema, schema_hint, options)
+        client = build_client(options.fetch(:provider, :openai), options.fetch(:provider_config, {}))
+        max_retries = options.fetch(:max_retries, 3)
+        sleep_time = options.fetch(:sleep_time, 0)
+        run_with_retries_helper(
+          client: client, schema: schema, prompt: prompt,
+          max_retries: max_retries,
+          sleep_time: sleep_time,
+          schema_hint: schema_hint
+        )
+      end
+
+      # Builds the provider client
+      def build_client(provider, config)
+        Provider::Factory.build(provider: provider, config: config)
+      end
+
+      # Runs AutoCorrection wrapper (max 5 parameters)
+      def run_with_retries_helper(options = {})
+        client      = options[:client]
+        schema      = options[:schema]
+        prompt      = options[:prompt]
+        max_retries = options[:max_retries] || 3
+        sleep_time  = options[:sleep_time] || 0
+        schema_hint = options[:schema_hint]
+
+        auto = AutoCorrection.new(
+          provider: client, schema: schema, max_retries: max_retries, sleep_time: sleep_time
+        )
+        auto.call(prompt: prompt, schema_hint: schema_hint)
+      end
+
+      # Applies blocklist filtering when needed
+      def apply_auto_fix(result, schema, hooks)
+        AiGuardrails::AutoFix.fix(result, schema: schema, hooks: hooks)
+      end
+
+      # Runs safety filter when needed.
+      def check_safety(result, blocklist)
+        return if blocklist.empty?
+
+        content = normalize_result(result)
+        check_blocklist(content, blocklist)
+      end
+
+      # Normalizes result into a simple string for safety scanning.
+      def normalize_result(result)
+        case result
+        when Hash
+          result.values.join(" ")
+        when String
+          parse_json_string(result)
+        else
+          result.to_s
+        end
+      end
+
+      # Attempt to parse string as JSON; fallback to original string if parsing fails
+      def parse_json_string(str)
+        parsed = JSON.parse(str)
+        parsed.is_a?(Hash) ? parsed.values.join(" ") : str
+      rescue JSON::ParserError
+        str
+      end
+
+      # Perform case-insensitive safety check using SafetyFilter
+      def check_blocklist(content, blocklist)
+        content_down = content.downcase
+        blocklist_down = blocklist.map(&:downcase)
+        SafetyFilter.new(blocklist: blocklist_down).check!(content_down)
+      end
+    end
+  end
+end

data/lib/ai_guardrails/json_repair.rb

@@ -0,0 +1,234 @@
+# frozen_string_literal: true
+
+require "json"
+
+module AiGuardrails
+  # Repairs malformed JSON strings
+  # rubocop:disable Metrics/ClassLength
+  class JsonRepair
+    class RepairError < StandardError; end
+
+    # Class method entrypoint
+    def self.repair(raw)
+      new(raw).repair
+    end
+
+    def initialize(raw)
+      @raw = raw.to_s.strip
+    end
+
+    # Main repair
+    def repair
+      raw_sanitized = sanitize_llm_output(@raw) # move here
+      return JSON.parse(raw_sanitized) if valid_json?(raw_sanitized)
+
+      repaired = run_full_repair(raw_sanitized)
+      raise RepairError, "Unable to repair JSON" unless valid_json?(repaired)
+
+      JSON.parse(repaired)
+    end
+
+    private
+
+    # --------------------------
+    # Full repair workflow extracted from original repair
+    # --------------------------
+    def run_full_repair(str)
+      str = preprocess(str)
+      str = normalize_structure(str)
+      str = balance_braces(str)
+      str = run_recursive_fixes(str)
+      str = remove_trailing_commas(str)
+      str.gsub(/\s+/, " ").strip
+    end
+
+    # --------------------------
+    # JSON validation
+    # --------------------------
+    def valid_json?(str)
+      JSON.parse(str)
+      true
+    rescue JSON::ParserError
+      false
+    end
+
+    # --------------------------
+    # Preprocessing
+    # --------------------------
+    def preprocess(str)
+      str = str.strip
+      str.gsub!("'", '"')
+      str = quote_all_keys(str)
+      str = insert_missing_commas_regex(str)
+      remove_trailing_commas(str)
+    end
+
+    # --------------------------
+    # Quote keys
+    # --------------------------
+    def quote_all_keys(str)
+      prev = nil
+      current = str.dup
+      while current != prev
+        prev = current
+        current.gsub!(/([{\s,])([a-zA-Z0-9_-]+)\s*:/, '\1"\2":')
+      end
+      current
+    end
+
+    def insert_missing_commas_regex(str)
+      str.gsub(/([}\]"0-9a-zA-Z])\s+("?[\w-]+"?\s*:)/, '\1, \2')
+    end
+
+    # --------------------------
+    # Normalization
+    # --------------------------
+    def normalize_structure(input)
+      repaired = input.dup
+      repaired = fix_double_braces(repaired)
+      repaired = fix_object_brace_spacing(repaired)
+      repaired = insert_missing_commas_by_scanner(repaired)
+      repaired.gsub!(/([}\]])\s*(?=([A-Za-z0-9_"-]+\s*:))/, '\1, ')
+      repaired.gsub!(/([}\]])\s*(?=(\{|\[|"|\d|true|false|null))/i, '\1, ')
+      repaired.gsub!(/,+/, ",")
+      repaired.gsub!(/\s+/, " ")
+      repaired.strip
+    end
+
+    def fix_double_braces(str)
+      prev = nil
+      current = str.dup
+      while current != prev
+        prev = current
+        current.gsub!(/(\[|,)\s*\{\s*\{/, '\1 {')
+      end
+      current
+    end
+
+    def fix_object_brace_spacing(str)
+      str.gsub(/}\s*{/, "}, {")
+         .gsub(/]\s*{/, "], {")
+         .gsub(/}\s*\]\s*\{/, "}], {")
+    end
+
+    # --------------------------
+    # Recursive fixes runner
+    # --------------------------
+    def run_recursive_fixes(str)
+      str = quote_all_keys(str)
+      str = insert_commas_recursively(str)
+      str = insert_final_commas(str)
+      str = insert_commas_recursive_nested(str)
+      str = fix_consecutive_objects_in_arrays(str)
+      str = fix_double_object_braces(str)
+      fix_adjacent_arrays(str)
+    end
+
+    def fix_adjacent_arrays(str)
+      str.gsub(/\]\s*\[/, "], [")
+    end
+
+    # --------------------------
+    # Scanner-based comma insertion
+    # --------------------------
+    def insert_missing_commas_by_scanner(str)
+      s = str.dup
+      out_chars = []
+      i = 0
+      while i < s.length
+        char = s[i]
+        out_chars << char
+        insert_comma_after_close_brace?(char, s, i, out_chars)
+        i += 1
+      end
+      out_chars.join.gsub(/,+/, ",").gsub(/\s+/, " ").strip
+    end
+
+    # rubocop:disable Metrics/CyclomaticComplexity
+    def insert_comma_after_close_brace?(char, string, index, output_chars)
+      return unless ["}", "]"].include?(char)
+
+      j = index + 1
+      j += 1 while j < string.length && string[j] =~ /\s/
+      next_char = j < string.length ? string[j] : nil
+      return unless next_char && ![",", "}", "]", ":"].include?(next_char)
+
+      output_chars << "," if next_char =~ /[\[\]"0-9A-Za-z_-]/
+    end
+    # rubocop:enable Metrics/CyclomaticComplexity
+
+    # --------------------------
+    # Recursive comma insertion
+    # --------------------------
+    def insert_commas_recursively(str)
+      loop do
+        prev = str.dup
+        str.gsub!(/([}\]"0-9a-zA-Z])\s+(?=(\{|"[^"]*"|\d+|true|false|null|\[))/i, '\1, ')
+        str.gsub!(/(\})\s+(?=\{)/, '\1, ')
+        str.gsub!(/(\])\s+(?=\[)/, '\1, ')
+        break if str == prev
+      end
+      str
+    end
+
+    def insert_final_commas(str)
+      loop do
+        prev = str.dup
+        str.gsub!(/([}\]])\s+(?=("[a-zA-Z_][a-zA-Z0-9_]*"\s*:))/, '\1, ')
+        str.gsub!(/([}\]])\s+(?=[{\[])/, '\1, ')
+        break if str == prev
+      end
+      str
+    end
+
+    def insert_commas_recursive_nested(str)
+      loop do
+        prev = str.dup
+        str.gsub!(/}\s*(?=\{)/, "}, {")
+        str.gsub!(/]\s*(?=\[)/, "], [")
+        str.gsub!(/([}\]])\s+(?=("[^"]+"\s*:))/, '\1, ')
+        break if str == prev
+      end
+      str
+    end
+
+    def fix_consecutive_objects_in_arrays(str)
+      loop do
+        prev = str.dup
+        str.gsub!(/({[^{}]*})\s*(?=\{)/, '\1, ')
+        break if str == prev
+      end
+      str
+    end
+
+    def fix_double_object_braces(str)
+      fix_double_braces(str)
+    end
+
+    def remove_trailing_commas(str)
+      str.gsub(/,(\s*[}\]])/, '\1')
+    end
+
+    def balance_braces(str)
+      open_braces = str.count("{")
+      close_braces = str.count("}")
+      open_brackets = str.count("[")
+      close_brackets = str.count("]")
+
+      str + "}" * [open_braces - close_braces, 0].max + "]" * [open_brackets - close_brackets, 0].max
+    end
+
+    def sanitize_llm_output(str)
+      return str unless str.is_a?(String)
+
+      # Remove everything before the first ```json or ```
+      sanitized = str.sub(/\A.*?```(?:json)?\s*/m, "")
+
+      # Remove trailing ```
+      sanitized = sanitized.sub(/```\s*\z/, "")
+
+      sanitized.strip
+    end
+  end
+  # rubocop:enable Metrics/ClassLength
+end

data/lib/ai_guardrails/logger.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+module AiGuardrails
+  # Simple wrapper for logging inside the gem.
+  # Allows the user to pass any logger (Rails.logger, Logger.new, etc.)
+  module Logger
+    class << self
+      attr_accessor :logger, :debug_mode
+
+      # Logs normal information
+      def info(message)
+        safe_logger.info("[AiGuardrails] #{message}")
+      end
+
+      # Logs errors only
+      def error(message)
+        safe_logger.error("[AiGuardrails ERROR] #{message}")
+      end
+
+      # Logs extra details when debug_mode is enabled
+      def debug(message)
+        return unless debug_mode
+
+        safe_logger.debug("[AiGuardrails DEBUG] #{message}")
+      end
+
+      private
+
+      # Uses null logger if no logger is configured
+      def safe_logger
+        logger || NullLogger.new
+      end
+    end
+
+    # Basic fallback logger that ignores messages.
+    # Prevents NoMethodError when users don't set a logger.
+    class NullLogger
+      def info(_msg); end
+
+      def error(_msg); end
+
+      def debug(_msg); end
+    end
+  end
+end

data/lib/ai_guardrails/mock_model_client.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+module AiGuardrails
+  # MockModelClient simulates AI LLM responses for tests
+  class MockModelClient
+    class MockError < StandardError; end
+
+    # Initialize with a hash of prompt => response
+    def initialize(responses = {})
+      @responses = responses.transform_keys(&:to_s)
+    end
+
+    # Simulates a call to the model
+    # Options can include:
+    # - prompt: string
+    # - raise_error: boolean to simulate API failure
+    def call(prompt:, raise_error: false, default_fallback: nil)
+      return default_fallback if raise_error == false && !@responses.key?(prompt.to_s)
+
+      raise MockError, "Simulated model error" if raise_error
+
+      response = @responses[prompt.to_s]
+
+      raise MockError, "No mock response defined for prompt: #{prompt}" unless response
+
+      response
+    end
+
+    # Add or update mock responses dynamically
+    def add_response(prompt, response)
+      @responses[prompt.to_s] = response
+    end
+  end
+end

data/lib/ai_guardrails/provider/base_client.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module AiGuardrails
+  module Provider
+    # BaseClient defines a common interface for all providers
+    class BaseClient
+      # Initialize with optional config hash
+      def initialize(config = {})
+        @config = config
+      end
+
+      # Call AI model with a prompt
+      # Must be implemented by subclasses
+      def call_model(prompt:)
+        raise NotImplementedError, "Subclasses must implement call_model"
+      end
+    end
+  end
+end

data/lib/ai_guardrails/provider/factory.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+module AiGuardrails
+  module Provider
+    # Factory returns the right provider client
+    class Factory
+      PROVIDERS = {
+        openai: OpenAIClient
+        # add :anthropic => AnthropicClient later
+      }.freeze
+
+      def self.build(provider:, config: {})
+        klass = PROVIDERS[provider.to_sym]
+        raise ArgumentError, "Unknown provider: #{provider}" unless klass
+
+        klass.new(config)
+      end
+    end
+  end
+end

data/lib/ai_guardrails/provider/openai_client.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+module AiGuardrails
+  module Provider
+    # Handles actual OpenAI API calls.
+    # The ruby-openai gem is only loaded when call_model is used.
+    class OpenAIClient < BaseClient
+      def initialize(config = {})
+        super
+        @client = nil
+      end
+
+      # Actual API call method
+      def call_model(prompt:)
+        ensure_provider_loaded
+
+        @client ||= ::OpenAI::Client.new(access_token: @config[:api_key])
+
+        response = @client.chat(
+          parameters: {
+            model: @config[:model] || "gpt-4o-mini",
+            messages: [{ role: "user", content: prompt }],
+            temperature: @config[:temperature] || 0.7
+          }
+        )
+
+        response.dig("choices", 0, "message", "content")
+      end
+
+      private
+
+      # Load ruby-openai only when needed
+      def ensure_provider_loaded
+        require "ruby/openai"
+      rescue LoadError
+        raise LoadError,
+              "ruby-openai gem is not installed. Add:\n" \
+              "  gem 'ruby-openai', require: false\n" \
+              "to your Gemfile if using OpenAI provider."
+      end
+    end
+  end
+end

data/lib/ai_guardrails/runner.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+module AiGuardrails
+  # Coordinates the full validation and repair flow.
+  class Runner
+    def initialize(prompt:, provider:, schema:, options: {})
+      @prompt = prompt
+      @provider = provider
+      @schema = schema
+      @options = options
+    end
+
+    # rubocop:disable Metrics/MethodLength
+    def run
+      Logger.info("Starting run")
+      Logger.debug("Prompt: #{@prompt}")
+
+      raw = @provider.call_model(prompt: @prompt)
+
+      Logger.debug("Raw model output: #{raw.inspect}")
+
+      repaired_json = JsonRepair.repair(raw)
+      Logger.debug("Repaired JSON: #{repaired_json.inspect}")
+
+      valid, result = SchemaValidator.new(@schema).validate(repaired_json)
+
+      unless valid
+        Logger.error("Schema validation failed: #{result}")
+        return { ok: false, errors: result }
+      end
+
+      Logger.info("Run completed successfully")
+      { ok: true, result: result }
+    rescue StandardError => e
+      Logger.error("Unhandled exception: #{e.class} - #{e.message}")
+      raise e
+    end
+    # rubocop:enable Metrics/MethodLength
+  end
+end