152 lines
5.7 KiB
Python
152 lines
5.7 KiB
Python
"""媒体类工具共享原语(从 seedance/seedream/gpt_image/look_at_image/read_document
|
|
五处同构析出,2026-07-23)。
|
|
|
|
收进来的判据:**逐字或参数化后逐字**的重复 —— 每日配额闸、`<ts>-<rand6>` 落盘命名、
|
|
meta.json 写入、记账 try/except 兜底、Ark 超时透明重试、chat 答案提取、响应递归找 URL。
|
|
各工具的首行 banner 格式 / 请求 body 组装 / seedance 轮询等有意各自不同,**不**收。
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import secrets
|
|
import time
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Optional
|
|
from uuid import UUID
|
|
|
|
from core.ark_client import ArkClient, ArkConfig, ArkError, ArkTimeoutError
|
|
from core.storage.usage import check_daily_quota
|
|
|
|
|
|
def quota_gate(user_id: UUID, *, kind: str, limit: int, what: str, noun: str) -> Optional[str]:
|
|
"""每账号每日配额闸(yaml quotas.*_per_day)。超额返回 [Error] 文案,未超返 None。
|
|
|
|
失败 retry 不计 —— record_*_usage 只在成功后才落库。tool 返串会进 LLM 上下文,
|
|
模型据此向用户解释,所以**只暴露用户该看的部分**(已用/上限 + 重置时间),
|
|
内部 yaml 路径不进对话。limit <= 0 = 不限。
|
|
"""
|
|
if limit <= 0:
|
|
return None
|
|
used, over = check_daily_quota(user_id=user_id, kind=kind, limit=limit)
|
|
if over:
|
|
return (
|
|
f"[Error] 已达每日{what}上限({used}/{limit} {noun}),"
|
|
f"次日 00:00 重置。"
|
|
)
|
|
return None
|
|
|
|
|
|
def stamped_path(directory: Path, ext: str) -> Path:
|
|
"""产物落盘命名:`<dir>/<YYYYMMDD-HHMMSS>-<rand6><ext>`,顺手建目录。"""
|
|
ts = datetime.now().strftime("%Y%m%d-%H%M%S")
|
|
short = secrets.token_hex(3)
|
|
directory.mkdir(parents=True, exist_ok=True)
|
|
return directory / f"{ts}-{short}{ext}"
|
|
|
|
|
|
def write_meta(dest: Path, meta: dict) -> None:
|
|
"""产物旁写同名 `.meta.json`(prompt/model/cost 等溯源信息)。"""
|
|
dest.with_suffix(".meta.json").write_text(
|
|
json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8"
|
|
)
|
|
|
|
|
|
def record_usage_safe(tool_name: str, fn: Callable, **kwargs) -> Any:
|
|
"""记账兜底:失败不阻塞 tool 返回(print 留痕),返回 fn 结果(失败返 None)。
|
|
|
|
tool 层没 sink 引用,print 是现阶段的兜底;后续可改成 sink 注入。
|
|
"""
|
|
try:
|
|
return fn(**kwargs)
|
|
except Exception as e:
|
|
print(f"[{tool_name}] {getattr(fn, '__name__', 'record_usage')} failed: "
|
|
f"{type(e).__name__}: {e}", flush=True)
|
|
return None
|
|
|
|
|
|
def ark_chat_with_retry(
|
|
ark_cfg: ArkConfig,
|
|
endpoint: str,
|
|
body: dict,
|
|
*,
|
|
timeout_s: float,
|
|
retries: int,
|
|
tool_name: str,
|
|
) -> tuple[Optional[dict], str]:
|
|
"""带超时透明重试的 Ark 单次 chat 调用。返回 (resp, err):成功 (dict, "");失败 (None, "[Error] ...")。
|
|
|
|
重试只吃 ArkTimeoutError(超时/网络抖动),tool 内消化掉不把 [Error] 抛给主模型 ——
|
|
否则主模型会重发整个 tool call(base64 重传、输入 token 再付一次)。HTTP 业务错误不重试。
|
|
退避 2**attempt 秒;retries=0 即只调一次。
|
|
"""
|
|
max_attempts = retries + 1
|
|
for attempt in range(max_attempts):
|
|
try:
|
|
with ArkClient(ark_cfg, timeout_s=timeout_s) as client:
|
|
return client.post_json(endpoint, body, timeout_s=timeout_s), ""
|
|
except ArkTimeoutError as e:
|
|
if attempt == max_attempts - 1:
|
|
return None, f"[Error] {tool_name} API: {e}(已重试 {attempt} 次仍超时)"
|
|
print(
|
|
f"[{tool_name}] timeout, retrying ({attempt + 1}/{max_attempts - 1}): {e}",
|
|
flush=True,
|
|
)
|
|
time.sleep(2 ** attempt)
|
|
except ArkError as e:
|
|
return None, f"[Error] {tool_name} API: {e}"
|
|
return None, f"[Error] {tool_name} API: unreachable" # pragma: no cover
|
|
|
|
|
|
def extract_chat_answer(resp: Optional[dict]) -> tuple[str, bool]:
|
|
"""OpenAI 兼容 chat 响应取文本 + 是否被输出上限截断(finish_reason=length)。
|
|
|
|
content 可能是 str,也可能是 list[{type:text,text:...}](多模态返回形态),都兜住。
|
|
"""
|
|
if not isinstance(resp, dict):
|
|
return "", False
|
|
choices = resp.get("choices")
|
|
if not (isinstance(choices, list) and choices and isinstance(choices[0], dict)):
|
|
return "", False
|
|
truncated = choices[0].get("finish_reason") == "length"
|
|
msg = choices[0].get("message")
|
|
if not isinstance(msg, dict):
|
|
return "", truncated
|
|
content = msg.get("content")
|
|
if isinstance(content, str):
|
|
return content.strip(), truncated
|
|
if isinstance(content, list):
|
|
parts = [
|
|
c.get("text", "")
|
|
for c in content
|
|
if isinstance(c, dict) and c.get("type") == "text"
|
|
]
|
|
return "\n".join(p for p in parts if p).strip(), truncated
|
|
return "", truncated
|
|
|
|
|
|
def find_first_url(
|
|
obj: Any,
|
|
*,
|
|
keys: tuple[str, ...],
|
|
accept: Optional[Callable[[str, str], bool]] = None,
|
|
) -> str:
|
|
"""递归搜响应里第一个 http URL(key ∈ keys)。accept(key, url) 可再过滤(如
|
|
seedance 只认 video_url 或视频扩展名);None = 任意 http 值都收。找不到返 ""。"""
|
|
def _walk(o: Any) -> Optional[str]:
|
|
if isinstance(o, dict):
|
|
for k, v in o.items():
|
|
if k in keys and isinstance(v, str) and v.startswith("http"):
|
|
if accept is None or accept(k, v):
|
|
return v
|
|
r = _walk(v)
|
|
if r:
|
|
return r
|
|
elif isinstance(o, list):
|
|
for x in o:
|
|
r = _walk(x)
|
|
if r:
|
|
return r
|
|
return None
|
|
return _walk(obj) or ""
|