#!/usr/bin/env python3 """ Step 2: Map data to template and generate HTML report. """ import os import sys import json import time import argparse import urllib.request import re 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") # Create output directory if it doesn't exist os.makedirs(OUTPUT_DIR, exist_ok=True) 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") # --------------------------------------------------------------------------- # Token usage tracking # --------------------------------------------------------------------------- import atexit, datetime as _dt _token_usage = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0, "calls": 0} def _flush_token_usage(): os.makedirs(OUTPUT_DIR, exist_ok=True) log_path = os.path.join(OUTPUT_DIR, "token_usage.jsonl") record = { "time": _dt.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "script": "step2_render", **_token_usage } with open(log_path, "a", encoding="utf-8") as f: f.write(json.dumps(record, ensure_ascii=False) + "\n") atexit.register(_flush_token_usage) 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) result = json.loads(resp.read().decode("utf-8")) usage = result.get("usage", {}) _token_usage["input_tokens"] += usage.get("input_tokens", 0) _token_usage["output_tokens"] += usage.get("output_tokens", 0) _token_usage["total_tokens"] += usage.get("total_tokens", 0) _token_usage["calls"] += 1 return result["output"]["choices"][0]["message"]["content"] except Exception as e: print(json.dumps({"status": "error", "message": f"LLM error: {e}"})) sys.exit(1) 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 main(): parser = argparse.ArgumentParser() parser.add_argument("--query", required=True) parser.add_argument("--template", required=True) parser.add_argument("--data", required=True) args = parser.parse_args() template_name = args.template # Parse input data try: rows = json.loads(args.data) except Exception as e: print(json.dumps({"status": "error", "message": f"Failed to parse data: {e}"})) sys.exit(1) result_str = json.dumps(rows, ensure_ascii=False) template_doc = read_file(os.path.join(TEMPLATES_DIR, f"{template_name}.md")) if not template_doc: print(json.dumps({"status": "error", "message": f"Template doc not found: {template_name}.md"})) sys.exit(1) map_resp = call_llm( f"USER QUERY: {args.query}\n\nQUERY RESULTS ({len(rows)} rows — ALL rows shown below):\n{result_str[:50000]}\n\nTEMPLATE DOCUMENTATION ({template_name}.md):\n{template_doc}\n\nCRITICAL: The data contains {len(rows)} rows. You MUST include ALL {len(rows)} data points in CHART_X_AXIS_DATA and CHART_SERIES_DATA (or equivalent) arrays. Do not skip any rows.\n\nReturn ONLY valid JSON. The JSON must include a 'summary' field (2-3 sentences in Chinese describing key findings from the actual data) in addition to the template data fields.", system="You are a data visualization engineer and analyst. 1. Map ALL data rows to the exact JSON structure required by the template. 2. Add a 'summary' field with 2-3 sentences of Chinese analysis based ONLY on the actual data. Do not skip or truncate any data rows. Do not fabricate data not present in the results." ) if not map_resp: print(json.dumps({"status": "error", "message": "Failed to map data"})) sys.exit(1) template_data = extract_json(map_resp) if not template_data: print(json.dumps({"status": "error", "message": "Failed to parse mapped data", "raw": map_resp})) sys.exit(1) summary = template_data.pop("summary", "数据分析报告已生成。") 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) # Replace text placeholders text_values = template_data.get("data", {}).get("text_values", {}) for key, val in text_values.items(): html = html.replace(f"{{{{{key}}}}}", str(val)) # Inject JS variables js_vars = template_data.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()] html = html.replace("// {{INJECT_DATA}}", "\n ".join(js_lines)) # Check for unreplaced placeholders remaining = re.findall(r'\{\{[A-Z_]+\}\}', html) if remaining: print(f" Warning: unreplaced placeholders: {set(remaining)}", flush=True) timestamp = int(time.time()) out_file = f"report_{timestamp}.html" out_path = os.path.join(OUTPUT_DIR, out_file) with open(out_path, 'w', encoding='utf-8') as f: f.write(html) # Final output print(json.dumps({ "status": "success", "message": "Report generated successfully", "template": template_name, "address": f"/skills/report-generator/output/{out_file}", "summary": summary.strip() if isinstance(summary, str) else str(summary) }, ensure_ascii=False)) if __name__ == "__main__": main()