Update sql_generator.py

sql_generator.py

@@ -1,67 +1,40 @@
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import requests
-from config import ACCESS_TOKEN, SHOP_NAME
+import sqlite3
+from typing import List, Dict, Any
+import logging
 
 class SQLGenerator:
-    def __init__(self):
-        self.model_name = "premai-io/prem-1B-SQL"
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
-        self.model = AutoModelForCausalLM.from_pretrained(self.model_name)
+    def __init__(self, db_path: str = "shopify.db"):
+        self.db_path = db_path
+        self.setup_logging()
         
-    def generate_query(self, natural_language_query):
-        schema_info = """
-        CREATE TABLE products (
-          id DECIMAL(8,2) PRIMARY KEY,
-          title VARCHAR(255),
-          body_html VARCHAR(255),
-          vendor VARCHAR(255),
-          product_type VARCHAR(255),
-          created_at VARCHAR(255),
-          handle VARCHAR(255),
-          updated_at DATE,
-          published_at VARCHAR(255),
-          template_suffix VARCHAR(255),
-          published_scope VARCHAR(255),
-          tags VARCHAR(255),
-          status VARCHAR(255),
-          admin_graphql_api_id DECIMAL(8,2),
-          variants VARCHAR(255),
-          options VARCHAR(255),
-          images VARCHAR(255),
-          image VARCHAR(255)
-        );
-        """
-        
-        prompt = f"""### Task: Generate a SQL query to answer the following question.
-        ### Database Schema:
-        {schema_info}
-        ### Question: {natural_language_query}
-        ### SQL Query:"""
-        
-        inputs = self.tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(self.model.device)
-        outputs = self.model.generate(
-            inputs["input_ids"],
-            max_length=256,
-            do_sample=False,
-            num_return_sequences=1,
-            eos_token_id=self.tokenizer.eos_token_id,
-            pad_token_id=self.tokenizer.pad_token_id
-        )
+    def setup_logging(self):
+        logging.basicConfig(level=logging.INFO)
+        self.logger = logging.getLogger(__name__)
         
-        return self.tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
-    
-    def fetch_shopify_data(self, endpoint):
-        headers = {
-            'X-Shopify-Access-Token': ACCESS_TOKEN,
-            'Content-Type': 'application/json'
-        }
-        url = f"https://{SHOP_NAME}/admin/api/2023-10/{endpoint}.json"
-        response = requests.get(url, headers=headers)
-        
-        if response.status_code == 200:
-            return response.json()
-        else:
-            print(f"Error fetching {endpoint}: {response.status_code} - {response.text}")
-            return None
+    def execute_query(self, query: str) -> List[Dict[str, Any]]:
+        """
+        Execute SQL query and return results as a list of dictionaries
+        """
+        try:
+            with sqlite3.connect(self.db_path) as conn:
+                conn.row_factory = sqlite3.Row
+                cursor = conn.cursor()
+                cursor.execute(query)
+                results = [dict(row) for row in cursor.fetchall()]
+                self.logger.info(f"Successfully executed query: {query[:100]}...")
+                return results
+        except sqlite3.Error as e:
+            self.logger.error(f"Database error: {e}")
+            raise
+        except Exception as e:
+            self.logger.error(f"Error executing query: {e}")
+            raise
+
+    def validate_query(self, query: str) -> bool:
+        """
+        Validate SQL query before execution
+        """
+        # Basic validation - you might want to add more sophisticated validation
+        dangerous_keywords = ["DROP", "DELETE", "TRUNCATE", "UPDATE", "INSERT"]
+        return not any(keyword in query.upper() for keyword in dangerous_keywords)