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",
    cache_folder=os.getenv("TRANSFORMERS_CACHE")
)
code_emb = model.encode(CODE, convert_to_tensor=True)

@app.route('/search', methods=['POST'])
def search():
    try:
        query = request.json['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'])
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500

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