rails-nl2sql 0.1.4 → 0.1.8

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: c8acace727f6a1308018979554b87e003ca3ce427600a2c1822c4c3bcdbbeaeb
-  data.tar.gz: 33cfaabfe6e5f069d78b274481a3c7a27d20a5bd634d37160f56fbc1383277a0
+  metadata.gz: 4c85af69e2b9d5c60665676f6b1aa853220bb44e98484bd0968b9e659528fa31
+  data.tar.gz: 8423aaec96729c329dad2fe601ba11ef84545054f84ab1f3e0fc32ac181a7a8f
 SHA512:
-  metadata.gz: d69f8f85025572b8aed9d0c991fc25108964a1a7dbc27c7e6f858f1d0c7a5c9775264e68b896a80dcb430f2df51447fe7cb1ad4aa4e46cb20fc912934e8f2512
-  data.tar.gz: e4f40efeec6ac18ce06d4d2875bec8aef20f1dae7201962b3050f0b29dcc0e736f45c5c583752b7eec55e1f39056688f86ece31ae94790f88a8a6e9dedab4459
+  metadata.gz: 91739a9e2c4456f0232da6f614a5f9b7647b77267bc8d9eb49b3bb36a7d64542cca9ebf5e02897424bd8dc68b20a71a2d336ffd7eededd3b4b05cc2879d25fd4
+  data.tar.gz: 4e4540b9939bbb22bb093cf15482215016817b5cee95dbefaf6713cce048e5464136c62aeef0116659f378f020443e409b92d927bd4755a22c2c3829d6c9d022

data/lib/rails/nl2sql/query_generator.rb

@@ -3,23 +3,178 @@ 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
+        
+        # Clean up the response to remove markdown formatting
+        generated_query = clean_sql_response(generated_query)
+        
+        # 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, NO MARKDOWN FORMATTING, NO CODE BLOCKS, NO 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 clean_sql_response(query)
+        return query unless query
+
+        # Remove markdown code blocks
+        query = query.gsub(/```sql\n?/, '')
+        query = query.gsub(/```\n?/, '')
+        
+        # Remove any leading/trailing whitespace
+        query = query.strip
+        
+        # Remove any explanatory text before or after the query
+        # Look for common patterns like "Here's the SQL query:" or "The query is:"
+        query = query.gsub(/^.*?(SELECT|WITH|INSERT|UPDATE|DELETE|CREATE|DROP|ALTER)/i, '\1')
+        
+        # Remove any trailing explanatory text after the query
+        # Split by newlines and take only the SQL part
+        lines = query.split("\n")
+        sql_lines = []
+        
+        lines.each do |line|
+          line = line.strip
+          # Skip empty lines or lines that look like explanations
+          next if line.empty?
+          next if line.match(/^(here|this|the query|explanation|note)/i)
+          
+          sql_lines << line
+        end
+        
+        # Rejoin the SQL lines
+        cleaned_query = sql_lines.join("\n").strip
+        
+        # Ensure it ends with a semicolon if it's a complete query
+        if cleaned_query.match(/^(SELECT|WITH)/i) && !cleaned_query.end_with?(';')
+          cleaned_query += ';'
+        end
+        
+        cleaned_query
+      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/query_validator.rb

@@ -2,17 +2,36 @@ module Rails
   module Nl2sql
     class QueryValidator
       def self.validate(query)
+        return false unless query && !query.strip.empty?
+        
+        # Clean the query first
+        query = query.strip
+        
+        # Check if query is malformed (contains markdown or other formatting)
+        if query.include?('```') || query.include?('```sql')
+          raise Rails::Nl2sql::Error, "Query contains markdown formatting and could not be cleaned properly"
+        end
+        
         # Basic validation: prevent destructive commands
-        disallowed_keywords = %w(DROP DELETE UPDATE INSERT TRUNCATE ALTER CREATE)
-        if disallowed_keywords.any? { |keyword| query.upcase.include?(keyword) }
-          raise "Query contains disallowed keywords."
+        disallowed_keywords = %w(DROP DELETE UPDATE INSERT TRUNCATE ALTER CREATE EXEC EXECUTE MERGE REPLACE)
+        query_upper = query.upcase
+        
+        if disallowed_keywords.any? { |keyword| query_upper.include?(keyword) }
+          raise Rails::Nl2sql::Error, "Query contains disallowed keywords."
+        end
+
+        # Ensure it's a SELECT query
+        unless query_upper.strip.start_with?('SELECT', 'WITH')
+          raise Rails::Nl2sql::Error, "Only SELECT queries are allowed."
         end
 
-        # Use Rails' built-in sanitization to be safe
+        # Use Rails' built-in validation with EXPLAIN
         begin
-          ActiveRecord::Base.connection.execute("EXPLAIN #{query}")
+          # Remove trailing semicolon for EXPLAIN
+          explain_query = query.gsub(/;\s*$/, '')
+          ActiveRecord::Base.connection.execute("EXPLAIN #{explain_query}")
         rescue ActiveRecord::StatementInvalid => e
-          raise "Invalid SQL query: #{e.message}"
+          raise Rails::Nl2sql::Error, "Invalid SQL query: #{e.message}"
         end
 
         true

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.4"
+    VERSION = "0.1.8"
   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.4
+  version: 0.1.8
 platform: ruby
 authors:
 - Russell Van Curen