Update sql_generator.py

sql_generator.py

@@ -1,40 +1,66 @@
-import sqlite3
-from typing import List, Dict, Any
-import logging
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import requests
+from config import ACCESS_TOKEN, SHOP_NAME
 
 class SQLGenerator:
-    def __init__(self, db_path: str = "shopify.db"):
-        self.db_path = db_path
-        self.setup_logging()
+    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 setup_logging(self):
-        logging.basicConfig(level=logging.INFO)
-        self.logger = logging.getLogger(__name__)
-        
-    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
+    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)
+        );
         """
-        # 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)
-
+        
+        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
+        )
+        
+        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