jimi
c39840fe15
主要变更: - 新增 tests/test_ai_chat.py: AI Agent 对话测试工具 - 优化 core/pydantic_ai_agent.py 和 db/chat_log_db.py - 清理归档文件，更新文档 Made-with: Cursor
239 lines
8.6 KiB
Python
Executable File
239 lines
8.6 KiB
Python
Executable File
# -*- coding: utf-8 -*-
|
||
"""
|
||
设计师派单数据库（SQLite）
|
||
|
||
同一设计师在不同店铺对应不同 group_id，派单时从在线设计师中轮询。
|
||
企微群「上线」/「下线」通过 update_online(wechat_user_id, is_online) 更新。
|
||
"""
|
||
import sqlite3
|
||
import os
|
||
from typing import Optional
|
||
|
||
_DB_PATH = os.path.join(os.path.dirname(__file__), "designer_roster_db", "roster.db")
|
||
_DB_TYPE = os.getenv("DB_TYPE", "sqlite").lower()
|
||
_MYSQL_HOST = os.getenv("MYSQL_HOST", "127.0.0.1")
|
||
_MYSQL_PORT = int(os.getenv("MYSQL_PORT", "3306"))
|
||
_MYSQL_USER = os.getenv("MYSQL_USER", "root")
|
||
_MYSQL_PASSWORD = os.getenv("MYSQL_PASSWORD", "")
|
||
_MYSQL_DATABASE = os.getenv("MYSQL_DATABASE", "ai_cs")
|
||
|
||
def _is_mysql() -> bool:
|
||
return _DB_TYPE in ("mysql", "mariadb")
|
||
|
||
def _sql(query: str) -> str:
|
||
return query.replace("?", "%s") if _is_mysql() else query
|
||
|
||
|
||
def _get_conn() -> sqlite3.Connection:
|
||
if _is_mysql():
|
||
import pymysql
|
||
return pymysql.connect(
|
||
host=_MYSQL_HOST,
|
||
port=_MYSQL_PORT,
|
||
user=_MYSQL_USER,
|
||
password=_MYSQL_PASSWORD,
|
||
database=_MYSQL_DATABASE,
|
||
charset="utf8mb4",
|
||
cursorclass=pymysql.cursors.DictCursor,
|
||
autocommit=False,
|
||
)
|
||
os.makedirs(os.path.dirname(_DB_PATH), exist_ok=True)
|
||
conn = sqlite3.connect(_DB_PATH)
|
||
conn.row_factory = sqlite3.Row
|
||
return conn
|
||
|
||
|
||
def init_db():
|
||
with _get_conn() as conn:
|
||
if _is_mysql():
|
||
conn.execute("""
|
||
CREATE TABLE IF NOT EXISTS designers (
|
||
id INTEGER PRIMARY KEY AUTO_INCREMENT,
|
||
name VARCHAR(255) NOT NULL,
|
||
wechat_user_id VARCHAR(128) UNIQUE NOT NULL
|
||
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
|
||
""")
|
||
conn.execute("""
|
||
CREATE TABLE IF NOT EXISTS designer_shops (
|
||
designer_id INTEGER NOT NULL,
|
||
shop_id VARCHAR(128) NOT NULL,
|
||
group_id VARCHAR(128) NOT NULL,
|
||
PRIMARY KEY (designer_id, shop_id),
|
||
FOREIGN KEY (designer_id) REFERENCES designers(id)
|
||
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
|
||
""")
|
||
conn.execute("""
|
||
CREATE TABLE IF NOT EXISTS designer_online (
|
||
wechat_user_id VARCHAR(128) PRIMARY KEY,
|
||
is_online INTEGER NOT NULL DEFAULT 0,
|
||
updated_at DATETIME
|
||
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
|
||
""")
|
||
conn.execute("""
|
||
CREATE TABLE IF NOT EXISTS round_robin (
|
||
shop_id VARCHAR(128) PRIMARY KEY,
|
||
last_index INTEGER NOT NULL DEFAULT 0
|
||
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
|
||
""")
|
||
else:
|
||
conn.execute("""
|
||
CREATE TABLE IF NOT EXISTS designers (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
name TEXT NOT NULL,
|
||
wechat_user_id TEXT UNIQUE NOT NULL
|
||
)
|
||
""")
|
||
conn.execute("""
|
||
CREATE TABLE IF NOT EXISTS designer_shops (
|
||
designer_id INTEGER NOT NULL,
|
||
shop_id TEXT NOT NULL,
|
||
group_id TEXT NOT NULL,
|
||
PRIMARY KEY (designer_id, shop_id),
|
||
FOREIGN KEY (designer_id) REFERENCES designers(id)
|
||
)
|
||
""")
|
||
conn.execute("""
|
||
CREATE TABLE IF NOT EXISTS designer_online (
|
||
wechat_user_id TEXT PRIMARY KEY,
|
||
is_online INTEGER NOT NULL DEFAULT 0,
|
||
updated_at TEXT
|
||
)
|
||
""")
|
||
conn.execute("""
|
||
CREATE TABLE IF NOT EXISTS round_robin (
|
||
shop_id TEXT PRIMARY KEY,
|
||
last_index INTEGER NOT NULL DEFAULT 0
|
||
)
|
||
""")
|
||
conn.commit()
|
||
|
||
|
||
init_db()
|
||
|
||
|
||
# ========== 设计师管理 ==========
|
||
|
||
def add_designer(name: str, wechat_user_id: str) -> int:
|
||
"""添加设计师，返回 id"""
|
||
with _get_conn() as conn:
|
||
if _is_mysql():
|
||
conn.execute(
|
||
"INSERT IGNORE INTO designers (name, wechat_user_id) VALUES (%s, %s)",
|
||
(name, wechat_user_id),
|
||
)
|
||
else:
|
||
conn.execute(
|
||
"INSERT OR IGNORE INTO designers (name, wechat_user_id) VALUES (?, ?)",
|
||
(name, wechat_user_id),
|
||
)
|
||
conn.commit()
|
||
row = conn.execute(_sql("SELECT id FROM designers WHERE wechat_user_id = ?"), (wechat_user_id,)).fetchone()
|
||
return row["id"] if row else 0
|
||
|
||
|
||
def set_designer_shop(designer_id: int, shop_id: str, group_id: str):
|
||
"""设置设计师在某店铺的分组 ID（同一设计师不同店铺不同 group_id）"""
|
||
with _get_conn() as conn:
|
||
if _is_mysql():
|
||
conn.execute(
|
||
"REPLACE INTO designer_shops (designer_id, shop_id, group_id) VALUES (%s, %s, %s)",
|
||
(designer_id, shop_id, group_id),
|
||
)
|
||
else:
|
||
conn.execute(
|
||
"INSERT OR REPLACE INTO designer_shops (designer_id, shop_id, group_id) VALUES (?, ?, ?)",
|
||
(designer_id, shop_id, group_id),
|
||
)
|
||
conn.commit()
|
||
|
||
|
||
def update_online(wechat_user_id: str, is_online: bool):
|
||
"""更新设计师在线状态（企微群「上线」/「下线」解析后调用）"""
|
||
from datetime import datetime
|
||
ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
with _get_conn() as conn:
|
||
if _is_mysql():
|
||
conn.execute(
|
||
"REPLACE INTO designer_online (wechat_user_id, is_online, updated_at) VALUES (%s, %s, %s)",
|
||
(wechat_user_id, 1 if is_online else 0, ts),
|
||
)
|
||
else:
|
||
conn.execute(
|
||
"INSERT OR REPLACE INTO designer_online (wechat_user_id, is_online, updated_at) VALUES (?, ?, ?)",
|
||
(wechat_user_id, 1 if is_online else 0, ts),
|
||
)
|
||
conn.commit()
|
||
|
||
|
||
# ========== 派单 ==========
|
||
|
||
def get_transfer_group_for_shop(shop_id: str) -> Optional[str]:
|
||
"""
|
||
为店铺轮询派单，返回分组 ID。
|
||
从该店铺的在线设计师中轮询选一个，返回其在该店铺的 group_id。
|
||
无人在线则返回 None。
|
||
"""
|
||
with _get_conn() as conn:
|
||
rows = conn.execute(_sql("""
|
||
SELECT d.wechat_user_id, ds.group_id
|
||
FROM designer_shops ds
|
||
JOIN designers d ON d.id = ds.designer_id
|
||
JOIN designer_online o ON o.wechat_user_id = d.wechat_user_id AND o.is_online = 1
|
||
WHERE ds.shop_id = ?
|
||
"""), (shop_id,)).fetchall()
|
||
|
||
if not rows:
|
||
return None
|
||
|
||
with _get_conn() as conn:
|
||
rr = conn.execute(_sql("SELECT last_index FROM round_robin WHERE shop_id = ?"), (shop_id,)).fetchone()
|
||
last = rr["last_index"] if rr else 0
|
||
idx = last % len(rows)
|
||
chosen = rows[idx]
|
||
if _is_mysql():
|
||
conn.execute(
|
||
"REPLACE INTO round_robin (shop_id, last_index) VALUES (%s, %s)",
|
||
(shop_id, idx + 1),
|
||
)
|
||
else:
|
||
conn.execute(
|
||
"INSERT OR REPLACE INTO round_robin (shop_id, last_index) VALUES (?, ?)",
|
||
(shop_id, idx + 1),
|
||
)
|
||
conn.commit()
|
||
|
||
return chosen["group_id"]
|
||
|
||
|
||
# ========== 查询 ==========
|
||
|
||
def get_all_wechat_user_ids() -> list:
|
||
"""获取所有设计师的 wechat_user_id（用于同步在线状态）"""
|
||
with _get_conn() as conn:
|
||
rows = conn.execute("SELECT wechat_user_id FROM designers").fetchall()
|
||
return [r["wechat_user_id"] for r in rows]
|
||
|
||
|
||
def list_designers():
|
||
"""列出所有设计师及其店铺分组"""
|
||
with _get_conn() as conn:
|
||
designers = conn.execute("SELECT id, name, wechat_user_id FROM designers").fetchall()
|
||
result = []
|
||
for d in designers:
|
||
shops = conn.execute(
|
||
_sql("SELECT shop_id, group_id FROM designer_shops WHERE designer_id = ?"),
|
||
(d["id"],),
|
||
).fetchall()
|
||
online = conn.execute(
|
||
_sql("SELECT is_online FROM designer_online WHERE wechat_user_id = ?"),
|
||
(d["wechat_user_id"],),
|
||
).fetchone()
|
||
result.append({
|
||
"id": d["id"],
|
||
"name": d["name"],
|
||
"wechat_user_id": d["wechat_user_id"],
|
||
"shops": {s["shop_id"]: s["group_id"] for s in shops},
|
||
"is_online": bool(online and online["is_online"]),
|
||
})
|
||
return result
|