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

app = Flask(__name__)

# 预定义代码片段
CODE = [
    """def sort_list(x): return sorted(x)""",
    """def count_above_threshold(elements, threshold=0):...""",
    """def find_min_max(elements):..."""
]

# 初始化模型（启动时自动下载）
model = SentenceTransformer("flax-sentence-embeddings/st-codesearch-distilroberta-base")
code_emb = model.encode(CODE, convert_to_tensor=True)

@app.route('/search', methods=['POST'])
def search():
    query = request.json.get('query', '')
    query_emb = model.encode(query, convert_to_tensor=True)
    hits = util.semantic_search(query_emb, code_emb)[0]
    best = hits[0]
    return jsonify({
        'code': CODE[best['corpus_id']],
        'score': float(best['score'])
    })

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)