import streamlit as st
import os
from azure.ai.inference import ChatCompletionsClient
from azure.ai.inference.models import SystemMessage
from azure.ai.inference.models import UserMessage
from azure.core.credentials import AzureKeyCredential
from azure.core.exceptions import ServiceResponseError

# To authenticate with the model you will need to generate a personal access token (PAT) in your GitHub settings. 
# Create your PAT token by following instructions here: https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/managing-your-personal-access-tokens
client = ChatCompletionsClient(
    endpoint="https://models.inference.ai.azure.com",
    credential=AzureKeyCredential(os.getenv("GITHUB_TOKEN")),
)

def main():
    st.title("Azure AI Inference")
    
    # Add a selection box for AI models
    models = ["Mistral-Nemo", "DeepSeek-R1", "Llama-3.3-70B-Instruct"]
    selected_model = st.selectbox("Select AI Model", models)
    st.write("The Deepseek R1 is too popular at the moment. The request might time out. Please try again later.")

    prompt = st.text_area("Enter your prompt here:", 
    value="Create a streamlit app that loads a csv file and runs the naive \
bayes model of scikit-learn on the data. Generate a synthetic data \
with two features and a target variable. The three classes in the \
data are non-linearly separable. Save to CSV. Load the CSV \
in the main app and display the dataframe, visualization, confusion \
matrix and classification report.".strip(), height=150)

    if st.button("Submit"):
        try:
            with st.spinner("Generating response..."):
                response = client.complete(
                    messages=[
                        SystemMessage(content="""You are an expert code assistant."""),
                        UserMessage(content=prompt),
                    ],            
                    model=selected_model,
                    max_tokens=2048,
                )
                st.markdown(response.choices[0].message.content)
        except ServiceResponseError as e:
            st.error(f"Error getting response: {e}")
        except Exception as e:
            st.error(f"Error: {e}") 

if __name__ == "__main__":
    main()