refactor: migrate workflow to v2 core and archive legacy modules
This commit is contained in:
@@ -1,480 +1,113 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
图片任务数据库管理
|
||||
支持客户后续增加需求细节
|
||||
"""
|
||||
import sqlite3
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
import os
|
||||
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):
|
||||
"""任务状态"""
|
||||
PENDING = "pending" # 待付款
|
||||
PAID = "paid" # 已付款，待处理
|
||||
PROCESSING = "processing" # 处理中
|
||||
AWAITING_CONFIRM = "awaiting_confirm" # 已完成，待客户确认
|
||||
COMPLETED = "completed" # 已完成
|
||||
FAILED = "failed" # 失败
|
||||
CANCELLED = "cancelled" # 已取消
|
||||
PENDING = "pending"
|
||||
PROCESSING = "processing"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
|
||||
class ImageTaskManager:
|
||||
"""图片任务管理器"""
|
||||
|
||||
def __init__(self, db_path: str = None):
|
||||
if db_path is None:
|
||||
db_path = Path(__file__).parent / "image_tasks.db"
|
||||
|
||||
self.db_path = db_path
|
||||
self._init_db()
|
||||
logger.info(f"图片任务管理器初始化完成，数据库：{self.db_path}")
|
||||
|
||||
|
||||
def _init_db(self):
|
||||
"""初始化数据库"""
|
||||
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("SHOW INDEX FROM image_tasks")
|
||||
exists = {str(r.get("Key_name", "")) for r in cursor.fetchall()}
|
||||
if "idx_customer" not in exists:
|
||||
cursor.execute('CREATE INDEX idx_customer ON image_tasks(customer_id)')
|
||||
if "idx_status" not in exists:
|
||||
cursor.execute('CREATE INDEX idx_status ON image_tasks(status)')
|
||||
if "idx_created" not in exists:
|
||||
cursor.execute('CREATE INDEX 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,
|
||||
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,
|
||||
platform TEXT DEFAULT 'qianniu',
|
||||
original_image TEXT NOT NULL,
|
||||
operation TEXT DEFAULT 'enhance',
|
||||
requirements TEXT,
|
||||
status TEXT DEFAULT 'pending',
|
||||
result_image TEXT,
|
||||
price REAL DEFAULT 0.0,
|
||||
outcome TEXT DEFAULT 'pending',
|
||||
created_at TEXT,
|
||||
updated_at TEXT
|
||||
)
|
||||
''')
|
||||
cursor.execute('CREATE INDEX IF NOT EXISTS idx_status ON image_tasks(status)')
|
||||
cursor.execute('CREATE INDEX IF NOT EXISTS idx_cust_plat ON image_tasks(customer_id, platform)')
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
def _get_conn(self):
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
conn.row_factory = sqlite3.Row
|
||||
return conn
|
||||
|
||||
def create_task(self, task_id: str, customer_id: str, customer_name: str,
|
||||
original_image: str, operation: str = 'enhance',
|
||||
requirements: dict = None, acc_id: str = '', acc_type: str = 'AliWorkbench') -> bool:
|
||||
"""创建图片任务"""
|
||||
|
||||
def add_task(self, customer_id: str, platform: str, original_image: str, operation: str, requirements: str = "", status: str = "pending") -> str:
|
||||
task_id = str(uuid.uuid4())
|
||||
now = datetime.now().isoformat()
|
||||
try:
|
||||
conn = self._get_conn()
|
||||
cursor = conn.cursor()
|
||||
|
||||
requirements_json = json.dumps(requirements) if requirements else None
|
||||
|
||||
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,
|
||||
original_image,
|
||||
operation,
|
||||
requirements_json,
|
||||
'', # 初始备注为空
|
||||
TaskStatus.PENDING.value,
|
||||
_now_str(),
|
||||
acc_id,
|
||||
acc_type
|
||||
))
|
||||
|
||||
cursor.execute('''
|
||||
INSERT INTO image_tasks (task_id, customer_id, platform, original_image, operation, requirements, status, created_at, updated_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
''', (task_id, customer_id, platform, original_image, operation, requirements, status, now, now))
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
logger.info(f"图片任务创建成功：{task_id}")
|
||||
return True
|
||||
|
||||
return task_id
|
||||
except Exception as e:
|
||||
logger.error(f"创建图片任务失败：{e}")
|
||||
return False
|
||||
|
||||
def get_task(self, task_id: str) -> Optional[dict]:
|
||||
"""查询任务"""
|
||||
logger.error(f"Failed to add task: {e}")
|
||||
return ""
|
||||
|
||||
def update_status(self, task_id: str, status: str, result_image: str = ""):
|
||||
now = datetime.now().isoformat()
|
||||
try:
|
||||
conn = self._get_conn()
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute(_sql('SELECT * FROM image_tasks WHERE task_id = ?'), (task_id,))
|
||||
row = cursor.fetchone()
|
||||
conn.close()
|
||||
|
||||
if row:
|
||||
task = dict(row)
|
||||
# 解析 JSON 字段
|
||||
if task.get('requirements'):
|
||||
task['requirements'] = json.loads(task['requirements'])
|
||||
return task
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"查询任务失败：{e}")
|
||||
return None
|
||||
|
||||
def get_customer_tasks(self, customer_id: str, status: str = None) -> List[dict]:
|
||||
"""查询客户的任务列表"""
|
||||
try:
|
||||
conn = self._get_conn()
|
||||
cursor = conn.cursor()
|
||||
|
||||
if status:
|
||||
cursor.execute(_sql('''
|
||||
SELECT * FROM image_tasks
|
||||
WHERE customer_id = ? AND status = ?
|
||||
ORDER BY created_at DESC
|
||||
'''), (customer_id, status))
|
||||
if result_image:
|
||||
cursor.execute('UPDATE image_tasks SET status = ?, result_image = ?, updated_at = ? WHERE task_id = ?',
|
||||
(status, result_image, now, task_id))
|
||||
else:
|
||||
cursor.execute(_sql('''
|
||||
SELECT * FROM image_tasks
|
||||
WHERE customer_id = ?
|
||||
ORDER BY created_at DESC
|
||||
'''), (customer_id,))
|
||||
|
||||
rows = cursor.fetchall()
|
||||
conn.close()
|
||||
|
||||
tasks = []
|
||||
for row in rows:
|
||||
task = dict(row)
|
||||
if task.get('requirements'):
|
||||
task['requirements'] = json.loads(task['requirements'])
|
||||
tasks.append(task)
|
||||
|
||||
return tasks
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"查询客户任务失败：{e}")
|
||||
return []
|
||||
|
||||
def update_status(self, task_id: str, status: TaskStatus):
|
||||
"""更新任务状态"""
|
||||
try:
|
||||
conn = self._get_conn()
|
||||
cursor = conn.cursor()
|
||||
|
||||
placeholder = "%s" if _is_mysql() else "?"
|
||||
updates = [f'status = {placeholder}']
|
||||
params = [status.value]
|
||||
|
||||
# 根据状态设置时间
|
||||
if status == TaskStatus.PAID:
|
||||
updates.append(f'paid_at = {placeholder}')
|
||||
params.append(_now_str())
|
||||
elif status == TaskStatus.PROCESSING:
|
||||
updates.append(f'started_at = {placeholder}')
|
||||
params.append(_now_str())
|
||||
elif status in [TaskStatus.COMPLETED, TaskStatus.FAILED]:
|
||||
updates.append(f'completed_at = {placeholder}')
|
||||
params.append(_now_str())
|
||||
|
||||
params.append(task_id)
|
||||
|
||||
cursor.execute(_sql(f'''
|
||||
UPDATE image_tasks
|
||||
SET {', '.join(updates)}
|
||||
WHERE task_id = ?
|
||||
'''), params)
|
||||
|
||||
cursor.execute('UPDATE image_tasks SET status = ?, updated_at = ? WHERE task_id = ?',
|
||||
(status, now, task_id))
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
logger.info(f"任务状态更新：{task_id} -> {status.value}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"更新任务状态失败：{e}")
|
||||
|
||||
def update_result(self, task_id: str, result_image: str, error_message: str = None):
|
||||
"""更新处理结果"""
|
||||
logger.error(f"Failed to update task status: {e}")
|
||||
|
||||
def update_outcome(self, customer_id: str, platform: str, outcome: str):
|
||||
"""记录任务的最终结局（用于训练样本分类）"""
|
||||
now = datetime.now().isoformat()
|
||||
try:
|
||||
conn = self._get_conn()
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute(_sql('''
|
||||
cursor.execute('''
|
||||
UPDATE image_tasks
|
||||
SET result_image = ?, error_message = ?
|
||||
WHERE task_id = ?
|
||||
'''), (result_image, error_message, task_id))
|
||||
|
||||
SET outcome = ?, updated_at = ?
|
||||
WHERE task_id = (
|
||||
SELECT task_id FROM image_tasks
|
||||
WHERE customer_id = ? AND platform = ?
|
||||
ORDER BY created_at DESC LIMIT 1
|
||||
)
|
||||
''', (outcome, now, customer_id, platform))
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
logger.info(f"任务结果更新：{task_id}")
|
||||
|
||||
logger.info(f"[DB] 客户 {customer_id} 任务结局更新为: {outcome}")
|
||||
except Exception as e:
|
||||
logger.error(f"更新任务结果失败：{e}")
|
||||
|
||||
def add_customer_note(self, task_id: str, note: str, changed_by: str = 'customer') -> bool:
|
||||
"""
|
||||
客户添加需求备注（支持后续增加细节）
|
||||
|
||||
Args:
|
||||
task_id: 任务 ID
|
||||
note: 备注内容
|
||||
changed_by: 修改者（customer/staff）
|
||||
|
||||
Returns:
|
||||
bool: 是否成功
|
||||
"""
|
||||
try:
|
||||
conn = self._get_conn()
|
||||
cursor = conn.cursor()
|
||||
|
||||
# 获取旧备注
|
||||
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(_sql('''
|
||||
UPDATE image_tasks
|
||||
SET customer_notes = ?
|
||||
WHERE task_id = ?
|
||||
'''), (new_note, task_id))
|
||||
|
||||
# 记录变更历史
|
||||
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,
|
||||
_now_str(),
|
||||
changed_by
|
||||
))
|
||||
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
logger.info(f"客户添加备注成功：{task_id}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"添加客户备注失败：{e}")
|
||||
return False
|
||||
|
||||
def modify_operation(self, task_id: str, new_operation: str, changed_by: str = 'customer') -> bool:
|
||||
"""
|
||||
修改操作类型（客户后续修改需求）
|
||||
|
||||
Args:
|
||||
task_id: 任务 ID
|
||||
new_operation: 新操作类型
|
||||
changed_by: 修改者
|
||||
|
||||
Returns:
|
||||
bool: 是否成功
|
||||
"""
|
||||
try:
|
||||
conn = self._get_conn()
|
||||
cursor = conn.cursor()
|
||||
|
||||
# 获取旧操作
|
||||
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(_sql('''
|
||||
UPDATE image_tasks
|
||||
SET operation = ?
|
||||
WHERE task_id = ?
|
||||
'''), (new_operation, task_id))
|
||||
|
||||
# 记录变更历史
|
||||
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,
|
||||
_now_str(),
|
||||
changed_by
|
||||
))
|
||||
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
logger.info(f"修改操作类型成功：{task_id} -> {new_operation}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"修改操作类型失败：{e}")
|
||||
return False
|
||||
|
||||
def get_requirement_history(self, task_id: str) -> List[dict]:
|
||||
"""获取需求变更历史"""
|
||||
try:
|
||||
conn = self._get_conn()
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute(_sql('''
|
||||
SELECT * FROM task_requirement_changes
|
||||
WHERE task_id = ?
|
||||
ORDER BY changed_at DESC
|
||||
'''), (task_id,))
|
||||
|
||||
rows = cursor.fetchall()
|
||||
conn.close()
|
||||
|
||||
return [dict(row) for row in rows]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"查询需求历史失败：{e}")
|
||||
return []
|
||||
|
||||
def get_pending_tasks(self) -> List[dict]:
|
||||
"""获取所有待处理任务"""
|
||||
return self.get_customer_tasks('', 'pending')
|
||||
|
||||
def increment_retry(self, task_id: str) -> int:
|
||||
"""增加重试次数"""
|
||||
try:
|
||||
conn = self._get_conn()
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute(_sql('''
|
||||
UPDATE image_tasks
|
||||
SET retry_count = retry_count + 1
|
||||
WHERE task_id = ?
|
||||
'''), (task_id,))
|
||||
|
||||
cursor.execute(_sql('SELECT retry_count FROM image_tasks WHERE task_id = ?'), (task_id,))
|
||||
row = cursor.fetchone()
|
||||
conn.close()
|
||||
|
||||
return row['retry_count'] if row else 0
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"增加重试次数失败：{e}")
|
||||
return 999
|
||||
logger.error(f"Failed to update outcome: {e}")
|
||||
|
||||
# 单例
|
||||
_task_manager: Optional[ImageTaskManager] = None
|
||||
|
||||
def get_image_task_manager() -> ImageTaskManager:
|
||||
"""获取图片任务管理器单例"""
|
||||
global _task_manager
|
||||
if _task_manager is None:
|
||||
_task_manager = ImageTaskManager()
|
||||
return _task_manager
|
||||
db = ImageTaskManager()
|
||||
|
||||
Reference in New Issue
Block a user