Update app.py

app.py

@@ -1,35 +1,41 @@
 from flask import Flask, request, jsonify
 from sentence_transformers import SentenceTransformer, util
-import os
 
 app = Flask(__name__)
 
-# 预定义代码片段
-CODE = [
-    """def sort_list(x): return sorted(x)""",
-    # 其他代码片段...
-]
+# 预加载模型（线上部署关键）
+model = SentenceTransformer("flax-sentence-embeddings/st-codesearch-distilroberta-base")
+code_snippets = [...]  # 你的代码片段
+code_embeddings = model.encode(code_snippets, convert_to_tensor=True)
 
-# 初始化模型（显示指定缓存路径）
-model = SentenceTransformer(
-    "flax-sentence-embeddings/st-codesearch-distilroberta-base",
-    cache_folder=os.getenv("TRANSFORMERS_CACHE")
-)
-code_emb = model.encode(CODE, convert_to_tensor=True)
-
-@app.route('/search', methods=['POST'])
-def search():
+@app.route('/v1/search', methods=['POST'])
+def api_handler():
+    """生产级API端点"""
     try:
-        query = request.json['query']
+        # 请求验证
+        if not request.is_json:
+            return jsonify({"error": "Invalid Content-Type"}), 415
+            
+        data = request.get_json()
+        if 'query' not in data:
+            return jsonify({"error": "Missing query parameter"}), 400
+            
+        # 执行语义搜索
+        query = data['query']
         query_emb = model.encode(query, convert_to_tensor=True)
-        hits = util.semantic_search(query_emb, code_emb)[0]
-        best = hits[0]
+        results = util.semantic_search(query_emb, code_embeddings, top_k=1)[0]
+        
+        # 构建响应
         return jsonify({
-            'code': CODE[best['corpus_id']],
-            'score': float(best['score'])
+            "data": {
+                "best_match": code_snippets[results[0]['corpus_id']],
+                "similarity": float(results[0]['score'])
+            }
         })
+        
     except Exception as e:
-        return jsonify({"error": str(e)}), 500
+        app.logger.error(f"API Error: {str(e)}")
+        return jsonify({"error": "Internal Server Error"}), 500
 
-if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=5000)
+if __name__ == "__main__":
+    app.run(host='0.0.0.0', port=8080)