feat: improve context memory and fix auto-draw gemini/upload chain

qingjian_cs/app/auto_draw.py

@@ -1,105 +1,13 @@
 from __future__ import annotations
 
-import asyncio
 import os
-import sys
 import tempfile
 import uuid
-from pathlib import Path
 from typing import Any
 
 import requests
-from dotenv import load_dotenv
 
-from .config import AUTO_DRAW_ENDPOINT, AUTO_DRAW_TIMEOUT_SECONDS
-
-
-def _add_legacy_tw_path() -> None:
-    root = os.getenv("LEGACY_TW_ROOT", r"D:\main\sandbox\tw_terminator").strip()
-    if not root:
-        return
-    p = Path(root)
-    # 先加载 legacy 项目的 .env，确保 service_tuhui_upload 在 import 时拿到正确账号配置
-    legacy_env = p / ".env"
-    if legacy_env.exists():
-        load_dotenv(legacy_env, override=True)
-    if p.exists() and str(p) not in sys.path:
-        sys.path.insert(0, str(p))
-
-
-async def _draw_via_legacy_tw(
-    image_url: str,
-    customer_id: str,
-    requirement: str = "",
-) -> dict[str, Any]:
-    # 优先使用当前项目中拷贝过来的 service_gemini
-    from services.service_gemini import GeminiExtractV2Service  # type: ignore
-
-    # 上传模块暂时仍走 legacy（你后续可替换为新项目本地上传实现）
-    _add_legacy_tw_path()
-    from services.service_tuhui_upload import upload_to_tuhui  # type: ignore
-
-    prompt = requirement.strip() or "按原图做高清修复，保留主体细节，输出清晰可用版本"
-
-    # 1) 下载原图到本地临时文件
-    input_path = os.path.join(tempfile.gettempdir(), f"qjcs_in_{uuid.uuid4().hex}.jpg")
-    output_path = os.path.join(tempfile.gettempdir(), f"qjcs_out_{uuid.uuid4().hex}.jpg")
-    headers = {
-        "User-Agent": (
-            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
-            "AppleWebKit/537.36 (KHTML, like Gecko) "
-            "Chrome/122.0.0.0 Safari/537.36"
-        ),
-        "Referer": "https://www.taobao.com/",
-        "Accept": "image/avif,image/webp,image/apng,image/*,*/*;q=0.8",
-    }
-    resp = requests.get(image_url, headers=headers, timeout=AUTO_DRAW_TIMEOUT_SECONDS)
-    if resp.status_code != 200:
-        return {"ok": False, "error": f"download_http_{resp.status_code}"}
-    with open(input_path, "wb") as f:
-        f.write(resp.content)
-
-    # 2) 直调你原来的 service_gemini 作图 API
-    service = GeminiExtractV2Service()
-    ok_extract, msg_extract, _ = await service.extract_pattern(
-        input_path=input_path,
-        output_path=output_path,
-        custom_prompt=prompt,
-        aspect_ratio="1:1",
-    )
-    if not ok_extract:
-        return {"ok": False, "error": f"extract_failed:{msg_extract}"}
-    if not os.path.exists(output_path):
-        return {"ok": False, "error": "extract_no_output_file"}
-
-    # 3) 上传图绘，返回可外发 URL
-    ok, link, _ = await upload_to_tuhui(
-        output_path,
-        title=f"客户{customer_id[-4:]}-预览图" if customer_id else "预览图",
-        description="AI自动作图预览",
-        price=1,
-    )
-    if not ok:
-        return {"ok": False, "error": str(link)}
-    return {"ok": True, "url": str(link)}
-
-
-def _draw_via_http_endpoint(image_url: str, customer_id: str, requirement: str = "") -> dict[str, Any]:
-    if not AUTO_DRAW_ENDPOINT:
-        return {"ok": False, "error": "AUTO_DRAW_ENDPOINT not configured"}
-    payload = {
-        "image_url": image_url,
-        "customer_id": customer_id,
-        "requirement": requirement,
-    }
-    resp = requests.post(AUTO_DRAW_ENDPOINT, json=payload, timeout=AUTO_DRAW_TIMEOUT_SECONDS)
-    if resp.status_code != 200:
-        return {"ok": False, "error": f"http_{resp.status_code}:{resp.text[:200]}"}
-    data = resp.json() if resp.text else {}
-    url = str(data.get("url", "") or data.get("preview_url", "") or "")
-    if not url:
-        return {"ok": False, "error": "missing_preview_url"}
-    return {"ok": True, "url": url}
+from .config import AUTO_DRAW_TIMEOUT_SECONDS
 
 
 async def auto_draw_preview(
@@ -108,24 +16,68 @@ async def auto_draw_preview(
     requirement: str = "",
 ) -> dict[str, Any]:
     """
-    统一自动作图入口：
-    1) 优先走 tw_terminator 的 service_gemini 直调链路
-    2) 失败时回退 AUTO_DRAW_ENDPOINT
+    统一自动作图入口（直调本地链路）：
+    1) 下载客户图
+    2) 调 Gemini 生成
+    3) 上传图绘，返回可外发 URL
     """
     try:
-        return await _draw_via_legacy_tw(image_url=image_url, customer_id=customer_id, requirement=requirement)
+        from services.service_gemini import GeminiExtractV2Service  # type: ignore
+        from services.service_tuhui_upload import upload_to_tuhui  # type: ignore
     except Exception as e:
-        legacy_error = str(e)
+        return {"ok": False, "error": f"import_failed:{e}"}
+
+    prompt = requirement.strip() or "按原图做高清修复，保留主体细节，输出清晰可用版本"
+    input_path = os.path.join(tempfile.gettempdir(), f"qjcs_in_{uuid.uuid4().hex}.jpg")
+    output_path = os.path.join(tempfile.gettempdir(), f"qjcs_out_{uuid.uuid4().hex}.jpg")
 
     try:
-        data = await asyncio.to_thread(
-            _draw_via_http_endpoint,
-            image_url,
-            customer_id,
-            requirement,
+        headers = {
+            "User-Agent": (
+                "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
+                "AppleWebKit/537.36 (KHTML, like Gecko) "
+                "Chrome/122.0.0.0 Safari/537.36"
+            ),
+            "Referer": "https://www.taobao.com/",
+            "Accept": "image/avif,image/webp,image/apng,image/*,*/*;q=0.8",
+        }
+        resp = requests.get(image_url, headers=headers, timeout=AUTO_DRAW_TIMEOUT_SECONDS)
+        if resp.status_code != 200:
+            return {"ok": False, "error": f"download_http_{resp.status_code}"}
+        with open(input_path, "wb") as f:
+            f.write(resp.content)
+
+        service = GeminiExtractV2Service()
+        ok_extract, msg_extract, _ = await service.extract_pattern(
+            input_path=input_path,
+            output_path=output_path,
+            custom_prompt=prompt,
+            aspect_ratio="1:1",
         )
-        if data.get("ok"):
-            return data
-        return {"ok": False, "error": f"legacy:{legacy_error}; endpoint:{data.get('error','unknown')}"}
+        if not ok_extract:
+            return {"ok": False, "error": f"extract_failed:{msg_extract}"}
+        if not os.path.exists(output_path):
+            return {"ok": False, "error": "extract_no_output_file"}
+
+        ok_upload, link, _ = await upload_to_tuhui(
+            output_path,
+            title=f"客户{customer_id[-4:]}-预览图" if customer_id else "预览图",
+            description="AI自动作图预览",
+            price=1,
+        )
+        if not ok_upload:
+            return {"ok": False, "error": f"upload_failed:{link}"}
+        return {"ok": True, "url": str(link)}
     except Exception as e:
-        return {"ok": False, "error": f"legacy:{legacy_error}; endpoint:{e}"}
+        return {"ok": False, "error": str(e)}
+    finally:
+        try:
+            if os.path.exists(input_path):
+                os.remove(input_path)
+        except Exception:
+            pass
+        try:
+            if os.path.exists(output_path):
+                os.remove(output_path)
+        except Exception:
+            pass