114 lines
4.2 KiB
Python
114 lines
4.2 KiB
Python
# -*- 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
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class TaskStatus(Enum):
|
|
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()
|
|
|
|
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,
|
|
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 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()
|
|
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()
|
|
return task_id
|
|
except Exception as e:
|
|
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()
|
|
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('UPDATE image_tasks SET status = ?, updated_at = ? WHERE task_id = ?',
|
|
(status, now, task_id))
|
|
conn.commit()
|
|
conn.close()
|
|
except Exception as e:
|
|
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('''
|
|
UPDATE image_tasks
|
|
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"[DB] 客户 {customer_id} 任务结局更新为: {outcome}")
|
|
except Exception as e:
|
|
logger.error(f"Failed to update outcome: {e}")
|
|
|
|
# 单例
|
|
db = ImageTaskManager()
|