# main.py

import streamlit as st
import aiohttp
import asyncio
import json

# Function to asynchronously generate text using the OpenAI API
async def generate_text(api_key, prompt, model, temperature, max_tokens, top_p, frequency_penalty, presence_penalty):
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    data = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": temperature,
        "max_tokens": max_tokens,
        "top_p": top_p,
        "frequency_penalty": frequency_penalty,
        "presence_penalty": presence_penalty,
        "stream": True
    }

    async with aiohttp.ClientSession() as session:
        async with session.post("https://api.openai.com/v1/chat/completions", headers=headers, json=data) as response:
            if response.status == 200:
                result_area = st.empty()
                generated_text = ""
                async for chunk in response.content.iter_chunked(1024):
                    chunk = chunk.decode('utf-8')
                    lines = chunk.split('\n')
                    for line in lines:
                        if line.strip() == "":
                            continue
                        if line.startswith("data:"):
                            try:
                                data = json.loads(line[5:])
                                if data["choices"][0]["finish_reason"] is None:
                                    delta = data["choices"][0]["delta"].get("content", "")
                                    generated_text += delta
                                    result_area.markdown(generated_text)
                                else:
                                    break
                            except json.JSONDecodeError as e:
                                print(f"JSONDecodeError: {e}")
                                continue
                return generated_text
            else:
                raise Exception(f"API request failed with status code {response.status}: {await response.text()}")

# Main function to setup Streamlit UI and handle user inputs
def main():
    st.title("GPT-4 Text Generation App")
    api_key = st.text_input("OpenAI API Key", type="password")
    system_prompt = st.text_area("System Prompt", value="You are a helpful assistant.")
    user_prompts = st.text_area("User Prompts (one per line)", value="What is the capital of France?\nExplain the concept of machine learning.").split("\n")
    assistant_prompt = st.text_area("Assistant Prompt", value="Here is the information you requested:")

    model = "gpt-4-turbo-preview"
    temperature = st.slider("Temperature", 0.0, 1.0, 1.0, 0.1)
    max_tokens = st.number_input("Max Tokens", 1, 4000, 256)
    top_p = st.slider("Top P", 0.0, 1.0, 1.0, 0.1)
    frequency_penalty = st.slider("Frequency Penalty", 0.0, 2.0, 0.0, 0.1)
    presence_penalty = st.slider("Presence Penalty", 0.0, 2.0, 0.0, 0.1)

    if st.button("Generate"):
        if not api_key:
            st.error("Please enter your OpenAI API key.")
        else:
            with st.spinner("Generating text..."):
                asyncio.run(generate_text(api_key, f"{system_prompt}\nUser: {user_prompts[0]}\nAssistant: {assistant_prompt}", model, temperature, max_tokens, top_p, frequency_penalty, presence_penalty))

if __name__ == "__main__":
    main()