Class: Aoororachain::Chains::RetrievalQA

Inherits:
Object
  • Object
show all
Defined in:
lib/aoororachain/chains/retrieval_qa.rb

Instance Method Summary collapse

Constructor Details

#initialize(llm, retriever, type: :stuff) ⇒ RetrievalQA

Returns a new instance of RetrievalQA.



6
7
8
9
10
 
# File 'lib/aoororachain/chains/retrieval_qa.rb', line 6

def initialize(llm, retriever, type: :stuff)
  @llm = llm
  @retriever = retriever
  @type = type
end

Instance Method Details

#complete(prompt:, prompt_template:, additional_context: "") ⇒ Object 



12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
 
# File 'lib/aoororachain/chains/retrieval_qa.rb', line 12

def complete(prompt:, prompt_template:, additional_context: "")
  context = @retriever.search(prompt)

  mapped_context = context.map(&:document)
  mapped_context << additional_context if !additional_context.nil? || additional_context != ""

  stuff_prompt = prompt_template % {context: mapped_context.join(" ").tr("\n", " "), prompt:}

  success, response = @llm.complete(prompt: stuff_prompt)

  if success
    completion = {
      "sources" => context.map(&:metadata)
    }.merge(response)
  else
    completion = {
      "response" => "Sorry we had a problem with the LLM",
      "sources" => [],
      "model" => ""
    }
    Aoororachain::Util.log_error("Failed to complete", message: response)
  end

  completion
end