import streamlit as st
import os
from ibm_watsonx_ai.foundation_models import Model

# Function to get credentials from user input
def get_credentials():
    st.sidebar.title("IBM Cloud Credentials")
    url = "https://us-south.ml.cloud.ibm.com"
    apikey = st.sidebar.text_input("Please enter your API key", type="password")
    if apikey:
        return {"url": url, "apikey": apikey}
    return None

# Streamlit app
def main():
    st.title("Granite Chat with IBM")

    # Get credentials from the user
    credentials = get_credentials()
    if not credentials:
        st.error("Please enter your API key.")
        return

    # Model and parameters
    model_id = "ibm/granite-13b-chat-v2"
    parameters = {
        "decoding_method": "greedy",
        "max_new_tokens": 900,
        "repetition_penalty": 1.05
    }
    project_id = os.getenv("PROJECT_ID")

    # Initialize the model
    model = Model(
        model_id=model_id,
        params=parameters,
        credentials=credentials,
        project_id=project_id
    )

    # User input
    prompt_input = """
    You are Granite Chat, an AI language model developed by IBM. You are a cautious assistant. You carefully follow instructions. You are helpful and harmless and you follow ethical guidelines and promote positive behavior. You always respond to greetings (for example, hi, hello, g'day, morning, afternoon, evening, night, what's up, nice to meet you, sup, etc) with "Hello! I am Granite Chat, created by IBM. How can I help you today?". Please do not say anything else and do not start a conversation.

    Write a poem about cats
    """

    question = st.text_area("Enter your question or prompt:")
    if st.button("Submit"):
        if question:
            formatted_question = f"{question}"
            prompt = f"{prompt_input}\n{formatted_question}"
            generated_response = model.generate_text(prompt=prompt, guardrails=False)
            st.write(f"**AI Response:** {generated_response}")
        else:
            st.error("Please enter a question or prompt.")

if __name__ == "__main__":
    main()