Update app.py

app.py

@@ -4,34 +4,60 @@ from langchain_huggingface import HuggingFaceEndpoint, HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
 from langchain_community.document_loaders import TextLoader
 from langchain_text_splitters import CharacterTextSplitter
-# 修改这里：使用更稳定的导入路径
-from langchain.chains.retrieval_qa.base import RetrievalQA
+from langchain.chains import RetrievalQA
 
-# 1. 引擎初始化
+# 1. 链接顶级大脑 Qwen 2.5
+# 系统会自动读取你刚才设置的 Secret: HF_TOKEN
 llm = HuggingFaceEndpoint(
     repo_id="Qwen/Qwen2.5-7B-Instruct",
-    huggingfacehub_api_token=os.getenv("HF_TOKEN")
+    huggingfacehub_api_token=os.getenv("HF_TOKEN"),
+    timeout=300
 )
 
-# 2. 知识库加载
-if not os.path.exists("knowledge.txt"):
-    with open("knowledge.txt", "w", encoding="utf-8") as f:
-        f.write("这里是你的私有大脑知识库。")
+# 2. 初始化与加载私有知识库
+def init_knowledge_base():
+    # 如果没有 knowledge.txt，先创建一个初始模版
+    if not os.path.exists("knowledge.txt"):
+        with open("knowledge.txt", "w", encoding="utf-8") as f:
+            f.write("大脑初始化成功。请在 knowledge.txt 文件中输入你的私有知识。")
+    
+    # 读取知识库
+    loader = TextLoader("knowledge.txt", encoding="utf-8")
+    documents = loader.load()
+    
+    # 文本切分：让 AI 更好找重点
+    text_splitter = CharacterTextSplitter(chunk_size=600, chunk_overlap=100)
+    docs = text_splitter.split_documents(documents)
+    
+    # 向量化处理：将文字坐标化 (使用中文优化模型)
+    embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-small-zh-v1.5")
+    vectorstore = FAISS.from_documents(docs, embeddings)
+    return vectorstore
 
-loader = TextLoader("knowledge.txt", encoding="utf-8")
-docs = CharacterTextSplitter(chunk_size=500, chunk_overlap=50).split_documents(loader.load())
-embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-small-zh-v1.5")
-vectorstore = FAISS.from_documents(docs, embeddings)
-
-# 3. 构建问答链
-qa_chain = RetrievalQA.from_chain_type(llm=llm, retriever=vectorstore.as_retriever())
+# 启动时先加载一次知识库
+vectorstore = init_knowledge_base()
+qa_chain = RetrievalQA.from_chain_type(
+    llm=llm, 
+    retriever=vectorstore.as_retriever(search_kwargs={"k": 3})
+)
 
-# 4. 界面逻辑
-def chat(msg, history):
+# 3. 聊天交互逻辑
+def predict(message, history):
     try:
-        res = qa_chain.invoke({"query": msg})
-        return res["result"]
+        # AI 会先搜索 knowledge.txt，再结合大模型给出回答
+        response = qa_chain.invoke({"query": message})
+        return response["result"]
     except Exception as e:
-        return f"运行出错：{str(e)}"
+        return f"大脑思考时遇到了一点问题，请检查 Token 或网络。详情: {str(e)}"
+
+# 4. 搭建前端界面
+demo = gr.ChatInterface(
+    predict,
+    title="我的全能私有大脑",
+    description="我已经读取了你的知识库。你可以问我关于你自己的事，也可以和我聊任何话题。",
+    examples=["你是谁？", "解释一下知识库里的核心内容", "帮我写一个 Python 脚本"],
+    theme="soft"
+)
 
-gr.ChatInterface(chat).launch()
+if __name__ == "__main__":
+    demo.launch()