Module: Ollama::Testing::ClientExtension

Defined in:
lib/ollama/testing.rb

Overview

Client extension to intercept requests and return stubbed responses during testing

Instance Method Summary collapse

Instance Method Details

#chat(messages:, **args) ⇒ Object



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

def chat(messages:, **args)
  stub = Ollama::Testing.pop_stub(:chat)
  return super if stub.nil?

  if args[:hooks] && stub["stream_chunks"]
    stub["stream_chunks"].each do |chunk|
      args[:hooks][:on_token]&.call(chunk)
    end
  end

  Response.new(stub)
end

#generate(prompt:, **args) ⇒ Object



44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
# File 'lib/ollama/testing.rb', line 44

def generate(prompt:, **args)
  stub = Ollama::Testing.pop_stub(:generate)
  return super if stub.nil?

  if args[:hooks] && stub["stream_chunks"]
    stub["stream_chunks"].each do |chunk|
      args[:hooks][:on_token]&.call(chunk)
    end
  end

  # When a schema is specified, generate returns a parsed Hash
  if args[:schema] && stub["response"].is_a?(String)
    begin
      JSON.parse(stub["response"])
    rescue JSON::ParserError
      stub["response"]
    end
  else
    stub["response"]
  end
end