"""豆包大模型 API 客户端(共享层). API 和 Worker 两边共用。基于火山引擎方舟平台的 OpenAI 兼容接口。 使用方式: from packages.shared.ai_client import get_doubao_client client = get_doubao_client() if client.is_available: result = client.chat_completion(messages=[...]) """ from __future__ import annotations import logging import os import time import uuid from typing import Any, Optional import httpx from packages.shared.config import get_shared_settings logger = logging.getLogger(__name__) class DoubaoClient: """豆包大模型 API 客户端. 封装 OpenAI 兼容的 Chat Completion 接口,支持自动重试。 未配置 API Key 时 is_available 为 False,调用方应降级处理。 """ def __init__(self) -> None: settings = get_shared_settings() self.api_key: str = settings.doubao_api_key self.model: str = settings.doubao_model self.base_url: str = settings.doubao_base_url.rstrip("/") self.timeout: int = settings.doubao_timeout self.max_retries: int = settings.doubao_max_retries self.vision_model: str = settings.doubao_vision_model def embed_text(self, text: str, timeout: int | None = None) -> list[float] | None: """调用豆包文本 Embedding API,返回浮点向量;失败返回 None。""" if not self.is_available or not text or not text.strip(): return None url = f"{self.base_url}/embeddings" headers = { "Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json", } payload: dict[str, Any] = { "model": getattr(self, "embedding_model", None) or "doubao-embedding-large-text-240915", "input": text.strip(), "encoding_format": "float", } req_timeout = timeout or self.timeout last_error: Exception | None = None for attempt in range(self.max_retries + 1): try: resp = httpx.post(url, headers=headers, json=payload, timeout=req_timeout) resp.raise_for_status() data = resp.json() emb_list = data.get("data") or [] if emb_list and isinstance(emb_list, list): vec = emb_list[0].get("embedding") if isinstance(vec, list) and vec: return [float(x) for x in vec] logger.warning("embedding 返回结构异常: %s", str(data)[:200]) return None except Exception as e: last_error = e if attempt < self.max_retries: wait = 0.5 * (2**attempt) logger.warning( "豆包 Embedding 调用失败,%.1fs 后重试 (%d/%d): %s", wait, attempt + 1, self.max_retries + 1, e ) time.sleep(wait) logger.error("豆包 Embedding 调用最终失败: %s", last_error) return None @property def is_available(self) -> bool: """是否可用(配置了 API Key).""" return bool(self.api_key) def chat_completion( self, messages: list[dict[str, str]], temperature: float = 0.7, max_tokens: int = 1024, ) -> Optional[str]: """调用 Chat Completion 接口. Args: messages: 对话消息列表,[{"role": "user"/"system"/"assistant", "content": "..."}] temperature: 采样温度,0-2,默认0.7 max_tokens: 最大生成token数,默认1024 Returns: 模型返回的文本内容,失败返回 None """ if not self.is_available: return None url = f"{self.base_url}/chat/completions" headers = { "Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json", } payload: dict[str, Any] = { "model": self.model, "messages": messages, "temperature": temperature, "max_tokens": max_tokens, } last_error: Optional[Exception] = None for attempt in range(self.max_retries + 1): try: response = httpx.post( url, headers=headers, json=payload, timeout=self.timeout, ) response.raise_for_status() data = response.json() content = data["choices"][0]["message"]["content"] return content.strip() except Exception as e: last_error = e if attempt < self.max_retries: wait = 0.5 * (2**attempt) logger.warning( "豆包API调用失败,%.1fs后重试 (第%d/%d次): %s", wait, attempt + 1, self.max_retries + 1, e, ) time.sleep(wait) logger.error("豆包API调用最终失败: %s", last_error) return None def vision_completion( self, messages: list[dict], images: list[str] | None = None, max_tokens: int = 2048, temperature: float = 0.3, timeout: int | None = None, ) -> Optional[str]: """调用豆包视觉理解 API(OpenAI 兼容多模态格式). 将 images 附加到最后一条 user message 的 content 中, 使用 vision_model(默认 doubao-1-5-vision-pro-250915)。 Args: messages: 对话消息列表。最后一条 user message 会被注入图片内容。 images: 图片列表,支持 base64 data URI 或 HTTP(S) URL。 max_tokens: 最大生成 token 数,默认 2048。 temperature: 采样温度,默认 0.3(视觉任务偏低更稳定)。 timeout: 单次请求超时秒数,不传则使用默认 self.timeout。 Returns: 模型返回的文本内容,失败返回 None。 """ if not self.is_available: return None # 构造多模态 content:先追加文本,再追加图片 vision_messages = [] for msg in messages: vision_messages.append(dict(msg)) # 将图片注入最后一条 user message if images and vision_messages: # 找到最后一条 user message for i in range(len(vision_messages) - 1, -1, -1): if vision_messages[i].get("role") == "user": text_content = vision_messages[i].get("content", "") multi_content: list[dict[str, Any]] = [] if text_content: multi_content.append({"type": "text", "text": text_content}) for img in images: if img.startswith("data:") or img.startswith("http://") or img.startswith("https://"): multi_content.append({"type": "image_url", "image_url": {"url": img}}) else: # 当作 base64 编码 multi_content.append( {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{img}"}} ) vision_messages[i]["content"] = multi_content break url = f"{self.base_url}/chat/completions" headers = { "Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json", } payload: dict[str, Any] = { "model": self.vision_model, "messages": vision_messages, "temperature": temperature, "max_tokens": max_tokens, } req_timeout = timeout or self.timeout last_error: Optional[Exception] = None for attempt in range(self.max_retries + 1): try: response = httpx.post( url, headers=headers, json=payload, timeout=req_timeout, ) response.raise_for_status() data = response.json() content = data["choices"][0]["message"]["content"] return content.strip() except Exception as e: last_error = e if attempt < self.max_retries: wait = 0.5 * (2**attempt) logger.warning( "豆包视觉API调用失败,%.1fs后重试 (第%d/%d次): %s", wait, attempt + 1, self.max_retries + 1, e, ) time.sleep(wait) logger.error("豆包视觉API调用最终失败: %s", last_error) return None # ── 视频生成(Seedance 2.5,异步任务)──────────────────────────── def video_generation( self, prompt: str, *, image_url: str | None = None, duration: int = 5, ratio: str = "9:16", resolution: str = "720p", generate_audio: bool = False, watermark: bool = False, output_dir: str | None = None, ) -> str | None: """调用 Seedance 2.5 文生/图生视频(异步任务→轮询→下载),返回本地 MP4 路径;失败返回 None。 Args: prompt: 文本提示词 image_url: 首帧参考图 URL(可选,提供则走图生视频) duration: 视频时长 2~30 秒,默认 5 ratio: 宽高比 16:9/9:16/1:1/4:3/3:4/21:9/adaptive resolution: 480p/720p/1080p generate_audio: 是否生成模型自带音效(默认 False,我们自己混 TTS) watermark: 是否加水印 output_dir: 下载目录,默认 /tmp Returns: 本地 MP4 文件路径,失败返回 None。 """ if not self.is_available: return None if not prompt or not prompt.strip(): return None settings = get_shared_settings() poll_interval = getattr(settings, "doubao_video_poll_interval", 10) or 10 total_timeout = getattr(settings, "doubao_video_timeout", 600) or 600 video_model = getattr(settings, "doubao_video_model", None) or "doubao-seedance-2-5-260628" content: list[dict[str, Any]] = [{"type": "text", "text": prompt.strip()}] if image_url: content.append({"type": "image_url", "image_url": {"url": image_url}}) create_payload: dict[str, Any] = { "model": video_model, "content": content, "generate_audio": generate_audio, "ratio": ratio, "duration": int(duration), "resolution": resolution, "watermark": watermark, } headers = { "Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json", } create_url = f"{self.base_url}/contents/generations/tasks" # 1) 创建任务(带重试) task_id: str | None = None last_error: Exception | None = None for attempt in range(self.max_retries + 1): try: resp = httpx.post(create_url, headers=headers, json=create_payload, timeout=self.timeout) resp.raise_for_status() data = resp.json() task_id = data.get("id") if task_id: break last_error = RuntimeError(f"create task returned no id: {str(data)[:200]}") except Exception as e: last_error = e if attempt < self.max_retries: wait = 0.5 * (2**attempt) logger.warning( "Seedance 创建任务失败,%.1fs 后重试 (%d/%d): %s", wait, attempt + 1, self.max_retries + 1, e ) time.sleep(wait) if not task_id: logger.error("Seedance 创建任务最终失败: %s", last_error) return None logger.info("Seedance 任务已创建: task_id=%s model=%s duration=%ds", task_id, video_model, duration) # 2) 轮询状态 poll_url = f"{create_url}/{task_id}" deadline = time.time() + total_timeout video_url: str | None = None last_status: str = "queued" while time.time() < deadline: try: resp = httpx.get(poll_url, headers=headers, timeout=self.timeout) resp.raise_for_status() data = resp.json() status = data.get("status", "") last_status = status if status == "succeeded": content_obj = data.get("content") or {} video_url = content_obj.get("video_url") if video_url: break last_error = RuntimeError(f"task succeeded but no video_url: {str(data)[:300]}") break if status == "failed": err = data.get("error") or {} last_error = RuntimeError(f"task failed: {err.get('code','')} {err.get('message','')}") break if status in ("expired", "cancelled"): last_error = RuntimeError(f"task {status}") break # queued / running: 继续轮询 except Exception as e: last_error = e logger.debug("Seedance 轮询异常: %s", e) time.sleep(poll_interval) if not video_url: logger.error("Seedance 任务未成功: task_id=%s status=%s err=%s", task_id, last_status, last_error) return None # 3) 下载到本地 try: out_dir = output_dir or "/tmp" os.makedirs(out_dir, exist_ok=True) local_path = f"{out_dir}/seedance_{task_id}_{uuid.uuid4().hex[:8]}.mp4" with httpx.stream("GET", video_url, timeout=300) as r: r.raise_for_status() with open(local_path, "wb") as f: for chunk in r.iter_bytes(chunk_size=1024 * 256): if chunk: f.write(chunk) logger.info("Seedance 视频下载完成: %s (%d bytes)", local_path, os.path.getsize(local_path)) return local_path except Exception as e: logger.error("Seedance 视频下载失败: %s", e) return None # ── 单例 ───────────────────────────────────────────────────────────────────── _client: Optional[DoubaoClient] = None def get_doubao_client() -> DoubaoClient: """获取豆包客户端单例.""" global _client if _client is None: _client = DoubaoClient() return _client