import streamlit as st
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.settings import Settings
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.openai_like import OpenAILike
import os
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 设置 API 参数
base_url = "https://internlm-chat.intern-ai.org.cn/puyu/api/v1/"
api_key = os.getenv("MY_API_KEY")
model = "internlm2.5-latest"

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

# 初始化一个 HuggingFaceEmbedding 对象，用于将文本转换为向量表示
embed_model = HuggingFaceEmbedding(
    model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
)

# 将创建的嵌入模型赋值给全局设置的 embed_model 属性
Settings.embed_model = embed_model

# 将创建的 LLM 赋值给全局设置的 llm 属性
Settings.llm = llm

# 从指定目录读取所有文档，并加载数据到内存中
documents = SimpleDirectoryReader("./data").load_data()

# 创建一个 VectorStoreIndex，并使用之前加载的文档来构建索引
index = VectorStoreIndex.from_documents(documents)

# 创建一个查询引擎，这个引擎可以接收查询并返回相关文档的响应
query_engine = index.as_query_engine()

# 设置页面标题
st.title("LlamaIndex Chat")

# 创建一个文本输入框供用户输入问题
user_input = st.text_input("请输入你的问题:")

# 按钮用于提交问题
if st.button("发送"):
    # 使用查询引擎获取回复
    response = query_engine.query(user_input)

    # 显示回复
    st.write(f"回复: {response}")