Module: PWN::AI::Context

Defined in:
lib/pwn/ai/context.rb,
lib/pwn/ai/context_ingestion.rb

Overview

Local evidence ingestion; unavailable embeddings never become synthetic vectors.

Constant Summary collapse

WINDOW =
32_768
CHUNK =
8_192

Class Method Summary collapse

Class Method Details

.attach_disasm(opts = {}) ⇒ Object



44
45
46
47
48
49
50
51
52
53
54
55
56
57
# File 'lib/pwn/ai/context.rb', line 44

public_class_method def self.attach_disasm(opts = {})
  path = opts[:path].to_s
  raise 'ERROR: path is required' if path.empty?

  fn = opts[:function].to_s
  text = if defined?(PWN::Plugins::Radare2) && File.file?(path)
           sid = PWN::Plugins::Radare2.open(path: path)
           addr = fn.empty? ? 'main' : fn
           PWN::Plugins::Radare2.disasm(session: sid, addr: addr, n: 64)
         else
           File.binread(path).to_s[0, CHUNK]
         end
  persist_and_chunk(bytes: text.to_s, path: path, kind: 'disasm', session_id: opts[:session_id], extra: { function: fn })
end

.attach_file(opts = {}) ⇒ Object



17
18
19
20
21
22
23
24
25
26
27
28
29
# File 'lib/pwn/ai/context.rb', line 17

public_class_method def self.attach_file(opts = {})
  path = opts[:path].to_s
  raise 'ERROR: path is required' if path.empty?
  raise "ERROR: file not found: #{path}" unless File.file?(path)

  range = opts[:range]
  data = if range.is_a?(Range)
           File.binread(path, range.end - range.begin, range.begin).to_s
         else
           File.binread(path)
         end
  persist_and_chunk(bytes: data, path: path, kind: 'file', session_id: opts[:session_id], mime: mime_of(path: path))
end

.attach_hexdump(opts = {}) ⇒ Object



31
32
33
34
35
36
37
38
39
40
41
42
# File 'lib/pwn/ai/context.rb', line 31

public_class_method def self.attach_hexdump(opts = {})
  path = opts[:path].to_s
  raise 'ERROR: path is required' if path.empty?
  raise "ERROR: file not found: #{path}" unless File.file?(path)

  offset = opts[:offset].to_i
  length = (opts[:length] || 512).to_i
  length = 512 if length <= 0
  slice = File.binread(path, length, offset).to_s
  hex = slice.unpack1('H*')
  persist_and_chunk(bytes: hex, path: path, kind: 'hexdump', session_id: opts[:session_id], extra: { offset: offset, length: slice.bytesize })
end

.attach_http_transcript(opts = {}) ⇒ Object



59
60
61
62
63
64
65
66
67
68
69
# File 'lib/pwn/ai/context.rb', line 59

public_class_method def self.attach_http_transcript(opts = {})
  src = opts[:har_or_raw] || opts[:har] || opts[:raw] || opts[:path]
  bytes = if src.to_s.empty?
            raise 'ERROR: har_or_raw is required'
          elsif File.file?(src.to_s)
            File.binread(src.to_s)
          else
            src.to_s
          end
  persist_and_chunk(bytes: bytes, path: src.to_s, kind: 'http_transcript', session_id: opts[:session_id])
end

.authorsObject



71
72
73
# File 'lib/pwn/ai/context.rb', line 71

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

.helpObject



75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
# File 'lib/pwn/ai/context.rb', line 75

public_class_method def self.help
  puts "USAGE:
    # Display module authors.
    #{self}.authors

    # Ingest local evidence into the session sqlite database with real local embeddings or explicit lexical degradation.
    #{self}.ingest(
      model: 'optional - installed Ollama embedding model name; defaults to nomic-embed-text',
      path: 'required - filesystem path to the local artifact or binary'
    )

    # Retrieve top-k chunks with source citations from the session evidence database.
    #{self}.retrieve(
      model: 'optional - installed Ollama embedding model name; defaults to nomic-embed-text',
      query: 'required - prompt text used to retrieve matching evidence chunks',
      top_k: 'optional - maximum retrieved chunks; clamped to 1 through 50'
    )
  "
  constants.sort
end

.ingest(opts = {}) ⇒ Object



16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
# File 'lib/pwn/ai/context_ingestion.rb', line 16

public_class_method def self.ingest(opts = {})
  opts = { path: opts } unless opts.is_a?(Hash)
  path = opts[:path]
  source = File.realpath(File.expand_path(path.to_s))
  records = ingestion_records(opts.merge(source: source))
  db, warnings = ingestion_db(opts)
  warnings.concat(records.filter_map { |row| row[:text].lines.first&.strip if row[:text].start_with?('DEGRADED') })
  count = 0
  embedding_error = nil
  db.transaction do
    prefix = "#{source}/"
    db.execute('DELETE FROM chunks WHERE source = ? OR substr(source,1,?) = ?', [source, prefix.length, prefix])
    records.group_by { |row| row[:source] }.each_value do |rows|
      rows.each do |row|
        file = row[:source]
        row[:text].encode('UTF-8', invalid: :replace, undef: :replace).scan(/.{1,2048}/m).each_with_index do |text, index|
          vector, error = embedding_error ? [nil, embedding_error] : ingestion_embedding(opts.merge(text: text))
          embedding_error = error if error
          warnings << error if error
          citation = "#{file}##{row[:locator]}:chunk-#{index + 1}"
          db.execute('INSERT INTO chunks(source,sha256,locator,text,vector,model) VALUES(?,?,?,?,?,?)', [file, row[:sha256], citation, text, vector && JSON.generate(vector), opts[:model] || 'nomic-embed-text'])
          count += 1
        end
      end
    end
  end
  { status: warnings.empty? ? 'ok' : 'degraded', chunks: count, warnings: warnings.uniq, database: ingestion_db_path(opts) }
ensure
  db&.close
end

.retrieve(opts = {}) ⇒ Object



47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
# File 'lib/pwn/ai/context_ingestion.rb', line 47

public_class_method def self.retrieve(opts = {})
  opts = { query: opts } unless opts.is_a?(Hash)
  query = opts[:query]
  db, warnings = ingestion_db(opts)
  vector, error = ingestion_embedding(opts.merge(text: query.to_s))
  warnings << error if error
  semantic = vector && warnings.empty?
  terms = query.to_s.downcase.scan(/[[:alnum:]_]+/).uniq
  rows = db.execute('SELECT source,sha256,locator,text,vector,model FROM chunks')
  if semantic && rows.any? { |row| row[4].nil? || row[5] != (opts[:model] || 'nomic-embed-text') }
    semantic = false
    warnings << 'stored embeddings missing or model mismatch; using lexical retrieval'
  end
  hits = rows.filter_map do |row|
    source, sha, citation, text, stored, model = row
    score = if semantic && stored && model == (opts[:model] || 'nomic-embed-text')
              candidate = JSON.parse(stored)
              next unless candidate.length == vector.length

              db.get_first_value('SELECT vec_distance_cosine(?, ?)', [JSON.generate(vector), stored]).to_f.then { |distance| 1.0 - distance }
            else
              next if semantic

              terms.sum { |term| text.downcase.scan(Regexp.new(Regexp.escape(term))).length }.to_f
            end
    next if !semantic && !score.positive?

    { source: source, sha256: sha, citation: citation, text: text, score: score }
  end
  hits = hits.sort_by { |row| -row[:score] }.first(opts.fetch(:top_k, 5).to_i.clamp(1, 50))
  { status: semantic ? 'ok' : 'degraded', backend: semantic ? 'sqlite-vec' : 'sqlite-lexical', warnings: warnings.uniq, chunks: hits, context: hits.map { |row| "[#{row[:citation]} sha256=#{row[:sha256]}]\n#{row[:text]}" }.join("\n\n") }
ensure
  db&.close
end