Class: AsktiveRecord::Query

Inherits:
Object
  • Object
show all
Defined in:
lib/asktive_record/query.rb

Overview

The Query class encapsulates a natural language question, its corresponding SQL, and provides methods for sanitization, execution, and generating answers using LLMs.

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(natural_question, raw_sql, model_class) ⇒ Query

Returns a new instance of Query.



10
11
12
13
14
15
 
# File 'lib/asktive_record/query.rb', line 10

def initialize(natural_question, raw_sql, model_class)
  @raw_sql = raw_sql
  @model_class = model_class
  @natural_question = natural_question
  @sanitized_sql = raw_sql # Initially, sanitized SQL is the same as raw SQL
end

Instance Attribute Details

#model_class ⇒ Object (readonly)

Returns the value of attribute model_class.



7
8
9
 
# File 'lib/asktive_record/query.rb', line 7

def model_class
  @model_class
end

#natural_question ⇒ Object (readonly)

Returns the value of attribute natural_question.



7
8
9
 
# File 'lib/asktive_record/query.rb', line 7

def natural_question
  @natural_question
end

#raw_sql ⇒ Object (readonly)

Returns the value of attribute raw_sql.



7
8
9
 
# File 'lib/asktive_record/query.rb', line 7

def raw_sql
  @raw_sql
end

#sanitized_sql ⇒ Object

Returns the value of attribute sanitized_sql.



8
9
10
 
# File 'lib/asktive_record/query.rb', line 8

def sanitized_sql
  @sanitized_sql
end

Instance Method Details

#answer ⇒ Object 



29
30
31
32
33
34
 
# File 'lib/asktive_record/query.rb', line 29

def answer
  response = execute
  llm = AsktiveRecord::LlmService.new(AsktiveRecord.configuration)
  response = response.inspect if response.respond_to?(:inspect)
  llm.answer(@natural_question, @sanitized_sql, response)
end

#execute ⇒ Object 



36
37
38
39
40
41
42
43
44
45
46
 
# File 'lib/asktive_record/query.rb', line 36

def execute
  unless @sanitized_sql
    raise QueryExecutionError,
          "Cannot execute raw SQL. Call sanitize! first or work with sanitized_sql."
  end

  result = execute_query
  extract_count_if_present(result)
rescue StandardError => e
  raise QueryExecutionError, "Failed to execute SQL query: #{e.message}"
end

#sanitize!(allow_only_select: true) ⇒ Object

Placeholder for sanitization logic In a real scenario, this would involve more sophisticated checks, potentially allowing only SELECT statements or using a whitelist of allowed SQL patterns.



20
21
22
23
24
25
26
27
 
# File 'lib/asktive_record/query.rb', line 20

def sanitize!(allow_only_select: true)
  if allow_only_select && !@sanitized_sql.strip.downcase.start_with?("select")
    raise SanitizationError, "Query sanitization failed: Only SELECT statements are allowed by default."
  end

  # Add more sanitization rules here as needed
  self # Return self for chaining
end

#to_s ⇒ Object 



48
49
50
 
# File 'lib/asktive_record/query.rb', line 48

def to_s
  @sanitized_sql || @raw_sql
end