hong-red/sql-optimization-agent

1import gradio as gr2import requests3import json4import os5import re6import urllib37from datetime import datetime8import config9 10urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)11 12# --- 使用 config.py 中的统一配置 ---13KIMI_API_KEY = config.KIMI_API_KEY14BASE_URL = config.BASE_URL15MODEL = config.MODEL16SYSTEM_PROMPT = config.SYSTEM_PROMPT17 18def ask_kimi(messages, temperature=0.3, max_tokens=2048):19    headers = {20        "Content-Type": "application/json",21        "Authorization": f"Bearer {KIMI_API_KEY}"22    }23    payload = {24        "model": MODEL,25        "messages": messages,26        "temperature": temperature,27        "max_tokens": max_tokens,28        "stream": True29    }30    try:31        response = requests.post(BASE_URL, headers=headers, json=payload, timeout=60, stream=True)32        33        # 如果报错，获取详细错误信息以供诊断34        if response.status_code != 200:35            error_msg = response.text36            yield f"❌ API 调用异常 (HTTP {response.status_code}): {error_msg}"37            return38        39        full_response = ""40        for line in response.iter_lines():41            if line:42                line_str = line.decode('utf-8')43                if line_str.startswith("data: "):44                    data_str = line_str[6:].strip()45                    if data_str == "[DONE]": break46                    try:47                        data_json = json.loads(data_str)48                        delta = data_json["choices"][0].get("delta", {})49                        if "content" in delta:50                            full_response += delta["content"]51                            yield full_response52                    except: continue53    except Exception as e:54        yield f"❌ 网络连接失败：{str(e)}"55 56def process_input(input_text, file):57    if file is not None:58        try:59            with open(file.name, "r", encoding="utf-8") as f:60                content = f.read()61        except:62            try:63                with open(file.name, "r", encoding="gbk") as f:64                    content = f.read()65            except Exception as e:66                content = f"文件读取失败：{str(e)}"67        input_content = f"请审计以下慢查询日志内容：\n\n{content}"68    else:69        input_content = input_text.strip()70 71    if not input_content:72        yield "请输入内容或上传文件"73        return74 75    messages = [76        {"role": "system", "content": SYSTEM_PROMPT},77        {"role": "user", "content": input_content}78    ]79 80    for partial_reply in ask_kimi(messages):81        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")82        yield f"### 🛡️ DBA 诊断报告 - {timestamp}\n\n{partial_reply}"83 84# Gradio 界面85with gr.Blocks(title="DBA-Agent - 慢查询优化助手", theme=gr.themes.Soft()) as demo:86    gr.Markdown("""87    # 🛡️ DBA-Agent - 慢查询优化助手88    89    上传慢查询日志 / SQL 文件，或直接粘贴 SQL / DDL，AI 将自动诊断并给出优化建议。90    """)91 92    with gr.Row():93        with gr.Column(scale=4):94            input_text = gr.Textbox(95                label="直接粘贴 SQL / DDL / 慢查询日志",96                placeholder="SELECT ... 或粘贴 EXPLAIN 输出...",97                lines=8,98                show_label=True99            )100        with gr.Column(scale=2):101            file_upload = gr.File(102                label="上传 .sql / .log / .txt 文件",103                file_types=[".sql", ".log", ".txt"]104            )105 106    submit_btn = gr.Button("🚀 开始诊断", variant="primary", scale=0)107 108    output_report = gr.Markdown(label="诊断报告")109 110    submit_btn.click(111        fn=process_input,112        inputs=[input_text, file_upload],113        outputs=output_report114    )115 116    gr.Markdown("""117    **提示**：118    - 建议上传包含 EXPLAIN 输出或慢查询日志的文件119    - 当前为只读分析模式，不执行任何 SQL120    - 如需真实执行计划对比，请提供数据库连接信息（后续版本支持）121    """)122 123if __name__ == "__main__":124    demo.queue().launch(125        server_name="0.0.0.0",126        server_port=7860,127        share=False,  # 改成 True 可临时生成公网链接128        debug=True129    )