feat: automate image pipeline and simplify gemini flow
This commit is contained in:
232
services/service_auto_image_pipeline.py
Normal file
232
services/service_auto_image_pipeline.py
Normal file
@@ -0,0 +1,232 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from db.customer_db import CustomerDatabase
|
||||
from db.image_tasks_db import TaskStatus, db as task_db
|
||||
from services.service_gemini import GeminiExtractV2Service
|
||||
from services.service_tuhui_upload import upload_to_tuhui
|
||||
from services.service_wecom_bot import wecom_bot_service
|
||||
|
||||
load_dotenv()
|
||||
logger = logging.getLogger("cs_agent")
|
||||
|
||||
AUTO_PROCESS_PRICE = int(os.getenv("AUTO_PROCESS_DEFAULT_PRICE", "12"))
|
||||
AUTO_PROCESS_CATEGORY = os.getenv("AUTO_PROCESS_CATEGORY", "设计素材")
|
||||
AUTO_PROCESS_ROOT = Path(
|
||||
os.getenv("AUTO_PROCESS_ROOT", str(Path(__file__).resolve().parents[1] / "runtime" / "auto_processed"))
|
||||
)
|
||||
|
||||
|
||||
def _safe_name(text: str, fallback: str = "image") -> str:
|
||||
cleaned = re.sub(r"[^0-9A-Za-z\u4e00-\u9fa5_-]+", "_", str(text or "").strip())
|
||||
cleaned = cleaned.strip("_")
|
||||
return cleaned[:40] or fallback
|
||||
|
||||
|
||||
def _suffix_from_url(url: str) -> str:
|
||||
path = urlparse(str(url or "")).path
|
||||
suffix = Path(path).suffix.lower()
|
||||
if suffix in {".png", ".jpg", ".jpeg", ".webp"}:
|
||||
return suffix
|
||||
return ".png"
|
||||
|
||||
|
||||
def _build_processing_prompt(intent: str, requirement_text: str, analysis: Dict) -> str:
|
||||
base_prompt = str((analysis or {}).get("gemini_prompt") or "").strip()
|
||||
req = str(requirement_text or "").strip()
|
||||
if base_prompt:
|
||||
return base_prompt
|
||||
if intent == "repair":
|
||||
return f"根据客户需求“{req or '高清修复'}”，保留主体和构图，做高清修复并补足细节。"
|
||||
return f"根据客户需求“{req or '找原图'}”，严格参考原图元素与构图，生成完整干净的高质量素材图。"
|
||||
|
||||
|
||||
class AutoImagePipelineService:
|
||||
def __init__(self):
|
||||
self.customer_db = CustomerDatabase()
|
||||
|
||||
async def _download_image(self, image_url: str, dest_path: Path) -> Path:
|
||||
dest_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
async with httpx.AsyncClient(timeout=60.0, follow_redirects=True) as client:
|
||||
response = await client.get(image_url)
|
||||
response.raise_for_status()
|
||||
dest_path.write_bytes(response.content)
|
||||
return dest_path
|
||||
|
||||
@staticmethod
|
||||
def _format_transfer_notice(
|
||||
customer_id: str,
|
||||
acc_id: str,
|
||||
designer_name: str,
|
||||
requirement_text: str,
|
||||
intent: str,
|
||||
image_urls: List[str],
|
||||
) -> str:
|
||||
lines = [
|
||||
"【AI自动转设计师】",
|
||||
f"店铺：{acc_id or '-'}",
|
||||
f"客户：{customer_id or '-'}",
|
||||
f"设计师：{designer_name or '-'}",
|
||||
f"需求：{requirement_text or '-'}",
|
||||
f"类型：{'高清修复' if intent == 'repair' else '找原图'}",
|
||||
f"默认价格：{AUTO_PROCESS_PRICE}元",
|
||||
]
|
||||
if image_urls:
|
||||
lines.append("原图URL：")
|
||||
lines.extend(image_urls[:5])
|
||||
return "\n".join(lines)
|
||||
|
||||
@staticmethod
|
||||
def _format_finish_notice(
|
||||
customer_id: str,
|
||||
acc_id: str,
|
||||
designer_name: str,
|
||||
links: List[Dict[str, str]],
|
||||
failures: List[str],
|
||||
) -> str:
|
||||
lines = [
|
||||
"【AI处理完成】",
|
||||
f"店铺：{acc_id or '-'}",
|
||||
f"客户：{customer_id or '-'}",
|
||||
f"设计师：{designer_name or '-'}",
|
||||
f"默认价格：{AUTO_PROCESS_PRICE}元",
|
||||
]
|
||||
if links:
|
||||
lines.append("处理结果：")
|
||||
for idx, item in enumerate(links, 1):
|
||||
lines.append(f"{idx}. 图绘链接：{item.get('download_url') or '-'}")
|
||||
lines.append(f" 原图URL：{item.get('source_url') or '-'}")
|
||||
if failures:
|
||||
lines.append("失败项：")
|
||||
lines.extend(failures[:5])
|
||||
return "\n".join(lines)
|
||||
|
||||
async def process_and_notify(
|
||||
self,
|
||||
*,
|
||||
session_key: str,
|
||||
customer_id: str,
|
||||
acc_id: str,
|
||||
designer_name: str,
|
||||
requirement_text: str,
|
||||
image_urls: List[str],
|
||||
intent: str = "",
|
||||
) -> Dict:
|
||||
image_urls = [str(url).strip() for url in (image_urls or []) if str(url).strip()]
|
||||
if not image_urls:
|
||||
return {"success": False, "message": "no_images"}
|
||||
image_urls = image_urls[:1]
|
||||
|
||||
profile = self.customer_db.get_customer(session_key)
|
||||
analysis = {}
|
||||
if getattr(profile, "last_image_analysis", ""):
|
||||
try:
|
||||
analysis = json.loads(profile.last_image_analysis)
|
||||
except Exception:
|
||||
analysis = {}
|
||||
|
||||
if not intent:
|
||||
intent = "repair" if "修复" in requirement_text else "find_original"
|
||||
|
||||
await wecom_bot_service.send_text(
|
||||
self._format_transfer_notice(
|
||||
customer_id=customer_id,
|
||||
acc_id=acc_id,
|
||||
designer_name=designer_name,
|
||||
requirement_text=requirement_text,
|
||||
intent=intent,
|
||||
image_urls=image_urls,
|
||||
)
|
||||
)
|
||||
|
||||
pipeline_root = AUTO_PROCESS_ROOT / _safe_name(customer_id, "customer")
|
||||
pipeline_root.mkdir(parents=True, exist_ok=True)
|
||||
gemini_service = GeminiExtractV2Service()
|
||||
uploaded_links: List[Dict[str, str]] = []
|
||||
failures: List[str] = []
|
||||
|
||||
for idx, image_url in enumerate(image_urls, 1):
|
||||
digest = hashlib.md5(f"{customer_id}|{acc_id}|{image_url}".encode("utf-8")).hexdigest()[:10]
|
||||
input_path = pipeline_root / f"{digest}_src{_suffix_from_url(image_url)}"
|
||||
output_path = pipeline_root / f"{digest}_out.png"
|
||||
title = f"{_safe_name(customer_id, '客户')}_{'修复' if intent == 'repair' else '原图'}_{idx}"
|
||||
prompt = _build_processing_prompt(intent, requirement_text, analysis)
|
||||
task_id = task_db.add_task(
|
||||
customer_id=customer_id,
|
||||
platform="qianniu",
|
||||
original_image=image_url,
|
||||
operation=intent or "auto_process",
|
||||
requirements=requirement_text,
|
||||
status=TaskStatus.PROCESSING.value,
|
||||
)
|
||||
try:
|
||||
await self._download_image(image_url, input_path)
|
||||
success, message, data = await gemini_service.extract_pattern(
|
||||
str(input_path),
|
||||
str(output_path),
|
||||
custom_prompt=prompt,
|
||||
aspect_ratio=str((analysis or {}).get("aspect_ratio") or "1:1"),
|
||||
)
|
||||
if not success or not output_path.exists():
|
||||
if task_id:
|
||||
task_db.update_status(task_id, TaskStatus.FAILED.value)
|
||||
failures.append(f"{idx}. Gemini失败：{message}")
|
||||
continue
|
||||
|
||||
upload_result = await upload_to_tuhui(
|
||||
image_path=str(output_path),
|
||||
title=title,
|
||||
description=requirement_text or prompt[:120],
|
||||
price=AUTO_PROCESS_PRICE,
|
||||
category=AUTO_PROCESS_CATEGORY,
|
||||
tags="AI处理,自动转接",
|
||||
designer_name=designer_name,
|
||||
)
|
||||
if not upload_result.success:
|
||||
if task_id:
|
||||
task_db.update_status(task_id, TaskStatus.FAILED.value)
|
||||
failures.append(f"{idx}. 图绘上传失败：{upload_result.message}")
|
||||
continue
|
||||
|
||||
if task_id:
|
||||
task_db.update_status(task_id, TaskStatus.COMPLETED.value, upload_result.download_url)
|
||||
uploaded_links.append(
|
||||
{
|
||||
"download_url": upload_result.download_url,
|
||||
"source_url": image_url,
|
||||
"work_id": str(upload_result.work_id),
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
if task_id:
|
||||
task_db.update_status(task_id, TaskStatus.FAILED.value)
|
||||
failures.append(f"{idx}. 处理异常：{e}")
|
||||
|
||||
await wecom_bot_service.send_text(
|
||||
self._format_finish_notice(
|
||||
customer_id=customer_id,
|
||||
acc_id=acc_id,
|
||||
designer_name=designer_name,
|
||||
links=uploaded_links,
|
||||
failures=failures,
|
||||
)
|
||||
)
|
||||
|
||||
return {
|
||||
"success": bool(uploaded_links),
|
||||
"uploaded": uploaded_links,
|
||||
"failures": failures,
|
||||
}
|
||||
|
||||
|
||||
auto_image_pipeline_service = AutoImagePipelineService()
|
||||
Reference in New Issue
Block a user