7f419f4c8b
- 015-017迁移：image_candidate 文案复审/AI视觉分/重生标记 - constants C7素人感约束(反电商摆拍对齐真实笔记)+C3叠字口子 - celery visibility_timeout=2h 防长任务被误判重投重复烧钱(task75教训) - image_scorer 评图分(只筛选+展示,真实信号才进飞轮权重) - storyboard/image_gen/pipeline_io 生图存量 - task_actions/tasks/task_service/flywheel 审核回路+飞轮存量 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
253 lines
13 KiB
Python
253 lines
13 KiB
Python
"""
|
||
gemini_factory.py — 每任务构建独立的 AI client 实例
|
||
解决全局单例问题（扒 banana gemini_service.py __init__，改造为每任务局部实例）
|
||
|
||
铁律（基石B）：
|
||
- 调用方只传 task_id，不传 key
|
||
- 本模块在 worker 内部查库 → Fernet 解密 → 构建 client
|
||
- 解密结果只活在局部变量，函数返回后即销毁
|
||
- 绝不打印 / 记录 / 传递明文 key
|
||
"""
|
||
from __future__ import annotations
|
||
import asyncio
|
||
import logging
|
||
import os
|
||
from dataclasses import dataclass, field
|
||
from typing import Any
|
||
|
||
import httpx
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@dataclass
|
||
class AIClients:
|
||
"""
|
||
一个任务专用的 AI client 集合。
|
||
worker 在任务开始时构建，任务结束后释放（局部变量，不存 Redis/DB）。
|
||
"""
|
||
# httpx AsyncClient 懒加载 + 按事件循环缓存：Celery 每任务多次 asyncio.run，
|
||
# 持久 client 会绑死到首个已关闭的 loop → "Event loop is closed"。
|
||
# 故只存 token/base，按当前运行 loop 缓存 client，loop 变了就重建。
|
||
_gpt_token: str | None = field(default=None, repr=False)
|
||
_gpt_base: str | None = field(default=None, repr=False)
|
||
_gpt_client: httpx.AsyncClient | None = field(default=None, repr=False)
|
||
_gpt_client_loop_id: int | None = field(default=None, repr=False)
|
||
# 备用 OpenAI 兼容中转站（codeproxy）：apiports 503 时真正切过去（独立 base+key）
|
||
_alt_token: str | None = field(default=None, repr=False)
|
||
_alt_base: str | None = field(default=None, repr=False)
|
||
# 多 base/token 的 client 池：key=(base,token的id)，按 loop 失效重建
|
||
_client_pool: dict = field(default_factory=dict, repr=False)
|
||
_pool_loop_id: int | None = field(default=None, repr=False)
|
||
_gemini_key: str | None = field(default=None, repr=False) # 局部变量不打印
|
||
_model_image: str = "gpt-image-2"
|
||
_model_text: str = "claude-opus-4-8" # 最强档(倩倩姐红线):Claude系一律4.8,绝不降级
|
||
|
||
def _client(self) -> httpx.AsyncClient:
|
||
"""主通道(apiports) client，按当前事件循环缓存"""
|
||
return self._client_for(self._gpt_base, self._gpt_token)
|
||
|
||
def _client_for(self, base: str | None, token: str | None) -> httpx.AsyncClient:
|
||
"""按 (base, token) 返回可用 client；loop 变化则整池重建（避免跨 loop 复用）"""
|
||
if not token:
|
||
raise RuntimeError("GPT client 未初始化（缺 token）")
|
||
loop_id = id(asyncio.get_running_loop())
|
||
if self._pool_loop_id != loop_id:
|
||
self._client_pool = {}
|
||
self._pool_loop_id = loop_id
|
||
ck = (base or "", token)
|
||
if ck not in self._client_pool:
|
||
self._client_pool[ck] = httpx.AsyncClient(
|
||
headers={"Authorization": f"Bearer {token}"},
|
||
base_url=base or None,
|
||
timeout=120.0,
|
||
)
|
||
return self._client_pool[ck]
|
||
|
||
# ── ImageClient 协议实现（供 image_gen.py 使用）────────
|
||
|
||
def _gpt_target(self, provider: str | None) -> tuple[str, str | None, httpx.AsyncClient]:
|
||
"""按 provider 选 (base, token, client)；codeproxy 走备用站独立 base+key"""
|
||
if provider == "codeproxy" and self._alt_token:
|
||
base = (self._alt_base or os.environ.get("CODEPROXY_BASE_URL") or "").rstrip("/")
|
||
return base, self._alt_token, self._client_for(base, self._alt_token)
|
||
base = (os.environ.get("IMAGE_API_BASE") or os.environ.get("APIPORTS_BASE_URL") or "").rstrip("/")
|
||
return base, self._gpt_token, self._client_for(self._gpt_base, self._gpt_token)
|
||
|
||
async def gpt_edits(self, prompt: str, reference_images: list[bytes], size: str, provider: str | None = None) -> bytes:
|
||
"""GPT edits endpoint（带产品参考图，禁纯文生图）"""
|
||
import io
|
||
files: list[tuple] = [("prompt", (None, prompt))]
|
||
for i, img in enumerate(reference_images):
|
||
files.append(("image[]", (f"ref_{i}.png", io.BytesIO(img), "image/png")))
|
||
files.append(("size", (None, size)))
|
||
files.append(("model", (None, self._model_image)))
|
||
base, _, client = self._gpt_target(provider)
|
||
resp = await client.post(f"{base}/images/edits", files=files, timeout=120.0)
|
||
resp.raise_for_status()
|
||
return _extract_image_bytes(resp.json())
|
||
|
||
async def gpt_generate(self, prompt: str, size: str, provider: str | None = None) -> bytes:
|
||
"""GPT 纯文生图（仅 ALLOW_TEXT_ONLY_IMAGE=true 时用）"""
|
||
base, _, client = self._gpt_target(provider)
|
||
payload = {"model": self._model_image, "prompt": prompt, "n": 1, "size": size}
|
||
resp = await client.post(f"{base}/images/generations", json=payload, timeout=120.0)
|
||
resp.raise_for_status()
|
||
return _extract_image_bytes(resp.json())
|
||
|
||
async def gemini_generate(self, prompt: str, reference_images: list[bytes], model: str) -> bytes:
|
||
"""Gemini 生图（备用通道）"""
|
||
if not self._gemini_key:
|
||
raise RuntimeError("Gemini key 未初始化")
|
||
gemini_base = os.environ.get("GEMINI_API_URL", "https://generativelanguage.googleapis.com/v1beta")
|
||
url = f"{gemini_base}/models/{model}:generateContent?key={self._gemini_key}"
|
||
parts: list[dict] = [{"text": prompt}]
|
||
for img in reference_images:
|
||
import base64
|
||
parts.append({"inline_data": {"mime_type": "image/png", "data": base64.b64encode(img).decode()}})
|
||
payload = {"contents": [{"role": "user", "parts": parts}], "generationConfig": {"responseModalities": ["IMAGE", "TEXT"]}}
|
||
async with httpx.AsyncClient() as client:
|
||
resp = await client.post(url, json=payload, timeout=120.0)
|
||
resp.raise_for_status()
|
||
return _extract_gemini_image(resp.json())
|
||
|
||
async def _chat_post_failover(self, payload: dict, timeout: float) -> dict:
|
||
"""
|
||
chat/completions 发送器，带 codeproxy 回落。
|
||
主通道(apiports, claude-opus-4-8)若 5xx/超时/连接错，自动切 codeproxy 重试一次。
|
||
⚠ codeproxy 账号池只支持 gpt 系(gpt-5.5)，无 claude，故回落时模型换成 gpt-5.5。
|
||
这是强档↔强档切换：红线"Claude系4.8/GPT系5.5"本就是两个平级最强档、互为兜底，
|
||
非降级(降级=落 sonnet/弱档)。codeproxy 回落档由 CODEPROXY_CHAT_MODEL 配置(默认 gpt-5.5)。
|
||
codeproxy 未配置(_alt_token 为空)时不回落，原样抛错。
|
||
"""
|
||
main_base = (os.environ.get("IMAGE_API_BASE") or os.environ.get("APIPORTS_BASE_URL") or "").rstrip("/")
|
||
try:
|
||
resp = await self._client().post(f"{main_base}/chat/completions", json=payload, timeout=timeout)
|
||
resp.raise_for_status()
|
||
return resp.json()
|
||
except (httpx.HTTPStatusError, httpx.TransportError, httpx.TimeoutException) as exc:
|
||
# 仅 5xx(服务端过载)或网络层错误才回落；4xx(参数/鉴权)回落也没用，直接抛。
|
||
status = getattr(getattr(exc, "response", None), "status_code", None)
|
||
retryable = status is None or status >= 500
|
||
if not (retryable and self._alt_token):
|
||
raise
|
||
alt_base = (self._alt_base or os.environ.get("CODEPROXY_BASE_URL") or "").rstrip("/")
|
||
alt_payload = dict(payload)
|
||
alt_payload["model"] = os.environ.get("CODEPROXY_CHAT_MODEL", "gpt-5.5")
|
||
logger.warning("apiports chat 失败(%s)，回落 codeproxy 用 %s 重试（强档兜底，非降级）",
|
||
status or type(exc).__name__, alt_payload["model"])
|
||
client = self._client_for(alt_base, self._alt_token)
|
||
resp = await client.post(f"{alt_base}/chat/completions", json=alt_payload, timeout=timeout)
|
||
resp.raise_for_status()
|
||
return resp.json()
|
||
|
||
async def chat_complete(self, messages: list[dict], model: str | None = None, max_tokens: int = 4096, temperature: float = 0.75) -> str:
|
||
"""文字生成（文案生成用）。apiports 503 时自动回落 codeproxy。"""
|
||
payload = {"model": model or self._model_text, "messages": messages, "max_tokens": max_tokens, "temperature": temperature}
|
||
# 单批≤4条文案正常 40-55s 返回；apiports 网关 ~60s 上限。客户端超时设 75s：
|
||
# 略高于网关上限即可，过长（如180s）会在 apiports 卡顿时干等，拖慢整体。
|
||
timeout = float(os.environ.get("TEXT_LLM_TIMEOUT", "75"))
|
||
data = await self._chat_post_failover(payload, timeout)
|
||
return data["choices"][0]["message"]["content"] or ""
|
||
|
||
async def gpt_vision_analyze(self, prompt: str, images: list[bytes], model: str | None = None) -> str:
|
||
"""
|
||
GPT/Claude vision 读产品图，返回 JSON 字符串。
|
||
messages content 混合 text + image_url(base64)，OpenAI vision 格式。
|
||
model 默认最强档（claude-opus-4-8），绝不偷降级。
|
||
最多传 4 张图，避免超 token。
|
||
"""
|
||
import base64
|
||
content: list[dict] = [{"type": "text", "text": prompt}]
|
||
for img in images[:4]:
|
||
b64 = base64.b64encode(img).decode()
|
||
content.append({
|
||
"type": "image_url",
|
||
"image_url": {"url": f"data:image/jpeg;base64,{b64}"},
|
||
})
|
||
used_model = model or os.environ.get("MODEL_TEXT", "claude-opus-4-8")
|
||
messages = [{"role": "user", "content": content}]
|
||
payload = {
|
||
"model": used_model,
|
||
"messages": messages,
|
||
"max_tokens": 2048,
|
||
"temperature": 0.2,
|
||
}
|
||
# 评图也走 codeproxy 回落：apiports 503 时切备用站，模型档不变（守红线）。
|
||
data = await self._chat_post_failover(payload, 90.0)
|
||
return data["choices"][0]["message"]["content"] or ""
|
||
|
||
# duck-type: text_variants._call_llm 用的属性
|
||
@property
|
||
def _model(self) -> str:
|
||
return self._model_text
|
||
|
||
async def aclose(self) -> None:
|
||
# client 可能绑在已关闭的 loop（Celery 多次 asyncio.run），aclose 也可能报
|
||
# "Event loop is closed"，吞掉即可——进程级连接随 loop 关闭自然释放。
|
||
for c in list(self._client_pool.values()):
|
||
try:
|
||
await c.aclose()
|
||
except Exception:
|
||
pass
|
||
self._client_pool = {}
|
||
self._pool_loop_id = None
|
||
self._gpt_client = None
|
||
self._gpt_client_loop_id = None
|
||
|
||
|
||
def build_ai_clients(plain_key: str, gemini_key: str | None = None) -> AIClients:
|
||
"""
|
||
用解密后的明文 key 构建 AIClients。
|
||
只在 Celery worker 函数体内调用，plain_key 是局部变量。
|
||
httpx client 不在此预创建（避免绑死到调用方 loop），首次 await 时按 loop 懒建。
|
||
调用完成后 caller 负责 await clients.aclose()。
|
||
"""
|
||
gpt_base = (
|
||
os.environ.get("IMAGE_API_BASE") # 旧变量名
|
||
or os.environ.get("APIPORTS_BASE_URL") # .env 实际变量名
|
||
or ""
|
||
).rstrip("/")
|
||
# 备用站 codeproxy：系统级 key（非用户录入），apiports 503 时切过去保生图成功
|
||
alt_base = (os.environ.get("CODEPROXY_BASE_URL") or "").rstrip("/")
|
||
alt_token = os.environ.get("CODEPROXY_KEY") or None
|
||
return AIClients(
|
||
_gpt_token=plain_key,
|
||
_gpt_base=gpt_base or None,
|
||
_alt_base=alt_base or None,
|
||
_alt_token=alt_token,
|
||
_gemini_key=gemini_key,
|
||
_model_image=os.environ.get("IMAGE_MODEL") or os.environ.get("MODEL_IMAGE", "gpt-image-2"),
|
||
_model_text=os.environ.get("MODEL_TEXT", "claude-opus-4-8"),
|
||
)
|
||
|
||
|
||
# ── 图片响应解析工具 ─────────────────────────────────────
|
||
|
||
def _extract_image_bytes(resp_json: dict) -> bytes:
|
||
"""从 OpenAI images API 响应提取图片 bytes（b64 或 url）"""
|
||
import base64
|
||
data = resp_json.get("data", [{}])
|
||
if not data:
|
||
raise ValueError("图片 API 返回空 data")
|
||
item = data[0]
|
||
if "b64_json" in item:
|
||
return base64.b64decode(item["b64_json"])
|
||
if "url" in item:
|
||
resp = httpx.get(item["url"], timeout=30.0)
|
||
resp.raise_for_status()
|
||
return resp.content
|
||
raise ValueError(f"无法解析图片响应：{list(item.keys())}")
|
||
|
||
|
||
def _extract_gemini_image(resp_json: dict) -> bytes:
|
||
"""从 Gemini generateContent 响应提取图片 bytes"""
|
||
import base64
|
||
candidates = resp_json.get("candidates", [])
|
||
for cand in candidates:
|
||
parts = cand.get("content", {}).get("parts", [])
|
||
for part in parts:
|
||
if "inlineData" in part:
|
||
return base64.b64decode(part["inlineData"]["data"])
|
||
raise ValueError("Gemini 响应中未找到图片数据")
|