Module: PWN::AI::Router

Defined in:
lib/pwn/ai/router.rb

Overview

Task-class model routing from pwn.yaml ai_router / model_routes.

Constant Summary collapse

TASKS = 
i[summarize plan exploit_dev triage].freeze
FRONTIER = 
%w[openai anthropic grok gemini].freeze
LOCAL = 
%w[ollama openwebui].freeze
ENGINE_MOD = 
{
  'ollama' => 'PWN::AI::Ollama',
  'openwebui' => 'PWN::AI::OpenWebUI',
  'openai' => 'PWN::AI::OpenAI',
  'anthropic' => 'PWN::AI::Anthropic',
  'grok' => 'PWN::AI::Grok',
  'gemini' => 'PWN::AI::Gemini'
}.freeze

Class Method Summary collapse

Class Method Details

.authors ⇒ Object 



67
68
69
 
# File 'lib/pwn/ai/router.rb', line 67

public_class_method def self.authors
  "AUTHOR(S):\n  0day Inc. <[email protected]>\n"
end

.help ⇒ Object 



71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
 
# File 'lib/pwn/ai/router.rb', line 71

public_class_method def self.help
  puts "USAGE:
    # List host binaries this module expects to be installed.
    #{self}.required_bins

    # Resolve engine/model fallbacks for a task class.
    #{self}.resolve(
      task: 'optional - summarize, plan, exploit_dev, or triage',
      class: 'optional - alias for task'
    )

    # Summarize text on the cheap local chain; never calls frontier models.
    #{self}.summarize(
      text: 'optional - artifact or page text to summarize',
      content: 'optional - alias for text',
      reset_usage: 'optional - true clears the usage ledger first',
      llm: 'optional - true tries local-engine chat after skipping frontier endpoints'
    )

    # Return per-task token accounting (frontier_tokens is always 0 for summarize).
    #{self}.usage(
      n: 'optional - unused placeholder so the method reads opts'
    )

    # Clear the in-process usage ledger.
    #{self}.reset_usage(
      n: 'optional - unused placeholder so the method reads opts'
    )

    # Print the AUTHOR(S) string for this module.
    #{self}.authors
  "
  constants.sort
end

.required_bins ⇒ Object 



21
22
23
 
# File 'lib/pwn/ai/router.rb', line 21

public_class_method def self.required_bins
  []
end

.reset_usage(opts = {}) ⇒ Object

Clear the in-process usage ledger.



62
63
64
65
 
# File 'lib/pwn/ai/router.rb', line 62

public_class_method def self.reset_usage(opts = {})
  _n = opts[:n]
  @usage = {}
end

.resolve(opts = {}) ⇒ Object

Resolve engine/model fallbacks for a task class.



26
27
28
29
30
31
32
33
 
# File 'lib/pwn/ai/router.rb', line 26

public_class_method def self.resolve(opts = {})
  task = (opts[:task] || opts[:class] || :summarize).to_s.to_sym
  entry = config_entry(task: task)
  endpoints = expand_endpoints(entry: entry)
  allow_frontier = !i[summarize plan].include?(task)
  allow_frontier &&= entry[:allow_frontier] != false && entry['allow_frontier'] != false
  { task: task, endpoints: endpoints, allow_frontier: allow_frontier }
end

.summarize(opts = {}) ⇒ Object

Summarize text on the cheap local chain; never calls frontier models.



36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
 
# File 'lib/pwn/ai/router.rb', line 36

public_class_method def self.summarize(opts = {})
  text = (opts[:text] || opts[:content]).to_s
  reset_task = opts[:reset_usage]
  reset_usage if reset_task
  note_usage(task: :summarize, frontier_tokens: 0)
  if opts[:llm]
    resolve(task: :summarize)[:endpoints].each do |endpoint|
      next if frontier?(engine: endpoint[:engine])

      out = try_chat(endpoint: endpoint, request: "Summarize this artifact for paging. Be terse.\n#{text.byteslice(0, 8_192)}")
      return out if out.to_s.strip != ''
    end
  end
  extractive(text: text)
rescue StandardError
  extractive(text: text)
end

.usage(opts = {}) ⇒ Object

Return per-task token accounting (frontier_tokens is always 0 for summarize).



55
56
57
58
59
 
# File 'lib/pwn/ai/router.rb', line 55

public_class_method def self.usage(opts = {})
  _n = opts[:n]
  @usage ||= {}
  @usage
end