rails-nl2sql 0.1.3 → 0.1.7

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: a0b8da6d6865e8ecdfd36ff345b61d5746871334b99131783660ff20ca1eb923
-  data.tar.gz: 3f8d16f25a5cdb09b4a6fafa2a5b297624c3d44335e7461f27b82df790c8167c
+  metadata.gz: 9723feb23a2bc5f1100279427d18b106964382689ade9c37b379cda0274e0a77
+  data.tar.gz: 9e4a429bc94b4f5d865e9efde2c19da5b6d9a775ffaf982e7ab1fe3070aa75f0
 SHA512:
-  metadata.gz: f670c26de092ee22d94026884b13166f4e94254c1a6796a36ab27670e810c8e82e32257e3aa9ec7d287a09bb22b1b8cd251a651ca45e2103ae68b52b2dd5372f
-  data.tar.gz: 81d13b15eed1cf42631c3740a6c729d3a68cbad2797deb3d72fde9d8f6539eb4438280088e76494f80dff9d31f2c182c2ea32800fa673af5baeafde196e6422b
+  metadata.gz: 24b563aeed1df87e3afc4683230471e18351e364794989791dd28e2057d8db42cd3a083b678640ac74dfff6bb4d4c28573ea6efa955ee1faae29d1dd2fc3984a
+  data.tar.gz: 74aa67f890208d440baaae2f863e6d92cc04d00c9ffc0012e14239e5f5bc0d2807cc170b460cba6c95c487812c9960209b6e1a7bf76ec238c8d63df342855e2b

data/lib/generators/rails/nl2sql/install_generator.rb

@@ -1,15 +1,9 @@
 require 'rails/generators'
 
-module Rails
-  module Nl2sql
-    module Generators
-      class InstallGenerator < Rails::Generators::Base
-        source_root File.expand_path('../templates', __FILE__)
+class Rails::Nl2sql::InstallGenerator < Rails::Generators::Base
+  source_root File.expand_path('../templates', __FILE__)
 
-        def copy_initializer
-          template 'rails_nl2sql.rb', 'config/initializers/rails_nl2sql.rb'
-        end
-      end
-    end
+  def copy_initializer
+    template 'rails_nl2sql.rb', 'config/initializers/rails_nl2sql.rb'
   end
 end

data/lib/rails/nl2sql/query_generator.rb

@@ -3,23 +3,136 @@ require "openai"
 module Rails
   module Nl2sql
     class QueryGenerator
-      def initialize(api_key, model = "text-davinci-003")
-        @client = OpenAI::Client.new(api_key: api_key)
+      def initialize(api_key, model = "gpt-3.5-turbo-instruct")
+        @client = OpenAI::Client.new(access_token: api_key)
         @model = model
       end
 
-      def generate_query(prompt, schema)
-        full_prompt = "Given the following schema:\n\n#{schema}\n\nGenerate a SQL query for the following request:\n\n#{prompt}"
+      def generate_query(prompt, schema, db_server = "PostgreSQL", tables = nil)
+        retrieved_context = build_context(schema, tables)
+        
+        system_prompt = build_system_prompt(db_server, retrieved_context)
+        user_prompt = build_user_prompt(prompt)
+        
+        full_prompt = "#{system_prompt}\n\n#{user_prompt}"
 
         response = @client.completions(
           parameters: {
             model: @model,
             prompt: full_prompt,
-            max_tokens: 150
+            max_tokens: 300,
+            temperature: 0.1
           }
         )
 
-        response.choices.first.text.strip
+        generated_query = response.dig("choices", 0, "text")&.strip
+        
+        # Safety check
+        validate_query_safety(generated_query)
+        
+        generated_query
+      end
+
+      private
+
+      def build_context(schema, tables)
+        if tables&.any?
+          # Filter schema to only include requested tables
+          filtered_schema = filter_schema_by_tables(schema, tables)
+          filtered_schema
+        else
+          schema
+        end
+      end
+
+      def filter_schema_by_tables(schema, tables)
+        # Simple filtering - in a real implementation, this would be more sophisticated
+        lines = schema.split("\n")
+        filtered_lines = []
+        current_table = nil
+        include_current = false
+        
+        lines.each do |line|
+          if line.match(/CREATE TABLE (\w+)/)
+            current_table = $1
+            include_current = tables.include?(current_table)
+          end
+          
+          if include_current || line.strip.empty?
+            filtered_lines << line
+          end
+        end
+        
+        filtered_lines.join("\n")
+      end
+
+      def build_system_prompt(db_server, retrieved_context)
+        <<~PROMPT
+          You are an expert SQL assistant specializing in generating dynamic queries based on natural language.
+          Your primary goal is to generate **correct, safe, and executable #{db_server} SQL queries** based on user questions.
+
+          ---
+          **DATABASE CONTEXT (SCHEMA):**
+          You are provided with relevant schema details from the database, retrieved to help you.
+          **STRICTLY adhere to this provided schema context.** Do not use any tables or columns not explicitly listed here.
+          #{retrieved_context}
+
+          ---
+          **SQL GENERATION RULES:**
+          1. **SQL Dialect:** All generated SQL must be valid **#{db_server} syntax**.
+             * For limiting results, use `LIMIT` (e.g., `LIMIT 10`) instead of `TOP`.
+             * Be mindful of #{db_server}'s specific function names (e.g., `COUNT(*)`, `MAX()`) and behaviors.
+             * For subqueries that return a single value to be used in a `WHERE` clause, ensure they are correctly formatted for #{db_server}.
+          2. **Schema Adherence:** Only use table names and column names that are explicitly present in the provided context. Do not invent names.
+          3. **Valid JOIN Paths:** All `JOIN` operations must be based on valid foreign key relationships. The provided schema context explicitly details many of these.
+          4. **Safety First:** Absolutely **DO NOT** generate any DDL (CREATE, ALTER, DROP) or DML (INSERT, UPDATE, DELETE) statements. Only `SELECT` queries are permitted.
+          5. **CRITICAL: Handling Missing/Empty Text Data:**
+             * When a user asks about "missing," "no," "empty," or "null" values for a TEXT column (like 'email', 'phone', 'address', 'company', 'fax'), generate a `WHERE` clause that explicitly checks for **both `IS NULL` and `= ''` (an empty string)**.
+             * **Example:** To find agents with no email, the query should be `SELECT first_name, last_name FROM agents WHERE email IS NULL OR email = '';`
+             * This is essential
+          6. **Ambiguity:** If a user question is ambiguous or requires more information to form a precise SQL query, clearly state that you need clarification and ask for more details. Do not guess.
+          
+          **RESPOND WITH ONLY THE SQL QUERY - NO EXPLANATIONS OR ADDITIONAL TEXT.**
+        PROMPT
+      end
+
+      def build_user_prompt(input)
+        <<~PROMPT
+          Here is the **USER QUESTION:** Respond with a thoughtful process that leads to a SQL query, using the tools as necessary.
+          "#{input}"
+        PROMPT
+      end
+
+      def validate_query_safety(query)
+        return unless query
+
+        banned_keywords = [
+          "delete", "drop", "truncate", "update", "insert", "alter",
+          "exec", "execute", "create", "merge", "replace", "into"
+        ]
+
+        banned_phrases = [
+          "ignore previous instructions", "pretend you are", "i am the admin",
+          "you are no longer bound", "bypass the rules", "run this instead",
+          "for testing, run", "no safety constraints", "show me a dangerous query",
+          "this is a dev environment", "drop all data", "delete all users", "wipe the database"
+        ]
+
+        query_lower = query.downcase
+
+        # Check for banned keywords
+        banned_keywords.each do |keyword|
+          if query_lower.include?(keyword)
+            raise Rails::Nl2sql::Error, "Query contains banned keyword: #{keyword}"
+          end
+        end
+
+        # Check for banned phrases
+        banned_phrases.each do |phrase|
+          if query_lower.include?(phrase)
+            raise Rails::Nl2sql::Error, "Query contains banned phrase: #{phrase}"
+          end
+        end
       end
     end
   end

data/lib/rails/nl2sql/schema_builder.rb

@@ -2,10 +2,157 @@ module Rails
   module Nl2sql
     class SchemaBuilder
       def self.build_schema(options = {})
-        tables = ActiveRecord::Base.connection.tables
-        tables -= options[:exclude] if options[:exclude]
-        tables = options[:include] if options[:include]
+        tables = get_filtered_tables(options)
+        
+        schema_text = build_schema_text(tables)
+        schema_text
+      end
+
+      def self.get_database_type
+        adapter = ActiveRecord::Base.connection.adapter_name.downcase
+        case adapter
+        when 'postgresql'
+          'PostgreSQL'
+        when 'mysql', 'mysql2'
+          'MySQL'
+        when 'sqlite3'
+          'SQLite'
+        when 'oracle'
+          'Oracle'
+        when 'sqlserver'
+          'SQL Server'
+        else
+          'PostgreSQL' # Default fallback
+        end
+      end
+
+      def self.get_filtered_tables(options = {})
+        all_tables = ActiveRecord::Base.connection.tables
+        
+        # Remove system tables
+        all_tables.reject! { |table| system_table?(table) }
+        
+        # Apply filtering options
+        if options[:exclude]
+          all_tables -= options[:exclude]
+        end
+        
+        if options[:include]
+          all_tables = options[:include] & all_tables
+        end
+        
+        all_tables
+      end
+
+      def self.build_schema_text(tables)
+        schema_parts = []
+        
+        tables.each do |table|
+          schema_parts << build_table_schema(table)
+        end
+        
+        schema_parts.join("\n\n")
+      end
+
+      def self.build_table_schema(table)
+        columns = ActiveRecord::Base.connection.columns(table)
+        
+        schema = "CREATE TABLE #{table} (\n"
+        
+        column_definitions = columns.map do |column|
+          type_info = get_column_type_info(column)
+          nullable = column.null ? "" : " NOT NULL"
+          default = column.default ? " DEFAULT #{column.default}" : ""
+          
+          "  #{column.name} #{type_info}#{nullable}#{default}"
+        end
+        
+        schema += column_definitions.join(",\n")
+        schema += "\n);"
+        
+        # Add indexes and foreign keys if available
+        indexes = get_table_indexes(table)
+        if indexes.any?
+          schema += "\n\n-- Indexes for #{table}:"
+          indexes.each do |index|
+            schema += "\n-- #{index[:type]}: #{index[:columns].join(', ')}"
+          end
+        end
+        
+        schema
+      end
+
+      def self.get_column_type_info(column)
+        case column.type
+        when :string
+          "VARCHAR(#{column.limit || 255})"
+        when :text
+          "TEXT"
+        when :integer
+          "INTEGER"
+        when :bigint
+          "BIGINT"
+        when :float
+          "FLOAT"
+        when :decimal
+          precision = column.precision || 10
+          scale = column.scale || 0
+          "DECIMAL(#{precision},#{scale})"
+        when :datetime
+          "TIMESTAMP"
+        when :date
+          "DATE"
+        when :time
+          "TIME"
+        when :boolean
+          "BOOLEAN"
+        when :json
+          "JSON"
+        else
+          column.sql_type || "TEXT"
+        end
+      end
+
+      def self.get_table_indexes(table)
+        indexes = []
+        
+        begin
+          connection = ActiveRecord::Base.connection
+          if connection.respond_to?(:indexes)
+            table_indexes = connection.indexes(table)
+            table_indexes.each do |index|
+              indexes << {
+                type: index.unique? ? "UNIQUE INDEX" : "INDEX",
+                columns: index.columns,
+                name: index.name
+              }
+            end
+          end
+        rescue => e
+          # Skip if indexes can't be retrieved
+        end
+        
+        indexes
+      end
+
+      def self.system_table?(table)
+        system_tables = [
+          'schema_migrations',
+          'ar_internal_metadata',
+          'sqlite_sequence',
+          'information_schema',
+          'performance_schema',
+          'mysql',
+          'sys'
+        ]
+        
+        system_tables.any? { |sys_table| table.include?(sys_table) }
+      end
 
+      # Legacy method for backward compatibility
+      def self.build_hash_schema(options = {})
+        tables = get_filtered_tables(options)
+        
         schema = {}
         tables.each do |table|
           schema[table] = ActiveRecord::Base.connection.columns(table).map(&:name)

data/lib/rails/nl2sql/version.rb

@@ -1,5 +1,5 @@
 module Rails
   module Nl2sql
-    VERSION = "0.1.3"
+    VERSION = "0.1.7"
   end
 end

data/lib/rails/nl2sql.rb

@@ -12,7 +12,7 @@ module Rails
       attr_accessor :api_key
       attr_accessor :model
     end
-    @@model = "text-davinci-003"
+    @@model = "gpt-3.5-turbo-instruct"
 
     def self.configure
       yield self
@@ -20,22 +20,67 @@ module Rails
 
     class Processor
       def self.execute(natural_language_query, options = {})
+        # Get database type
+        db_server = SchemaBuilder.get_database_type
+        
+        # Build schema with optional table filtering
         schema = SchemaBuilder.build_schema(options)
+        
+        # Extract tables for filtering if specified
+        tables = options[:tables]
+        
+        # Generate query with enhanced prompt
         query_generator = QueryGenerator.new(Rails::Nl2sql.api_key, Rails::Nl2sql.model)
-        generated_query = query_generator.generate_query(natural_language_query, schema)
+        generated_query = query_generator.generate_query(
+          natural_language_query, 
+          schema, 
+          db_server, 
+          tables
+        )
 
+        # Validate the generated query
         QueryValidator.validate(generated_query)
 
+        # Execute the query
         ActiveRecord::Base.connection.execute(generated_query)
       end
 
+      def self.generate_query_only(natural_language_query, options = {})
+        # Get database type
+        db_server = SchemaBuilder.get_database_type
+        
+        # Build schema with optional table filtering
+        schema = SchemaBuilder.build_schema(options)
+        
+        # Extract tables for filtering if specified
+        tables = options[:tables]
+        
+        # Generate query with enhanced prompt
+        query_generator = QueryGenerator.new(Rails::Nl2sql.api_key, Rails::Nl2sql.model)
+        generated_query = query_generator.generate_query(
+          natural_language_query, 
+          schema, 
+          db_server, 
+          tables
+        )
+
+        # Validate the generated query
+        QueryValidator.validate(generated_query)
+
+        generated_query
+      end
+
       def self.get_tables(options = {})
-        ActiveRecord::Base.connection.tables
+        SchemaBuilder.get_filtered_tables(options)
       end
 
       def self.get_schema(options = {})
         SchemaBuilder.build_schema(options)
       end
+
+      def self.get_database_type
+        SchemaBuilder.get_database_type
+      end
     end
   end
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: rails-nl2sql
 version: !ruby/object:Gem::Version
-  version: 0.1.3
+  version: 0.1.7
 platform: ruby
 authors:
 - Russell Van Curen