QuerySUTRA 0.5.2__py3-none-any.whl → 0.6.0__py3-none-any.whl

sutra/nlp_processor.py

@@ -1,144 +1,176 @@
-"""NLP to SQL query processor with relevancy checking"""
-
-import pandas as pd
-from typing import Optional, Tuple
-from tabulate import tabulate
-from sutra.cache_manager import CacheManager
-import openai
-import config
-from sutra.feedback import SimpleFeedback
-from sutra.schema_embeddings import SchemaEmbeddings
-from sutra.feedback_matcher import FeedbackMatcher
-
-class NLPProcessor:
-    """Process natural language questions to SQL queries"""
-    
-    def __init__(self, db_manager, openai_client=None):
-        self.db = db_manager
-        self.cache = CacheManager() if config.CACHE_ENABLED else None
-        self.model_name = config.MODEL_NAME
-        
-        # Set the API key directly for openai 0.28.1
-        openai.api_key = config.OPENAI_API_KEY
-
-        # Added for feedback handling and tracking
-        self.feedback = SimpleFeedback()
-        self.last_question = None
-        self.last_sql = None
-        
-        # ✅ NEW: Auto-load schema embeddings
-        self.relevancy_checker = SchemaEmbeddings(db_manager)
-        
-        # ✅ NEW: Smart feedback matcher
-        self.feedback_matcher = FeedbackMatcher()
-    
-    def nlp_to_sql(self, question: str) -> str:
-        """Convert natural language question to SQL"""
-        
-        # ✅ NEW: Check feedback for similar queries first
-        similar_sql, similarity = self.feedback_matcher.find_similar_query(question)
-        if similar_sql:
-            print(f"🎯 Found similar query in feedback (similarity: {similarity:.2f})")
-            return similar_sql
-        
-        # Check cache next
-        if self.cache:
-            cached_sql = self.cache.get_cached_query(question)
-            if cached_sql:
-                print("⚡ Using cached query")
-                return cached_sql
-        
-        # Only call API if no feedback match and no cache
-        print("🤖 Calling OpenAI API...")
-        
-        # Get schema context
-        schema = self.db.get_schema_context()
-        
-        prompt = f"""
-Convert this question to a SQLite query:
-
-Question: {question}
-
-Database schema:
-{schema}
-
-Return ONLY the SELECT statement. No explanations, no markdown.
-"""
-        
-        # Use openai.ChatCompletion directly for version 0.28.1
-        response = openai.ChatCompletion.create(
-            model=self.model_name,
-            messages=[{"role": "user", "content": prompt}],
-            temperature=0
-        )
-        
-        sql_query = response['choices'][0]['message']['content'].strip()
-        sql_query = sql_query.replace('```sql', '').replace('```', '').strip()
-        
-        # Cache the result
-        if self.cache:
-            self.cache.add_to_cache(question, sql_query)
-        
-        return sql_query
-    
-    def process_question(self, question: str) -> Tuple[Optional[pd.DataFrame], str]:
-        """Process a natural language question and return results"""
-        
-        # ✅ NEW: Check relevancy FIRST - BEFORE any API calls
-        is_relevant, similarity, info = self.relevancy_checker.is_relevant(question)
-        
-        if not is_relevant:
-            print(f"\n❌ Question not relevant to database (similarity: {similarity:.2f})")
-            for item in info:
-                print(f"   {item}")
-            return None, ""
-        
-        print(f"✅ Relevant question (similarity: {similarity:.2f})")
-        
-        try:
-            # Convert to SQL (only if relevant)
-            sql_query = self.nlp_to_sql(question)
-            print(f"\n🔍 Generated SQL Query:")
-            print(f"   {sql_query}")
-            
-            # Track for feedback
-            self.last_question = question
-            self.last_sql = sql_query
-            
-            # Execute query
-            result_df = self.db.execute_query(sql_query)
-            
-            return result_df, sql_query
-            
-        except Exception as e:
-            print(f"❌ Error processing question: {e}")
-            return None, ""
-    
-    def display_results(self, df: pd.DataFrame, max_rows: int = 15):
-        """Display query results in a formatted table"""
-        if df is None or df.empty:
-            print("   No results found")
-            return  # Exit early if no results
-        
-        # Show the table
-        display_df = df.head(max_rows) if len(df) > max_rows else df
-        print(tabulate(display_df, headers='keys', tablefmt='grid', showindex=False))
-        
-        if len(df) > max_rows:
-            print(f"   ... showing first {max_rows} of {len(df)} rows")
-        
-        # ✅ UPDATED: Only ask for feedback for relevant questions with results
-        # (Irrelevant questions never reach here due to early return)
-        feedback = input("\n👍 or 👎? (y/n): ").lower()
-        if feedback == 'y':
-            self.feedback.save(self.last_question, self.last_sql, True)
-            print("✅ Saved as good")
-            # Reload feedback matcher with new data
-            self.feedback_matcher.reload_feedback()
-        elif feedback == 'n':
-            correct = input("Correct SQL: ").strip()
-            self.feedback.save(self.last_question, self.last_sql, False, correct)
-            if correct:
-                print("✅ Learned correction")
-                # Reload feedback matcher with new data
+"""NLP to SQL query processor with relevancy checking"""
+
+import pandas as pd
+from typing import Optional, Tuple
+from tabulate import tabulate
+from sutra.cache_manager import CacheManager
+import openai
+import config
+from sutra.feedback import SimpleFeedback
+from sutra.schema_embeddings import SchemaEmbeddings
+from sutra.feedback_matcher import FeedbackMatcher
+
+class NLPProcessor:
+    """Process natural language questions to SQL queries"""
+    
+    def __init__(self, db_manager, openai_client=None):
+        self.db = db_manager
+        self.cache = CacheManager() if config.CACHE_ENABLED else None
+        self.model_name = config.MODEL_NAME
+        
+        # Set the API key directly for openai 0.28.1
+        openai.api_key = config.OPENAI_API_KEY
+
+        # Added for feedback handling and tracking
+        self.feedback = SimpleFeedback()
+        self.last_question = None
+        self.last_sql = None
+        
+        # ✅ NEW: Auto-load schema embeddings
+        self.relevancy_checker = SchemaEmbeddings(db_manager)
+        
+        # ✅ NEW: Smart feedback matcher
+        self.feedback_matcher = FeedbackMatcher()
+    
+    def nlp_to_sql(self, question: str) -> str:
+        """Convert natural language question to SQL"""
+        
+        # ✅ NEW: Check feedback for similar queries first
+        similar_sql, similarity = self.feedback_matcher.find_similar_query(question)
+        if similar_sql:
+            print(f"🎯 Found similar query in feedback (similarity: {similarity:.2f})")
+            return similar_sql
+        
+        # Check cache next
+        if self.cache:
+            cached_sql = self.cache.get_cached_query(question)
+            if cached_sql:
+                print("⚡ Using cached query")
+                return cached_sql
+        
+        # Only call API if no feedback match and no cache
+        print("🤖 Calling OpenAI API...")
+        
+        # Get schema context
+        schema = self.db.get_schema_context()
+        
+        prompt = f"""
+Convert this question to a SQLite query.
+
+Question: {question}
+
+Database schema:
+{schema}
+
+CRITICAL INSTRUCTIONS FOR MULTI-TABLE QUERIES:
+
+**STEP 1: CHECK TABLE RELATIONSHIPS FIRST**
+Look at the "=== TABLE RELATIONSHIPS ===" section at the top of the schema.
+These show you exactly how tables are connected via foreign keys.
+Format: table1.column → table2.column means table1.column references table2.column
+
+**STEP 2: IDENTIFY REQUIRED TABLES**
+Analyze which tables contain the data needed to answer the question.
+If information is spread across multiple tables, you MUST join them.
+
+**STEP 3: USE THE RELATIONSHIPS TO JOIN**
+When you need data from multiple tables:
+- Use the foreign key relationships shown in the TABLE RELATIONSHIPS section
+- Join table1 to table2 using: JOIN table2 ON table1.fk_column = table2.pk_column
+- Use INNER JOIN when both tables must have matching data
+- Use LEFT JOIN when you need all rows from the first table regardless of matches
+
+**STEP 4: WRITE THE QUERY**
+- Use table aliases (t1, t2, etc.) for readability
+- Qualify all column names with table aliases to avoid ambiguity
+- Include all necessary columns from all joined tables in SELECT
+
+EXAMPLES:
+❌ WRONG: SELECT name FROM customers WHERE city = 'NYC'
+   (if you need order information too)
+
+✅ CORRECT: SELECT c.name, o.order_date, o.total 
+             FROM customers c 
+             JOIN orders o ON c.customer_id = o.customer_id 
+             WHERE c.city = 'NYC'
+
+Return ONLY the executable SELECT statement. No explanations, no markdown, no code blocks.
+"""
+        
+        # Use openai.ChatCompletion directly for version 0.28.1
+        response = openai.ChatCompletion.create(
+            model=self.model_name,
+            messages=[{"role": "user", "content": prompt}],
+            temperature=0
+        )
+        
+        sql_query = response['choices'][0]['message']['content'].strip()
+        sql_query = sql_query.replace('```sql', '').replace('```', '').strip()
+        
+        # Cache the result
+        if self.cache:
+            self.cache.add_to_cache(question, sql_query)
+        
+        return sql_query
+    
+    def process_question(self, question: str) -> Tuple[Optional[pd.DataFrame], str]:
+        """Process a natural language question and return results"""
+        
+        # ✅ NEW: Check relevancy FIRST - BEFORE any API calls
+        is_relevant, similarity, info = self.relevancy_checker.is_relevant(question)
+        
+        if not is_relevant:
+            print(f"\n❌ Question not relevant to database (similarity: {similarity:.2f})")
+            for item in info:
+                print(f"   {item}")
+            return None, ""
+        
+        print(f"✅ Relevant question (similarity: {similarity:.2f})")
+        
+        try:
+            # Convert to SQL (only if relevant)
+            sql_query = self.nlp_to_sql(question)
+            print(f"\n🔍 Generated SQL Query:")
+            print(f"   {sql_query}")
+            
+            # Track for feedback
+            self.last_question = question
+            self.last_sql = sql_query
+            
+            # Execute query
+            result_df = self.db.execute_query(sql_query)
+            
+            return result_df, sql_query
+            
+        except Exception as e:
+            print(f"❌ Error processing question: {e}")
+            return None, ""
+    
+    def display_results(self, df: pd.DataFrame, max_rows: int = 15):
+        """Display query results in a formatted table"""
+        if df is None or df.empty:
+            print("   No results found")
+            return  # Exit early if no results
+        
+        # Show the table
+        display_df = df.head(max_rows) if len(df) > max_rows else df
+        print(tabulate(display_df, headers='keys', tablefmt='grid', showindex=False))
+        
+        if len(df) > max_rows:
+            print(f"   ... showing first {max_rows} of {len(df)} rows")
+        
+        # ✅ UPDATED: Only ask for feedback for relevant questions with results
+        # (Irrelevant questions never reach here due to early return)
+        feedback = input("\n👍 or 👎? (y/n): ").lower()
+        if feedback == 'y':
+            self.feedback.save(self.last_question, self.last_sql, True)
+            print("✅ Saved as good")
+            # Reload feedback matcher with new data
+            self.feedback_matcher.reload_feedback()
+        elif feedback == 'n':
+            correct = input("Correct SQL: ").strip()
+            self.feedback.save(self.last_question, self.last_sql, False, correct)
+            if correct:
+                print("✅ Learned correction")
+                # Reload feedback matcher with new data
                 self.feedback_matcher.reload_feedback()

sutra/schema_generator.py

@@ -1,53 +1,57 @@
-"""SQL schema generation from unstructured text using AI"""
-
-import openai
-import config
-
-class SchemaGenerator:
-    """Generate SQL schema from unstructured data using OpenAI"""
-    
-    def __init__(self, api_key: str, model_name: str = "gpt-3.5-turbo"):
-        openai.api_key = api_key
-        self.model_name = model_name
-        self.temperature = config.TEMPERATURE
-    
-    def generate_schema(self, unstructured_data: str) -> str:
-        """Generate SQL schema from unstructured text"""
-        
-        # Truncate if too long
-        if len(unstructured_data) > config.MAX_TEXT_LENGTH:
-            unstructured_data = unstructured_data[:config.MAX_TEXT_LENGTH]
-            print(f"⚠️ Data truncated to {config.MAX_TEXT_LENGTH} characters")
-        
-        prompt = f"""
-Convert this unstructured text into a SQLite database:
-
-{unstructured_data}
-
-Requirements:
-1. Create tables based on what entities you find in the text
-2. Add foreign keys to connect related tables  
-3. Extract ALL data from the text - don't add anything not in the text
-4. Use INTEGER PRIMARY KEY AUTOINCREMENT for IDs
-
-Return ONLY executable SQLite statements:
-- DROP TABLE IF EXISTS statements
-- CREATE TABLE statements with PRIMARY KEY and FOREIGN KEY
-- INSERT statements with the actual data from the text
-
-No markdown, no code blocks, just SQL.
-"""
-        
-        print("🔄 Generating schema via OpenAI API...")
-        
-        response = openai.ChatCompletion.create(
-            model=self.model_name,
-            messages=[{"role": "user", "content": prompt}],
-            temperature=self.temperature
-        )
-        
-        generated_schema = response['choices'][0]['message']['content'].strip()
-        generated_schema = generated_schema.replace('```sql', '').replace('```', '').strip()
-        
-        print("✅ Schema generated!")
+"""SQL schema generation from unstructured text using AI"""
+
+import openai
+import config
+
+class SchemaGenerator:
+    """Generate SQL schema from unstructured data using OpenAI"""
+    
+    def __init__(self, api_key: str, model_name: str = "gpt-3.5-turbo"):
+        openai.api_key = api_key
+        self.model_name = model_name
+        self.temperature = config.TEMPERATURE
+    
+    def generate_schema(self, unstructured_data: str) -> str:
+        """Generate SQL schema from unstructured text"""
+        
+        # Truncate if too long
+        if len(unstructured_data) > config.MAX_TEXT_LENGTH:
+            unstructured_data = unstructured_data[:config.MAX_TEXT_LENGTH]
+            print(f"⚠️ Data truncated to {config.MAX_TEXT_LENGTH} characters")
+        
+        prompt = f"""
+Convert this unstructured text into a SQLite database:
+
+{unstructured_data}
+
+CRITICAL Requirements:
+1. Identify all entities in the text and create a table for each
+2. **MANDATORY:** Add FOREIGN KEY constraints to connect related tables
+   - If table A references table B, add: FOREIGN KEY (column_name) REFERENCES table_b(id)
+   - Example: If employees belong to departments, employees table must have:
+     department_id INTEGER, FOREIGN KEY (department_id) REFERENCES departments(id)
+3. Extract ALL data from the text - don't add anything not in the text
+4. Use INTEGER PRIMARY KEY AUTOINCREMENT for all ID columns
+5. Ensure parent tables (referenced tables) are created BEFORE child tables
+
+Return ONLY executable SQLite statements in this order:
+1. DROP TABLE IF EXISTS statements (child tables first, parent tables last)
+2. CREATE TABLE statements (parent tables first, child tables last)
+3. INSERT statements (parent tables first, child tables last)
+
+No markdown, no code blocks, no explanations - just SQL statements.
+"""
+        
+        print("🔄 Generating schema via OpenAI API...")
+        
+        response = openai.ChatCompletion.create(
+            model=self.model_name,
+            messages=[{"role": "user", "content": prompt}],
+            temperature=self.temperature
+        )
+        
+        generated_schema = response['choices'][0]['message']['content'].strip()
+        generated_schema = generated_schema.replace('```sql', '').replace('```', '').strip()
+        
+        print("✅ Schema generated!")
         return generated_schema