From 8038fb7491f578e932a8b41ea0a59a13ebcd28f1 Mon Sep 17 00:00:00 2001 From: caoqianming Date: Thu, 23 Jul 2026 13:10:32 +0800 Subject: [PATCH] =?UTF-8?q?refactor(tools):=20=E5=AA=92=E4=BD=93=E5=B7=A5?= =?UTF-8?q?=E5=85=B7=E5=90=8C=E6=9E=84=E6=9E=90=E5=87=BA=20tools/media=5Fc?= =?UTF-8?q?ommon.py=20+=2017=20=E4=B8=AA=E5=8D=95=E6=B5=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 审查 P1#6:seedance/seedream/gpt_image/look_at_image/read_document 五处 同构(逐字或参数化后逐字)收敛为共享原语: - quota_gate(每日配额闸,文案逐字保持)/ stamped_path(- 落盘 命名+建目录)/ write_meta(.meta.json)/ record_usage_safe(记账失败不 阻塞)/ ark_chat_with_retry(超时透明重试,业务错误不重试)/ extract_chat_answer(chat 文本+截断标志)/ find_first_url(递归找 URL, accept 谓词覆盖 seedance 的 video_url 优先语义) - 五工具各自的 banner 格式 / body 组装 / seedance 轮询刻意保持不动 - tests/test_media_common.py:17 个纯函数/打桩单测(配额文案、重试次序、 业务错误不重试、视频 URL 谓词等) 318 测试全过;真实配置冒烟 32 工具挂载不变、schema 完整。 Co-Authored-By: Claude Fable 5 --- tests/test_media_common.py | 168 +++++++++++++++++++++++++++++++++++++ tools/gpt_image.py | 49 +++++------ tools/look_at_image.py | 88 ++++++------------- tools/media_common.py | 151 +++++++++++++++++++++++++++++++++ tools/read_document.py | 82 ++++++------------ tools/seedance.py | 106 ++++++++++------------- tools/seedream.py | 79 +++++++---------- 7 files changed, 460 insertions(+), 263 deletions(-) create mode 100644 tests/test_media_common.py create mode 100644 tools/media_common.py diff --git a/tests/test_media_common.py b/tests/test_media_common.py new file mode 100644 index 0000000..8c7b9e5 --- /dev/null +++ b/tests/test_media_common.py @@ -0,0 +1,168 @@ +"""tools/media_common.py 单测 —— 五个媒体工具同构析出后的共享原语。 + +全部纯函数/打桩,不碰网络与 DB(quota_gate 的 check_daily_quota、 +ark_chat_with_retry 的 ArkClient 均 monkeypatch)。 +""" +from __future__ import annotations + +import re +import json +import unittest +import uuid +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest import mock + +from core.ark_client import ArkError, ArkTimeoutError +from tools import media_common as mc + + +class StampedPathTests(unittest.TestCase): + def test_creates_dir_and_names_with_ts_rand(self): + with TemporaryDirectory() as td: + d = Path(td) / "figures" / "sub" + p = mc.stamped_path(d, ".png") + self.assertTrue(d.is_dir()) # 顺手建目录 + self.assertEqual(p.parent, d) + self.assertRegex(p.name, r"^\d{8}-\d{6}-[0-9a-f]{6}\.png$") + + def test_write_meta_alongside(self): + with TemporaryDirectory() as td: + dest = Path(td) / "a.mp4" + mc.write_meta(dest, {"prompt": "早安", "cost_cny": 0.22}) + meta = json.loads((Path(td) / "a.meta.json").read_text(encoding="utf-8")) + self.assertEqual(meta["prompt"], "早安") + self.assertEqual(meta["cost_cny"], 0.22) + + +class QuotaGateTests(unittest.TestCase): + def test_no_limit_short_circuits(self): + # limit<=0 不查 DB(不打桩也不该被调用) + with mock.patch.object(mc, "check_daily_quota", side_effect=AssertionError("不该调")): + self.assertIsNone(mc.quota_gate(uuid.uuid4(), kind="image", limit=0, what="图片生成", noun="张")) + + def test_over_and_under(self): + uid = uuid.uuid4() + with mock.patch.object(mc, "check_daily_quota", return_value=(30, True)): + msg = mc.quota_gate(uid, kind="image", limit=30, what="图片生成", noun="张") + self.assertEqual(msg, "[Error] 已达每日图片生成上限(30/30 张),次日 00:00 重置。") + with mock.patch.object(mc, "check_daily_quota", return_value=(3, False)): + self.assertIsNone(mc.quota_gate(uid, kind="video", limit=10, what="视频生成", noun="个")) + + +class RecordUsageSafeTests(unittest.TestCase): + def test_returns_value(self): + self.assertEqual(mc.record_usage_safe("t", lambda **kw: 42, a=1), 42) + + def test_swallow_exception_returns_none(self): + def boom(**kw): + raise RuntimeError("db down") + self.assertIsNone(mc.record_usage_safe("t", boom)) + + +class ExtractChatAnswerTests(unittest.TestCase): + def test_str_content(self): + resp = {"choices": [{"finish_reason": "stop", "message": {"content": " 你好 "}}]} + self.assertEqual(mc.extract_chat_answer(resp), ("你好", False)) + + def test_list_content_and_truncated(self): + resp = {"choices": [{ + "finish_reason": "length", + "message": {"content": [ + {"type": "text", "text": "第一段"}, + {"type": "image_url", "image_url": {}}, + {"type": "text", "text": "第二段"}, + ]}, + }]} + self.assertEqual(mc.extract_chat_answer(resp), ("第一段\n第二段", True)) + + def test_degenerate_shapes(self): + self.assertEqual(mc.extract_chat_answer(None), ("", False)) + self.assertEqual(mc.extract_chat_answer({}), ("", False)) + self.assertEqual(mc.extract_chat_answer({"choices": []}), ("", False)) + self.assertEqual( + mc.extract_chat_answer({"choices": [{"finish_reason": "length"}]}), ("", True) + ) + + +class FindFirstUrlTests(unittest.TestCase): + def test_nested_image_url(self): + resp = {"data": {"images": [{"url": "http://x/img.png"}]}} + self.assertEqual(mc.find_first_url(resp, keys=("url", "image_url")), "http://x/img.png") + + def test_non_http_ignored(self): + resp = {"url": "figures/local.png", "data": [{"image_url": "https://y/z.png"}]} + self.assertEqual(mc.find_first_url(resp, keys=("url", "image_url")), "https://y/z.png") + + def test_accept_predicate_video(self): + accept = lambda k, v: k == "video_url" or v.lower().endswith((".mp4", ".webm", ".mov")) + # 裸 url 非视频扩展 → 拒;video_url → 收 + resp = {"url": "http://x/preview.png", "content": {"video_url": "http://x/v"}} + self.assertEqual( + mc.find_first_url(resp, keys=("video_url", "url"), accept=accept), "http://x/v" + ) + # 只有裸视频扩展 url → 收 + resp2 = {"data": [{"url": "http://x/out.MP4"}]} + self.assertEqual( + mc.find_first_url(resp2, keys=("video_url", "url"), accept=accept), "http://x/out.MP4" + ) + + def test_not_found(self): + self.assertEqual(mc.find_first_url({"a": [1, {"b": "no"}]}, keys=("url",)), "") + + +class _FakeClient: + """打桩 ArkClient:按脚本依次返回/抛出。""" + script: list = [] + + def __init__(self, *a, **kw): + pass + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def post_json(self, endpoint, body, timeout_s=None): + act = _FakeClient.script.pop(0) + if isinstance(act, Exception): + raise act + return act + + +class ArkChatWithRetryTests(unittest.TestCase): + def _run(self, script, retries=1): + _FakeClient.script = list(script) + with mock.patch.object(mc, "ArkClient", _FakeClient), \ + mock.patch.object(mc.time, "sleep"): + return mc.ark_chat_with_retry( + object(), "/chat/completions", {}, timeout_s=1.0, + retries=retries, tool_name="look_at_image", + ) + + def test_success_first_try(self): + resp, err = self._run([{"ok": 1}]) + self.assertEqual(resp, {"ok": 1}) + self.assertEqual(err, "") + + def test_timeout_then_success(self): + resp, err = self._run([ArkTimeoutError("t"), {"ok": 2}], retries=1) + self.assertEqual(resp, {"ok": 2}) + self.assertEqual(err, "") + + def test_timeout_exhausted(self): + resp, err = self._run([ArkTimeoutError("t1"), ArkTimeoutError("t2")], retries=1) + self.assertIsNone(resp) + self.assertIn("look_at_image API", err) + self.assertIn("已重试 1 次仍超时", err) + + def test_business_error_no_retry(self): + resp, err = self._run([ArkError("HTTP 400 bad"), {"ok": 3}], retries=3) + self.assertIsNone(resp) + self.assertIn("HTTP 400 bad", err) + self.assertEqual(len(_FakeClient.script), 1) # 第二个动作没被消费 = 未重试 + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/gpt_image.py b/tools/gpt_image.py index 733a861..f7d9d50 100644 --- a/tools/gpt_image.py +++ b/tools/gpt_image.py @@ -15,7 +15,6 @@ from __future__ import annotations import base64 import json -import secrets import time from datetime import datetime from pathlib import Path @@ -23,9 +22,10 @@ from typing import Optional from uuid import UUID from core.ark_client import ArkClient, ArkConfig, ArkError -from core.storage.usage import check_daily_quota, record_image_usage +from core.storage.usage import record_image_usage from .base import Tool +from .media_common import quota_gate, record_usage_safe, stamped_path, write_meta class GptImageTool(Tool): @@ -76,14 +76,12 @@ class GptImageTool(Tool): if not (prompt or "").strip(): return "[Error] prompt 不能为空" - # 每账号每日配额(kind="image" 与 seedream 同口径合计;失败不落库故 retry 不计) - if self.daily_limit > 0: - used, over = check_daily_quota(user_id=self.user_id, kind="image", limit=self.daily_limit) - if over: - return ( - f"[Error] 已达每日图片生成上限({used}/{self.daily_limit} 张)," - f"次日 00:00 重置。" - ) + # 每账号每日配额(kind="image" 与 seedream 同口径合计) + quota_err = quota_gate( + self.user_id, kind="image", limit=self.daily_limit, what="图片生成", noun="张", + ) + if quota_err: + return quota_err cfg = self.cfg model_id = cfg["model_id"] @@ -111,10 +109,7 @@ class GptImageTool(Tool): except Exception as e: return f"[Error] gpt_image b64 解码失败: {e}" - ts = datetime.now().strftime("%Y%m%d-%H%M%S") - short = secrets.token_hex(3) - dest_png = self.working_dir / "figures" / f"{ts}-{short}.png" - dest_png.parent.mkdir(parents=True, exist_ok=True) + dest_png = stamped_path(self.working_dir / "figures", ".png") dest_png.write_bytes(img_bytes) elapsed = time.monotonic() - t0 @@ -135,22 +130,18 @@ class GptImageTool(Tool): "elapsed_s": round(elapsed, 2), "ts": datetime.now().isoformat(timespec="seconds"), } - dest_png.with_suffix(".meta.json").write_text( - json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8" - ) + write_meta(dest_png, meta) - try: - record_image_usage( - task_id=self.task_id, - user_id=self.user_id, - model_profile=f"unifyllm.{self.variant_key}", - n_images=1, - size=actual_size, - price_cny_per_image=price, - extra_units={"output_tokens": output_tokens, "quality": quality}, - ) - except Exception as e: - print(f"[gpt_image] record_image_usage failed: {type(e).__name__}: {e}", flush=True) + record_usage_safe( + "gpt_image", record_image_usage, + task_id=self.task_id, + user_id=self.user_id, + model_profile=f"unifyllm.{self.variant_key}", + n_images=1, + size=actual_size, + price_cny_per_image=price, + extra_units={"output_tokens": output_tokens, "quality": quality}, + ) disp = self._display(dest_png) # 首行 banner 协议同 seedream(`key=value · ` 分隔,前端 extractMediaBanner 解析); diff --git a/tools/look_at_image.py b/tools/look_at_image.py index 42b371b..bf784d4 100644 --- a/tools/look_at_image.py +++ b/tools/look_at_image.py @@ -13,13 +13,12 @@ from pathlib import Path from typing import Any, Optional from uuid import UUID -import time - -from core.ark_client import ArkClient, ArkConfig, ArkError, ArkTimeoutError +from core.ark_client import ArkConfig from core.storage.usage import record_vision_usage from .base import Tool, compact_tool_output from .image_ref import load_image_as_data_url +from .media_common import ark_chat_with_retry, extract_chat_answer, record_usage_safe # 不带 question 时的默认提问:全覆盖(描述 + OCR + 图表读数),让模型一次把图里能用的信息都吐出来 _DEFAULT_QUESTION = ( @@ -117,28 +116,17 @@ class LookAtImageTool(Tool): ], } - # 透明重试:Seed 2.0 Lite 非流式,长 OCR 偶发超时/网络抖动。tool 内消化掉, - # 不把 [Error] 抛给主模型 —— 否则主模型会重发整个 tool call(图 base64 重传、 - # 输入 token 再付一次)。仅 ArkTimeoutError(超时/网络)重试;HTTP 业务错误不重试。 - max_attempts = int(cfg.get("timeout_retries", 1)) + 1 - resp = None - for attempt in range(max_attempts): - try: - with ArkClient(self.ark_cfg, timeout_s=timeout_s) as client: - resp = client.post_json(endpoint, body, timeout_s=timeout_s) - break - except ArkTimeoutError as e: - if attempt == max_attempts - 1: - return f"[Error] look_at_image API: {e}(已重试 {attempt} 次仍超时)" - print( - f"[look_at_image] timeout, retrying ({attempt + 1}/{max_attempts - 1}): {e}", - flush=True, - ) - time.sleep(2 ** attempt) - except ArkError as e: - return f"[Error] look_at_image API: {e}" + # 透明重试(超时/网络抖动 tool 内消化,细节见 media_common.ark_chat_with_retry) + resp, api_err = ark_chat_with_retry( + self.ark_cfg, endpoint, body, + timeout_s=timeout_s, + retries=int(cfg.get("timeout_retries", 1)), + tool_name="look_at_image", + ) + if api_err: + return api_err - answer = self._extract_answer(resp) + answer, _truncated = extract_chat_answer(resp) if not answer: return ( "[Error] vision 响应缺内容(模型未返回文本)。" @@ -149,22 +137,18 @@ class LookAtImageTool(Tool): tin = int(usage.get("prompt_tokens", 0) or 0) tout = int(usage.get("completion_tokens", 0) or 0) - # 记账;失败不阻塞 tool 返回(沿用 seedream 兜底) - cost_cny = 0.0 - try: - cost = record_vision_usage( - task_id=self.task_id, - user_id=self.user_id, - model_profile=f"doubao.{self.variant_key}", - prompt_tokens=tin, - completion_tokens=tout, - input_cny_per_mtoken=float(cfg.get("price_cny_per_mtoken_input", 0)), - output_cny_per_mtoken=float(cfg.get("price_cny_per_mtoken_output", 0)), - extra_units={"image": disp}, - ) - cost_cny = float(cost) - except Exception as e: - print(f"[look_at_image] record_vision_usage failed: {type(e).__name__}: {e}", flush=True) + cost = record_usage_safe( + "look_at_image", record_vision_usage, + task_id=self.task_id, + user_id=self.user_id, + model_profile=f"doubao.{self.variant_key}", + prompt_tokens=tin, + completion_tokens=tout, + input_cny_per_mtoken=float(cfg.get("price_cny_per_mtoken_input", 0)), + output_cny_per_mtoken=float(cfg.get("price_cny_per_mtoken_output", 0)), + extra_units={"image": disp}, + ) + cost_cny = float(cost or 0) # 第一行 banner(key=value · 分隔,与 seedream/seedance 同协议,便于前端/对账) banner = ( @@ -173,27 +157,3 @@ class LookAtImageTool(Tool): ) # 图片解读正文可能很长(整页 OCR),压一下防爆上下文(保头尾) return f"{banner}\n\n{compact_tool_output(answer)}" - - @staticmethod - def _extract_answer(resp: dict) -> str: - """OpenAI 兼容 chat 响应取文本: choices[0].message.content。 - - content 可能是 str,也可能是 list[{type:text,text:...}](多模态返回形态),都兜住。 - """ - choices = resp.get("choices") - if not (isinstance(choices, list) and choices): - return "" - msg = choices[0].get("message") if isinstance(choices[0], dict) else None - if not isinstance(msg, dict): - return "" - content = msg.get("content") - if isinstance(content, str): - return content.strip() - 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() - return "" diff --git a/tools/media_common.py b/tools/media_common.py new file mode 100644 index 0000000..40463f6 --- /dev/null +++ b/tools/media_common.py @@ -0,0 +1,151 @@ +"""媒体类工具共享原语(从 seedance/seedream/gpt_image/look_at_image/read_document +五处同构析出,2026-07-23)。 + +收进来的判据:**逐字或参数化后逐字**的重复 —— 每日配额闸、`-` 落盘命名、 +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: + """产物落盘命名:`/-`,顺手建目录。""" + 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 "" diff --git a/tools/read_document.py b/tools/read_document.py index 3d0437d..d61aad2 100644 --- a/tools/read_document.py +++ b/tools/read_document.py @@ -11,16 +11,16 @@ file 内容块 + `data:application/pdf;base64,` 前缀;单页栅格化 3600 万 """ from __future__ import annotations -import time from pathlib import Path from typing import Any, Optional from uuid import UUID -from core.ark_client import ArkClient, ArkConfig, ArkError, ArkTimeoutError +from core.ark_client import ArkConfig from core.storage.usage import record_vision_usage from .base import Tool, compact_tool_output from .image_ref import _CONTAINER_ROOT, load_pdf_as_data_url, resolve_in_root +from .media_common import ark_chat_with_retry, extract_chat_answer, record_usage_safe _DEFAULT_QUESTION = ( "这是一份多页 PDF 文档。请逐页把其中的文字完整 OCR 成 markdown:" @@ -158,26 +158,17 @@ class ReadDocumentTool(Tool): ], } - # 超时透明重试,理由同 look_at_image(避免主模型整调用重发、base64 重传) - max_attempts = int(cfg.get("timeout_retries", 1)) + 1 - resp = None - for attempt in range(max_attempts): - try: - with ArkClient(self.ark_cfg, timeout_s=timeout_s) as client: - resp = client.post_json(endpoint, body, timeout_s=timeout_s) - break - except ArkTimeoutError as e: - if attempt == max_attempts - 1: - return f"[Error] read_document API: {e}(已重试 {attempt} 次仍超时)" - print( - f"[read_document] timeout, retrying ({attempt + 1}/{max_attempts - 1}): {e}", - flush=True, - ) - time.sleep(2 ** attempt) - except ArkError as e: - return f"[Error] read_document API: {e}" + # 超时透明重试(细节见 media_common.ark_chat_with_retry) + resp, api_err = ark_chat_with_retry( + self.ark_cfg, endpoint, body, + timeout_s=timeout_s, + retries=int(cfg.get("timeout_retries", 1)), + tool_name="read_document", + ) + if api_err: + return api_err - answer, truncated = self._extract_answer(resp) + answer, truncated = extract_chat_answer(resp) if not answer: return ( "[Error] 文档理解响应缺内容(模型未返回文本)。" @@ -188,21 +179,18 @@ class ReadDocumentTool(Tool): tin = int(usage.get("prompt_tokens", 0) or 0) tout = int(usage.get("completion_tokens", 0) or 0) - cost_cny = 0.0 - try: - cost = record_vision_usage( - task_id=self.task_id, - user_id=self.user_id, - model_profile=f"doubao.{self.variant_key}", - prompt_tokens=tin, - completion_tokens=tout, - input_cny_per_mtoken=float(cfg.get("price_cny_per_mtoken_input", 0)), - output_cny_per_mtoken=float(cfg.get("price_cny_per_mtoken_output", 0)), - extra_units={"document": disp}, - ) - cost_cny = float(cost) - except Exception as e: - print(f"[read_document] record_vision_usage failed: {type(e).__name__}: {e}", flush=True) + cost = record_usage_safe( + "read_document", record_vision_usage, + task_id=self.task_id, + user_id=self.user_id, + model_profile=f"doubao.{self.variant_key}", + prompt_tokens=tin, + completion_tokens=tout, + input_cny_per_mtoken=float(cfg.get("price_cny_per_mtoken_input", 0)), + output_cny_per_mtoken=float(cfg.get("price_cny_per_mtoken_output", 0)), + extra_units={"document": disp}, + ) + cost_cny = float(cost or 0) banner = ( f"[read_document] model={model_id} · document={disp}" @@ -251,25 +239,3 @@ class ReadDocumentTool(Tool): except OSError as e: return "", f"[Error] 写入 {rel} 失败: {type(e).__name__}: {e}" return self._display(target), "" - - @staticmethod - def _extract_answer(resp: dict) -> tuple[str, bool]: - """取 choices[0].message.content 文本 + 是否被输出上限截断(finish_reason=length)。""" - 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 diff --git a/tools/seedance.py b/tools/seedance.py index e5a3adc..d2b1060 100644 --- a/tools/seedance.py +++ b/tools/seedance.py @@ -18,7 +18,6 @@ W×H 由 resolution + ratio 推算(横版 height=resolution_num,竖版 width=res from __future__ import annotations import json -import secrets import time from datetime import datetime from pathlib import Path @@ -26,9 +25,16 @@ from typing import Any, Callable, Optional from uuid import UUID from core.ark_client import ArkClient, ArkConfig, ArkError -from core.storage.usage import check_daily_quota, record_video_usage +from core.storage.usage import record_video_usage from .base import Tool +from .media_common import ( + find_first_url, + quota_gate, + record_usage_safe, + stamped_path, + write_meta, +) # resolution → 短边像素;W/H 实际由 ratio 决定(横版短边=H,竖版短边=W) @@ -159,18 +165,14 @@ class SeedanceTool(Tool): return "[Error] prompt 不能为空" resume_id = (resume_task_id or "").strip() - # 每账号每日配额(yaml quotas.videos_per_day)。失败 / cancel 不计,因为 - # record_video_usage 只在 succeeded+下载完才落库。tool 返串会进 LLM 上下文 - # → 模型据此向用户解释,所以**只暴露用户该看的部分**(已用/上限 + 重置时间), - # 内部 yaml 路径不进对话(管理员要改的地方读代码/yaml 自己找)。 + # 每账号每日配额(yaml quotas.videos_per_day,细节见 media_common.quota_gate)。 # resume 不过配额闸:原次提交已经占过额度,续查不是新生成。 - if self.daily_limit > 0 and not resume_id: - used, over = check_daily_quota(user_id=self.user_id, kind="video", limit=self.daily_limit) - if over: - return ( - f"[Error] 已达每日视频生成上限({used}/{self.daily_limit} 个)," - f"次日 00:00 重置。" - ) + if not resume_id: + quota_err = quota_gate( + self.user_id, kind="video", limit=self.daily_limit, what="视频生成", noun="个", + ) + if quota_err: + return quota_err cfg = self.cfg model_id = cfg["model_id"] @@ -249,10 +251,7 @@ class SeedanceTool(Tool): return f"[Error] seedance succeeded 但响应缺 video url: {json.dumps(final_resp, ensure_ascii=False)[:400]}" # 3. download - ts = datetime.now().strftime("%Y%m%d-%H%M%S") - short = secrets.token_hex(3) - videos_dir = self.working_dir / "videos" - dest_mp4 = videos_dir / f"{ts}-{short}.mp4" + dest_mp4 = stamped_path(self.working_dir / "videos", ".mp4") client.download(video_url, dest_mp4, timeout_s=300.0) except ArkError as e: return f"[Error] seedance API: {e}" @@ -284,32 +283,29 @@ class SeedanceTool(Tool): "cgt_id": cgt_id, "ts": datetime.now().isoformat(timespec="seconds"), } - meta_path = dest_mp4.with_suffix(".meta.json") - meta_path.write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8") + write_meta(dest_mp4, meta) - try: - record_video_usage( - task_id=self.task_id, - user_id=self.user_id, - model_profile=f"doubao.{self.variant_key}", - resolution=chosen_resolution, - ratio=chosen_ratio, - duration_s=chosen_duration, - fps=fps, - width=width, - height=height, - tokens=tokens_actual, - price_cny_per_mtoken=price_t2v, - has_video_input=False, # phase 1 仅 t2v;i2v 接入后这里读 body 判断 - watermark=chosen_watermark, - extra_units={ - "cgt_id": cgt_id, - "elapsed_s": round(elapsed, 1), - "generate_audio": chosen_generate_audio, - }, - ) - except Exception as e: - print(f"[seedance] record_video_usage failed: {type(e).__name__}: {e}", flush=True) + record_usage_safe( + "seedance", record_video_usage, + task_id=self.task_id, + user_id=self.user_id, + model_profile=f"doubao.{self.variant_key}", + resolution=chosen_resolution, + ratio=chosen_ratio, + duration_s=chosen_duration, + fps=fps, + width=width, + height=height, + tokens=tokens_actual, + price_cny_per_mtoken=price_t2v, + has_video_input=False, # phase 1 仅 t2v;i2v 接入后这里读 body 判断 + watermark=chosen_watermark, + extra_units={ + "cgt_id": cgt_id, + "elapsed_s": round(elapsed, 1), + "generate_audio": chosen_generate_audio, + }, + ) disp = self._display(dest_mp4) # banner 协议与 seedream 一致:首行 `[tool] key=value · key=value ...` @@ -358,26 +354,12 @@ class SeedanceTool(Tool): if isinstance(v, str) and v.startswith("http"): return v # output[0].content[0].text → json parse(代理路径,极少用到) - # 递归兜底 - def _find_url(o: Any) -> Optional[str]: - if isinstance(o, dict): - for k, v in o.items(): - if k in ("video_url", "url") and isinstance(v, str) and v.startswith("http"): - # 过滤明显非视频的(image_url 等),video_url 优先级最高 - if k == "video_url": - return v - if v.lower().endswith((".mp4", ".webm", ".mov")): - return v - r = _find_url(v) - if r: - return r - elif isinstance(o, list): - for x in o: - r = _find_url(x) - if r: - return r - return None - return _find_url(resp) or "" + # 递归兜底(media_common.find_first_url):video_url 直收;裸 url 须是视频扩展名 + return find_first_url( + resp, + keys=("video_url", "url"), + accept=lambda k, v: k == "video_url" or v.lower().endswith((".mp4", ".webm", ".mov")), + ) @staticmethod def _extract_tokens(resp: dict) -> Optional[int]: diff --git a/tools/seedream.py b/tools/seedream.py index 4efcff7..00eeed6 100644 --- a/tools/seedream.py +++ b/tools/seedream.py @@ -9,7 +9,6 @@ from __future__ import annotations import json -import secrets import time from datetime import datetime from pathlib import Path @@ -17,10 +16,17 @@ from typing import Any, Optional from uuid import UUID from core.ark_client import ArkClient, ArkConfig, ArkError -from core.storage.usage import check_daily_quota, record_image_usage +from core.storage.usage import record_image_usage from .base import Tool from .image_ref import load_image_as_data_url +from .media_common import ( + find_first_url, + quota_gate, + record_usage_safe, + stamped_path, + write_meta, +) class SeedreamTool(Tool): @@ -124,17 +130,12 @@ class SeedreamTool(Tool): ref_disp.append(disp) is_i2i = bool(ref_data_urls) - # 每账号每日配额(yaml quotas.images_per_day)。失败 retry 不计,因为 - # record_image_usage 只在成功+下载完才落库。tool 返串会进 LLM 上下文, - # 模型据此向用户解释,所以**只暴露用户该看的部分**(已用/上限 + 重置时间), - # 内部 yaml 路径不进对话(管理员要改的地方读代码/yaml 自己找)。 - if self.daily_limit > 0: - used, over = check_daily_quota(user_id=self.user_id, kind="image", limit=self.daily_limit) - if over: - return ( - f"[Error] 已达每日图片生成上限({used}/{self.daily_limit} 张)," - f"次日 00:00 重置。" - ) + # 每账号每日配额(yaml quotas.images_per_day),细节见 media_common.quota_gate + quota_err = quota_gate( + self.user_id, kind="image", limit=self.daily_limit, what="图片生成", noun="张", + ) + if quota_err: + return quota_err cfg = self.cfg model_id = cfg["model_id"] @@ -176,10 +177,7 @@ class SeedreamTool(Tool): return f"[Error] seedream response 缺 image url: {json.dumps(resp, ensure_ascii=False)[:300]}" # 落盘 figures/-.png + .meta.json - ts = datetime.now().strftime("%Y%m%d-%H%M%S") - short = secrets.token_hex(3) - figures_dir = self.working_dir / "figures" - dest_png = figures_dir / f"{ts}-{short}.png" + dest_png = stamped_path(self.working_dir / "figures", ".png") client.download(image_url, dest_png, timeout_s=120.0) except ArkError as e: return f"[Error] seedream API: {e}" @@ -203,24 +201,19 @@ class SeedreamTool(Tool): "response_id": response_id, "ts": datetime.now().isoformat(timespec="seconds"), } - meta_path = dest_png.with_suffix(".meta.json") - meta_path.write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8") + write_meta(dest_png, meta) - # usage_events 记账;失败不阻塞 tool 返回,但 emit 一条 warn 给 sink 的事走不到这里 - # (tool 层没 sink 引用),先 print 兜底;后续可改成 sink 注入。 - try: - record_image_usage( - task_id=self.task_id, - user_id=self.user_id, - model_profile=f"doubao.{self.variant_key}", - n_images=1, - size=chosen_size, - price_cny_per_image=float(price), - search=chosen_search, - extra_units={"search_extra_cny": extra_cny} if chosen_search else None, - ) - except Exception as e: - print(f"[seedream] record_image_usage failed: {type(e).__name__}: {e}", flush=True) + record_usage_safe( + "seedream", record_image_usage, + task_id=self.task_id, + user_id=self.user_id, + model_profile=f"doubao.{self.variant_key}", + n_images=1, + size=chosen_size, + price_cny_per_image=float(price), + search=chosen_search, + extra_units={"search_extra_cny": extra_cny} if chosen_search else None, + ) disp = self._display(dest_png) # 第一行 banner:前端 SPA 把这行(name===seedream 时)单独提到 details summary @@ -318,19 +311,5 @@ class SeedreamTool(Tool): u = imgs[0].get("url") if isinstance(imgs[0], dict) else None if isinstance(u, str): return u, rid - # 兜底:递归搜 - def _find_url(o: Any) -> Optional[str]: - if isinstance(o, dict): - for k, v in o.items(): - if k in ("url", "image_url") and isinstance(v, str) and v.startswith("http"): - return v - r = _find_url(v) - if r: - return r - elif isinstance(o, list): - for x in o: - r = _find_url(x) - if r: - return r - return None - return (_find_url(resp) or ""), rid + # 兜底:递归搜(media_common.find_first_url) + return find_first_url(resp, keys=("url", "image_url")), rid