Module: PWN::AI::Ollama

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

Overview

This plugin is used for interacting w/ Ollama's REST API using the 'rest' browser type of PWN::Plugins::TransparentBrowser. This is based on the following Ollama API Specification: https://api.openai.com/v1

Class Method Summary collapse

Class Method Details

.authors ⇒ Object

Author(s)

0day Inc. support@0dayinc.com



313
314
315
316
317
 
# File 'lib/pwn/ai/ollama.rb', line 313

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

.chat(opts = {}) ⇒ Object

Supported Method Parameters

response = PWN::AI::Ollama.chat( request: 'required - message to Ollama' model: 'optional - model to use for text generation (defaults to PWN::Env[:ollama][:model])', temp: 'optional - creative response float (deafults to PWN::Env[:ollama][:temp])', system_role_content: 'optional - context to set up the model behavior for conversation (Default: PWN::Env[:ollama][:system_role_content])', response_history: 'optional - pass response back in to have a conversation', speak_answer: 'optional speak answer using PWN::Plugins::Voice.text_to_speech (Default: nil)', timeout: 'optional timeout in seconds (defaults to 900)', spinner: 'optional - display spinner (defaults to false)' )



229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
 
# File 'lib/pwn/ai/ollama.rb', line 229

public_class_method def self.chat(opts = {})
  engine = PWN::Env[:ai][:ollama]
  request = opts[:request]
  max_prompt_length = engine[:max_prompt_length] ||= 1_000_000
  request_trunc_idx = ((max_prompt_length - 1) / 3.36).floor
  request = request[0..request_trunc_idx]

  model = opts[:model] ||= engine[:model]
  raise 'ERROR: Model is required.  Call #get_models method for details' if model.nil?

  temp = opts[:temp].to_f ||= engine[:temp].to_f
  temp = 1 if temp.zero?

  rest_call = 'ollama/v1/chat/completions'

  response_history = opts[:response_history]

  max_tokens = response_history[:usage][:total_tokens] unless response_history.nil?

  system_role_content = opts[:system_role_content] ||= engine[:system_role_content]

  system_role = {
    role: 'system',
    content: system_role_content
  }

  user_role = {
    role: 'user',
    content: request
  }

  response_history ||= { choices: [system_role] }
  choices_len = response_history[:choices].length

  http_body = {
    model: model,
    messages: [system_role],
    temperature: temp,
    stream: false
  }

  if response_history[:choices].length > 1
    response_history[:choices][1..-1].each do |message|
      http_body[:messages].push(message)
    end
  end

  http_body[:messages].push(user_role)

  timeout = opts[:timeout]
   = opts[:spinner]

  response = ollama_rest_call(
    http_method: :post,
    rest_call: rest_call,
    http_body: http_body,
    timeout: timeout,
    spinner: 
  )

  json_resp = JSON.parse(response, symbolize_names: true)
  assistant_resp = json_resp[:choices].first[:message]
  json_resp[:choices] = http_body[:messages]
  json_resp[:choices].push(assistant_resp)

  speak_answer = true if opts[:speak_answer]

  if speak_answer
    answer = assistant_resp[:content]
    text_path = "/tmp/#{SecureRandom.hex}.pwn_voice"
    # answer = json_resp[:choices].last[:text]
    # answer = json_resp[:choices].last[:content] if gpt
    File.write(text_path, answer)
    PWN::Plugins::Voice.text_to_speech(text_path: text_path)
    File.unlink(text_path)
  end

  json_resp
rescue StandardError => e
  raise e
end

.chat_with_tools(opts = {}) ⇒ Object

Supported Method Parameters

response = PWN::AI::Ollama.chat_with_tools( messages: 'required - full OpenAI-format messages array (system/user/assistant/tool)', tools: 'optional - OpenAI tools array [function:{...}]', tool_choice: 'optional - "auto" | "none" | function:{name:..}', model: 'optional - overrides PWN::Env[:ollama][:model]', temp: 'optional - temperature (defaults to PWN::Env[:ollama][:temp] || 1)', timeout: 'optional - seconds (default 900)', spinner: 'optional - display spinner (default false)' )

Returns a Hash with :choices / :assistant_message intact (including :message) — used by PWN::AI::Agent::Loop.

LOCAL-MODEL SCAFFOLDING

This hits Ollama's NATIVE /api/chat (not the OpenAI-compat shim) so the following actually take effect:

options.num_ctx     - Ollama defaults to 2048; the pwn-ai system
                    prompt alone blows that. Defaults here to
                    PWN::Env[:ai][:ollama][:num_ctx] || 32768.
options.temperature - forced to 0.1 on tool-bearing turns for
                    deterministic tool selection; engine[:temp]
                    (creative) on the final text-only turn.
format: 'json'      - constrains the sampler to valid JSON when
                    tools are present so Dispatch never sees
                    malformed function.arguments.
keep_alive: '30m'   - avoids reload latency between iterations.

tool_calls come back with function.arguments as a Hash (not a JSON string), which PWN::AI::Agent::Dispatch.parse_args handles.



166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
 
# File 'lib/pwn/ai/ollama.rb', line 166

public_class_method def self.chat_with_tools(opts = {})
  engine   = PWN::Env[:ai][:ollama]
  messages = opts[:messages]
  raise 'ERROR: messages array is required' if messages.nil? || messages.empty?

  model = opts[:model] ||= engine[:model]
  raise 'ERROR: Model is required.  Call #get_models method for details' if model.nil?

  temp = opts[:temp].to_f
  temp = engine[:temp].to_f.nonzero? || 1 if temp.zero?

  tools_present = opts[:tools] && !opts[:tools].empty?
  tool_temp     = (engine[:tool_temp] || 0.1).to_f
  num_ctx       = (engine[:num_ctx]   || 32_768).to_i
  keep_alive    = engine[:keep_alive] || '30m'

  http_body = {
    model: model,
    messages: messages,
    stream: false,
    keep_alive: keep_alive,
    options: {
      num_ctx: num_ctx,
      temperature: tools_present ? tool_temp : temp
    }
  }
  if tools_present
    http_body[:tools]  = opts[:tools]
    http_body[:format] = engine[:format] || 'json'
  end
  http_body[:tool_choice] = opts[:tool_choice] if opts[:tool_choice]

  response = ollama_rest_call(
    http_method: :post,
    rest_call: 'ollama/api/chat',
    http_body: http_body,
    timeout: opts[:timeout],
    spinner: opts[:spinner]
  )
  return nil if response.nil?

  json_resp = JSON.parse(response, symbolize_names: true)
  # Normalise native /api/chat shape to what Loop.normalize_llm expects.
  msg = json_resp[:message] || json_resp.dig(:choices, 0, :message)
  json_resp[:choices] = [{ message: msg }] if msg && !json_resp.key?(:choices)
  json_resp[:assistant_message] = msg
  json_resp
rescue StandardError => e
  raise e
end

.get_models ⇒ Object

Supported Method Parameters

response = PWN::AI::Ollama.get_models



127
128
129
130
131
132
133
 
# File 'lib/pwn/ai/ollama.rb', line 127

public_class_method def self.get_models
  models = ollama_rest_call(rest_call: 'ollama/api/tags')

  JSON.parse(models, symbolize_names: true)[:models]
rescue StandardError => e
  raise e
end

.help ⇒ Object

Display Usage for this Module



321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
 
# File 'lib/pwn/ai/ollama.rb', line 321

public_class_method def self.help
  puts "USAGE:
    models = #{self}.get_models

    response = #{self}.chat(
      request: 'required - message to Ollama',
      model: 'optional - model to use for text generation (defaults to PWN::Env[:ai][:ollama][:model])',
      temp: 'optional - creative response float (defaults to PWN::Env[:ai][:ollama][:temp])',
      system_role_content: 'optional - context to set up the model behavior for conversation (Default: PWN::Env[:ai][:ollama][:system_role_content])',
      response_history: 'optional - pass response back in to have a conversation',
      speak_answer: 'optional speak answer using PWN::Plugins::Voice.text_to_speech (Default: nil)',
      timeout: 'optional - timeout in seconds (defaults to 900)',
      spinner: 'optional - display spinner (defaults to false)'
    )

    #{self}.authors
  "
end