import streamlit as st
from llama_index 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

# Callback manager
callback_manager = CallbackManager()

# API 信息
api_base_url = "https://internlm-chat.intern-ai.org.cn/puyu/api/v1/"
model = "internlm2.5-latest"
api_key = "eyJ0eXBlIjoiSldUIiwiYWxnIjoiSFM1MTIifQ.eyJqdGkiOiI4MDAwNzM0OCIsInJvbCI6IlJPTEVfUkVHSVNURVIiLCJpc3MiOiJPcGVuWExhYiIsImlhdCI6MTczNzczMjUwMCwiY2xpZW50SWQiOiJlYm1ydm9kNnlvMG5semFlazF5cCIsInBob25lIjoiMTk4MjEyMTUyNzEiLCJ1dWlkIjoiZjc2MDM3NTctMTU2Yy00MTM3LWE1YmEtNTk3MDljODRiNDRkIiwiZW1haWwiOiIiLCJleHAiOjE3NTMyODQ1MDB9.4zD9ixAv1JdP_AJLQhNW3tCgzCGquW6eFcbV0XNqmqCZ0pL5A4hIPVA0zeFleg-n04O1IsyIZZ0rmkATZ1V6_A"  # 替换为你的真实 API Key

# 初始化 LLM
llm = OpenAILike(
    model=model,
    api_base=api_base_url,
    api_key=api_key,
    is_chat_model=True,
    callback_manager=callback_manager,
)

# Streamlit 页面配置
st.set_page_config(page_title="Llama Index Demo", page_icon="🦜🔗")
st.title("Llama Index Demo")

# 初始化模型
@st.cache_resource
def init_models():
    try:
        # 使用 Hugging Face Hub 的公开模型
        embed_model = HuggingFaceEmbedding(model_name="sentence-transformers/all-MiniLM-L6-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
    except Exception as e:
        st.error(f"模型初始化失败：{e}")
        return None

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

# 模型查询函数
def greet2(question):
    query_engine = st.session_state['query_engine']
    if query_engine:
        response = query_engine.query(question)
        return response.response  # 确保返回的内容是字符串
    else:
        return "模型初始化失败，请检查环境配置。"

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

# 初始化聊天记录
if "messages" not in st.session_state:
    st.session_state.messages = [{"role": "assistant", "content": "你好，我是你的助手，有什么我可以帮助你的吗？"}]

# 显示聊天记录
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.write(message["content"])

# 清空聊天记录按钮
st.sidebar.button('Clear Chat History', on_click=clear_chat_history)

# 用户输入处理
if prompt := st.chat_input("请输入您的问题："):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.write(prompt)

    # Assistant 响应
    with st.chat_message("assistant"):
        with st.spinner("思考中..."):
            response = greet2(prompt)  # 调用查询函数
            if response:
                st.write(response)
                st.session_state.messages.append({"role": "assistant", "content": response})
            else:
                st.error("未能生成响应，请检查模型状态。")