#!/usr/bin/env python3 """ Step 1: Select views and template based on user query. For 园区概览 and 企业概览 queries, use LLM to extract dynamic parameters (date_type, date, park_name, enterprise_name) and return fixed template identifiers. For all other queries, use LLM to select views and template normally. """ import os import sys import json import argparse import urllib.request import re import datetime SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) SKILL_DIR = os.path.dirname(SCRIPT_DIR) TEMPLATES_DIR = os.path.join(SKILL_DIR, "templates") REFERENCES_DIR = os.path.join(SKILL_DIR, "references") VIEWS_DIR = os.path.join(REFERENCES_DIR, "views") VIEWS_INDEX = os.path.join(REFERENCES_DIR, "views_index.md") 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(): output_dir = os.path.join(SKILL_DIR, "output") 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": "step1_select", **_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 get_template_list(): templates = [] for fname in sorted(os.listdir(TEMPLATES_DIR)): if fname.endswith(".md"): name = fname[:-3] if os.path.exists(os.path.join(TEMPLATES_DIR, f"{name}.html")): doc = read_file(os.path.join(TEMPLATES_DIR, fname)) summary = "\n".join(doc.split("\n")[:5]) templates.append({"name": name, "summary": summary}) return templates def main(): parser = argparse.ArgumentParser() parser.add_argument("--query", required=True) args = parser.parse_args() today = datetime.date.today() last_month = today.replace(day=1) - datetime.timedelta(days=1) last_month_str = last_month.strftime("%Y-%m") current_year = str(today.year) # ----------------------------------------------------------------------- # Single LLM call: detect intent AND extract parameters simultaneously # ----------------------------------------------------------------------- intent_resp = call_llm( f"""用户查询: {args.query} 今天的日期是: {today.strftime("%Y年%m月%d日")} 请判断用户的查询意图，并返回 JSON。 --- 如果用户明确要求【园区概览】（必须是要求生成园区的「综合概览/总览/全面报告」，包含能耗、碳排、设备、价格等多维度汇总），返回： {{ "intent": "park_overview", "date_type": "year 或 month（用户只说了某年如'2025年'则为year；说了某年某月如'2025年3月'则为month；未提及时间则为month）", "date": "若date_type=year，返回4位年份字符串如 '2025'；若date_type=month，返回YYYY-MM格式如 '2025-03'；若未提及时间，返回上个月 '{last_month_str}'", "park_name": "园区名称，若未提及则返回 null", "enterprise_name": null }} 如果用户明确要求【企业概览】（必须是要求生成某企业的「综合概览/总览/全面报告」），返回： {{ "intent": "enterprise_overview", "date_type": "year 或 month", "date": "若date_type=year，返回4位年份字符串如 '2025'；若date_type=month，返回YYYY-MM格式如 '2025-03'；若未提及时间，返回上个月 '{last_month_str}'", "park_name": "所在园区名称，若未提及则返回 null", "enterprise_name": "企业名称，若未提及则返回 null" }} 如果是其他类型的数据查询（查询某个具体指标、趋势、对比、明细等），返回： {{ "intent": "general" }} --- 关键判断规则： 【园区概览】必须同时满足：(1)明确说出“概览”“总览”“报告”等关键词，(2)是对园区整体的综合查询。 如果查询只涉及某个具体指标（如电力能耗、用电量、碳排放、光伏发电量等），即使提到了园区，也应判断为 general。 判断示例： - "园区概览" → park_overview（默认上个月） - "2025年的园区概览" → park_overview, year - "2025年3月园区概览" → park_overview, month - "生成园区总览报告" → park_overview - "查询园区的电力能耗情况" → general（具体指标） - "查询园区用电量" → general（具体指标） - "2025年3月园区的电力能耗情况" → general（具体指标，即使有时间和园区） - "各企业碳排放对比" → general - "查询上月用电量" → general - "各园区光伏发电量排名" → general 只返回 JSON，不要其他内容。""", system="你是一个数据分析助手，负责理解用户的查询意图并提取关键参数。只返回 JSON，不要解释。", model="qwen-turbo", max_tokens=200 ) intent_data = extract_json(intent_resp) if intent_resp else None # ----------------------------------------------------------------------- # Fixed template branch: 园区概览 or 企业概览 # ----------------------------------------------------------------------- if intent_data and intent_data.get("intent") in ("park_overview", "enterprise_overview"): intent = intent_data["intent"] date_type = intent_data.get("date_type", "month") date_val = intent_data.get("date") or last_month_str park_name = intent_data.get("park_name") enterprise_name = intent_data.get("enterprise_name") # Normalize: if date_type is year but date looks like YYYY-MM, strip month if date_type == "year" and date_val and len(date_val) > 4: date_val = date_val[:4] # If date_type is month but date looks like YYYY (4 chars), treat as year if date_type == "month" and date_val and len(date_val) == 4: date_type = "year" if intent == "park_overview": template_name = "dashboard_overview" reason = f"园区概览使用固定SQL模板查询（{date_type}维度，date={date_val}）" else: template_name = "dashboard_enterprise" reason = f"企业概览使用固定SQL模板查询（{date_type}维度，date={date_val}）" print(json.dumps({ "status": "success", "selected_views": [], "template_name": template_name, "is_fixed_template": True, "fixed_template_type": template_name, "date_type": date_type, "date": date_val, "park_name": park_name, "enterprise_name": enterprise_name, "reason": reason }, ensure_ascii=False)) return # ----------------------------------------------------------------------- # Normal LLM-based flow: select views and template # ----------------------------------------------------------------------- views_index = read_file(VIEWS_INDEX) templates = get_template_list() templates_summary = "\n\n".join([f"[{t['name']}]\n{t['summary']}" for t in templates]) step1_resp = call_llm( f"USER QUERY: {args.query}\n\nAVAILABLE DATABASE VIEWS INDEX:\n{views_index}\n\nAVAILABLE TEMPLATES:\n{templates_summary}\n\nReturn ONLY valid JSON.", system="""You are a data analysis expert. Given the user's query, do TWO things simultaneously: 1. Select 1-3 most relevant database views from the index. 2. Select the single best HTML template for visualizing the result. Template guidelines: - Time-series trend (每天/每月/每年趋势) → lineChart - Comparison by category/entity (各企业/各园区对比) → barChart - Distribution / proportion (占比/分布/比例) → pieChart - Detailed tabular data (列表/明细/详情) → dataTable - KPI summary metrics (总量/汇总/概览) → infoCard or dashboardKpiBar Return ONLY valid JSON in this exact format: {"selected_views": ["视图名1", "视图名2"], "template_name": "templateName", "reason": "one sentence explanation"} Use EXACT view names from the index. Use EXACT template names from the list.""", max_tokens=512 ) if not step1_resp: print(json.dumps({"status": "error", "message": "Failed to get response from LLM"})) sys.exit(1) step1 = extract_json(step1_resp) if not step1 or "selected_views" not in step1 or "template_name" not in step1: print(json.dumps({"status": "error", "message": "Failed to parse step 1 JSON", "raw": step1_resp})) sys.exit(1) print(json.dumps({ "status": "success", "selected_views": step1["selected_views"], "template_name": step1["template_name"], "is_fixed_template": False, "reason": step1.get("reason", "") }, ensure_ascii=False)) if __name__ == "__main__": main()