Update app.py

app.py

@@ -1,124 +1,191 @@
-import streamlit as st
-import pandas as pd
 import os
+import logging
+import requests
+import json
+from typing import Dict, Any, List
+from dataclasses import dataclass
 from dotenv import load_dotenv
+import streamlit as st
+import pandas as pd
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import logging
 
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+@dataclass
+class GraphQLSchemaType:
+    """Store GraphQL type information including fields and relationships"""
+    name: str
+    fields: List[Dict[str, Any]]
+    relationships: List[Dict[str, str]]
+
+class ShopifyGraphQLConverter:
+    def __init__(self, shop_url: str, access_token: str, api_key: str, model_name: str):
+        """
+        Initialize Shopify GraphQL converter
+        
+        :param shop_url: Shopify store URL
+        :param access_token: Shopify Admin API access token
+        :param api_key: LLM service API key
+        :param model_name: Model name for Hugging Face
+        """
+        load_dotenv()
+
+        # Ensure shop URL has https:// scheme
+        if not shop_url.startswith(('http://', 'https://')):
+            shop_url = f'https://{shop_url}'
+        
+        # Shopify GraphQL endpoint configuration
+        self.shop_url = shop_url
+        self.graphql_endpoint = f"{shop_url}/admin/api/2024-04/graphql.json"
+        self.access_token = access_token
+        
+        # LLM API configuration
+        self.api_key = api_key
+        self.llm_api_url = "https://api.groq.com/openai/v1/chat/completions"
+        self.llm_headers = {
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json"
+        }
+
+        # Load model directly for natural language processing
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(model_name)
+
+        # Predefined schema for Shopify resources
+        self.schema = {
+            "Product": GraphQLSchemaType(
+                name="Product",
+                fields=[
+                    {"name": "id", "type": "ID", "required": False},
+                    {"name": "title", "type": "String", "required": False},
+                    {"name": "description", "type": "String", "required": False},
+                    {"name": "productType", "type": "String", "required": False},
+                    {"name": "vendor", "type": "String", "required": False},
+                    {"name": "priceRangeV2", "type": "ProductPriceRangeV2", "required": False}
+                ],
+                relationships=[
+                    {"from_field": "variants", "to_type": "ProductVariant"},
+                    {"from_field": "collections", "to_type": "Collection"}
+                ]
+            ),
+        }
+
+        # Setup logging
+        logging.basicConfig(level=logging.INFO)
+        self.logger = logging.getLogger(__name__)
+
+    def generate_graphql_query(self, natural_query: str) -> str:
+        """
+        Generate GraphQL query from natural language using Llama model
+        
+        :param natural_query: The query in natural language
+        :return: GraphQL query as a string
+        """
+        inputs = self.tokenizer(natural_query, return_tensors="pt")
+        outputs = self.model.generate(**inputs, max_length=500)
+        query = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        return query
+
+    def convert_to_graphql_query(self, natural_query: str) -> Dict[str, Any]:
+        """
+        Convert natural language to Shopify GraphQL query
+        
+        :param natural_query: Natural language query string
+        :return: Dictionary containing GraphQL query or error
+        """
+        try:
+            query = self.generate_graphql_query(natural_query)
 
-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")
+            # Basic query validation
+            if query.startswith("query") and "products" in query:
+                return {"success": True, "query": query}
 
-    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}
+            return {"success": False, "error": "Failed to generate valid GraphQL query"}
+            
         except Exception as e:
-            logger.error(f"Error generating SQL query: {e}")
+            self.logger.error(f"Query generation error: {str(e)}")
             return {"success": False, "error": str(e)}
 
-    def execute_query(self, sql_query):
+    def execute_query(self, graphql_query: str) -> Dict[str, Any]:
+        """
+        Execute the GraphQL query against Shopify Admin API
+        
+        :param graphql_query: GraphQL query to execute
+        :return: Dictionary containing query results or error
+        """
         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}")
+            payload = {"query": graphql_query}
+            response = requests.post(
+                self.graphql_endpoint,
+                headers={
+                    "Content-Type": "application/json",
+                    "X-Shopify-Access-Token": self.access_token
+                },
+                json=payload
+            )
+            response.raise_for_status()
+            
+            result = response.json()
+            return {"success": True, "data": result.get('data', {}), "errors": result.get('errors', [])}
+                
+        except requests.exceptions.RequestException as e:
+            self.logger.error(f"Shopify GraphQL query execution error: {str(e)}")
             return {"success": False, "error": str(e)}
 
 def main():
-    st.title("Natural Language to SQL Query Converter")
-    st.write("Enter your question about the database in natural language.")
-    
-    # Load environment variables
+    st.title("Shopify GraphQL Natural Language Query Converter")
+
     load_dotenv()
-    db_path = os.getenv("DB_PATH")
-    
-    if not db_path:
-        st.error("Missing database path in environment variables.")
-        logger.error("DB path not found in environment variables.")
+
+    shop_url = os.getenv("SHOPIFY_STORE_URL", "https://agkd0n-fa.myshopify.com")
+    access_token = os.getenv("SHOPIFY_ACCESS_TOKEN")
+    groq_api_key = os.getenv("GROQ_API_KEY")
+    model_name = "Qwen/Qwen2.5-72B-Instruct"  # Modify this for Llama3 if needed
+
+    if not all([shop_url, access_token, groq_api_key]):
+        st.error("Missing environment variables. Please set SHOPIFY_STORE_URL, SHOPIFY_ACCESS_TOKEN, and GROQ_API_KEY")
         return
     
-    # Initialize LLM Service
     try:
-        llm_service = LLMService(db_path=db_path)
+        graphql_converter = ShopifyGraphQLConverter(shop_url, access_token, groq_api_key, model_name)
     except Exception as e:
         st.error(f"Error initializing service: {str(e)}")
         return
 
-    # Input for natural language query
-    natural_query = st.text_area("Enter your query", "Show me all albums by artist 'Queen'", height=100)
+    natural_query = st.text_area("Enter your Shopify query in natural language", "Find shirt with red color", height=100)
 
-    if st.button("Generate and Execute Query"):
+    if st.button("Generate and Execute GraphQL Query"):
         if not natural_query.strip():
             st.warning("Please enter a valid query.")
             return
 
-        # Convert to SQL
-        with st.spinner("Generating SQL query..."):
-            sql_result = llm_service.convert_to_sql_query(natural_query)
+        with st.spinner("Generating GraphQL query..."):
+            graphql_result = graphql_converter.convert_to_graphql_query(natural_query)
 
-        if not sql_result["success"]:
-            st.error(f"Error generating SQL query: {sql_result['error']}")
+        if not graphql_result["success"]:
+            st.error(f"Error generating GraphQL query: {graphql_result['error']}")
             return
 
-        # Display generated SQL
-        st.subheader("Generated SQL Query:")
-        st.code(sql_result["query"], language="sql")
+        st.subheader("Generated GraphQL Query:")
+        st.code(graphql_result["query"], language="graphql")
 
-        # Execute query
         with st.spinner("Executing query..."):
-            query_result = llm_service.execute_query(sql_result["query"])
+            query_result = graphql_converter.execute_query(graphql_result["query"])
 
         if not query_result["success"]:
             st.error(f"Error executing query: {query_result['error']}")
             return
 
-        # 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")
-                    
-                    # 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
+        st.subheader("Query Results:")
+        if query_result["errors"]:
+            st.error(f"GraphQL Errors: {query_result['errors']}")
+
+        if query_result["data"]:
+            products = query_result["data"].get("products", {}).get("edges", [])
+            if products:
+                product_list = [{"Title": p["node"]["title"], "Vendor": p["node"]["vendor"]} for p in products]
+                st.dataframe(pd.DataFrame(product_list))
+            else:
+                st.info("No products found.")
         else:
             st.info("No results found.")