Module: PWN::AI::Agent::Registry

Defined in:
lib/pwn/ai/agent/registry.rb

Overview

Central registry for pwn-ai agent tools.

Each file under lib/pwn/ai/agent/tools/*.rb calls PWN::AI::Agent::Registry.register(...) at load time to declare a JSON-Schema (what the LLM sees) and a handler lambda (what pwn runs).

Registry.definitions(...) returns the OpenAI-format tools: array; Registry.lookup(name:) returns the entry for dispatch.

Import chain (circular-import safe):

agent/registry.rb        (no deps on tool files)
     ^
agent/tools/*.rb         (require registry, call .register at top level)
     ^
agent/loop.rb            (calls Registry.discover then .definitions)

DYNAMIC TOOL-SET SLIMMING (local-model scaffolding)

Shipping all ~47 tool schemas on every call overwhelms a 35B local model — it mis-routes (extro_rf_tune for a git question) because the choice space is huge. When PWN::Env[:agent][:tool_router] is truthy AND definitions(relevance: request) is passed, the pool is reduced to CORE_TOOLS + the top-K keyword-ranked matches. Routing accuracy is fed back into Metrics under name:'tool_router' so the router itself becomes a learned component.

Defined Under Namespace

Classes: Entry

Constant Summary collapse

CORE_TOOLS = 
%w[shell pwn_eval memory_remember memory_recall
mistakes_record mistakes_resolve learning_note_outcome].freeze

Class Method Summary collapse

Class Method Details

.all ⇒ Object

Supported Method Parameters

entries = PWN::AI::Agent::Registry.all



87
88
89
 
# File 'lib/pwn/ai/agent/registry.rb', line 87

public_class_method def self.all
  @entries.values
end

.authors ⇒ Object

Author(s)

0day Inc. support@0dayinc.com



191
192
193
 
# File 'lib/pwn/ai/agent/registry.rb', line 191

public_class_method def self.authors
  "AUTHOR(S):\n  0day Inc. <support@0dayinc.com>\n"
end

.definitions(opts = {}) ⇒ Object

Supported Method Parameters

tools = PWN::AI::Agent::Registry.definitions( enabled: 'optional - Array of toolset names to include; nil = all whose check passes', relevance: 'optional - user request; when set AND :tool_router is enabled, slim to CORE + top-K keyword matches', top_k: 'optional - keyword-ranked tools to keep beyond CORE (default 10)' )



105
106
107
108
109
110
111
112
113
114
115
116
117
 
# File 'lib/pwn/ai/agent/registry.rb', line 105

public_class_method def self.definitions(opts = {})
  enabled = opts[:enabled]
  enabled = enabled.map(&:to_s) if enabled
  pool = @entries.values.select { |e| (enabled.nil? || enabled.include?(e.toolset)) && safe_check(entry: e) }

  if opts[:relevance] && router_enabled?
    keep  = rank(query: opts[:relevance], entries: pool).first(opts[:top_k] || 10).map(&:name)
    names = (CORE_TOOLS + keep).uniq
    pool  = pool.select { |e| names.include?(e.name) }
  end

  pool.map { |e| { type: 'function', function: e.schema } }
end

.discover(opts = {}) ⇒ Object

Supported Method Parameters

names = PWN::AI::Agent::Registry.discover( force: 'optional - re-require tool files even if already discovered (default false)' )



152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
 
# File 'lib/pwn/ai/agent/registry.rb', line 152

public_class_method def self.discover(opts = {})
  force = opts[:force] ||= false
  return @entries.keys if @discovered && !force

  tools_dir = File.join(__dir__, 'tools')
  if Dir.exist?(tools_dir)
    Dir[File.join(tools_dir, '*.rb')].each do |f|
      require f
    rescue StandardError, LoadError => e
      warn "[pwn-ai] failed to load tool #{File.basename(f)}: #{e.class}: #{e.message}"
    end
  end
  @discovered = true
  @entries.keys
end

.help ⇒ Object

Display Usage for this Module



197
198
199
200
201
202
203
204
205
206
207
208
209
210
 
# File 'lib/pwn/ai/agent/registry.rb', line 197

public_class_method def self.help
  puts <<~USAGE
    USAGE:
      PWN::AI::Agent::Registry.discover
      PWN::AI::Agent::Registry.definitions(enabled: %w[terminal pwn])
      PWN::AI::Agent::Registry.definitions(relevance: 'nmap sweep 10/8')  # slim (needs :tool_router)
      PWN::AI::Agent::Registry.rank(query: 'run a shell command')
      PWN::AI::Agent::Registry.lookup(name: 'shell')   # => Entry
      PWN::AI::Agent::Registry.toolsets                # => ["memory","pwn","skills","terminal"]
      PWN::AI::Agent::Registry.register(name:, toolset:, schema:, handler:)

      #{self}.authors
  USAGE
end

.lookup(opts = {}) ⇒ Object

Supported Method Parameters

entry = PWN::AI::Agent::Registry.lookup( name: 'required - registered tool name' )



79
80
81
82
 
# File 'lib/pwn/ai/agent/registry.rb', line 79

public_class_method def self.lookup(opts = {})
  name = opts[:name]
  @entries[name.to_s]
end

.rank(opts = {}) ⇒ Object

Supported Method Parameters

ranked = PWN::AI::Agent::Registry.rank( query: 'required - user request text', entries: 'optional - Entry pool to rank (default .all)' )

Lightweight keyword TF router: scores each tool by the number of query tokens (≥ 3 chars, downcased) appearing in its name / description / property names. Zero-score entries drop off; ties break on historical Metrics success_rate so the router learns.



130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
 
# File 'lib/pwn/ai/agent/registry.rb', line 130

public_class_method def self.rank(opts = {})
  query   = opts[:query].to_s.downcase
  entries = opts[:entries] || all
  return entries if query.strip.empty?

  tokens = query.scan(/[a-z0-9_]{3,}/).uniq
  rates  = metrics_rates
  scored = entries.map do |e|
    hay   = "#{e.name} #{e.toolset} #{e.schema[:description]} #{Array(e.schema.dig(:parameters, :properties)&.keys).join(' ')}".downcase
    score = tokens.count { |t| hay.include?(t) }
    [e, score, rates[e.name] || 0.0]
  end
  scored.reject { |_, s, _| s.zero? }
        .sort_by { |_, s, r| [-s, -r] }
        .map(&:first)
end

.register(opts = {}) ⇒ Object

Supported Method Parameters

PWN::AI::Agent::Registry.register( name: 'required - tool name exposed to the model', toolset: 'required - grouping for enable/disable (terminal, file, pwn, memory…)', schema: 'required - OpenAI function schema description:, parameters:', handler: 'required - ->(args_hash) { ... } returning a JSON-serialisable object', check: 'optional - -> { bool } gate; tool only advertised when truthy', max_chars: 'optional - cap on serialised result (default 24_000)' )



58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
 
# File 'lib/pwn/ai/agent/registry.rb', line 58

public_class_method def self.register(opts = {})
  name = opts[:name].to_s
  raise 'ERROR: name is required' if name.empty?
  raise 'ERROR: schema is required' unless opts[:schema]
  raise 'ERROR: handler is required' unless opts[:handler].respond_to?(:call)

  @entries[name] = Entry.new(
    name: name,
    toolset: opts[:toolset].to_s,
    schema: opts[:schema],
    handler: opts[:handler],
    check: opts[:check] ||= -> { true },
    max_chars: opts[:max_chars] ||= 24_000
  )
end

.toolsets ⇒ Object

Supported Method Parameters

names = PWN::AI::Agent::Registry.toolsets



94
95
96
 
# File 'lib/pwn/ai/agent/registry.rb', line 94

public_class_method def self.toolsets
  @entries.values.map(&:toolset).uniq.sort
end