#!/usr/bin/env python3 """ Step 3: Map query results to template and generate HTML report. For 园区概览 (dashboard_overview) and 企业概览 (dashboard_enterprise), step2_query.py already used LLM to produce a flat dict of placeholder_values (including INJECT_DATA). This script just substitutes them into the HTML. For all other templates, uses LLM to map data to template format as before. v2 Optimizations: - [优化1] 数据映射任务切换为 qwen-turbo（结构化输出，速度快2-3倍） - [优化2] 精简模板文档：仅提取 JSON 结构示例段落，减少 LLM 输入上下文 - [优化3] 并发调用：数据映射（qwen-turbo）与摘要生成（qwen-turbo）同时发起， 总耗时取决于较慢的那个，而非两者之和 """ import os import sys import json import time import argparse import urllib.request import re from concurrent.futures import ThreadPoolExecutor, as_completed SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) SKILL_DIR = os.path.dirname(SCRIPT_DIR) TEMPLATES_DIR = os.path.join(SKILL_DIR, "templates") OUTPUT_DIR = os.path.join(SKILL_DIR, "output") LLM_API_URL = os.environ.get("LLM_API_URL", "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation") LLM_API_KEY = os.environ.get("LLM_API_KEY", "sk-b36adcdd25d64b7c847e7c941b5c1e03") LLM_MODEL = os.environ.get("LLM_MODEL", "qwen-plus") # [优化1] 数据映射专用模型：qwen-turbo 速度快，结构化输出能力足够 LLM_MODEL_FAST = os.environ.get("LLM_MODEL_FAST", "qwen-turbo") # 按模板类型动态设置 max_tokens 上限，避免过度生成 # lineChart/barChart 需要填充完整数组，给足空间 TEMPLATE_MAX_TOKENS = { "infoCard": 800, "pieChart": 1000, "barChart": 2500, "lineChart": 2500, "lineChartCompare": 3000, "dataTable": 2500, "dashboardKpiBar": 1500, "dashboardKpiBarCompare":2000, } DEFAULT_MAX_TOKENS = 2500 def call_llm(prompt, system="You are a helpful assistant.", model=None, max_tokens=None): use_model = model or LLM_MODEL headers = {"Authorization": f"Bearer {LLM_API_KEY}", "Content-Type": "application/json"} parameters = {"result_format": "message"} if max_tokens: parameters["max_tokens"] = max_tokens payload = {"model": use_model, "input": {"messages": [ {"role": "system", "content": system}, {"role": "user", "content": prompt} ]}, "parameters": parameters} req = urllib.request.Request(LLM_API_URL, data=json.dumps(payload).encode("utf-8"), headers=headers, method="POST") try: resp = urllib.request.urlopen(req, timeout=60) return json.loads(resp.read().decode("utf-8"))["output"]["choices"][0]["message"]["content"] except Exception as e: print(json.dumps({"status": "error", "message": f"LLM error: {e}"})) return None def extract_json(text): for pattern in [r'```json\s*(.*?)\s*```', r'```\s*([\[{].*?)\s*```']: m = re.search(pattern, text, re.DOTALL) if m: try: return json.loads(m.group(1)) except: pass try: return json.loads(text.strip()) except: pass brace_match = re.search(r'(\{.*\})', text, re.DOTALL) if brace_match: try: return json.loads(brace_match.group(1)) except: pass return None def read_file(path): if os.path.exists(path): with open(path, 'r', encoding='utf-8') as f: return f.read() return "" def inject_placeholders(html, placeholder_values): """ Replace all {{KEY}} placeholders in html with the corresponding values from placeholder_values dict. Special handling for INJECT_DATA: the template uses {{INJECT_DATA}} as a standalone placeholder (not prefixed with // ), so we replace it directly with the full `const mockData = {...};` JS string. """ for key, value in placeholder_values.items(): if value is None: value = "" html = html.replace("{{" + key + "}}", str(value)) return html def slim_template_doc(template_doc, max_chars=3000): """ [优化2] 精简模板文档：从完整的 .md 文档中提取关键的 JSON 结构示例， 去除冗长的背景说明，减少传入 LLM 的上下文长度。 提取策略： 1. 优先提取包含 JSON 代码块（```json ... ```）的段落，完整保留不截断。 JSON 示例是 LLM 理解数据结构的关键，截断会导致 LLM 无法正确填充真实数据。 2. 如果没有 JSON 代码块，则提取包含 "js_variables"、"text_values"、"INJECT_DATA" 等关键字的段落。 3. 最终截断到 max_chars 字符以内（仅作为最后兜底）。 """ if not template_doc or len(template_doc) <= max_chars: return template_doc # 优先提取 JSON 代码块，完整保留（不截断单个代码块） json_blocks = re.findall(r'```json.*?```', template_doc, re.DOTALL) if json_blocks: # 完整保留所有 JSON 代码块，不截断 slim = "\n\n".join(json_blocks) return slim # 不截断，保证 LLM 能看到完整结构 # 次选：提取含关键字的段落 key_terms = ["js_variables", "text_values", "INJECT_DATA", "CHART_", "SERIES_", "xAxis", "series"] paragraphs = template_doc.split("\n\n") selected = [] for para in paragraphs: if any(term in para for term in key_terms): selected.append(para) if selected: slim = "\n\n".join(selected) return slim[:max_chars] # 兜底：直接截断 return template_doc[:max_chars] def build_summary_context(rows, max_chars=800): """ [优化3] 为摘要生成构建精简的数据上下文： 只传入聚合统计信息（总行数、前5行、数值列的最大/最小/均值）， 而非完整的原始数据，大幅减少摘要 LLM 调用的输入长度。 """ if not rows: return "无数据" total = len(rows) preview = rows[:5] preview_str = json.dumps(preview, ensure_ascii=False, separators=(',', ':')) # 计算数值列的简单统计 stats = {} if rows: numeric_keys = [k for k, v in rows[0].items() if isinstance(v, (int, float))] for k in numeric_keys[:6]: # 最多统计6个数值列 vals = [r[k] for r in rows if isinstance(r.get(k), (int, float))] if vals: stats[k] = { "sum": round(sum(vals), 3), "max": round(max(vals), 3), "min": round(min(vals), 3), "avg": round(sum(vals) / len(vals), 3) } context = f"共 {total} 条数据。前5条: {preview_str}" if stats: context += f"\n数值统计: {json.dumps(stats, ensure_ascii=False, separators=(',', ':'))}" return context[:max_chars] def main(): parser = argparse.ArgumentParser() parser.add_argument("--query", required=True) parser.add_argument("--template", required=True) parser.add_argument("--result_file", required=True) args = parser.parse_args() template_name = args.template # Read result file written by step2_query.py try: with open(args.result_file, 'r', encoding='utf-8') as f: result_data = json.load(f) except Exception as e: print(json.dumps({"status": "error", "message": f"Failed to read result file: {e}"})) sys.exit(1) # Load HTML template html_path = os.path.join(TEMPLATES_DIR, f"{template_name}.html") if not os.path.exists(html_path): print(json.dumps({"status": "error", "message": f"Template HTML not found: {html_path}"})) sys.exit(1) html = read_file(html_path) # ----------------------------------------------------------------------- # Fixed template branch: placeholder_values already prepared by step2 # ----------------------------------------------------------------------- if isinstance(result_data, dict) and result_data.get("is_fixed_template"): placeholder_values = result_data.get("placeholder_values", {}) summary = result_data.get("summary", "数据分析报告已生成。") if not placeholder_values: # Fallback: if LLM in step2 failed to produce values, error out print(json.dumps({"status": "error", "message": "Fixed template: placeholder_values is empty. Check step2 output."})) sys.exit(1) html = inject_placeholders(html, placeholder_values) # ----------------------------------------------------------------------- # Normal flow: call LLM to map raw rows to template format # [优化1] 使用 qwen-turbo 替代 qwen-plus # [优化2] 精简 template_doc 后再传入 LLM # [优化3] 数据映射与摘要生成并发执行 # ----------------------------------------------------------------------- else: rows = result_data.get("rows", []) template_doc = result_data.get("template_doc", "") summary = "数据分析报告已生成。" if not rows: print(json.dumps({"status": "error", "message": "Normal template: rows is empty. Check step2 output."})) sys.exit(1) # 截断原始数据，避免超出 token 限制 rows_str = json.dumps(rows, ensure_ascii=False, separators=(',', ':')) if len(rows_str) > 12000: rows_str = rows_str[:12000] + "...(truncated)" # [优化2] 精简模板文档，只保留关键 JSON 结构示例 slim_doc = slim_template_doc(template_doc, max_chars=1500) # [优化1] 按模板类型动态确定 max_tokens map_max_tokens = TEMPLATE_MAX_TOKENS.get(template_name, DEFAULT_MAX_TOKENS) # [优化3] 并发发起两个 LLM 调用 # - 任务A（数据映射）：qwen-turbo，使用精简后的模板文档，生成 JSON 结构 # - 任务B（摘要生成）：qwen-turbo，使用压缩后的统计上下文，生成中文摘要 summary_context = build_summary_context(rows, max_chars=800) def task_map(): return call_llm( f"USER QUERY: {args.query}\n\n" f"QUERY RESULTS ({len(rows)} rows — these are the REAL data you MUST use):\n{rows_str}\n\n" f"TEMPLATE SPEC (the JSON example below shows the STRUCTURE only; all numeric values and labels in the example are PLACEHOLDERS — replace them entirely with values from QUERY RESULTS above):\n{slim_doc}\n\n" f"CRITICAL RULES:\n" f"1. CHART_X_AXIS_DATA must contain ALL {len(rows)} x-axis labels from QUERY RESULTS (e.g. dates or categories).\n" f"2. CHART_SERIES_DATA (and any other series arrays) must contain ALL {len(rows)} numeric values from QUERY RESULTS.\n" f"3. Do NOT copy example values from the template spec. Use ONLY real data from QUERY RESULTS.\n" f"4. Return ONLY valid JSON. Do NOT include a 'summary' field.", system="You are a data visualization engineer. Your ONLY job is to map the QUERY RESULTS data into the template JSON structure. The template spec shows structure only — never copy its example values. Always use real data from QUERY RESULTS. Return ONLY valid JSON.", model=LLM_MODEL_FAST, # [优化1] qwen-turbo max_tokens=map_max_tokens ) def task_summary(): return call_llm( f"用户查询：{args.query}\n\n数据概况：{summary_context}\n\n请用2-3句中文分析这份数据的关键发现，直接输出文字，不要 JSON 格式。", system="你是数据分析师，根据数据概况生成简洁的中文分析摘要。", model=LLM_MODEL_FAST, # [优化1] qwen-turbo max_tokens=300 ) # [优化3] 使用 ThreadPoolExecutor 并发执行，并记录各任务耗时 map_resp = None summary_raw = None t0 = time.time() with ThreadPoolExecutor(max_workers=2) as executor: future_map = executor.submit(task_map) future_summary = executor.submit(task_summary) # 记录各任务完成时间 import concurrent.futures done_times = {} for future in concurrent.futures.as_completed({future_map: 'task_map', future_summary: 'task_summary'}): label = {future_map: 'task_map', future_summary: 'task_summary'}[future] done_times[label] = round(time.time() - t0, 2) map_resp = future_map.result() summary_raw = future_summary.result() total_elapsed = round(time.time() - t0, 2) # 将计时结果输出到 stderr，方便日志排查 import sys as _sys _sys.stderr.write(json.dumps({ "timing": { "task_map_done_at": done_times.get('task_map'), "task_summary_done_at": done_times.get('task_summary'), "total_elapsed": total_elapsed, "bottleneck": "task_map" if done_times.get('task_map', 0) > done_times.get('task_summary', 0) else "task_summary" } }, ensure_ascii=False) + "\n") if not map_resp: print(json.dumps({"status": "error", "message": "Failed to map data with LLM"})) sys.exit(1) mapped = extract_json(map_resp) if not mapped: print(json.dumps({"status": "error", "message": "Failed to parse LLM mapped data", "raw": map_resp[:500]})) sys.exit(1) # 摘要：优先使用并发任务B的结果，回退到 mapped 中的 summary 字段 summary = (summary_raw or "").strip() or mapped.pop("summary", "数据分析报告已生成。") # Render: text_values + js_variables text_values = mapped.get("data", {}).get("text_values", {}) for key, val in text_values.items(): html = html.replace(f"{{{{{key}}}}}", str(val)) js_vars = mapped.get("data", {}).get("js_variables", {}) if js_vars: js_lines = [f"var {k} = {json.dumps(v, ensure_ascii=False)};" for k, v in js_vars.items()] injected = "\n ".join(js_lines) if "// {{INJECT_DATA}}" in html: html = html.replace("// {{INJECT_DATA}}", injected) elif "{{INJECT_DATA}}" in html: html = html.replace("{{INJECT_DATA}}", injected) # ----------------------------------------------------------------------- # Write output HTML # ----------------------------------------------------------------------- timestamp = int(time.time()) out_file = f"report_{timestamp}.html" out_path = os.path.join(OUTPUT_DIR, out_file) os.makedirs(OUTPUT_DIR, exist_ok=True) with open(out_path, 'w', encoding='utf-8') as f: f.write(html) # Clean up temp file try: os.remove(args.result_file) except: pass print(json.dumps({ "status": "success", "message": "Report generated successfully", "template": template_name, "address": f"/skill/sql-report-generator/output/{out_file}", "summary": summary.strip() if isinstance(summary, str) else str(summary) }, ensure_ascii=False)) if __name__ == "__main__": main()