接企业数据&RAG做知识型智能体：多表、多文档、多轮问答的代码套路

pip install -r requirements.txt

# 自动加载 .env（如果存在）；不覆盖已存在的系统环境变量

load_dotenv(find_dotenv(), override=False)

ZHIPU_API_KEY = os.environ.get(“ZHIPUAI_API_KEY”)

if not ZHIPU_API_KEY:

raise RuntimeError(“请在 .env 中设置 ZHIPUAI_API_KEY，或设置系统环境变量”)

CHAT_MODEL = os.environ.get(“ZHIPU_CHAT_MODEL”, “glm-4.6”) # 可改为 glm-4-flash

EMBED_MODEL = os.environ.get(“ZHIPU_EMBED_MODEL”, “embedding-3”) # 可改 embedding-2

EMBED_DIM = int(os.environ.get(“ZHIPU_EMBED_DIM”, “1024”)) # embedding-3 可自定义维度

— 客户信息表

CREATE TABLE customers(

id INTEGER PRIMARY KEY, — 客户主键（自增整数，便于 JOIN 与索引）

name TEXT, — 客户名称（示例使用中文企业名）

region TEXT — 客户所属区域（示例值：华东、华北、华南、西南）。

)

— 销售明细表

CREATE TABLE orders(

id INTEGER PRIMARY KEY, — 订单主键

customer_id INTEGER, — 客户ID，外键，指向customers.id

product TEXT, — 产品/服务名称

amount REAL, — 订单金额

date TEXT, — 订单日期（以 ISO 8601（YYYY-MM-DD）字符串存储）

FOREIGN KEY(customer_id) REFERENCES customers(id) — 关系约束

)

customers = [

(1, “星辰科技”, “华东”),

(2, “海纳能源”, “华北”),

(3, “南风智造”, “华南”),

(4, “川岳电气”, “西南”),

(5, “凌云数码”, “华东”),

……

]

cur.executemany(“INSERT INTO customers VALUES (?, ?, ?)”, customers)

# 生成最近 ~540 天订单

rng = np.random.default_rng(42)

products = [“A100”, “B200”, “服务包”]

base = date.today() – timedelta(days=540)

rows, oid = [], 1

for _ in range(30000):

cust = int(rng.integers(1, 50))

prod = str(rng.choice(products))

amount = float(np.round(rng.uniform(2000, 50000), 2))

d = base + timedelta(days=int(rng.integers(0, 540)))

rows.append((oid, cust, prod, amount, d.iso8601() if hasattr(d, “iso8601”) else d.isoformat()))

oid += 1

cur.executemany(“INSERT INTO orders VALUES (?, ?, ?, ?, ?)”, rows)

@dataclass

class Chunk:

# 文档分片元数据：所属文档 ID、片内序号、文本内容

doc_id: str

chunk_id: int

text: str

def chunk_text(text: str, max_len: int = 350) -> List[str]:

# 按段落拼接，保持语义完整度，直到接近 max_len 时换片

paras = [p.strip() for p in text.strip().split(“\n”) if p.strip()]

out, buf = [], “”

for p in paras:

if len(buf) + len(p) + 1 <= max_len:

buf = (buf + “\n” + p).strip()

else:

if buf:

out.append(buf)

buf = p

if buf:

out.append(buf)

return out

def build_corpus(docs: Dict[str, str]) -> List[Chunk]:

# 将多文档拆分为 Chunk 列表，保留文档 ID 与分片顺序，供后续向量化检索。

chunks: List[Chunk] = []

for doc_id, content in docs.items():

parts = chunk_text(content, 350)

for i, part in enumerate(parts):

# chunk_id 保留原文内顺序，方便引用

chunks.append(Chunk(doc_id=doc_id, chunk_id=i, text=part))

return chunks

ZHIPU_API_KEY = os.environ.get(“ZHIPUAI_API_KEY”)

if not ZHIPU_API_KEY:

raise RuntimeError(“请在 .env 中设置 ZHIPUAI_API_KEY，或设置系统环境变量”)

client = ZhipuAI(api_key=ZHIPU_API_KEY)

CHAT_MODEL = os.environ.get(“ZHIPU_CHAT_MODEL”, “glm-4.6”) # 可改为 glm-4-flash

EMBED_MODEL = os.environ.get(“ZHIPU_EMBED_MODEL”, “embedding-3”) # 可改 embedding-2

EMBED_DIM = int(os.environ.get(“ZHIPU_EMBED_DIM”, “1024”)) # embedding-3 可自定义维度

# EMBED_DIM 仅在 embedding-3 下生效，用于控制输出向量长度

def embed_texts(texts: List[str], batch_size: int = 64) -> np.ndarray:

# 调用 ZhipuAI 向量接口做批量向量化。embedding-3 支持自定义维度。

vecs: List[List[float]] = []

for i in range(0, len(texts), batch_size):

batch = texts[i:i+batch_size]

resp = client.embeddings.create(

model=EMBED_MODEL,

input=batch if len(batch) > 1 else batch[0],

**({“dimensions”: EMBED_DIM} if EMBED_MODEL == “embedding-3” else {})

)

# 接口输出格式：resp.data 每条对应一个 embedding

if hasattr(resp, “data”) and isinstance(resp.data, list):

for item in resp.data:

vecs.append(item.embedding)

else:

vecs.append(resp.data[0].embedding) # 兼容性处理

import numpy as _np

arr = _np.array(vecs, dtype=_np.float32)

norms = _np.linalg.norm(arr, axis=1, keepdims=True) + 1e-9

# 统一做 L2 归一化，便于余弦相似度检索

return arr / norms

def _month_range(reference: date, offset: int = 0) -> Tuple[date, date]:

# 计算 reference 所在月的前/后 offset 个月的起止日期

year = reference.year

month = reference.month + offset

while month <= 0:

month += 12

year -= 1

while month > 12:

month -= 12

year += 1

start = date(year, month, 1)

end = date(year, month, calendar.monthrange(year, month)[1])

return start, end

def _week_range(reference: date, offset: int = 0) -> Tuple[date, date]:

# 计算 reference 所在周（周一至周日）相对位移 offset 后的起止日期

start = reference – timedelta(days=reference.weekday()) + timedelta(weeks=offset)

end = start + timedelta(days=6)

return start, end

def augment_question_with_time_hint(question: str) -> Tuple[str, List[str]]:

“””为包含本月/上月/本周等相对时间的问题补充明确的日期范围，降低 LLM 误判。”””

today = date.today()

hints: Dict[str, Tuple[date, date]] = {}

patterns = [

({“本月”, “这个月”, “当月”}, lambda t: (“本月”, *_month_range(t, 0))),

({“上月”, “上个月”, “上一个月”}, lambda t: (“上月”, *_month_range(t, -1))),

({“本周”, “这周”, “本星期”}, lambda t: (“本周”, *_week_range(t, 0))),

({“上周”, “上星期”}, lambda t: (“上周”, *_week_range(t, -1))),

({“今年”, “本年”}, lambda t: (“今年”, date(t.year, 1, 1), date(t.year, 12, 31))),

({“去年”}, lambda t: (“去年”, date(t.year – 1, 1, 1), date(t.year – 1, 12, 31))),

({“今天”, “今日”}, lambda t: (“今天”, t, t)),

({“昨天”, “昨日”}, lambda t: (“昨天”, t – timedelta(days=1), t – timedelta(days=1))),

({“前天”}, lambda t: (“前天”, t – timedelta(days=2), t – timedelta(days=2))),

]

for keywords, fn in patterns:

if any(k in question for k in keywords):

label, start, end = fn(today)

hints.setdefault(label, (start, end))

if not hints:

return question, []

parts = []

for label, (start, end) in hints.items():

if start == end:

parts.append(f”{label}={start.isoformat()}”)

else:

parts.append(f”{label}={start.isoformat()}~{end.isoformat()}”)

hint_text = “；”.join(parts)

return f”{question}（时间范围提示：{hint_text}）”, parts

SCHEMA_TEXT = “””– 数据库表结构

TABLE customers(id INTEGER PRIMARY KEY, name TEXT, region TEXT);

TABLE orders(id INTEGER PRIMARY KEY, customer_id INTEGER, product TEXT, amount REAL, date TEXT);

— 字段含义：

— customers.id：客户主键；customers.name：客户名称；customers.region：客户所属大区，仅存储“华东/华南/华北/西南”等简洁大区名，不含“大区”“地区”字样。

— orders.id：订单主键；orders.customer_id：关联的客户 ID；orders.product：产品型号/类型；

— orders.amount：订单金额（元）；orders.date：订单日期（ISO 格式字符串）。

— 关系：

— orders.customer_id -> customers.id

— 仅允许只读 SELECT，不允许 INSERT/UPDATE/DELETE/DDL；不允许多条语句；不允许子查询以外的 ; 号。

“””

SQL_SYSTEM_PROMPT = f”””# 角色

你是一个严谨的企业数据分析助手。你的核心任务是**识别用户的查询意图**，并**仅在意图与数据库无关时拒绝查询**。

# 工作流程与判断标准

你必须严格遵循以下逻辑：

## 步骤 1：明确拒绝

首先，检查用户问题是否属于以下**完全无关**的类别。如果是，**必须且仅输出 `NO_SQL`**。

– **无关对话**：打招呼、问候、闲聊（如“你好”、“谢谢”、“在吗？”）。

– **元问题**：询问关于你自身、系统或规则的问题（如“你是谁？”、“你能做什么？”）。

– **纯知识问答**：与数据库数据无关的通用知识或技术问题（如“什么是SQL？”、“华东地区包括哪些省份？”）。

## 步骤 2：数据关联性判断

如果问题不属于上述类别，则进一步判断：**问题是否意图查询 `customers` 或 `orders` 表中的任何实体信息？**

– **实体信息包括**：客户、订单、产品、金额、日期、区域、数量、排名等。

– **关键原则**：**只要问题中提到了与数据库实体相关的关键词（如“最近一个订单”、“销售额最高的产品”、“华北的客户”），即使它询问的属性（如“质保期”、“联系方式”、“库存”）在数据库中不存在，也**必须**被视为与数据相关，并进入下一步生成 SQL。**

## 步骤 3：SQL 生成

如果步骤 2 的判断是“与数据相关”，则根据以下规则生成单条 SQL 语句。

### 生成规则

– **核心任务**：你的任务是查询数据库，以定位问题中提到的核心实体（如“最近一个订单”、“销售额最高的产品”）。

– **表与字段**：仅使用 `customers` 和 `orders` 表，且字段名必须来自 Schema。不得臆造字段名。

– **查询类型**：只能生成单条 `SELECT` 语句。严禁 `INSERT`, `UPDATE`, `DELETE`, `DDL`以及任何形式的注释。

– **语句格式**：禁止以分号 `;` 结尾。

– **字段选择**：显式列出所需字段，避免使用 `SELECT *`。

– **聚合与排序**：如需回答“最新/最近/最后/最多/最少”等问题，必须使用完整的 `ORDER BY` 和 `LIMIT` 子句。时间相关的排序基于 `orders.date` 字段。

– **时间范围**：如果用户消息包含“时间范围提示：xxx~yyy”，必须使用 `orders.date` 进行精确过滤。

– **数据特性**：`customers.region` 字段仅存储“华东”、“华南”等简洁大区名，不含“大区”等后缀。

# 上下文

当前日期：{date.today().isoformat()}

{SCHEMA_TEXT}

# 输出格式

– 如果问题完全无关，仅输出 `NO_SQL`。

– 如果问题与数据库实体相关，只输出 SQL 语句，并将其包裹在 “`sql … “` 代码块中。

“””

def llm_nl2sql(question: str) -> Optional[str]:

# 交给 GLM 生成 SQL，再从代码块中抽取裸 SQL 字符串

messages = [

{“role”: “system”, “content”: SQL_SYSTEM_PROMPT},

{“role”: “user”, “content”: question}

]

resp = client.chat.completions.create(

model=CHAT_MODEL,

messages=messages,

temperature=0.2,

)

text = resp.choices[0].message.content or “”

if “NO_SQL” in text.upper():

return None

m = re.search(r”“`sql\s*(.*?)\s*“`”, text, re.S | re.I)

sql = m.group(1).strip() if m else text.strip()

return sql or None

ALLOWED_TABLES = {“customers”, “orders”}

def audit_sql(sql: str) -> Optional[str]:

s = sql.strip().rstrip(“;”)

if not re.match(r”(?is)^\s*select\b”, s):

return “仅允许 SELECT 查询”

banned = [“insert”, “update”, “delete”, “drop”, “pragma”, “attach”, “alter”, “create”, “grant”, “reindex”, “;–“, “/*”, “*/”]

# 粗粒度黑名单 + 表名白名单拦截，避免误执行危险语句

if any(b in s.lower() for b in banned):

return “检测到危险关键字或注释”

tokens = re.findall(r”\b(from|join)\s+([a-zA-Z_][a-zA-Z0-9_]*)”, s, re.I)

for _, t in tokens:

if t.lower() not in ALLOWED_TABLES:

return f”不允许访问表：{t}”

return None # 通过

def run_sql(sql: str) -> List[Dict[str, Any]]:

# 仅执行 SELECT；返回 [{列: 值}, …]

conn = sqlite3.connect(DB_PATH)

conn.row_factory = sqlite3.Row

cur = conn.cursor()

cur.execute(sql)

rows = [dict(r) for r in cur.fetchall()]

conn.close()

return rows

def topk_by_embedding(

query: str,

chunk_vecs: np.ndarray,

chunks: List[Chunk],

k: int = 4,

) -> List[Tuple[float, Chunk]]:

# 对查询句向量化并与缓存向量矩阵做内积，返回相似度 Top-K

q = client.embeddings.create(

model=EMBED_MODEL,

input=query,

**({“dimensions”: EMBED_DIM} if EMBED_MODEL == “embedding-3” else {})

)

import numpy as _np

qv = _np.array(q.data[0].embedding, dtype=_np.float32)

qv = qv / (_np.linalg.norm(qv) + 1e-9)

sims = chunk_vecs @ qv

idx = _np.argsort(-sims)[:k]

return [(float(sims[i]), chunks[i]) for i in idx]

ANSWER_SYSTEM = “””你是企业知识问答助手。

请仅依据“SQL结果”和“文档片段”作答，必要时给出简短解释。

– 引用文档时使用 [doc_id#chunk_id] 形式。

– 如果证据不足，明确说明不足并给出需要的补充信息。”””

def synthesize_with_llm(question: str,

sql_rows: List[Dict[str, Any]],

doc_hits: List[Tuple[float, ‘Chunk’]],

generated_sql: Optional[str] = None

) -> str:

# 将 SQL 结果与文档片段打包成 JSON 传入 LLM，指导其引用证据作答

sql_preview = json.dumps(sql_rows[:10], ensure_ascii=False)

doc_preview = [

{“ref”: f”{c.doc_id}#{c.chunk_id}”, “text”: c.text.strip()} for _, c in doc_hits[:4]

]

payload = {

“question”: question,

“sql_rows”: json.loads(sql_preview),

“doc_snippets”: doc_preview

}

if generated_sql:

payload[“generated_sql”] = generated_sql

messages = [

{“role”: “system”, “content”: ANSWER_SYSTEM},

{“role”: “user”, “content”: json.dumps(payload, ensure_ascii=False)}

]

resp = client.chat.completions.create(

model=CHAT_MODEL,

messages=messages,

temperature=0.2,

)

return resp.choices[0].message.content or “”

# ====== 交互主流程 ======

def main():

print(“== 初始化数据与索引 ==”)

ensure_sqlite_db(DB_PATH)

print(“构建文档分片…”)

chunks = build_corpus(DOCS)

print(f”分片数：{len(chunks)}”)

print(f”向量化（{EMBED_MODEL}, dim={EMBED_DIM}）…”)

chunk_vecs = embed_texts([c.text for c in chunks])

print(“就绪。示例：”)

print(” – 星辰科技2024年订单总额”)

print(” – 华东地区上月销售额Top 3客户”)

print(” – 按地区统计今年销售额”)

print(” – A100的质保多久？”)

print(” – 出差住宿与餐补标准是什么？”)

print(“输入 exit 或 quit 结束。\n”)

while True:

q = input(“你：”).strip()

if q.lower() in (“exit”, “quit”, “q”):

print(“再见！”)

break

q_with_time, time_hints = augment_question_with_time_hint(q)

if time_hints:

print(f”[DEBUG] 解析相对时间：{‘；’.join(time_hints)}”)

# 1) 文档检索

doc_hits = topk_by_embedding(q_with_time, chunk_vecs, chunks, k=4)

# 2) NL2SQL + 审计 + 查询（失败则空列表）

sql_rows: List[Dict[str, Any]] = []

generated_sql: Optional[str] = None

try:

# 由 LLM 判断是否需要 SQL；若返回 None 表示无需结构化查询

generated_sql = llm_nl2sql(q_with_time)

if not generated_sql:

print(“[DEBUG] 略过 SQL 生成（判断为纯文档类问题）”)

else:

print(f”[DEBUG] LLM 生成 SQL: {generated_sql}”)

reason = audit_sql(generated_sql)

if reason:

print(f”[WARN] SQL 审计未通过：{reason}”)

else:

sql_rows = run_sql(generated_sql)

print(f”[DEBUG] SQL 结果条数：{len(sql_rows)}”)

except Exception as e:

print(f”[WARN] NL2SQL 或执行异常：{e}”)

# 3) 让 LLM 综合“结构化结果+文档片段”作答

answer = synthesize_with_llm(q_with_time, sql_rows, doc_hits, generated_sql)

# 4) 控制台展示

print(“\n【检索Top文档】”)

for s, ch in doc_hits:

print(f” – ({s:.3f}) [{ch.doc_id}#{ch.chunk_id}] {ch.text.strip()[:80]}…”)

print(“【结构化结果示例】”, json.dumps(sql_rows[:5], ensure_ascii=False))

print(“\n【最终回答】\n” + answer + “\n” + “-” * 60 + “\n”)