Update app.py

app.py

@@ -1,6 +1,7 @@
 from flask import Flask, request, jsonify
 from sentence_transformers import SentenceTransformer, util
 import logging
+import os
 
 app = Flask(__name__)
 
@@ -8,7 +9,7 @@ app = Flask(__name__)
 logging.basicConfig(level=logging.INFO)
 app.logger = logging.getLogger("CodeSearchAPI")
 
-# 预定义代码片段（已验证数据）
+# 预定义代码片段
 CODE_SNIPPETS = [
     """def sort_list(x): return sorted(x)""",
     """def count_above_threshold(elements, threshold=0):
@@ -17,66 +18,61 @@ CODE_SNIPPETS = [
     return min(elements), max(elements)"""
 ]
 
-# 输入数据验证
-def validate_snippets(snippets):
-    cleaned = []
-    for idx, s in enumerate(snippets):
-        if not isinstance(s, str):
-            app.logger.warning(f"索引 {idx} 类型错误，已转换为字符串")
-            s = str(s)
-        cleaned.append(s.replace("...", "").strip())
-    return [s for s in cleaned if len(s) > 0]
+# 初始化标记
+model_ready = False
 
-# 初始化模型和编码
 try:
-    model = SentenceTransformer("flax-sentence-embeddings/st-codesearch-distilroberta-base")
-    valid_snippets = validate_snippets(CODE_SNIPPETS)
-    code_emb = model.encode(valid_snippets, convert_to_tensor=True)
-    app.logger.info(f"成功加载模型，编码 {len(valid_snippets)} 个有效代码片段")
+    # 初始化模型（使用预下载的缓存）
+    model = SentenceTransformer(
+        "flax-sentence-embeddings/st-codesearch-distilroberta-base",
+        cache_folder=os.getenv("HF_HOME")
+    )
+    
+    # 预计算编码
+    code_emb = model.encode(CODE_SNIPPETS, convert_to_tensor=True)
+    model_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:
+        return jsonify({"status": "ready"}), 200
+    else:
+        return jsonify({"status": "initializing"}), 503
+
 @app.route('/search', methods=['POST'])
 def handle_search():
-    """API 处理端点"""
+    """搜索请求处理"""
+    if not model_ready:
+        return jsonify({"error": "服务正在初始化"}), 503
+        
     try:
         # 请求验证
         if not request.is_json:
-            app.logger.warning("无效的 Content-Type")
             return jsonify({"error": "需要 application/json"}), 415
             
         data = request.get_json()
         query = data.get('query', '').strip()
         
         if not query:
-            app.logger.warning("收到空查询")
             return jsonify({"error": "查询不能为空"}), 400
             
-        # 编码查询
-        try:
-            query_emb = model.encode(query, convert_to_tensor=True)
-        except Exception as e:
-            app.logger.error(f"编码失败: {str(e)}")
-            return jsonify({"error": "查询处理失败"}), 500
-            
-        # 语义搜索
-        try:
-            hits = util.semantic_search(query_emb, code_emb, top_k=1)[0]
-            best = hits[0]
-            result = {
-                "code": valid_snippets[best['corpus_id']],
-                "score": round(float(best['score']), 4)
-            }
-            app.logger.info(f"成功处理查询: '{query}'")
-            return jsonify(result)
-            
-        except IndexError:
-            app.logger.error("无匹配结果")
-            return jsonify({"error": "无可用匹配"}), 404
-            
+        # 处理查询
+        query_emb = model.encode(query, convert_to_tensor=True)
+        hits = util.semantic_search(query_emb, code_emb, top_k=1)[0]
+        best = hits[0]
+        
+        return jsonify({
+            "code": CODE_SNIPPETS[best['corpus_id']],
+            "score": round(float(best['score']), 4)
+        })
+        
     except Exception as e:
-        app.logger.error(f"未知错误: {str(e)}", exc_info=True)
+        app.logger.error(f"请求处理失败: {str(e)}")
         return jsonify({"error": "服务器内部错误"}), 500
 
 if __name__ == "__main__":