Class: Rails::Nl2sql::QueryGenerator

Inherits:
Object
  • Object
show all
Defined in:
lib/rails/nl2sql/query_generator.rb

Constant Summary collapse

DEFAULT_MODEL = 
'gpt-3.5-turbo-instruct'

Instance Method Summary collapse

Constructor Details

#initialize(provider: nil, model: DEFAULT_MODEL) ⇒ QueryGenerator

Initializes a new QueryGenerator.

Parameters:

  • provider (Object) (defaults to: nil) —

    The AI provider to use for generating queries.

  • model (String) (defaults to: DEFAULT_MODEL) —

    The name of the AI model to use.



14
15
16
17
 
# File 'lib/rails/nl2sql/query_generator.rb', line 14

def initialize(provider: nil, model: DEFAULT_MODEL)
  @provider = provider || Rails::Nl2sql.provider || default_provider(model)
  @model = model
end

Instance Method Details

#generate_query(prompt, schema, db_server = 'PostgreSQL', tables = nil) ⇒ String

Generates a SQL query from a natural language prompt.

Parameters:

  • prompt (String) —

    The natural language prompt.

  • schema (String) —

    The database schema.

  • db_server (String) (defaults to: 'PostgreSQL') —

    The type of database server.

  • tables (Array<String>) (defaults to: nil) —

    The tables to include in the schema.

Returns:

  • (String) —

    The generated SQL query.



26
27
28
29
30
31
32
33
34
35
36
37
38
 
# File 'lib/rails/nl2sql/query_generator.rb', line 26

def generate_query(prompt, schema, db_server = 'PostgreSQL', tables = nil)
  retrieved_context = build_context(schema, tables)

  full_prompt = PromptBuilder.build(prompt, db_server, retrieved_context)

  response = @provider.complete(prompt: full_prompt, max_tokens: 500, temperature: 0.1)
  generated_query = extract_text(response)

  generated_query = clean_sql_response(generated_query)
  validate_query_safety(generated_query)

  generated_query
end