import os
import streamlit as st
import requests

# Retrieve Hugging Face API token from environment variable
API_TOKEN = os.environ.get("HUGGING_FACE_API_TOKEN")

# Define the Hugging Face API URL
API_URL = "https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3-8B"
headers = {"Authorization": f"Bearer {API_TOKEN}"}

# Function to query the Hugging Face API
def query(payload):
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.json()

# Streamlit app
def main():
    st.title("SQL Query Generator")

    # User prompt input
    prompt = st.text_input("Enter your prompt:", "Please generate a SQL query to fetch data from the database.")

    # Button to generate SQL query
    if st.button("Generate SQL Query"):
        # Generate payload for Hugging Face API
        payload = {"inputs": prompt}
        
        # Query the Hugging Face API
        with st.spinner('Generating SQL query...'):
            output = query(payload)
        
        # Display the SQL query response
        if "generated_text" in output:
            st.write("Generated SQL Query:")
            st.code(output["generated_text"])
        else:
            st.error("Failed to generate SQL query. Please try again.")

if __name__ == "__main__":
    main()