zuowei/tw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 添加 AI Agent 对话测试工具 + 代码优化

Browse Source 
主要变更:

- 新增 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:
jimi
2026-02-28 16:19:35 +08:00
parent a6c42d505a
commit c39840fe15
49 changed files with 2453 additions and 8556 deletions
Download Patch File Download Diff File Expand all files Collapse all files

237
db/image_tasks_db.py
View File

@@ -10,8 +10,26 @@ from typing import Optional, List, Dict
from pathlib import Path
from datetime import datetime
from enum import Enum
import os
logger = logging.getLogger(__name__)
_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 _now_str() -> str:
if _is_mysql():
return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
return datetime.now().isoformat()
class TaskStatus(Enum):
"""任务状态"""
@@ -36,61 +54,105 @@ class ImageTaskManager:
def _init_db(self):
"""初始化数据库"""
self.db_path.parent.mkdir(parents=True, exist_ok=True)
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
# 创建图片任务表
cursor.execute('''
CREATE TABLE IF NOT EXISTS image_tasks (
task_id TEXT PRIMARY KEY,
customer_id TEXT NOT NULL,
customer_name TEXT,
original_image TEXT NOT NULL,
operation TEXT DEFAULT 'enhance',
requirements TEXT, -- JSON 格式：复杂度、比例、透视等
customer_notes TEXT, -- 客户备注/需求细节
status TEXT DEFAULT 'pending',
created_at TEXT,
paid_at TEXT,
started_at TEXT,
completed_at TEXT,
result_image TEXT,
error_message TEXT,
retry_count INTEGER DEFAULT 0,
-- 店铺信息
acc_id TEXT,
acc_type TEXT DEFAULT 'AliWorkbench'
)
''')
# 创建需求变更记录表（支持客户后续增加需求）
cursor.execute('''
CREATE TABLE IF NOT EXISTS task_requirement_changes (
id INTEGER PRIMARY KEY AUTOINCREMENT,
task_id TEXT NOT NULL,
change_type TEXT, -- add_note/modify_operation/add_requirement
old_value TEXT,
new_value TEXT,
changed_at TEXT,
changed_by TEXT, -- customer/staff
FOREIGN KEY (task_id) REFERENCES image_tasks(task_id)
)
''')
# 创建索引
cursor.execute('CREATE INDEX IF NOT EXISTS idx_customer ON image_tasks(customer_id)')
cursor.execute('CREATE INDEX IF NOT EXISTS idx_status ON image_tasks(status)')
cursor.execute('CREATE INDEX IF NOT EXISTS idx_created ON image_tasks(created_at)')
conn.commit()
conn.close()
if _is_mysql():
conn = self._get_conn()
cursor = conn.cursor()
cursor.execute('''
CREATE TABLE IF NOT EXISTS image_tasks (
task_id VARCHAR(128) PRIMARY KEY,
customer_id VARCHAR(128) NOT NULL,
customer_name VARCHAR(255),
original_image TEXT NOT NULL,
operation VARCHAR(64) DEFAULT 'enhance',
requirements TEXT,
customer_notes TEXT,
status VARCHAR(32) DEFAULT 'pending',
created_at DATETIME,
paid_at DATETIME,
started_at DATETIME,
completed_at DATETIME,
result_image TEXT,
error_message TEXT,
retry_count INT DEFAULT 0,
acc_id VARCHAR(128),
acc_type VARCHAR(64) DEFAULT 'AliWorkbench'
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
''')
cursor.execute('''
CREATE TABLE IF NOT EXISTS task_requirement_changes (
id INTEGER PRIMARY KEY AUTO_INCREMENT,
task_id VARCHAR(128) NOT NULL,
change_type VARCHAR(64),
old_value TEXT,
new_value TEXT,
changed_at DATETIME,
changed_by VARCHAR(32),
FOREIGN KEY (task_id) REFERENCES image_tasks(task_id)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
''')
cursor.execute('CREATE INDEX IF NOT EXISTS idx_customer ON image_tasks(customer_id)')
cursor.execute('CREATE INDEX IF NOT EXISTS idx_status ON image_tasks(status)')
cursor.execute('CREATE INDEX IF NOT EXISTS idx_created ON image_tasks(created_at)')
conn.commit()
conn.close()
else:
self.db_path.parent.mkdir(parents=True, exist_ok=True)
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
cursor.execute('''
CREATE TABLE IF NOT EXISTS image_tasks (
task_id TEXT PRIMARY KEY,
customer_id TEXT NOT NULL,
customer_name TEXT,
original_image TEXT NOT NULL,
operation TEXT DEFAULT 'enhance',
requirements TEXT,
customer_notes TEXT,
status TEXT DEFAULT 'pending',
created_at TEXT,
paid_at TEXT,
started_at TEXT,
completed_at TEXT,
result_image TEXT,
error_message TEXT,
retry_count INTEGER DEFAULT 0,
acc_id TEXT,
acc_type TEXT DEFAULT 'AliWorkbench'
)
''')
cursor.execute('''
CREATE TABLE IF NOT EXISTS task_requirement_changes (
id INTEGER PRIMARY KEY AUTOINCREMENT,
task_id TEXT NOT NULL,
change_type TEXT,
old_value TEXT,
new_value TEXT,
changed_at TEXT,
changed_by TEXT,
FOREIGN KEY (task_id) REFERENCES image_tasks(task_id)
)
''')
cursor.execute('CREATE INDEX IF NOT EXISTS idx_customer ON image_tasks(customer_id)')
cursor.execute('CREATE INDEX IF NOT EXISTS idx_status ON image_tasks(status)')
cursor.execute('CREATE INDEX IF NOT EXISTS idx_created ON image_tasks(created_at)')
conn.commit()
conn.close()
logger.info("数据库表初始化完成")
def _get_conn(self):
"""获取数据库连接"""
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,
)
conn = sqlite3.connect(self.db_path)
conn.row_factory = sqlite3.Row
return conn
@@ -105,13 +167,13 @@ class ImageTaskManager:
requirements_json = json.dumps(requirements) if requirements else None
cursor.execute('''
cursor.execute(_sql('''
INSERT INTO image_tasks (
task_id, customer_id, customer_name, original_image,
operation, requirements, customer_notes, status,
created_at, acc_id, acc_type
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
''', (
'''), (
task_id,
customer_id,
customer_name,
@@ -120,7 +182,7 @@ class ImageTaskManager:
requirements_json,
'', # 初始备注为空
TaskStatus.PENDING.value,
datetime.now().isoformat(),
_now_str(),
acc_id,
acc_type
))
@@ -141,7 +203,7 @@ class ImageTaskManager:
conn = self._get_conn()
cursor = conn.cursor()
cursor.execute('SELECT * FROM image_tasks WHERE task_id = ?', (task_id,))
cursor.execute(_sql('SELECT * FROM image_tasks WHERE task_id = ?'), (task_id,))
row = cursor.fetchone()
conn.close()
@@ -164,17 +226,17 @@ class ImageTaskManager:
cursor = conn.cursor()
if status:
cursor.execute('''
cursor.execute(_sql('''
SELECT * FROM image_tasks
WHERE customer_id = ? AND status = ?
ORDER BY created_at DESC
''', (customer_id, status))
'''), (customer_id, status))
else:
cursor.execute('''
cursor.execute(_sql('''
SELECT * FROM image_tasks
WHERE customer_id = ?
ORDER BY created_at DESC
''', (customer_id,))
'''), (customer_id,))
rows = cursor.fetchall()
conn.close()
@@ -198,27 +260,28 @@ class ImageTaskManager:
conn = self._get_conn()
cursor = conn.cursor()
updates = ['status = ?']
placeholder = "%s" if _is_mysql() else "?"
updates = [f'status = {placeholder}']
params = [status.value]
# 根据状态设置时间
if status == TaskStatus.PAID:
updates.append('paid_at = ?')
params.append(datetime.now().isoformat())
updates.append(f'paid_at = {placeholder}')
params.append(_now_str())
elif status == TaskStatus.PROCESSING:
updates.append('started_at = ?')
params.append(datetime.now().isoformat())
updates.append(f'started_at = {placeholder}')
params.append(_now_str())
elif status in [TaskStatus.COMPLETED, TaskStatus.FAILED]:
updates.append('completed_at = ?')
params.append(datetime.now().isoformat())
updates.append(f'completed_at = {placeholder}')
params.append(_now_str())
params.append(task_id)
cursor.execute(f'''
cursor.execute(_sql(f'''
UPDATE image_tasks
SET {', '.join(updates)}
WHERE task_id = ?
''', params)
'''), params)
conn.commit()
conn.close()
@@ -234,11 +297,11 @@ class ImageTaskManager:
conn = self._get_conn()
cursor = conn.cursor()
cursor.execute('''
cursor.execute(_sql('''
UPDATE image_tasks
SET result_image = ?, error_message = ?
WHERE task_id = ?
''', (result_image, error_message, task_id))
'''), (result_image, error_message, task_id))
conn.commit()
conn.close()
@@ -265,30 +328,30 @@ class ImageTaskManager:
cursor = conn.cursor()
# 获取旧备注
cursor.execute('SELECT customer_notes FROM image_tasks WHERE task_id = ?', (task_id,))
cursor.execute(_sql('SELECT customer_notes FROM image_tasks WHERE task_id = ?'), (task_id,))
row = cursor.fetchone()
old_note = row['customer_notes'] if row else ''
# 更新备注
new_note = f"{old_note}\n[{datetime.now().strftime('%m-%d %H:%M')}] {note}" if old_note else f"[{datetime.now().strftime('%m-%d %H:%M')}] {note}"
cursor.execute('''
cursor.execute(_sql('''
UPDATE image_tasks
SET customer_notes = ?
WHERE task_id = ?
''', (new_note, task_id))
'''), (new_note, task_id))
# 记录变更历史
cursor.execute('''
cursor.execute(_sql('''
INSERT INTO task_requirement_changes (
task_id, change_type, old_value, new_value, changed_at, changed_by
) VALUES (?, ?, ?, ?, ?, ?)
''', (
'''), (
task_id,
'add_note',
old_note or '无',
note,
datetime.now().isoformat(),
_now_str(),
changed_by
))
@@ -319,28 +382,28 @@ class ImageTaskManager:
cursor = conn.cursor()
# 获取旧操作
cursor.execute('SELECT operation FROM image_tasks WHERE task_id = ?', (task_id,))
cursor.execute(_sql('SELECT operation FROM image_tasks WHERE task_id = ?'), (task_id,))
row = cursor.fetchone()
old_operation = row['operation'] if row else ''
# 更新操作
cursor.execute('''
cursor.execute(_sql('''
UPDATE image_tasks
SET operation = ?
WHERE task_id = ?
''', (new_operation, task_id))
'''), (new_operation, task_id))
# 记录变更历史
cursor.execute('''
cursor.execute(_sql('''
INSERT INTO task_requirement_changes (
task_id, change_type, old_value, new_value, changed_at, changed_by
) VALUES (?, ?, ?, ?, ?, ?)
''', (
'''), (
task_id,
'modify_operation',
old_operation,
new_operation,
datetime.now().isoformat(),
_now_str(),
changed_by
))
@@ -360,11 +423,11 @@ class ImageTaskManager:
conn = self._get_conn()
cursor = conn.cursor()
cursor.execute('''
cursor.execute(_sql('''
SELECT * FROM task_requirement_changes
WHERE task_id = ?
ORDER BY changed_at DESC
''', (task_id,))
'''), (task_id,))
rows = cursor.fetchall()
conn.close()
@@ -385,13 +448,13 @@ class ImageTaskManager:
conn = self._get_conn()
cursor = conn.cursor()
cursor.execute('''
cursor.execute(_sql('''
UPDATE image_tasks
SET retry_count = retry_count + 1
WHERE task_id = ?
''', (task_id,))
'''), (task_id,))
cursor.execute('SELECT retry_count FROM image_tasks WHERE task_id = ?', (task_id,))
cursor.execute(_sql('SELECT retry_count FROM image_tasks WHERE task_id = ?'), (task_id,))
row = cursor.fetchone()
conn.close()