Update app.py

app.py

@@ -1,66 +1,54 @@
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import requests
-from config import ACCESS_TOKEN, SHOP_NAME
+import gradio as gr
+from sql_generator import SQLGenerator
+from intent_classifier import IntentClassifier
+from rag_system import RAGSystem
 
-class SQLGenerator:
+class UnifiedSystem:
     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)
+        self.sql_generator = SQLGenerator()
+        self.intent_classifier = IntentClassifier()
+        self.rag_system = RAGSystem()
         
-    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)
-        );
-        """
+    def process_query(self, query):
+        intent, confidence = self.intent_classifier.classify(query)
         
-        prompt = f"""### Task: Generate a SQL query to answer the following question.
-        ### Database Schema:
-        {schema_info}
-        ### Question: {natural_language_query}
-        ### SQL Query:"""
+        if intent == "database_query":
+            sql_query = self.sql_generator.generate_query(query)
+            products = self.sql_generator.fetch_shopify_data("products")
+            
+            if products and 'products' in products:
+                results = "\n".join([
+                    f"Title: {p['title']}, Vendor: {p['vendor']}" 
+                    for p in products['products']
+                ])
+                return f"Intent: Database Query (Confidence: {confidence:.2f})\n\n" \
+                       f"SQL Query: {sql_query}\n\nResults:\n{results}"
+            else:
+                return "No results found or error fetching data from Shopify."
+                
+        elif intent == "product_description":
+            rag_response = self.rag_system.process_query(query)
+            return f"Intent: Product Description (Confidence: {confidence:.2f})\n\n" \
+                   f"Response: {rag_response}"
         
-        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()
+        return "Intent not recognized."
+
+def create_interface():
+    system = UnifiedSystem()
     
-    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
+    iface = gr.Interface(
+        fn=system.process_query,
+        inputs=gr.Textbox(
+            label="Enter your query",
+            placeholder="e.g., 'Show me all T-shirts' or 'Describe the product features'"
+        ),
+        outputs=gr.Textbox(label="Response"),
+        title="Unified Query Processing System",
+        description="Enter a natural language query to search products or get descriptions."
+    )
+    
+    return iface
+
+if __name__ == "__main__":
+    iface = create_interface()
+    iface.launch()