Class: PromptWarden::Policy

Inherits:
Object
  • Object
show all
Includes:
Singleton
Defined in:
lib/prompt_warden/policy.rb

Instance Method Summary collapse

Constructor Details

#initialize ⇒ Policy

Returns a new instance of Policy.



12
13
14
15
16
 
# File 'lib/prompt_warden/policy.rb', line 12

def initialize
  path = ENV['PROMPT_WARDEN_POLICY'] || 'config/promptwarden.yml'
  @rules = File.exist?(path) ? YAML.load_file(path) : {}
  @rules = @rules.transform_keys(&:to_s)
end

Instance Method Details

#check!(prompt:, cost_estimate:) ⇒ Object 



59
60
61
62
63
64
65
66
67
68
69
70
71
 
# File 'lib/prompt_warden/policy.rb', line 59

def check!(prompt:, cost_estimate:)
  max_cost = @rules['max_cost_usd'].to_f
  if max_cost.positive? && cost_estimate > max_cost
    raise PolicyError, "Cost #{cost_estimate}$ exceeds max_cost_usd #{max_cost}"
  end

  Array(@rules['reject_if_regex']).each do |regex|
    r = parse_regex(regex)
    raise PolicyError, "Prompt matches reject regex #{regex}" if prompt.match?(r)
  end

  :ok
end

#check_alerts(prompt:, cost_estimate:) ⇒ Object 



39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
 
# File 'lib/prompt_warden/policy.rb', line 39

def check_alerts(prompt:, cost_estimate:)
  alerts = []

  # Check warn_if_regex patterns
  Array(@rules['warn_if_regex']).each do |regex|
    r = parse_regex(regex)
    if prompt.match?(r)
      alerts << { type: 'regex', rule: regex, level: 'warn' }
    end
  end

  # Check cost limits
  max_cost = @rules['max_cost_usd'].to_f
  if max_cost.positive? && cost_estimate > max_cost
    alerts << { type: 'cost', limit: max_cost, level: 'block' }
  end

  alerts
end

#parse_regex(str) ⇒ Object 



24
25
26
27
28
29
30
31
32
33
34
35
36
37
 
# File 'lib/prompt_warden/policy.rb', line 24

def parse_regex(str)
  # Handles strings like "/pattern/i" or "/pattern/"
  if str.is_a?(String) && str =~ %r{^/(.*?)/([imx]*)$}
    pattern = ::Regexp.last_match(1)
    flags = ::Regexp.last_match(2)
    options = 0
    options |= Regexp::IGNORECASE if flags.include?('i')
    options |= Regexp::MULTILINE  if flags.include?('m')
    options |= Regexp::EXTENDED   if flags.include?('x')
    Regexp.new(pattern, options)
  else
    Regexp.new(str.to_s)
  end
end

#reload! ⇒ Object 



18
19
20
21
22
 
# File 'lib/prompt_warden/policy.rb', line 18

def reload!
  path = ENV['PROMPT_WARDEN_POLICY'] || 'config/promptwarden.yml'
  @rules = File.exist?(path) ? YAML.load_file(path) : {}
  @rules = @rules.transform_keys(&:to_s)
end