Update app.py

app.py

@@ -2,57 +2,55 @@ import os
 import gradio as gr
 from langchain_huggingface import HuggingFaceEndpoint, HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
-from langchain_community.document_loaders import PyPDFLoader, TextLoader, Docx2txtLoader
-from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import TextLoader
+from langchain_text_splitters import CharacterTextSplitter
 from langchain.chains import RetrievalQA
 
-# 1. 接入顶级大脑 API
+# 1. 配置大模型动力源 (使用 Qwen 2.5)
+# 确保你在 Settings -> Secrets 中设置了 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. 向量化模型（负责把文字变成AI能懂的坐标）
+# 2. 初始化私有知识库
+# 如果 knowledge.txt 不存在，先创建一个简单的，防止报错
+if not os.path.exists("knowledge.txt"):
+    with open("knowledge.txt", "w", encoding="utf-8") as f:
+        f.write("欢迎使用全能 AI 大脑！请在 knowledge.txt 中输入你的私有知识。")
+
+# 加载并切分文档
+loader = TextLoader("knowledge.txt", encoding="utf-8")
+documents = loader.load()
+text_splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+docs = text_splitter.split_documents(documents)
+
+# 3. 创建向量检索系统 (使用中文优化的 Embedding 模型)
 embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-small-zh-v1.5")
+vectorstore = FAISS.from_documents(docs, embeddings)
+
+# 4. 构建问答链 (RAG 核心)
+qa_chain = RetrievalQA.from_chain_type(
+    llm=llm,
+    retriever=vectorstore.as_retriever(search_kwargs={"k": 3})
+)
+
+# 5. 定义界面交互
+def chat_fn(message, history):
+    try:
+        # 执行检索并生成回答
+        response = qa_chain.invoke({"query": message})
+        return response["result"]
+    except Exception as e:
+        return f"大脑连接超时或出错，请检查 Token 设置。错误详情: {str(e)}"
 
-def process_files(files):
-    all_docs = []
-    for file in files:
-        # 根据文件后缀选择加载器
-        if file.name.endswith('.pdf'):
-            loader = PyPDFLoader(file.name)
-        elif file.name.endswith('.docx'):
-            loader = Docx2txtLoader(file.name)
-        else:
-            loader = TextLoader(file.name)
-        all_docs.extend(loader.load())
-    
-    # 3. 智能切片（防止文档太长 AI 记不住）
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=600, chunk_overlap=100)
-    split_docs = text_splitter.split_documents(all_docs)
-    
-    # 4. 存入临时向量库
-    vectorstore = FAISS.from_documents(split_docs, embeddings)
-    return vectorstore
-
-# 5. 问答函数
-def predict(message, history, file_output):
-    if not file_output:
-        return "请先上传 PDF/Word/TXT 文档，我才能开启私有大脑模式。"
-    
-    # 处理文件并建立检索链
-    vectorstore = process_files(file_output)
-    qa_chain = RetrievalQA.from_chain_type(llm=llm, retriever=vectorstore.as_retriever())
-    
-    response = qa_chain.invoke(message)
-    return response["result"]
-
-# 6. 界面设计
+# 启动 Gradio 界面
 demo = gr.ChatInterface(
-    predict,
-    additional_inputs=[gr.File(file_count="multiple", label="上传私有文档 (.pdf, .docx, .txt)")],
+    chat_fn, 
     title="我的全能私有大脑",
-    description="上传文件后，我可以基于文档内容回答任何专业问题。"
+    description="基于 Qwen 2.5 + RAG 技术。它会先查阅你的 knowledge.txt 再回答。"
 )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()