Update app.py

app.py

@@ -1,8 +1,10 @@
 from flask import Flask, request, jsonify
 from sentence_transformers import SentenceTransformer, util
 import logging
-import os
+import sys
+import signal
 
+# 初始化Flask应用
 app = Flask(__name__)
 
 # 配置日志
@@ -18,51 +20,61 @@ CODE_SNIPPETS = [
     return min(elements), max(elements)"""
 ]
 
-# 初始化标记
-model_ready = False
+# 全局服务状态
+service_ready = False
 
+# 优雅关闭处理
+def handle_shutdown(signum, frame):
+    app.logger.info("收到终止信号，开始关闭...")
+    sys.exit(0)
+
+signal.signal(signal.SIGTERM, handle_shutdown)
+signal.signal(signal.SIGINT, handle_shutdown)
+
+# 初始化模型和编码
 try:
-    # 初始化模型（使用预下载的缓存）
+    # Hugging Face Spaces专用缓存路径
     model = SentenceTransformer(
         "flax-sentence-embeddings/st-codesearch-distilroberta-base",
-        cache_folder=os.getenv("HF_HOME")
+        cache_folder="/model-cache"
     )
     
-    # 预计算编码
-    code_emb = model.encode(CODE_SNIPPETS, convert_to_tensor=True)
-    model_ready = True
-    app.logger.info("模型加载完成，服务就绪")
+    # 预计算编码（强制使用CPU）
+    code_emb = model.encode(CODE_SNIPPETS, 
+                          convert_to_tensor=True,
+                          device="cpu")
+    
+    service_ready = True
+    app.logger.info("服务初始化完成")
 except Exception as e:
-    app.logger.error(f"模型初始化失败: {str(e)}")
+    app.logger.error(f"初始化失败: {str(e)}")
     raise
 
-@app.route('/health')
-def health_check():
-    """健康检查端点"""
-    if model_ready:
+# Hugging Face健康检查端点
+@app.route('/')
+def hf_health_check():
+    """必须响应根路径的健康检查"""
+    if service_ready:
         return jsonify({"status": "ready"}), 200
     else:
         return jsonify({"status": "initializing"}), 503
 
+# 搜索API端点
 @app.route('/search', methods=['POST'])
 def handle_search():
-    """搜索请求处理"""
-    if not model_ready:
+    if not service_ready:
         return jsonify({"error": "服务正在初始化"}), 503
-        
+    
     try:
-        # 请求验证
-        if not request.is_json:
-            return jsonify({"error": "需要 application/json"}), 415
-            
         data = request.get_json()
         query = data.get('query', '').strip()
         
         if not query:
             return jsonify({"error": "查询不能为空"}), 400
             
-        # 处理查询
-        query_emb = model.encode(query, convert_to_tensor=True)
+        query_emb = model.encode(query, 
+                               convert_to_tensor=True,
+                               device="cpu")
         hits = util.semantic_search(query_emb, code_emb, top_k=1)[0]
         best = hits[0]
         
@@ -76,4 +88,5 @@ def handle_search():
         return jsonify({"error": "服务器内部错误"}), 500
 
 if __name__ == "__main__":
-    app.run(host='0.0.0.0', port=8080)
+    # Hugging Face Spaces会通过gunicorn启动，此处仅为本地测试保留
+    app.run(host='0.0.0.0', port=7860)