...
llm = Ollama(model="deepseek-r1")
embed_model = OllamaEmbedding(model_name="milkey/dmeta-embedding-zh:f16")

Settings.llm = llm
Settings.embed_model = embed_model

...
# 上传文件（这里允许上传的类型为 txt/ PDF）
uploaded_file = st.file_uploader("", type=["pdf", "txt"])

if uploaded_file is not None:
    
    # 计算文件内容的哈希值
    file_content = uploaded_file.getvalue()
    new_file_hash = hashlib.md5(file_content).hexdigest()

    # 如果文件发生变化或第一次上传，则重新创建 query_engine
    if st.session_state.file_hash != new_file_hash:
        st.session_state.file_hash = new_file_hash

        with st.spinner("正在处理文档..."):

            st.session_state.messages = []

            # 将上传的文件保存到临时位置
            with open("temp.pdf", "wb") as f:
                f.write(file_content)
...

...
            # 加载并处理文档
            docs = SimpleDirectoryReader(input_files=["temp.pdf"]).load_data()
            
            # 创建向量存储索引
            index = VectorStoreIndex.from_documents(docs)

            # 使用自定义提示创建查询引擎
            query_prompt = """
            1. 使用以下上下文来回答最后的问题。
            2. 如果你不知道答案，请直接说"我不知道"，不要编造答案。
            3. 保持答案简洁，限制在3-4句话内。
            上下文：{context_str}
            问题：{query_str}
            答案："""
            
            st.session_state.query_engine = index.as_query_engine(
                text_qa_template=PromptTemplate(query_prompt),
                similarity_top_k=3,
            )

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

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

        # 生成回答
        with st.chat_message("assistant"):
            with st.spinner("思考中..."):
                response = st.session_state.query_engine.query(prompt)
                st.markdown(response.response)
                st.session_state.messages.append({"role": "assistant", "response": response.response})
else:
    st.write("上传 PDF 文件以开始对话。")

...
# 添加清除对话历史的按钮
if st.button("清除对话历史"):
    st.session_state.messages = []
    st.rerun()

...
                # 处理回答中的标签
                text = response.response
                text = re.sub(
                    r'([\s\S]*?)', 
                    lambda m: f'
{m.group(1)}
',
                    text,
                    flags=re.DOTALL
                )
                st.markdown(text, unsafe_allow_html=True)
                clean_text = re.sub(r'[\s\S]*?', '', response.response).strip()
                st.session_state.messages.append({"role": "assistant", "content": clean_text})

...
THINK_PROMPT = """
    请结合下面提供的上下文(context字段),仔细思考我的问题(question字段)，完善之前的推理过程(thoughts字段)。
    如果上下文与问题并不相关，请忽略context字段，直接思考问题。
    注意你只需思考该问题，不需要输出答案，答案部分直接输出"Reasoning Done"即可。
"""

# 初始化组件
reasoning_llm = Ollama(model="deepseek-r1:1.5b")

class DeepSeek:
    def __init__(self, reasoning_llm, reasoning_times=1):
        self.reasoning_llm = reasoning_llm
        self.reasoning_times = reasoning_times

    #检索相关上下文
    def retrieve(self, query):
        nodes = st.session_state.query_engine.retrieve(query["question"])
        return {
            "context": nodes,
            "question": query["question"]
        }

    #迭代推理，并返回最终COT
    def think(self, input: str) -> str:

        thoughts = input
        for _ in range(self.reasoning_times):

            retrieved_docs = self.retrieve({"question": thoughts})
            docs_content = "\n\n".join(
                node.text for node in retrieved_docs["context"]
            )

            prompt_json = {
                "context": docs_content,
                "question": input,
                "thoughts": thoughts
            }

            response = self.reasoning_llm.complete(THINK_PROMPT + '\n\n' + dumps(prompt_json, ensure_ascii=False))

            think = re.findall(r'(.*?)', response.text, re.DOTALL)
            if not think:
                return thoughts or "No reasoning available."
            thoughts = think[0]

            st.write(f"thoughts: {thoughts}")

        return thoughts