app.py

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

@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)

            # Basic query validation
            if query.startswith("query") and "products" in query:
                return {"success": True, "query": query}

            return {"success": False, "error": "Failed to generate valid GraphQL query"}
            
        except Exception as e:
            self.logger.error(f"Query generation error: {str(e)}")
            return {"success": False, "error": str(e)}

    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:
            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("Shopify GraphQL Natural Language Query Converter")

    load_dotenv()

    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
    
    try:
        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

    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 GraphQL Query"):
        if not natural_query.strip():
            st.warning("Please enter a valid query.")
            return

        with st.spinner("Generating GraphQL query..."):
            graphql_result = graphql_converter.convert_to_graphql_query(natural_query)

        if not graphql_result["success"]:
            st.error(f"Error generating GraphQL query: {graphql_result['error']}")
            return

        st.subheader("Generated GraphQL Query:")
        st.code(graphql_result["query"], language="graphql")

        with st.spinner("Executing 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

        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.")

if __name__ == "__main__":
    main()