hong-red/sql-optimization-agent

1import http.server2import socketserver3import json4import requests5import os6import urllib37import logging8from datetime import datetime9import config10 11# ====== 配置日志（生产友好） ======12logging.basicConfig(13    level=logging.INFO,14    format='%(asctime)s [%(levelname)s] %(message)s',15    handlers=[16        logging.StreamHandler(),17        logging.FileHandler("web_agent.log", encoding='utf-8')18    ]19)20logger = logging.getLogger(__name__)21 22# --- 使用 config.py 中的统一配置 ---23KIMI_API_KEY = config.KIMI_API_KEY24BASE_URL = config.BASE_URL25MODEL = config.MODEL26SYSTEM_PROMPT = config.SYSTEM_PROMPT27 28# ====== 全局上下文（简单版：每个会话独立） ======29sessions = {}  # {session_id: messages}  简单内存存储，生产可换 Redis 或文件30 31def ask_kimi(messages, temperature=0.3, max_tokens=4000):32    """调用 Kimi API"""33    headers = {34        "Content-Type": "application/json",35        "Authorization": f"Bearer {KIMI_API_KEY}"36    }37    payload = {38        "model": MODEL,39        "messages": messages,40        "temperature": temperature,41        "max_tokens": max_tokens42    }43    proxies = {"http": None, "https": None}44    45    try:46        logger.info(f"API 调用开始，消息数: {len(messages)}")47        response = requests.post(48            BASE_URL,49            headers=headers,50            json=payload,51            proxies=proxies,52            timeout=60,53            verify=False  # 生产建议移除或用 certifi 证书54        )55        response.raise_for_status()56        content = response.json()["choices"][0]["message"]["content"].strip()57        logger.info("API 调用成功")58        return content59    except Exception as e:60        logger.error(f"API 调用失败: {str(e)}")61        raise Exception(f"API Error: {str(e)}")62 63 64class SimpleHTTPRequestHandler(http.server.BaseHTTPRequestHandler):65    def _set_headers(self, status=200, content_type='text/html; charset=utf-8'):66        self.send_response(status)67        self.send_header('Content-type', content_type)68        self.end_headers()69 70    def do_GET(self):71        if self.path == '/':72            self._set_headers()73            self.wfile.write(HTML_CONTENT.encode('utf-8'))74        else:75            self.send_error(404, "Not Found")76 77    def do_POST(self):78        if self.path == '/chat':79            content_length = int(self.headers['Content-Length'])80            post_data = self.rfile.read(content_length)81            try:82                data = json.loads(post_data)83                session_id = data.get('session_id', 'default')  # 支持多会话（前端可传随机ID）84                user_message = data.get('message', '').strip()85 86                if not user_message:87                    self._set_headers(400, 'application/json')88                    self.wfile.write(json.dumps({'error': '消息不能为空'}).encode('utf-8'))89                    return90 91                # 获取或初始化会话92                if session_id not in sessions:93                    sessions[session_id] = [{"role": "system", "content": SYSTEM_PROMPT}]94                95                messages = sessions[session_id]96                messages.append({"role": "user", "content": user_message})97 98                # 调用 API99                reply = ask_kimi(messages)100                messages.append({"role": "assistant", "content": reply})101 102                # 限制上下文长度（防止 token 爆）103                if len(messages) > 20:104                    messages = [messages[0]] + messages[-19:]  # 保留 system + 最近19条105 106                self._set_headers(200, 'application/json')107                self.wfile.write(json.dumps({'reply': reply}).encode('utf-8'))108 109            except json.JSONDecodeError:110                self._set_headers(400, 'application/json')111                self.wfile.write(json.dumps({'error': '无效的 JSON'}).encode('utf-8'))112            except Exception as e:113                logger.error(f"POST 处理异常: {str(e)}")114                self._set_headers(500, 'application/json')115                self.wfile.write(json.dumps({'error': str(e)}).encode('utf-8'))116        else:117            self.send_error(404)118 119# ====== HTML 内容（稍作优化，增加 session_id 支持） ======120HTML_CONTENT = """121<!DOCTYPE html>122<html lang="zh">123<head>124    <meta charset="UTF-8">125    <meta name="viewport" content="width=device-width, initial-scale=1.0">126    <title>SQL 优化智能体 - 可视化终端</title>127    <script src="https://cdn.tailwindcss.com"></script>128    <script src="https://cdn.jsdelivr.net/npm/marked/marked.min.js"></script>129    <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/styles/github-dark.min.css">130    <script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/highlight.min.js"></script>131    <style>132        .chat-container { height: calc(100vh - 180px); }133        .markdown-body pre { background: #1e1e1e; padding: 1rem; border-radius: 0.5rem; overflow-x: auto; }134        .markdown-body code { font-family: 'Fira Code', monospace; }135        .custom-scrollbar::-webkit-scrollbar { width: 6px; }136        .custom-scrollbar::-webkit-scrollbar-thumb { background: #4b5563; border-radius: 3px; }137    </style>138</head>139<body class="bg-gray-900 text-gray-100 font-sans">140    <div class="max-w-5xl mx-auto p-4">141        <header class="flex items-center justify-between mb-6 border-b border-gray-700 pb-4">142            <div class="flex items-center gap-3">143                <div class="w-10 h-10 bg-blue-600 rounded-lg flex items-center justify-center shadow-lg shadow-blue-900/20">144                    <svg xmlns="http://www.w3.org/2000/svg" class="h-6 w-6" fill="none" viewBox="0 0 24 24" stroke="currentColor">145                        <path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M4 7v10c0 2.21 3.582 4 8 4s8-1.79 8-4V7M4 7c0 2.21 3.582 4 8 4s8-1.79 8-4M4 7c0-2.21 3.582-4 8-4s8 1.79 8 4m0 5c0 2.21-3.582 4-8 4s-8-1.79-8-4" />146                    </svg>147                </div>148                <div>149                    <h1 class="text-xl font-bold">SQL 优化智能体</h1>150                    <p class="text-xs text-gray-400">基于 Kimi API 的数据库专家</p>151                </div>152            </div>153            <div class="flex gap-2">154                <span class="px-3 py-1 bg-green-900/30 text-green-400 text-xs rounded-full border border-green-800">在线</span>155            </div>156        </header>157 158        <div id="chat-box" class="chat-container overflow-y-auto space-y-4 mb-4 pr-2 custom-scrollbar">159            <div class="flex justify-start">160                <div class="bg-gray-800 rounded-2xl px-4 py-3 max-w-[85%] border border-gray-700 shadow-sm">161                    <p class="text-sm">你好！我是 SQL 优化专家。描述业务需求或粘贴慢 SQL，我来帮你设计/优化。</p>162                </div>163            </div>164        </div>165 166        <div class="relative">167            <textarea id="user-input" 168                class="w-full bg-gray-800 border border-gray-700 rounded-xl px-4 py-3 pr-12 text-sm focus:outline-none focus:ring-2 focus:ring-blue-500 focus:border-transparent transition-all resize-none"169                placeholder="输入你的问题或粘贴 SQL..." rows="3"></textarea>170            <button id="send-btn" class="absolute right-3 bottom-3 p-2 bg-blue-600 hover:bg-blue-500 rounded-lg transition-colors">171                <svg xmlns="http://www.w3.org/2000/svg" class="h-5 w-5" viewBox="0 0 20 20" fill="currentColor">172                    <path d="M10.894 2.553a1 1 0 00-1.788 0l-7 14a1 1 0 001.169 1.409l5-1.429A1 1 0 009 15.571V11a1 1 0 112 0v4.571a1 1 0 00.725.962l5 1.428a1 1 0 001.17-1.408l-7-14z" />173                </svg>174            </button>175        </div>176    </div>177 178    <script>179        const chatBox = document.getElementById('chat-box');180        const userInput = document.getElementById('user-input');181        const sendBtn = document.getElementById('send-btn');182        const sessionId = 'sess_' + Math.random().toString(36).substr(2, 9);  // 简单随机 session183 184        marked.setOptions({185            highlight: function(code, lang) {186                return hljs.highlightAuto(code).value;187            },188            breaks: true189        });190 191        function addMessage(role, content) {192            const div = document.createElement('div');193            div.className = `flex ${role === 'user' ? 'justify-end' : 'justify-start'}`;194            const inner = document.createElement('div');195            inner.className = `${role === 'user' ? 'bg-blue-600' : 'bg-gray-800 border border-gray-700'} rounded-2xl px-4 py-3 max-w-[90%] shadow-sm markdown-body`;196            if (role === 'assistant') {197                inner.innerHTML = marked.parse(content);198            } else {199                inner.textContent = content;200            }201            div.appendChild(inner);202            chatBox.appendChild(div);203            chatBox.scrollTop = chatBox.scrollHeight;204            hljs.highlightAll();205        }206 207        async function sendMessage() {208            const text = userInput.value.trim();209            if (!text) return;210 211            addMessage('user', text);212            userInput.value = '';213 214            const loading = document.createElement('div');215            loading.className = 'flex justify-start text-gray-400 italic text-sm';216            loading.innerHTML = '<span class="animate-pulse">思考中...</span>';217            chatBox.appendChild(loading);218            chatBox.scrollTop = chatBox.scrollHeight;219 220            try {221                const res = await fetch('/chat', {222                    method: 'POST',223                    headers: {'Content-Type': 'application/json'},224                    body: JSON.stringify({session_id: sessionId, message: text})225                });226                const data = await res.json();227                loading.remove();228 229                if (data.reply) {230                    addMessage('assistant', data.reply);231                } else {232                    addMessage('assistant', '错误：' + (data.error || '未知问题'));233                }234            } catch (err) {235                loading.remove();236                addMessage('assistant', '网络错误，请重试');237            }238        }239 240        sendBtn.addEventListener('click', sendMessage);241        userInput.addEventListener('keydown', e => {242            if (e.key === 'Enter' && !e.shiftKey) {243                e.preventDefault();244                sendMessage();245            }246        });247    </script>248</body>249</html>250"""251 252def run_server(port=8000):253    handler = SimpleHTTPRequestHandler254    with socketserver.TCPServer(("", port), handler) as httpd:255        print(f"服务已启动 → 请访问: http://localhost:{port}")256        logger.info(f"服务器启动于端口 {port}")257        try:258            httpd.serve_forever()259        except KeyboardInterrupt:260            logger.info("服务器正常关闭")261            httpd.server_close()262 263if __name__ == "__main__":264    PORT = 8000265    run_server(PORT)