from flask import Flask, request, jsonify
from sentence_transformers import SentenceTransformer, util

app = Flask(__name__)

# 预加载模型（线上部署关键）
model = SentenceTransformer("flax-sentence-embeddings/st-codesearch-distilroberta-base")
code_snippets = [...]  # 你的代码片段
code_embeddings = model.encode(code_snippets, convert_to_tensor=True)

@app.route('/v1/search', methods=['POST'])
def api_handler():
    """生产级API端点"""
    try:
        # 请求验证
        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)
        results = util.semantic_search(query_emb, code_embeddings, top_k=1)[0]
        
        # 构建响应
        return jsonify({
            "data": {
                "best_match": code_snippets[results[0]['corpus_id']],
                "similarity": float(results[0]['score'])
            }
        })
        
    except Exception as e:
        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=8080)