import streamlit as st
from openai import OpenAI, AzureOpenAI
from typing import Iterator
import os
from phoenix.otel import register

PHOENIX_API_KEY=st.secrets['PHOENIX_API_KEY']
os.environ["PHOENIX_CLIENT_HEADERS"] = f"api_key={PHOENIX_API_KEY}"

tracer_provider = register(
  project_name=st.secrets['PHOENIX_PROJECT_NAME'],
  endpoint=os.getenv('PHOENIX_COLLECTOR_ENDPOINT'),
)

from openinference.instrumentation.openai import OpenAIInstrumentor
OpenAIInstrumentor().instrument(tracer_provider=tracer_provider)

st.set_page_config(
    page_title="Free GPT-4o Chat",
    page_icon="💬",
    layout="centered"
)

st.title("💬 GPT-4o Chat")
with st.expander("Notice"):
    st.write('''
        Please note that this app collects your conversation records. 
        Do not input any personal or sensitive information, such as:
        - National identification numbers
        - Contact information (phone numbers, email addresses)
        - Bank account or credit card details
        - Health or medical information
        - Any other data that can be used to identify you

        Use this app cautiously and avoid sharing sensitive data. 
        Do not use this app in inappropriate contexts. By using this app, you agree to these terms and conditions.
    ''')

client = AzureOpenAI(
    api_key=st.secrets['API_KEY'],
    api_version=st.secrets['API_VERSION'],
    azure_endpoint=st.secrets['ENDPOINT']
)

if "openai_model" not in st.session_state:
    st.session_state["openai_model"] = st.secrets['MODEL']

if "messages" not in st.session_state:
    st.session_state.messages = []

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

if prompt := st.chat_input("What is up?"):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        stream = client.chat.completions.create(
            model=st.session_state["openai_model"],
            messages=[
                {"role": m["role"], "content": m["content"]}
                for m in st.session_state.messages
            ],
            stream=True,
        )
        response = st.write_stream(stream)
    st.session_state.messages.append(
        {"role": "assistant", "content": response})