import streamlit as st
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.legacy.callbacks import CallbackManager
from llama_index.llms.openai_like import OpenAILike

# 设置页面配置必须是第一个 Streamlit 命令
st.set_page_config(page_title="llama_index_demo", page_icon="🦜🔗")

# Create an instance of CallbackManager
callback_manager = CallbackManager()

api_base_url = "https://internlm-chat.intern-ai.org.cn/puyu/api/v1/"
model = "internlm2.5-latest"

# 使用 Streamlit 侧边栏让用户输入 API Key
api_key = st.sidebar.text_input('API Key', value='', type='password')

# 确保 api_key 不为空
if not api_key:
    st.warning("请输入有效的 API Key.")
else:
    # 初始化 LLM 时使用用户提供的 API Key
    llm = OpenAILike(model=model, api_base=api_base_url, api_key=api_key, is_chat_model=True, callback_manager=callback_manager)

    st.title("llama_index_demo")

    # 初始化模型
    @st.cache_resource
    def init_models():
        embed_model = HuggingFaceEmbedding(
            model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
        )
        Settings.embed_model = embed_model

        # 用初始化llm
        Settings.llm = llm

        documents = SimpleDirectoryReader("./data").load_data()
        index = VectorStoreIndex.from_documents(documents)
        query_engine = index.as_query_engine()

        return query_engine

    # 检查是否需要初始化模型
    if 'query_engine' not in st.session_state and api_key:
        st.session_state['query_engine'] = init_models()

    def greet2(question):
        response = st.session_state['query_engine'].query(question)
        return response

    # Store LLM generated responses
    if "messages" not in st.session_state.keys():
        st.session_state.messages = [{"role": "assistant", "content": "你好，我是你的助手，有什么我可以帮助你的吗？"}]

    # Display or clear chat messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.write(message["content"])

    def clear_chat_history():
        st.session_state.messages = [{"role": "assistant", "content": "你好，我是你的助手，有什么我可以帮助你的吗？"}]

    st.sidebar.button('Clear Chat History', on_click=clear_chat_history)

    # Function for generating LLaMA2 response
    def generate_llama_index_response(prompt_input):
        return greet2(prompt_input)

    # User-provided prompt
    if prompt := st.chat_input():
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.write(prompt)

    # Generate response only if the last message is not from assistant
    if st.session_state.messages[-1]["role"] != "assistant":
        with st.chat_message("assistant"):
            with st.spinner("Thinking..."):
                response = generate_llama_index_response(prompt)
                placeholder = st.empty()
                placeholder.markdown(response)
        message = {"role": "assistant", "content": response}
        st.session_state.messages.append(message)