|
|
@@ -0,0 +1,300 @@
|
|
|
+"""Doris 自然语言查询助手 —— 基于 openai-agents-python + glm-5 (百炼) + Gradio WebUI。
|
|
|
+
|
|
|
+用法:
|
|
|
+ python app.py --build-schema # 重新抓取 ods 库的表/列/注释到 schema.json
|
|
|
+ python app.py --selftest # 不连 Doris,只验证 LLM 接线
|
|
|
+ python app.py # 启动 WebUI (默认 http://0.0.0.0:7860)
|
|
|
+"""
|
|
|
+import argparse
|
|
|
+import asyncio
|
|
|
+import json
|
|
|
+import os
|
|
|
+import re
|
|
|
+from functools import lru_cache
|
|
|
+
|
|
|
+import pymysql
|
|
|
+from dotenv import load_dotenv
|
|
|
+
|
|
|
+load_dotenv()
|
|
|
+
|
|
|
+# --- 配置 -----------------------------------------------------------------
|
|
|
+DB = dict(
|
|
|
+ host=os.getenv("DB_HOST"),
|
|
|
+ port=int(os.getenv("DB_PORT", "9030")),
|
|
|
+ user=os.getenv("DB_USER"),
|
|
|
+ password=os.getenv("DB_PASSWORD"),
|
|
|
+ database=os.getenv("DB_NAME", "ods"),
|
|
|
+ charset="utf8mb4",
|
|
|
+ connect_timeout=10,
|
|
|
+ read_timeout=30,
|
|
|
+ write_timeout=30,
|
|
|
+ autocommit=True,
|
|
|
+)
|
|
|
+SCHEMA_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "schema.json")
|
|
|
+
|
|
|
+
|
|
|
+def _conn(retries: int = 4, delay: float = 1.5):
|
|
|
+ """连 Doris。经 VPN 的链路会在 MySQL 握手阶段偶发掉线(2013),这里重试。"""
|
|
|
+ import time
|
|
|
+
|
|
|
+ last = None
|
|
|
+ for i in range(retries):
|
|
|
+ try:
|
|
|
+ return pymysql.connect(**DB)
|
|
|
+ except pymysql.err.OperationalError as e:
|
|
|
+ last = e
|
|
|
+ if i < retries - 1:
|
|
|
+ time.sleep(delay)
|
|
|
+ raise last
|
|
|
+
|
|
|
+
|
|
|
+# --- schema 缓存 ----------------------------------------------------------
|
|
|
+def build_schema():
|
|
|
+ """从 Doris INFORMATION_SCHEMA 抓取所有表/列/注释,写成 schema.json。"""
|
|
|
+ conn = _conn()
|
|
|
+ try:
|
|
|
+ with conn.cursor(pymysql.cursors.DictCursor) as cur:
|
|
|
+ cur.execute(
|
|
|
+ "SELECT TABLE_NAME, TABLE_COMMENT FROM INFORMATION_SCHEMA.TABLES "
|
|
|
+ "WHERE TABLE_SCHEMA=%s",
|
|
|
+ (DB["database"],),
|
|
|
+ )
|
|
|
+ tables = cur.fetchall()
|
|
|
+ cur.execute(
|
|
|
+ "SELECT TABLE_NAME, COLUMN_NAME, COLUMN_TYPE, COLUMN_COMMENT, IS_NULLABLE "
|
|
|
+ "FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_SCHEMA=%s "
|
|
|
+ "ORDER BY TABLE_NAME, ORDINAL_POSITION",
|
|
|
+ (DB["database"],),
|
|
|
+ )
|
|
|
+ cols = cur.fetchall()
|
|
|
+ finally:
|
|
|
+ conn.close()
|
|
|
+
|
|
|
+ schema = {t["TABLE_NAME"]: {"comment": t["TABLE_COMMENT"] or "", "columns": []}
|
|
|
+ for t in tables}
|
|
|
+ for c in cols:
|
|
|
+ schema.setdefault(c["TABLE_NAME"], {"comment": "", "columns": []})["columns"].append(
|
|
|
+ {
|
|
|
+ "name": c["COLUMN_NAME"],
|
|
|
+ "type": c["COLUMN_TYPE"],
|
|
|
+ "comment": c["COLUMN_COMMENT"] or "",
|
|
|
+ "nullable": c["IS_NULLABLE"],
|
|
|
+ }
|
|
|
+ )
|
|
|
+ with open(SCHEMA_PATH, "w", encoding="utf-8") as f:
|
|
|
+ json.dump(schema, f, ensure_ascii=False, indent=2)
|
|
|
+ print(f"schema.json 已生成,共 {len(schema)} 张表。")
|
|
|
+ return schema
|
|
|
+
|
|
|
+
|
|
|
+@lru_cache(maxsize=1)
|
|
|
+def load_schema():
|
|
|
+ if not os.path.exists(SCHEMA_PATH):
|
|
|
+ try:
|
|
|
+ return build_schema()
|
|
|
+ except Exception as e: # selftest 等场景下 Doris 不可达时也能继续
|
|
|
+ print(f"[警告] 无法连接 Doris 生成 schema,使用空 schema: {e}")
|
|
|
+ return {}
|
|
|
+ with open(SCHEMA_PATH, encoding="utf-8") as f:
|
|
|
+ return json.load(f)
|
|
|
+
|
|
|
+
|
|
|
+# --- Agent 设置 -----------------------------------------------------------
|
|
|
+from agents import ( # noqa: E402
|
|
|
+ Agent,
|
|
|
+ OpenAIChatCompletionsModel,
|
|
|
+ Runner,
|
|
|
+ function_tool,
|
|
|
+ set_tracing_disabled,
|
|
|
+)
|
|
|
+from openai import AsyncOpenAI # noqa: E402
|
|
|
+
|
|
|
+# 框架默认会把 trace 上报到 OpenAI,用国产模型时必须关掉,否则会报错/泄露。
|
|
|
+set_tracing_disabled(True)
|
|
|
+
|
|
|
+_client = AsyncOpenAI(
|
|
|
+ base_url=os.getenv("LLM_BASE_URL"),
|
|
|
+ api_key=os.getenv("LLM_API_KEY"),
|
|
|
+)
|
|
|
+_model = OpenAIChatCompletionsModel(
|
|
|
+ model=os.getenv("LLM_MODEL", "glm-5"),
|
|
|
+ openai_client=_client,
|
|
|
+)
|
|
|
+
|
|
|
+SCHEMA = load_schema()
|
|
|
+
|
|
|
+# 关键表的人工说明(平台字典为空,这些规则是探查数据得出的,agent 易出错的地方)。
|
|
|
+TABLE_NOTES = {
|
|
|
+ "ods_sa_device_cbdphoto_b": (
|
|
|
+ "测报灯图片识别【主表,约2000万行,查虫量默认用这个】。"
|
|
|
+ "indentify_result 格式 '虫码,数量#虫码,数量'(虫码=ods_base_pest_code.pest_yfkj_code,数量=头数),"
|
|
|
+ "例 '3,1#53,3#136,1'。device_id = ods_sa_device.id。时间字段 uptime(datetime)。"
|
|
|
+ "取某虫数量: regexp_extract(indentify_result,'(^|#)<虫码>,([0-9]+)',2),配 RLIKE '(^|#)<虫码>,' 过滤。"
|
|
|
+ ),
|
|
|
+ "ods_sa_device_lpsphoto": "性诱识别(约2.6万行)。indentify_result='虫码,数量'。device_id=ods_sa_device.id。时间 uptime。",
|
|
|
+ "ods_sa_device_lpsphoto_count": "性诱识别计数(约2.6万行),格式同 lpsphoto。",
|
|
|
+ "ods_sa_device": (
|
|
|
+ "设备主表。province/city/district 值带后缀(如'河南省''新乡市')。"
|
|
|
+ "device_id 是 IMEI;id 是数字主键,被各识别表的 device_id 引用。"
|
|
|
+ ),
|
|
|
+ "ods_base_pest_code": "害虫字典(1226种)。pest_yfkj_code=识别结果里的虫码;pest_name=虫名;pest_level=等级;pest_code 是另一套编码,别混用。",
|
|
|
+ "ods_sa_device_cbd_data": "测报灯设备遥测(JSON),只有温湿度/电池/灯状态等,【没有虫种和数量】,别用来查虫量。",
|
|
|
+}
|
|
|
+
|
|
|
+# --- tools ---------------------------------------------------------------
|
|
|
+_DML = re.compile(
|
|
|
+ r"\b(insert|update|delete|drop|alter|create|truncate|grant|revoke|load|merge)\b"
|
|
|
+)
|
|
|
+
|
|
|
+
|
|
|
+@function_tool
|
|
|
+def list_tables(keyword: str = "") -> str:
|
|
|
+ """列出 ods 库的表。可传 keyword 按表名或注释过滤。返回 '表名 - 注释'。"""
|
|
|
+ items = [(n, d["comment"]) for n, d in SCHEMA.items()]
|
|
|
+ if keyword:
|
|
|
+ k = keyword.lower()
|
|
|
+ items = [(n, c) for n, c in items if k in n.lower() or k in (c or "").lower()]
|
|
|
+ items = items[:50]
|
|
|
+ return "\n".join(f"{n} - {c}" for n, c in items) or "没有匹配的表。"
|
|
|
+
|
|
|
+
|
|
|
+@function_tool
|
|
|
+def get_table_schema(table_name: str) -> str:
|
|
|
+ """获取某张表的列定义、类型和注释。"""
|
|
|
+ t = SCHEMA.get(table_name) or SCHEMA.get(table_name.lower())
|
|
|
+ if not t:
|
|
|
+ cands = [n for n in SCHEMA if table_name.lower() in n.lower()][:10]
|
|
|
+ return f"表 {table_name} 不存在。相近的表: {cands}"
|
|
|
+ lines = [f"表 {table_name}: {t['comment']}"]
|
|
|
+ if TABLE_NOTES.get(table_name):
|
|
|
+ lines.append(f" ⚑ 数据说明: {TABLE_NOTES[table_name]}")
|
|
|
+ for col in t["columns"]:
|
|
|
+ lines.append(
|
|
|
+ f" - {col['name']} {col['type']} "
|
|
|
+ f"{'NULL' if col['nullable'] == 'YES' else 'NOT NULL'} -- {col['comment']}"
|
|
|
+ )
|
|
|
+ return "\n".join(lines)
|
|
|
+
|
|
|
+
|
|
|
+@function_tool
|
|
|
+def execute_sql(sql: str) -> str:
|
|
|
+ """对 Doris 执行只读 SELECT 查询,返回最多 200 行结果。只允许 SELECT。"""
|
|
|
+ s = sql.strip().rstrip(";").strip()
|
|
|
+ low = re.sub(r"/\*.*?\*/", " ", s, flags=re.S) # 去掉块注释
|
|
|
+ low = re.sub(r"--[^\n]*", " ", low) # 去掉行注释
|
|
|
+ if not low.strip().lower().startswith("select") or _DML.search(low.lower()):
|
|
|
+ return "错误: 只允许只读 SELECT 查询。"
|
|
|
+ conn = _conn()
|
|
|
+ try:
|
|
|
+ with conn.cursor() as cur:
|
|
|
+ # ponytail: 把用户 SQL 包进子查询再 LIMIT,是只读/行数上限的主要保障
|
|
|
+ # (只读账号 + 子查询无法 DML/DDL),上面的正则只是第二道防线。
|
|
|
+ cur.execute("SET query_timeout = 30")
|
|
|
+ cur.execute(f"SELECT * FROM ({s}) AS _q LIMIT 200")
|
|
|
+ cols = [d[0] for d in cur.description]
|
|
|
+ rows = cur.fetchall()
|
|
|
+ if not rows:
|
|
|
+ return "查询成功,但没有数据。"
|
|
|
+ body = "\n".join(" | ".join(str(v) for v in r) for r in rows)
|
|
|
+ return f"列: {cols}\n行数(已截断到200): {len(rows)}\n{body}"
|
|
|
+ except Exception as e:
|
|
|
+ return f"SQL 执行错误: {e}\n请根据错误修正 SQL 后重试。"
|
|
|
+ finally:
|
|
|
+ conn.close()
|
|
|
+
|
|
|
+
|
|
|
+@function_tool
|
|
|
+def search_pest(keyword: str) -> str:
|
|
|
+ """按虫名(或虫码)查害虫编码。返回 pest_yfkj_code | 虫名 | 等级。
|
|
|
+ 查询某虫的数量前【必须】先调用本工具拿到 pest_yfkj_code(共1226种,不要靠记忆猜)。"""
|
|
|
+ conn = _conn()
|
|
|
+ try:
|
|
|
+ with conn.cursor() as cur:
|
|
|
+ cur.execute(
|
|
|
+ "SELECT pest_yfkj_code, pest_name, pest_level FROM ods_base_pest_code "
|
|
|
+ "WHERE pest_name LIKE %s OR pest_yfkj_code = %s LIMIT 20",
|
|
|
+ (f"%{keyword}%", keyword),
|
|
|
+ )
|
|
|
+ rows = cur.fetchall()
|
|
|
+ finally:
|
|
|
+ conn.close()
|
|
|
+ if not rows:
|
|
|
+ return f"没找到含 '{keyword}' 的害虫,请换关键词(如'夜蛾''飞虱')。"
|
|
|
+ return "\n".join(f"{c} | {n} | {lv}" for c, n, lv in rows)
|
|
|
+
|
|
|
+
|
|
|
+INSTRUCTIONS = """你是农业病虫害监测数据平台的自然语言查询助手。底层是 Apache Doris,库 ods。用户用自然语言问"某地某时某虫监测到多少头",你转成 SQL 查询并用中文回答。
|
|
|
+
|
|
|
+==== 数据模型(务必遵守,这是查对虫量的关键)====
|
|
|
+"虫量(头数)"来自设备的**图片识别结果**,不是设备遥测。
|
|
|
+1. 主表 ods_sa_device_cbdphoto_b(测报灯,约2000万行,查虫量默认用它):每行=一张虫情照片。
|
|
|
+ - indentify_result 格式 `虫码,数量#虫码,数量`,例 "3,1#53,3#136,1" = 虫码3计1头、53计3头、136计1头。
|
|
|
+ - 虫码 = ods_base_pest_code.pest_yfkj_code。
|
|
|
+ - 时间字段 uptime(datetime)。
|
|
|
+ - 取某虫数量: regexp_extract(indentify_result, '(^|#)<虫码>,([0-9]+)', 2) ,并用 indentify_result RLIKE '(^|#)<虫码>,' 过滤后再 SUM。
|
|
|
+2. 地理: 识别表 device_id = ods_sa_device.id(数字主键,不是IMEI)。JOIN ods_sa_device d ON d.id = <识别表>.device_id,取 d.province/d.city/district(值带后缀:"河南省""新乡市")。
|
|
|
+3. 虫名→虫码: 先 search_pest(虫名) 拿 pest_yfkj_code。
|
|
|
+4. device_data 等 JSON 遥测字段没有虫,别用。
|
|
|
+
|
|
|
+==== 工作流程 ====
|
|
|
+1. search_pest 把虫名转 pest_yfkj_code(多个结果挑最匹配并说明)。
|
|
|
+2. 必要时 get_table_schema 看表结构(默认查虫量用 ods_sa_device_cbdphoto_b)。
|
|
|
+3. 写 SQL:识别表 JOIN ods_sa_device 过滤省市 + uptime 时间段 + regexp 取该虫数量求 SUM。
|
|
|
+4. 报错就修正重试(最多3次)。
|
|
|
+5. 中文回答:给总头数 + 覆盖地区/时间/虫名 + 末尾附 SQL 代码块。
|
|
|
+
|
|
|
+==== 注意 ====
|
|
|
+- "某月"按自然月;时间范围拿不准先 SELECT MIN/MAX(uptime) 确认。
|
|
|
+- 大表查询【必须】带时间范围过滤,避免全表扫描。
|
|
|
+- 若该虫查出来是 0 或没有:如实说明——这类设备(测报灯/性诱)主要监测趋光/趋性诱的蛾、飞虱、叶蝉等,不监测蝗虫等非趋光昆虫,不要编造数字。
|
|
|
+- Doris 与 MySQL 有差异,日期用 'YYYY-MM-DD'。
|
|
|
+"""
|
|
|
+
|
|
|
+agent = Agent(
|
|
|
+ name="Doris查询助手",
|
|
|
+ instructions=INSTRUCTIONS,
|
|
|
+ model=_model,
|
|
|
+ tools=[list_tables, get_table_schema, search_pest, execute_sql],
|
|
|
+)
|
|
|
+
|
|
|
+
|
|
|
+# --- WebUI ---------------------------------------------------------------
|
|
|
+def _run(user_message: str, history: list[dict]):
|
|
|
+ # 历史由 Gradio 管理(每用户每标签页独立),直接拼成 input 交给 Runner。
|
|
|
+ result = asyncio.run(
|
|
|
+ Runner.run(agent, history + [{"role": "user", "content": user_message}])
|
|
|
+ )
|
|
|
+ return result.final_output
|
|
|
+
|
|
|
+
|
|
|
+def build_ui():
|
|
|
+ import gradio as gr
|
|
|
+
|
|
|
+ # gradio 6 默认 history 即为 OpenAI 消息字典格式 {"role","content"}。
|
|
|
+ return gr.ChatInterface(
|
|
|
+ fn=_run,
|
|
|
+ title="Doris 数据查询助手",
|
|
|
+ description="用自然语言查询 ods 库。例:『上个月每天的订单量』",
|
|
|
+ textbox=gr.Textbox(placeholder="问点什么,比如:每个地区近 7 天的销售额"),
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+def main():
|
|
|
+ p = argparse.ArgumentParser()
|
|
|
+ p.add_argument("--build-schema", action="store_true")
|
|
|
+ p.add_argument("--selftest", action="store_true")
|
|
|
+ args = p.parse_args()
|
|
|
+
|
|
|
+ if args.build_schema:
|
|
|
+ build_schema()
|
|
|
+ return
|
|
|
+ if args.selftest:
|
|
|
+ r = asyncio.run(Runner.run(agent, "只回复 OK"))
|
|
|
+ print("模型自检返回:", repr(r.final_output))
|
|
|
+ return
|
|
|
+
|
|
|
+ build_ui().launch(server_name="0.0.0.0", server_port=7860)
|
|
|
+
|
|
|
+
|
|
|
+if __name__ == "__main__":
|
|
|
+ main()
|
YangLiyang
vor 4 Tagen