Create app.py

app.py

@@ -0,0 +1,192 @@
+# app.py
+import os
+import threading
+from functools import wraps
+import google.generativeai as genai
+from flask import Flask, request, jsonify, send_from_directory
+from flask_cors import CORS
+from dotenv import load_dotenv
+import chromadb
+
+# 从您的核心逻辑文件中导入类
+from app_chromadb import MarkdownKnowledgeBase
+
+# --- 初始化与配置 ---
+load_dotenv()
+app = Flask(__name__, static_folder='.', static_url_path='')
+CORS(app)
+
+# --- API 密钥认证配置 ---
+VALID_API_KEYS_STR = os.environ.get("KNOWLEDGE_BASE_API_KEYS", "")
+VALID_API_KEYS = {key.strip() for key in VALID_API_KEYS_STR.split(',') if key.strip()}
+if not VALID_API_KEYS:
+    print("⚠️ 警告: 未配置 KNOWLEDGE_BASE_API_KEYS。API 将对所有人开放！")
+
+# --- ChromaDB, Gemini, SiliconFlow 实例配置 ---
+try:
+    CHROMA_DATA_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), "chroma_db")
+    COLLECTION_NAME = "markdown_knowledge_base_m3"
+    chroma_client = chromadb.PersistentClient(path=CHROMA_DATA_PATH)
+    collection = chroma_client.get_or_create_collection(name=COLLECTION_NAME)
+    print(f"✅ ChromaDB 客户端已连接，数据存储在 '{CHROMA_DATA_PATH}'")
+except Exception as e:
+    chroma_client = None
+    collection = None
+    print(f"❌ 初始化 ChromaDB 失败: {e}")
+
+try:
+    GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
+    genai.configure(api_key=GEMINI_API_KEY)
+    gemini_model = genai.GenerativeModel('gemini-1.5-flash')
+    print("✅ Gemini API 已配置。")
+except Exception as e:
+    gemini_model = None
+    print(f"❌ Gemini API 配置失败: {e}")
+
+try:
+    SF_API_TOKEN = os.environ.get("SILICONFLOW_API_TOKEN")
+    kb_instance = MarkdownKnowledgeBase(api_token=SF_API_TOKEN, chroma_collection=collection)
+    print("✅ SiliconFlow 与知识库实例已配置。")
+except Exception as e:
+    kb_instance = None
+    print(f"❌ 知识库实例配置失败: {e}")
+
+kb_status = { "is_building": False }
+
+# --- API 密钥认证装饰器 ---
+def require_api_key(f):
+    @wraps(f)
+    def decorated_function(*args, **kwargs):
+        if not VALID_API_KEYS: return f(*args, **kwargs)
+        api_key = request.headers.get('X-API-Key')
+        if api_key and api_key in VALID_API_KEYS:
+            return f(*args, **kwargs)
+        else:
+            return jsonify({"error": "授权失败。请提供有效'X-API-Key'请求头。"}), 403
+    return decorated_function
+
+# --- 前端页面路由 ---
+@app.route('/')
+def serve_index():
+    return send_from_directory('.', 'index.html')
+
+# --- API 端点 ---
+@app.route('/status', methods=['GET'])
+def get_status():
+    if collection:
+        kb_status['total_items'] = collection.count()
+        kb_status['is_built'] = kb_status['total_items'] > 0
+        if not kb_status['is_building']:
+             kb_status['message'] = f"知识库已就绪，共有 {kb_status['total_items']} 个条目。"
+    else:
+        kb_status['message'] = "ChromaDB 未连接。"
+    return jsonify(kb_status)
+
+@app.route('/build', methods=['POST'])
+@require_api_key
+def build_knowledge_base():
+    if kb_status['is_building']:
+        return jsonify({"error": "知识库已在构建中，请稍后。"}), 409
+    
+    if not kb_instance:
+        return jsonify({"error": "知识库实例未初始化，无法构建。"}), 500
+        
+    data = request.get_json()
+    clear_existing = data.get('clear_existing', False)
+    
+    build_params = {
+        'folder_path': data.get('folder_path'),
+        'chunk_size': data.get('chunk_size', 4096),
+        'overlap': data.get('overlap', 400),
+        'max_files': data.get('max_files', 500),
+        'sample_mode': data.get('sample_mode', 'largest')
+    }
+    
+    def build_in_background():
+        global kb_status, collection
+        kb_status['is_building'] = True
+        kb_status['message'] = "构建任务开始..."
+        try:
+            if clear_existing and chroma_client:
+                print(f"正在清空现有集合: {COLLECTION_NAME}")
+                chroma_client.delete_collection(name=COLLECTION_NAME)
+                collection = chroma_client.get_or_create_collection(name=COLLECTION_NAME)
+                kb_instance.collection = collection
+                print("集合已清空并重建。")
+
+            kb_instance.build_knowledge_base(**build_params)
+            kb_status['message'] = f"构建完成！知识库现有 {collection.count()} 个条目。"
+        except Exception as e:
+            kb_status['message'] = f"构建时出错: {e}"
+            print(f"Error during build: {e}")
+        finally:
+            kb_status['is_building'] = False
+
+    thread = threading.Thread(target=build_in_background)
+    thread.start()
+    return jsonify({"message": "知识库构建任务已在后台启动。"}), 202
+
+@app.route('/search', methods=['GET'])
+@require_api_key
+def search_in_kb():
+    if not (collection and collection.count() > 0):
+        return jsonify({"error": "知识库为空，请先构建。"}), 400
+
+    if not kb_instance:
+        return jsonify({"error": "知识库实例未初始化，无法搜索。"}), 500
+
+    query = request.args.get('query')
+    top_k = request.args.get('top_k', default=5, type=int)
+    
+    if not query:
+        return jsonify({"error": "必须提供 'query' 参数"}), 400
+
+    try:
+        results = kb_instance.search(query, top_k=top_k)
+        return jsonify(results)
+    except Exception as e:
+        return jsonify({"error": f"搜索时发生错误: {e}"}), 500
+
+@app.route('/summarize', methods=['POST'])
+@require_api_key
+def summarize_results():
+    if not gemini_model:
+        return jsonify({"error": "Gemini API 未配置或初始化失败。"}), 500
+    
+    data = request.get_json()
+    query = data.get('query')
+    search_results = data.get('results')
+
+    if not query or not search_results:
+        return jsonify({"error": "必须提供查询和搜索结果。"}), 400
+
+    context = "\n\n---\n\n".join([item['content'] for item in search_results])
+    prompt = f"""
+    根据以下本地知识库中搜索到的内容，请用清晰、简洁的中文直接回答用户的问题。
+    如果内容不足以回答，请说明现有信息无法直接回答。
+
+    用户问题: "{query}"
+
+    搜索到的内容:
+    ---
+    {context}
+    ---
+
+    你的回答:
+    """
+
+    try:
+        print(f"正在向 Gemini 模型 '{gemini_model.model_name}' 发送请求...")
+        response = gemini_model.generate_content(prompt)
+        summary = response.text
+        return jsonify({"summary": summary})
+    except Exception as e:
+        print(f"调用 Gemini API 时出错: {e}")
+        return jsonify({"error": f"调用 AI 服务时出错: {e}"}), 500
+
+
+if __name__ == '__main__':
+    print("知识库后端服务 (最终版) 启动...")
+    print("✅ 服务已启动！请在浏览器中打开 http://127.0.0.1:5000")
+    # 使用生产级服务器时应移除 debug=False
+    app.run(host='0.0.0.0', port=5000, debug=False)