Forrest99/codesearchBase

1from flask import Flask, request, jsonify, render_template_string2from sentence_transformers import SentenceTransformer, util3import logging4import sys5import signal6 7# 初始化 Flask 应用8app = Flask(__name__)9 10# 配置日志，级别设为 INFO11logging.basicConfig(level=logging.INFO)12app.logger = logging.getLogger("CodeSearchAPI")13 14# 预定义代码片段15CODE_SNIPPETS = [16 17"print('Hello, World!')",18"def add(a, b): return a + b",19"import random; def generate_random(): return random.randint(1, 100)",20"def is_even(n): return n % 2 == 0",21"def string_length(s): return len(s)",22"from datetime import date; def get_current_date(): return date.today()",23"import os; def file_exists(path): return os.path.exists(path)",24"def read_file(path): return open(path, 'r').read()",25"def write_file(path, content): open(path, 'w').write(content)",26"from datetime import datetime; def get_current_time(): return datetime.now()",27"def to_upper(s): return s.upper()",28"def to_lower(s): return s.lower()",29"def reverse_string(s): return s[::-1]",30"def list_length(lst): return len(lst)",31"def list_max(lst): return max(lst)",32"def list_min(lst): return min(lst)",33"def sort_list(lst): return sorted(lst)",34"def merge_lists(lst1, lst2): return lst1 + lst2",35"def remove_element(lst, element): lst.remove(element)",36"def is_list_empty(lst): return len(lst) == 0",37"def count_char(s, char): return s.count(char)",38"def contains_substring(s, sub): return sub in s",39"def int_to_str(n): return str(n)",40"def str_to_int(s): return int(s)",41"def is_numeric(s): return s.isdigit()",42"def get_index(lst, element): return lst.index(element)",43"def clear_list(lst): lst.clear()",44"def reverse_list(lst): lst.reverse()",45"def remove_duplicates(lst): return list(set(lst))",46"def is_in_list(lst, value): return value in lst",47"def create_dict(): return {}",48"def add_to_dict(d, key, value): d[key] = value",49"def delete_key(d, key): del d[key]",50"def get_keys(d): return list(d.keys())",51"def get_values(d): return list(d.values())",52"def merge_dicts(d1, d2): return {**d1, **d2}",53"def is_dict_empty(d): return len(d) == 0",54"def get_value(d, key): return d[key]",55"def key_exists(d, key): return key in d",56"def clear_dict(d): d.clear()",57"def count_lines(path): return len(open(path).readlines())",58"def write_list_to_file(path, lst): open(path, 'w').write('\\n'.join(map(str, lst)))",59"def read_list_from_file(path): return open(path, 'r').read().splitlines()",60"def count_words(path): return len(open(path, 'r').read().split())",61"def is_leap_year(year): return year % 4 == 0 and (year % 100 != 0 or year % 400 == 0)",62"from datetime import datetime; def format_time(dt): return dt.strftime('%Y-%m-%d %H:%M:%S')",63"from datetime import date; def days_between(d1, d2): return (d2 - d1).days",64"import os; def get_current_dir(): return os.getcwd()",65"import os; def list_files(path): return os.listdir(path)",66"import os; def create_dir(path): os.mkdir(path)"67"import os; def remove_dir(path): os.rmdir(path)",68"import os; def is_file(path): return os.path.isfile(path)",69"import os; def is_dir(path): return os.path.isdir(path)",70"import os; def get_file_size(path): return os.path.getsize(path)",71"import os; def rename_file(src, dst): os.rename(src, dst)",72"import shutil; def copy_file(src, dst): shutil.copy(src, dst)",73"import shutil; def move_file(src, dst): shutil.move(src, dst)",74"import os; def delete_file(path): os.remove(path)",75"import os; def get_env_var(key): return os.getenv(key)",76"import os; def set_env_var(key, value): os.environ[key] = value",77"import webbrowser; def open_url(url): webbrowser.open(url)",78"import requests; def send_get_request(url): return requests.get(url).text",79"import json; def parse_json(data): return json.loads(data)",80"import json; def write_json(data, path): open(path, 'w').write(json.dumps(data))",81"import json; def read_json(path): return json.loads(open(path, 'r').read())",82"def list_to_string(lst): return ''.join(lst)",83"def string_to_list(s): return list(s)",84"def join_with_comma(lst): return ','.join(lst)",85"def join_with_newline(lst): return '\\n'.join(lst)",86"def split_by_space(s): return s.split()",87"def split_by_char(s, char): return s.split(char)",88"def split_to_chars(s): return list(s)",89"def replace_string(s, old, new): return s.replace(old, new)",90"def remove_spaces(s): return s.replace(' ', '')",91"import string; def remove_punctuation(s): return s.translate(str.maketrans('', '', string.punctuation))",92"def is_string_empty(s): return len(s) == 0",93"def is_palindrome(s): return s == s[::-1]",94"import csv; def write_csv(data, path): open(path, 'w', newline='').write('\\n'.join([','.join(map(str, row)) for row in data]))",95"import csv; def read_csv(path): return [row for row in csv.reader(open(path, 'r'))]",96"def count_csv_lines(path): return len(open(path).readlines())",97"import random; def shuffle_list(lst): random.shuffle(lst)",98"import random; def random_choice(lst): return random.choice(lst)",99"import random; def random_sample(lst, k): return random.sample(lst, k)",100"import random; def roll_dice(): return random.randint(1, 6)",101"import random; def flip_coin(): return random.choice(['Heads', 'Tails'])",102"import random; import string; def generate_password(length=8): return ''.join(random.choices(string.ascii_letters + string.digits, k=length))",103"import random; def generate_color(): return '#%06x' % random.randint(0, 0xFFFFFF)",104"import uuid; def generate_uuid(): return str(uuid.uuid4())",105"class MyClass: pass",106"def create_instance(): return MyClass()",107"class MyClass: def my_method(self): pass",108"class MyClass: def __init__(self): self.my_attr = None",109"class ChildClass(MyClass): pass",110"class ChildClass(MyClass): def my_method(self): pass",111"class MyClass: @classmethod def my_class_method(cls): pass",112"class MyClass: @staticmethod def my_static_method(): pass",113"def check_type(obj): return type(obj)",114"def get_attr(obj, attr): return getattr(obj, attr)",115"def set_attr(obj, attr, value): setattr(obj, attr, value)",116"def del_attr(obj, attr): delattr(obj, attr)",117"""try:118    x = 1 / 0119except ZeroDivisionError:120    pass""",121"""class CustomError(Exception): pass122def raise_custom_error(): raise CustomError('Error occurred')""",123"""try:124    x = 1 / 0125except Exception as e: return str(e)""",126"""import logging; logging.basicConfig(filename='error.log', level=logging.ERROR); logging.error('Error occurred')""",127"""import time; def timer(): start = time.time(); return lambda: time.time() - start""",128"""import time; def run_time(): start = time.time(); return lambda: time.time() - start""",129"""import sys; def print_progress(progress): sys.stdout.write(f'\\rProgress: {progress}%'); sys.stdout.flush()""",130"""import time; def delay(seconds): time.sleep(seconds)""",131"lambda x: x * 2",132"map(lambda x: x * 2, [1, 2, 3])",133"filter(lambda x: x > 2, [1, 2, 3])",134"from functools import reduce; reduce(lambda x, y: x + y, [1, 2, 3])",135"[x * 2 for x in [1, 2, 3]]",136"{x: x * 2 for x in [1, 2, 3]}",137"{x for x in [1, 2, 3]}",138"set1 & set2",139"set1 | set2",140"set1 - set2",141"[x for x in lst if x is not None]",142"""try:143    with open('file.txt', 'r') as f: pass144except IOError: pass""",145"type(var)",146"bool(s)",147"if condition: pass",148"while condition: pass",149"for item in lst: pass",150"for key, value in d.items(): pass",151"for char in s: pass",152"for item in lst: if condition: break",153"for item in lst: if condition: continue",154"def my_func(): pass",155"def my_func(param=1): pass",156"def my_func(): return 1, 2",157"def my_func(*args): pass",158"def my_func(**kwargs): pass",159"""import time; def timer(func):160    def wrapper(*args, **kwargs):161        start = time.time()162        result = func(*args, **kwargs)163        print(f'Time: {time.time() - start}'); return result164    return wrapper""",165"""def decorator(func):166    def wrapper(*args, **kwargs): return func(*args, **kwargs)167    return wrapper""",168"""from functools import lru_cache; @lru_cache(maxsize=None)169def my_func(): pass""",170"def my_generator(): yield 1",171"gen = my_generator(); next(gen)",172"class MyIterator: def __iter__(self): return self; def __next__(self): pass",173"it = iter([1, 2, 3]); next(it)",174"for i, val in enumerate(lst): pass",175"list(zip(lst1, lst2))",176"dict(zip(keys, values))",177"lst1 == lst2",178"dict1 == dict2",179"set1 == set2",180"set(lst)",181"set.clear()",182"len(set) == 0",183"set.add(item)",184"set.remove(item)",185"item in set",186"len(set)",187"set1 & set2",188"set(lst1).issubset(lst2)",189"sub in s",190"s[0]",191"s[-1]",192"import mimetypes; mimetypes.guess_type(path)[0] == 'text/plain'",193"import mimetypes; mimetypes.guess_type(path)[0].startswith('image/')",194"round(num)",195"import math; math.ceil(num)",196"import math; math.floor(num)",197"f'{num:.2f}'",198"import random; import string; ''.join(random.choices(string.ascii_letters + string.digits, k=8))",199"import os; os.path.exists(path)",200"import os; for root, dirs, files in os.walk(path): pass",201"import os; os.path.splitext(path)[1]",202"import os; os.path.basename(path)",203"import os; os.path.abspath(path)",204"import platform; platform.python_version()",205"import platform; platform.system()",206"import multiprocessing; multiprocessing.cpu_count()",207"import psutil; psutil.virtual_memory().total",208"import psutil; psutil.disk_usage('/')",209"import socket; socket.gethostbyname(socket.gethostname())",210"import requests; try: requests.get('http://www.google.com'); return True; except: return False",211"import requests; def download_file(url, path): with open(path, 'wb') as f: f.write(requests.get(url).content)",212"def upload_file(path): with open(path, 'rb') as f: requests.post('http://example.com/upload', files={'file': f})",213"import requests; requests.post(url, data={'key': 'value'})",214"import requests; requests.get(url, params={'key': 'value'})",215"import requests; requests.get(url, headers={'key': 'value'})",216"from bs4 import BeautifulSoup; BeautifulSoup(html, 'html.parser')",217"from bs4 import BeautifulSoup; soup.title.text",218"from bs4 import BeautifulSoup; [a['href'] for a in soup.find_all('a')]",219"from bs4 import BeautifulSoup; import requests; for img in soup.find_all('img'): requests.get(img['src']).content",220"from collections import Counter; Counter(text.split())",221"import requests; session = requests.Session(); session.post(login_url, data={'username': 'user', 'password': 'pass'})",222"from bs4 import BeautifulSoup; soup.get_text()",223"import re; re.findall(r'[\\w.-]+@[\\w.-]+', text)",224"import re; re.findall(r'\\+?\\d[\\d -]{8,12}\\d', text)",225"import re; re.findall(r'\\d+', text)",226"import re; re.sub(pattern, repl, text)",227"import re; re.match(pattern, text)",228"from bs4 import BeautifulSoup; soup.get_text()",229"import html; html.escape(text)",230"import html; html.unescape(text)",231"import tkinter as tk; root = tk.Tk(); root.mainloop()",232"import tkinter as tk; def add_button(window, text): return tk.Button(window, text=text)",233"""def bind_click(button, func): button.config(command=func)""",234"""import tkinter.messagebox; def show_alert(message): tkinter.messagebox.showinfo('Info', message)""",235"""def get_entry_text(entry): return entry.get()""",236"def set_title(window, title): window.title(title)",237"def set_size(window, width, height): window.geometry(f'{width}x{height}')",238"""def center_window(window):239    window.update_idletasks()240    width = window.winfo_width()241    height = window.winfo_height()242    x = (window.winfo_screenwidth() // 2) - (width // 2)243    y = (window.winfo_screenheight() // 2) - (height // 2)244    window.geometry(f'{width}x{height}+{x}+{y}')""",245"""def create_menu(window): return tk.Menu(window)""",246"def create_combobox(window): return ttk.Combobox(window)",247"def create_radiobutton(window, text): return tk.Radiobutton(window, text=text)",248"def create_checkbutton(window, text): return tk.Checkbutton(window, text=text)",249"""from PIL import ImageTk, Image; def show_image(window, path):250    img = Image.open(path)251    photo = ImageTk.PhotoImage(img)252    label = tk.Label(window, image=photo)253    label.image = photo254    return label""",255"import pygame; def play_audio(path): pygame.mixer.init(); pygame.mixer.music.load(path); pygame.mixer.music.play()",256"import cv2; def play_video(path): cap = cv2.VideoCapture(path); while cap.isOpened(): ret, frame = cap.read()",257"def get_playback_time(): return pygame.mixer.music.get_pos()",258"import pyautogui; def screenshot(): return pyautogui.screenshot()",259"import pyautogui; import time; def record_screen(duration): return [pyautogui.screenshot() for _ in range(duration)]",260"def get_mouse_pos(): return pyautogui.position()",261"import pyautogui; def type_text(text): pyautogui.write(text)",262"import pyautogui; def click_mouse(x, y): pyautogui.click(x, y)",263"import time; def get_timestamp(): return int(time.time())",264"import datetime; def timestamp_to_date(ts): return datetime.datetime.fromtimestamp(ts)",265"import time; def date_to_timestamp(dt): return int(time.mktime(dt.timetuple()))",266"def get_weekday(): return datetime.datetime.now().strftime('%A')",267"import calendar; def get_month_days(): return calendar.monthrange(datetime.datetime.now().year, datetime.datetime.now().month)[1]",268"def first_day_of_year(): return datetime.date(datetime.datetime.now().year, 1, 1)",269"def last_day_of_year(): return datetime.date(datetime.datetime.now().year, 12, 31)",270"def first_day_of_month(month): return datetime.date(datetime.datetime.now().year, month, 1)",271"import calendar; def last_day_of_month(month): return datetime.date(datetime.datetime.now().year, month, calendar.monthrange(datetime.datetime.now().year, month)[1])",272"def is_weekday(): return datetime.datetime.now().weekday() < 5",273"def is_weekend(): return datetime.datetime.now().weekday() >= 5",274"def current_hour(): return datetime.datetime.now().hour",275"def current_minute(): return datetime.datetime.now().minute",276"def current_second(): return datetime.datetime.now().second",277"import time; def delay_1s(): time.sleep(1)",278"import time; def millis_timestamp(): return int(time.time() * 1000)",279"def format_time(dt, fmt='%Y-%m-%d %H:%M:%S'): return dt.strftime(fmt)",280"from dateutil.parser import parse; def parse_time(s): return parse(s)",281"import threading; def create_thread(target): return threading.Thread(target=target)",282"import time; def thread_pause(seconds): time.sleep(seconds)",283"def run_threads(*threads): [t.start() for t in threads]",284"import threading; def current_thread_name(): return threading.current_thread().name",285"def set_daemon(thread): thread.daemon = True",286"import threading; lock = threading.Lock()",287"import multiprocessing; def create_process(target): return multiprocessing.Process(target=target)",288"import os; def get_pid(): return os.getpid()",289"import psutil; def is_process_alive(pid): return psutil.pid_exists(pid)",290"def run_processes(*procs): [p.start() for p in procs]",291"from queue import Queue; q = Queue()",292"from multiprocessing import Pipe; parent_conn, child_conn = Pipe()",293"import os; def limit_cpu_usage(percent): os.system(f'cpulimit -p {os.getpid()} -l {percent}')",294"import subprocess; def run_command(cmd): subprocess.run(cmd, shell=True)",295"import subprocess; def get_command_output(cmd): return subprocess.check_output(cmd, shell=True).decode()",296"def get_exit_code(cmd): return subprocess.call(cmd, shell=True)",297"def is_success(code): return code == 0",298"import os; def script_path(): return os.path.realpath(__file__)",299"import sys; def get_cli_args(): return sys.argv[1:]",300"import argparse; parser = argparse.ArgumentParser()",301"parser.print_help()",302"help('modules')",303"import pip; def install_pkg(pkg): pip.main(['install', pkg])",304"import pip; def uninstall_pkg(pkg): pip.main(['uninstall', pkg])",305"import pkg_resources; def get_pkg_version(pkg): return pkg_resources.get_distribution(pkg).version",306"import venv; def create_venv(path): venv.create(path)",307"import pip; def list_pkgs(): return pip.get_installed_distributions()",308"import pip; def upgrade_pkg(pkg): pip.main(['install', '--upgrade', pkg])",309"import sqlite3; conn = sqlite3.connect(':memory:')",310"def execute_query(conn, query): return conn.execute(query)",311"""def insert_record(conn, table, data): conn.execute(f'INSERT INTO {table} VALUES ({",".join("?"*len(data))})', data)""",312"def delete_record(conn, table, condition): conn.execute(f'DELETE FROM {table} WHERE {condition}')",313"def update_record(conn, table, set_clause, condition): conn.execute(f'UPDATE {table} SET {set_clause} WHERE {condition}')",314"def fetch_all(conn, query): return conn.execute(query).fetchall()",315"def safe_query(conn, query, params): return conn.execute(query, params)",316"def close_db(conn): conn.close()",317"def create_table(conn, name, columns): conn.execute(f'CREATE TABLE {name} ({columns})')",318"def drop_table(conn, name): conn.execute(f'DROP TABLE {name}')",319"def table_exists(conn, name): return conn.execute(f\"SELECT name FROM sqlite_master WHERE type='table' AND name='{name}'\").fetchone()",320"def list_tables(conn): return conn.execute(\"SELECT name FROM sqlite_master WHERE type='table'\").fetchall()",321"""from sqlalchemy import Column, Integer, String322class User(Base):323    __tablename__ = 'users'324    id = Column(Integer, primary_key=True)325    name = Column(String)""",326"session.add(User(name='John'))",327"session.query(User).filter_by(name='John')",328"session.query(User).filter_by(name='John').delete()",329"session.query(User).filter_by(name='John').update({'name': 'Bob'})",330"Base = declarative_base()",331"class Admin(User): pass",332"id = Column(Integer, primary_key=True)",333"name = Column(String, unique=True)",334"name = Column(String, default='Unknown')",335"import csv; def export_csv(data, path): open(path, 'w').write('\\n'.join([','.join(map(str, row)) for row in data]))",336"import pandas as pd; pd.DataFrame(data).to_excel(path)",337"import json; json.dump(data, open(path, 'w'))",338"pd.read_excel(path).values.tolist()",339"pd.concat([pd.read_excel(f) for f in files])",340"with pd.ExcelWriter(path, mode='a') as writer: df.to_excel(writer, sheet_name='New')",341"from openpyxl.styles import copy; copy.copy(style)",342"from openpyxl.styles import PatternFill; cell.fill = PatternFill(start_color='FFFF00', fill_type='solid')",343"from openpyxl.styles import Font; cell.font = Font(bold=True)",344"sheet['A1'].value",345"sheet['A1'] = value",346"from PIL import Image; Image.open(path).size",347"from PIL import Image; Image.open(path).resize((w, h))"348 349 350 351 352 353 354 355    356]357 358# 全局服务状态359service_ready = False360 361# 优雅关闭处理362def handle_shutdown(signum, frame):363    app.logger.info("收到终止信号，开始关闭...")364    sys.exit(0)365 366signal.signal(signal.SIGTERM, handle_shutdown)367signal.signal(signal.SIGINT, handle_shutdown)368 369# 初始化模型和预计算编码370try:371    app.logger.info("开始加载模型...")372    model = SentenceTransformer(373        "flax-sentence-embeddings/st-codesearch-distilroberta-base",374        cache_folder="/model-cache"375    )376    # 预计算代码片段的编码（强制使用 CPU）377    code_emb = model.encode(CODE_SNIPPETS, convert_to_tensor=True, device="cpu")378    service_ready = True379    app.logger.info("服务初始化完成")380except Exception as e:381    app.logger.error("初始化失败: %s", str(e))382    raise383 384# Hugging Face 健康检查端点，必须响应根路径385@app.route('/')386def hf_health_check():387    # 如果请求接受 HTML，则返回一个简单的 HTML 页面（包含测试链接）388    if request.accept_mimetypes.accept_html:389        html = """390        <h2>CodeSearch API</h2>391        <p>服务状态：{{ status }}</p>392        <p>你可以在地址栏输入 /search?query=你的查询 来测试接口</p>393        """394        status = "ready" if service_ready else "initializing"395        return render_template_string(html, status=status)396    # 否则返回 JSON 格式的健康检查397    if service_ready:398        return jsonify({"status": "ready"}), 200399    else:400        return jsonify({"status": "initializing"}), 503401 402# 搜索 API 端点，同时支持 GET 和 POST 请求403@app.route('/search', methods=['GET', 'POST'])404def handle_search():405    if not service_ready:406        app.logger.info("服务未就绪")407        return jsonify({"error": "服务正在初始化"}), 503408 409    try:410        # 根据请求方法提取查询内容411        if request.method == 'GET':412            query = request.args.get('query', '').strip()413        else:414            data = request.get_json() or {}415            query = data.get('query', '').strip()416 417        if not query:418            app.logger.info("收到空的查询请求")419            return jsonify({"error": "查询不能为空"}), 400420 421        # 记录接收到的查询422        app.logger.info("收到查询请求: %s", query)423 424        # 对查询进行编码，并进行语义搜索425        query_emb = model.encode(query, convert_to_tensor=True, device="cpu")426        hits = util.semantic_search(query_emb, code_emb, top_k=1)[0]427        best = hits[0]428 429        result = {430            "code": CODE_SNIPPETS[best['corpus_id']],431            "score": round(float(best['score']), 4)432        }433 434        # 记录返回结果435        app.logger.info("返回结果: %s", result)436        return jsonify(result)437 438    except Exception as e:439        app.logger.error("请求处理失败: %s", str(e))440        return jsonify({"error": "服务器内部错误"}), 500441 442if __name__ == "__main__":443    # 本地测试用，Hugging Face Spaces 通常通过 gunicorn 启动444    app.run(host='0.0.0.0', port=7860)445