Update app.py

app.py

@@ -1,13 +1,13 @@
-from flask import Flask, request, jsonify
+from flask import Flask, request, jsonify, render_template_string
 from sentence_transformers import SentenceTransformer, util
 import logging
 import sys
 import signal
 
-# 初始化Flask应用
+# 初始化 Flask 应用
 app = Flask(__name__)
 
-# 配置日志，日志级别设为INFO
+# 配置日志，级别设为 INFO
 logging.basicConfig(level=logging.INFO)
 app.logger = logging.getLogger("CodeSearchAPI")
 
@@ -31,42 +31,48 @@ def handle_shutdown(signum, frame):
 signal.signal(signal.SIGTERM, handle_shutdown)
 signal.signal(signal.SIGINT, handle_shutdown)
 
-# 初始化模型和编码
+# 初始化模型和预计算编码
 try:
-    # Hugging Face Spaces专用缓存路径
+    app.logger.info("开始加载模型...")
     model = SentenceTransformer(
         "flax-sentence-embeddings/st-codesearch-distilroberta-base",
         cache_folder="/model-cache"
     )
-    
-    # 预计算编码（强制使用CPU）
-    code_emb = model.encode(CODE_SNIPPETS, 
-                            convert_to_tensor=True,
-                            device="cpu")
-    
+    # 预计算代码片段的编码（强制使用 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("初始化失败: %s", str(e))
     raise
 
-# Hugging Face健康检查端点
+# Hugging Face 健康检查端点，必须响应根路径
 @app.route('/')
 def hf_health_check():
-    """必须响应根路径的健康检查"""
+    # 如果请求接受 HTML，则返回一个简单的 HTML 页面（包含测试链接）
+    if request.accept_mimetypes.accept_html:
+        html = """
+        <h2>CodeSearch API</h2>
+        <p>服务状态：{{ status }}</p>
+        <p>你可以在地址栏输入 /search?query=你的查询 来测试接口</p>
+        """
+        status = "ready" if service_ready else "initializing"
+        return render_template_string(html, status=status)
+    # 否则返回 JSON 格式的健康检查
     if service_ready:
         return jsonify({"status": "ready"}), 200
     else:
         return jsonify({"status": "initializing"}), 503
 
-# 支持GET和POST请求的搜索API端点
+# 搜索 API 端点，同时支持 GET 和 POST 请求
 @app.route('/search', methods=['GET', 'POST'])
 def handle_search():
     if not service_ready:
+        app.logger.info("服务未就绪")
         return jsonify({"error": "服务正在初始化"}), 503
 
     try:
-        # 区分GET和POST请求，GET从URL参数中获取query，POST从JSON体中获取
+        # 根据请求方法提取查询内容
         if request.method == 'GET':
             query = request.args.get('query', '').strip()
         else:
@@ -76,14 +82,12 @@ def handle_search():
         if not query:
             app.logger.info("收到空的查询请求")
             return jsonify({"error": "查询不能为空"}), 400
-        
+
         # 记录接收到的查询
         app.logger.info("收到查询请求: %s", query)
 
-        # 对查询进行编码，并搜索最匹配的代码片段
-        query_emb = model.encode(query, 
-                                 convert_to_tensor=True,
-                                 device="cpu")
+        # 对查询进行编码，并进行语义搜索
+        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]
 
@@ -91,15 +95,15 @@ def handle_search():
             "code": CODE_SNIPPETS[best['corpus_id']],
             "score": round(float(best['score']), 4)
         }
+
         # 记录返回结果
         app.logger.info("返回结果: %s", result)
-
         return jsonify(result)
-        
+
     except Exception as e:
         app.logger.error("请求处理失败: %s", str(e))
         return jsonify({"error": "服务器内部错误"}), 500
 
 if __name__ == "__main__":
-    # Hugging Face Spaces会通过gunicorn启动，此处仅为本地测试保留
+    # 本地测试用，Hugging Face Spaces 通常通过 gunicorn 启动
     app.run(host='0.0.0.0', port=7860)