feat: 添加 AI Agent 对话测试工具 + 代码优化
主要变更: - 新增 tests/test_ai_chat.py: AI Agent 对话测试工具 - 优化 core/pydantic_ai_agent.py 和 db/chat_log_db.py - 清理归档文件，更新文档 Made-with: Cursor
This commit is contained in:
@@ -10,9 +10,33 @@ 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
|
||||
@@ -21,35 +45,66 @@ def _get_conn() -> sqlite3.Connection:
|
||||
|
||||
def init_db():
|
||||
with _get_conn() as conn:
|
||||
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
|
||||
)
|
||||
""")
|
||||
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()
|
||||
|
||||
|
||||
@@ -61,22 +116,34 @@ init_db()
|
||||
def add_designer(name: str, wechat_user_id: str) -> int:
|
||||
"""添加设计师，返回 id"""
|
||||
with _get_conn() as conn:
|
||||
conn.execute(
|
||||
"INSERT OR IGNORE INTO designers (name, wechat_user_id) VALUES (?, ?)",
|
||||
(name, wechat_user_id),
|
||||
)
|
||||
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("SELECT id FROM designers WHERE wechat_user_id = ?", (wechat_user_id,)).fetchone()
|
||||
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:
|
||||
conn.execute(
|
||||
"INSERT OR REPLACE INTO designer_shops (designer_id, shop_id, group_id) VALUES (?, ?, ?)",
|
||||
(designer_id, shop_id, group_id),
|
||||
)
|
||||
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()
|
||||
|
||||
|
||||
@@ -85,10 +152,16 @@ 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:
|
||||
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),
|
||||
)
|
||||
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()
|
||||
|
||||
|
||||
@@ -101,26 +174,32 @@ def get_transfer_group_for_shop(shop_id: str) -> Optional[str]:
|
||||
无人在线则返回 None。
|
||||
"""
|
||||
with _get_conn() as conn:
|
||||
rows = conn.execute("""
|
||||
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()
|
||||
"""), (shop_id,)).fetchall()
|
||||
|
||||
if not rows:
|
||||
return None
|
||||
|
||||
with _get_conn() as conn:
|
||||
rr = conn.execute("SELECT last_index FROM round_robin WHERE shop_id = ?", (shop_id,)).fetchone()
|
||||
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]
|
||||
conn.execute(
|
||||
"INSERT OR REPLACE INTO round_robin (shop_id, last_index) VALUES (?, ?)",
|
||||
(shop_id, idx + 1),
|
||||
)
|
||||
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"]
|
||||
@@ -142,11 +221,11 @@ def list_designers():
|
||||
result = []
|
||||
for d in designers:
|
||||
shops = conn.execute(
|
||||
"SELECT shop_id, group_id FROM designer_shops WHERE designer_id = ?",
|
||||
_sql("SELECT shop_id, group_id FROM designer_shops WHERE designer_id = ?"),
|
||||
(d["id"],),
|
||||
).fetchall()
|
||||
online = conn.execute(
|
||||
"SELECT is_online FROM designer_online WHERE wechat_user_id = ?",
|
||||
_sql("SELECT is_online FROM designer_online WHERE wechat_user_id = ?"),
|
||||
(d["wechat_user_id"],),
|
||||
).fetchone()
|
||||
result.append({
|
||||
|
||||
Reference in New Issue
Block a user