Update app.py

app.py

@@ -1,307 +1,126 @@
 import streamlit as st
-import requests
-import json
-import re
+import pandas as pd
+import os
+from dotenv import load_dotenv
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import logging
-from typing import Dict, Any, Optional
 
-# Setup logging
+# Set up logging
 logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-# API Configuration
-GROQ_API_URL = "https://api.groq.com/openai/v1/chat/completions"
-GROQ_API_KEY = "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"
-GROQ_HEADERS = {
-    "Authorization": f"Bearer {GROQ_API_KEY}",
-    "Content-Type": "application/json",
-}
+class LLMService:
+    def __init__(self, db_path):
+        self.db_path = db_path
+        # Load tokenizer and model
+        self.tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-72B-Instruct")
+        self.model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-72B-Instruct")
 
-# Shopify Configuration
-STORE_NAME = "agkd0n-fa"  # Your store name
-API_VERSION = "2024-04"    # Current API version
-SHOPIFY_URL = f"https://{STORE_NAME}.myshopify.com/admin/api/{API_VERSION}/graphql.json"
-SHOPIFY_ACCESS_TOKEN = "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"
-SHOPIFY_HEADERS = {
-    "Content-Type": "application/json",
-    "X-Shopify-Access-Token": SHOPIFY_ACCESS_TOKEN,
-}
-
-def fetch_graphql_schema() -> Dict[str, Any]:
-    """Fetch the GraphQL schema from Shopify."""
-    introspection_query = """
-    {
-      "__schema" {
-        "types" {
-          "name"
-        }
-      }
-    }
-    """
-    try:
-        response = requests.post(
-            SHOPIFY_URL,
-            headers=SHOPIFY_HEADERS,
-            json={"query": introspection_query}
-        )
-        response.raise_for_status()
-        return response.json()
-    except Exception as e:
-        logging.error(f"Error fetching schema: {str(e)}")
-        return {"error": str(e)}
+    def convert_to_sql_query(self, natural_query):
+        try:
+            # Tokenize input
+            inputs = self.tokenizer(f"Translate this to SQL: {natural_query}", return_tensors="pt")
+            # Generate output
+            outputs = self.model.generate(**inputs, max_length=512, num_beams=5)
+            # Decode output
+            sql_query = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            return {"success": True, "query": sql_query}
+        except Exception as e:
+            logger.error(f"Error generating SQL query: {e}")
+            return {"success": False, "error": str(e)}
 
+    def execute_query(self, sql_query):
+        try:
+            import sqlite3
+            conn = sqlite3.connect(self.db_path)
+            cursor = conn.cursor()
+            cursor.execute(sql_query)
+            results = cursor.fetchall()
+            columns = [desc[0] for desc in cursor.description]
+            conn.close()
+            return {"success": True, "results": results, "columns": columns}
+        except Exception as e:
+            logger.error(f"Error executing SQL query: {e}")
+            return {"success": False, "error": str(e)}
 
-def convert_to_graphql_query(natural_language: str) -> str:
-    """Convert natural language to a GraphQL query using the Groq API."""
-    # Enhanced prompt with Shopify Admin API specific guidance
-    prompt = f"""
-    Create a Shopify Admin API GraphQL query for this request: "{natural_language}"
+def main():
+    st.title("Natural Language to SQL Query Converter")
+    st.write("Enter your question about the database in natural language.")
     
-    Follow these Shopify Admin API rules:
-    1. For text searches in product queries:
-       - Use the query parameter with these prefixes:
-         * title: for product titles
-         * description: for product descriptions
-         * tag: for product tags
-         * product_type: for product types
-       - Example: query: "title:shirt OR description:cotton"
-       
-    2. For price filters:
-       - Use variants.price for exact price
-       - Use variants.price:<number for maximum price
-       - Use variants.price:>number for minimum price
-       - Example: query: "variants.price:<50.00"
-       
-    3. Include these fields in the response:
-       - id
-       - title
-       - description
-       - productType
-       - priceRangeV2 with minVariantPrice (amount and currencyCode)
-       - images (first: 1) with url
-       
-    4. Use proper pagination with first: 250
+    # Load environment variables
+    load_dotenv()
+    db_path = os.getenv("DB_PATH")
     
-    Return only the complete GraphQL query, no explanations.
-    """
+    if not db_path:
+        st.error("Missing database path in environment variables.")
+        logger.error("DB path not found in environment variables.")
+        return
     
-    payload = {
-        "model": "llama3-8b-8192",
-        "messages": [{"role": "user", "content": prompt}],
-        "temperature": 0.1,
-        "max_tokens": 500,
-    }
-
+    # Initialize LLM Service
     try:
-        response = requests.post(GROQ_API_URL, headers=GROQ_HEADERS, json=payload)
-        response.raise_for_status()
-        
-        query = response.json()['choices'][0]['message']['content'].strip()
-        return clean_query(query)
+        llm_service = LLMService(db_path=db_path)
     except Exception as e:
-        logging.error(f"Query generation error: {str(e)}")
-        return get_default_query()
+        st.error(f"Error initializing service: {str(e)}")
+        return
 
-def clean_query(query: str) -> str:
-    """Clean and format the GraphQL query."""
-    # Remove code block markers and quotes
-    query = re.sub(r'```(?:graphql)?\s*|\s*```|^["\']\s*|\s*["\']$', '', query)
-    
-    # Extract the query content
-    query_match = re.search(r'({[\s\S]*})', query)
-    if not query_match:
-        return get_default_query()
-    
-    query = query_match.group(1)
-    
-    # Ensure the query has all required fields
-    if not all(field in query for field in ['id', 'title', 'description', 'productType', 'priceRangeV2', 'images']):
-        return get_default_query()
-        
-    # Format the query
-    try:
-        formatted_query = format_query(query)
-        return formatted_query
-    except Exception as e:
-        logging.error(f"Query formatting error: {str(e)}")
-        return get_default_query()
+    # Input for natural language query
+    natural_query = st.text_area("Enter your query", "Show me all albums by artist 'Queen'", height=100)
 
-def format_query(query: str) -> str:
-    """Format the GraphQL query with proper indentation."""
-    depth = 0
-    formatted = []
-    lines = query.split('\n')
-    
-    for line in lines:
-        line = line.strip()
-        
-        # Adjust depth based on brackets
-        if '}' in line:
-            depth -= 1
-        
-        # Add indentation
-        if line:
-            formatted.append('  ' * depth + line)
-            
-        # Increase depth for next line if needed
-        if '{' in line:
-            depth += 1
-            
-    return '\n'.join(formatted)
+    if st.button("Generate and Execute Query"):
+        if not natural_query.strip():
+            st.warning("Please enter a valid query.")
+            return
 
-def get_default_query() -> str:
-    """Return a default working query with all required fields."""
-    return """
-{
-  products(first: 250) {
-    edges {
-      node {
-        id
-        title
-        description
-        productType
-        priceRangeV2 {
-          minVariantPrice {
-            amount
-            currencyCode
-          }
-        }
-        images(first: 1) {
-          edges {
-            node {
-              url
-            }
-          }
-        }
-      }
-    }
-  }
-}
-""".strip()
+        # Convert to SQL
+        with st.spinner("Generating SQL query..."):
+            sql_result = llm_service.convert_to_sql_query(natural_query)
 
-def execute_query(query: str) -> Dict[str, Any]:
-    """Execute a GraphQL query against the Shopify Admin API."""
-    try:
-        response = requests.post(
-            SHOPIFY_URL,
-            headers=SHOPIFY_HEADERS,
-            json={"query": query}
-        )
-        response.raise_for_status()
-        return response.json()
-    except Exception as e:
-        logging.error(f"Query execution error: {str(e)}")
-        return {"error": str(e)}
+        if not sql_result["success"]:
+            st.error(f"Error generating SQL query: {sql_result['error']}")
+            return
 
-def format_response(response: Dict[str, Any]) -> Dict[str, Any]:
-    """Format the API response for display."""
-    if "errors" in response:
-        error_message = response["errors"][0]["message"]
-        logging.error(f"Shopify API returned an error: {error_message}")
-        return {"error": error_message}
-    
-    try:
-        products_data = response.get("data", {}).get("products", {}).get("edges", [])
-        if not products_data:
-            return {"error": "No products found in response"}
-            
-        formatted = {"products": []}
-        for edge in products_data:
-            node = edge.get("node", {})
-            product = {
-                "title": node.get("title", ""),
-                "description": node.get("description", ""),
-                "productType": node.get("productType", ""),
-                "price": None,
-                "image": None
-            }
-            price_range = node.get("priceRangeV2", {}).get("minVariantPrice", {})
-            if price_range:
-                amount = price_range.get("amount")
-                currency = price_range.get("currencyCode")
-                if amount and currency:
-                    product["price"] = f"{float(amount):.2f} {currency}"
-            
-            images = node.get("images", {}).get("edges", [])
-            if images:
-                product["image"] = images[0].get("node", {}).get("url")
-            
-            formatted["products"].append(product)
-        
-        return formatted
-    except Exception as e:
-        logging.error(f"Response formatting error: {str(e)}")
-        return {"error": str(e)}
+        # Display generated SQL
+        st.subheader("Generated SQL Query:")
+        st.code(sql_result["query"], language="sql")
 
-def main():
-    st.title("Shopify Product Search")
-    
-    # Fetch schema for dynamic query generation
-    schema = fetch_graphql_schema()
-    
-    if "error" in schema:
-        st.error(f"Failed to fetch schema: {schema['error']}")
-        return
-    
-    # Display the available types or fields for user to choose from
-    types = schema.get("data", {}).get("__schema", {}).get("types", [])
-    st.subheader("Available Types in Shopify GraphQL Schema:")
-    for type_info in types:
-        st.write(type_info["name"])
+        # Execute query
+        with st.spinner("Executing query..."):
+            query_result = llm_service.execute_query(sql_result["query"])
 
-    st.markdown("""
-    ### Example Queries:
-    - Show all products
-    - Find products under $50
-    - Search for products with "cotton" in description
-    - Find products tagged as "summer"
-    - Search for specific product types
-    """)
-    
-    query = st.text_input(
-        "What products are you looking for?",
-        placeholder="e.g., Find t-shirts under $50"
-    )
+        if not query_result["success"]:
+            st.error(f"Error executing query: {query_result['error']}")
+            return
 
-    if query and st.button("Search"):
-        try:
-            with st.spinner("Generating query..."):
-                graphql_query = convert_to_graphql_query(query)
-                
-                with st.expander("View GraphQL Query"):
-                    st.code(graphql_query, language='graphql')
-                
-            with st.spinner("Searching products..."):
-                response = execute_query(graphql_query)
-                
-                if "errors" in response:
-                    st.error("Query Error")
-                    st.json(response["errors"])
-                    return
-                
-                formatted = format_response(response)
-                
-                if "error" in formatted:
-                    st.error(formatted["error"])
-                    return
-                
-                if formatted["products"]:
-                    st.subheader(f"Found {len(formatted['products'])} products")
-                    for product in formatted["products"]:
-                        with st.container():
-                            cols = st.columns([1, 2])
-                            with cols[0]:
-                                if product["image"]:
-                                    st.image(product["image"])
-                            with cols[1]:
-                                st.markdown(f"**{product['title']}**")
-                                st.write(f"Type: {product['productType']}")
-                                if product["price"]:
-                                    st.write(f"Price: {product['price']}")
-                                st.write(product["description"])
-                            st.divider()
-                else:
-                    st.warning("No products found.")
+        # Check if there are results
+        if query_result["results"]:
+            df = pd.DataFrame(query_result["results"], columns=query_result["columns"])
+
+            # Create a collapsible DataFrame using Streamlit's expander
+            with st.expander("Click to view query results as a DataFrame"):
+                st.dataframe(df)
+
+            # Extract product details from the JSON result and display them
+            json_results = df.to_dict(orient='records')
+            if "title" in json_results[0] and "images" in json_results[0] and "price" in json_results[0]:
+                st.subheader("Product Details:")
+                for product in json_results:
+                    price = product.get("price", "N/A")
+                    title = product.get("handle", "N/A")
+                    src = product.get("src", "N/A")
                     
-        except Exception as e:
-            st.error(f"An error occurred: {str(e)}")
-            logging.exception("Main execution error")
+                    # Display product details in a neat format using columns for alignment
+                    with st.container():
+                        col1, col2, col3 = st.columns([1, 2, 3])  # Adjust column widths as needed
+                        
+                        with col1:
+                            st.image(src, use_container_width=True)  # Display product image with container width
+                        with col2:
+                            st.write(f"**Price:** {price}")  # Display price
+                            st.write(f"**Title:** {title}")  # Display title
+                        with col3:
+                            st.write(f"**Image Source:** [Link]( {src} )")  # Link to the image if needed
+        else:
+            st.info("No results found.")
+
+if __name__ == "__main__":
+    main()