refactor(tools): 媒体工具同构析出 tools/media_common.py + 17 个单测

审查 P1#6:seedance/seedream/gpt_image/look_at_image/read_document 五处
同构(逐字或参数化后逐字)收敛为共享原语:

- quota_gate(每日配额闸,文案逐字保持)/ stamped_path(<ts>-<rand6> 落盘
  命名+建目录)/ 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 <noreply@anthropic.com>
This commit is contained in:
caoqianming 2026-07-23 13:10:32 +08:00
parent 7568ceeb48
commit 8038fb7491
7 changed files with 460 additions and 263 deletions

168
tests/test_media_common.py Normal file
View File

@ -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()

View File

@ -15,7 +15,6 @@ from __future__ import annotations
import base64 import base64
import json import json
import secrets
import time import time
from datetime import datetime from datetime import datetime
from pathlib import Path from pathlib import Path
@ -23,9 +22,10 @@ from typing import Optional
from uuid import UUID from uuid import UUID
from core.ark_client import ArkClient, ArkConfig, ArkError 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 .base import Tool
from .media_common import quota_gate, record_usage_safe, stamped_path, write_meta
class GptImageTool(Tool): class GptImageTool(Tool):
@ -76,14 +76,12 @@ class GptImageTool(Tool):
if not (prompt or "").strip(): if not (prompt or "").strip():
return "[Error] prompt 不能为空" return "[Error] prompt 不能为空"
# 每账号每日配额(kind="image" 与 seedream 同口径合计;失败不落库故 retry 不计) # 每账号每日配额(kind="image" 与 seedream 同口径合计)
if self.daily_limit > 0: quota_err = quota_gate(
used, over = check_daily_quota(user_id=self.user_id, kind="image", limit=self.daily_limit) self.user_id, kind="image", limit=self.daily_limit, what="图片生成", noun="",
if over: )
return ( if quota_err:
f"[Error] 已达每日图片生成上限({used}/{self.daily_limit} 张)," return quota_err
f"次日 00:00 重置。"
)
cfg = self.cfg cfg = self.cfg
model_id = cfg["model_id"] model_id = cfg["model_id"]
@ -111,10 +109,7 @@ class GptImageTool(Tool):
except Exception as e: except Exception as e:
return f"[Error] gpt_image b64 解码失败: {e}" return f"[Error] gpt_image b64 解码失败: {e}"
ts = datetime.now().strftime("%Y%m%d-%H%M%S") dest_png = stamped_path(self.working_dir / "figures", ".png")
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.write_bytes(img_bytes) dest_png.write_bytes(img_bytes)
elapsed = time.monotonic() - t0 elapsed = time.monotonic() - t0
@ -135,22 +130,18 @@ class GptImageTool(Tool):
"elapsed_s": round(elapsed, 2), "elapsed_s": round(elapsed, 2),
"ts": datetime.now().isoformat(timespec="seconds"), "ts": datetime.now().isoformat(timespec="seconds"),
} }
dest_png.with_suffix(".meta.json").write_text( write_meta(dest_png, meta)
json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8"
)
try: record_usage_safe(
record_image_usage( "gpt_image", record_image_usage,
task_id=self.task_id, task_id=self.task_id,
user_id=self.user_id, user_id=self.user_id,
model_profile=f"unifyllm.{self.variant_key}", model_profile=f"unifyllm.{self.variant_key}",
n_images=1, n_images=1,
size=actual_size, size=actual_size,
price_cny_per_image=price, price_cny_per_image=price,
extra_units={"output_tokens": output_tokens, "quality": quality}, 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)
disp = self._display(dest_png) disp = self._display(dest_png)
# 首行 banner 协议同 seedream(`key=value · ` 分隔,前端 extractMediaBanner 解析); # 首行 banner 协议同 seedream(`key=value · ` 分隔,前端 extractMediaBanner 解析);

View File

@ -13,13 +13,12 @@ from pathlib import Path
from typing import Any, Optional from typing import Any, Optional
from uuid import UUID from uuid import UUID
import time from core.ark_client import ArkConfig
from core.ark_client import ArkClient, ArkConfig, ArkError, ArkTimeoutError
from core.storage.usage import record_vision_usage from core.storage.usage import record_vision_usage
from .base import Tool, compact_tool_output from .base import Tool, compact_tool_output
from .image_ref import load_image_as_data_url 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 + 图表读数),让模型一次把图里能用的信息都吐出来 # 不带 question 时的默认提问:全覆盖(描述 + OCR + 图表读数),让模型一次把图里能用的信息都吐出来
_DEFAULT_QUESTION = ( _DEFAULT_QUESTION = (
@ -117,28 +116,17 @@ class LookAtImageTool(Tool):
], ],
} }
# 透明重试:Seed 2.0 Lite 非流式,长 OCR 偶发超时/网络抖动。tool 内消化掉, # 透明重试(超时/网络抖动 tool 内消化,细节见 media_common.ark_chat_with_retry)
# 不把 [Error] 抛给主模型 —— 否则主模型会重发整个 tool call(图 base64 重传、 resp, api_err = ark_chat_with_retry(
# 输入 token 再付一次)。仅 ArkTimeoutError(超时/网络)重试;HTTP 业务错误不重试。 self.ark_cfg, endpoint, body,
max_attempts = int(cfg.get("timeout_retries", 1)) + 1 timeout_s=timeout_s,
resp = None retries=int(cfg.get("timeout_retries", 1)),
for attempt in range(max_attempts): tool_name="look_at_image",
try: )
with ArkClient(self.ark_cfg, timeout_s=timeout_s) as client: if api_err:
resp = client.post_json(endpoint, body, timeout_s=timeout_s) return api_err
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}"
answer = self._extract_answer(resp) answer, _truncated = extract_chat_answer(resp)
if not answer: if not answer:
return ( return (
"[Error] vision 响应缺内容(模型未返回文本)。" "[Error] vision 响应缺内容(模型未返回文本)。"
@ -149,22 +137,18 @@ class LookAtImageTool(Tool):
tin = int(usage.get("prompt_tokens", 0) or 0) tin = int(usage.get("prompt_tokens", 0) or 0)
tout = int(usage.get("completion_tokens", 0) or 0) tout = int(usage.get("completion_tokens", 0) or 0)
# 记账;失败不阻塞 tool 返回(沿用 seedream 兜底) cost = record_usage_safe(
cost_cny = 0.0 "look_at_image", record_vision_usage,
try: task_id=self.task_id,
cost = record_vision_usage( user_id=self.user_id,
task_id=self.task_id, model_profile=f"doubao.{self.variant_key}",
user_id=self.user_id, prompt_tokens=tin,
model_profile=f"doubao.{self.variant_key}", completion_tokens=tout,
prompt_tokens=tin, input_cny_per_mtoken=float(cfg.get("price_cny_per_mtoken_input", 0)),
completion_tokens=tout, output_cny_per_mtoken=float(cfg.get("price_cny_per_mtoken_output", 0)),
input_cny_per_mtoken=float(cfg.get("price_cny_per_mtoken_input", 0)), extra_units={"image": disp},
output_cny_per_mtoken=float(cfg.get("price_cny_per_mtoken_output", 0)), )
extra_units={"image": disp}, cost_cny = float(cost or 0)
)
cost_cny = float(cost)
except Exception as e:
print(f"[look_at_image] record_vision_usage failed: {type(e).__name__}: {e}", flush=True)
# 第一行 banner(key=value · 分隔,与 seedream/seedance 同协议,便于前端/对账) # 第一行 banner(key=value · 分隔,与 seedream/seedance 同协议,便于前端/对账)
banner = ( banner = (
@ -173,27 +157,3 @@ class LookAtImageTool(Tool):
) )
# 图片解读正文可能很长(整页 OCR),压一下防爆上下文(保头尾) # 图片解读正文可能很长(整页 OCR),压一下防爆上下文(保头尾)
return f"{banner}\n\n{compact_tool_output(answer)}" 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 ""

151
tools/media_common.py Normal file
View File

@ -0,0 +1,151 @@
"""媒体类工具共享原语(从 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 ""

View File

@ -11,16 +11,16 @@ file 内容块 + `data:application/pdf;base64,` 前缀;单页栅格化 3600 万
""" """
from __future__ import annotations from __future__ import annotations
import time
from pathlib import Path from pathlib import Path
from typing import Any, Optional from typing import Any, Optional
from uuid import UUID 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 core.storage.usage import record_vision_usage
from .base import Tool, compact_tool_output from .base import Tool, compact_tool_output
from .image_ref import _CONTAINER_ROOT, load_pdf_as_data_url, resolve_in_root 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 = ( _DEFAULT_QUESTION = (
"这是一份多页 PDF 文档。请逐页把其中的文字完整 OCR 成 markdown:" "这是一份多页 PDF 文档。请逐页把其中的文字完整 OCR 成 markdown:"
@ -158,26 +158,17 @@ class ReadDocumentTool(Tool):
], ],
} }
# 超时透明重试,理由同 look_at_image(避免主模型整调用重发、base64 重传) # 超时透明重试(细节见 media_common.ark_chat_with_retry)
max_attempts = int(cfg.get("timeout_retries", 1)) + 1 resp, api_err = ark_chat_with_retry(
resp = None self.ark_cfg, endpoint, body,
for attempt in range(max_attempts): timeout_s=timeout_s,
try: retries=int(cfg.get("timeout_retries", 1)),
with ArkClient(self.ark_cfg, timeout_s=timeout_s) as client: tool_name="read_document",
resp = client.post_json(endpoint, body, timeout_s=timeout_s) )
break if api_err:
except ArkTimeoutError as e: return api_err
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}"
answer, truncated = self._extract_answer(resp) answer, truncated = extract_chat_answer(resp)
if not answer: if not answer:
return ( return (
"[Error] 文档理解响应缺内容(模型未返回文本)。" "[Error] 文档理解响应缺内容(模型未返回文本)。"
@ -188,21 +179,18 @@ class ReadDocumentTool(Tool):
tin = int(usage.get("prompt_tokens", 0) or 0) tin = int(usage.get("prompt_tokens", 0) or 0)
tout = int(usage.get("completion_tokens", 0) or 0) tout = int(usage.get("completion_tokens", 0) or 0)
cost_cny = 0.0 cost = record_usage_safe(
try: "read_document", record_vision_usage,
cost = record_vision_usage( task_id=self.task_id,
task_id=self.task_id, user_id=self.user_id,
user_id=self.user_id, model_profile=f"doubao.{self.variant_key}",
model_profile=f"doubao.{self.variant_key}", prompt_tokens=tin,
prompt_tokens=tin, completion_tokens=tout,
completion_tokens=tout, input_cny_per_mtoken=float(cfg.get("price_cny_per_mtoken_input", 0)),
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)),
output_cny_per_mtoken=float(cfg.get("price_cny_per_mtoken_output", 0)), extra_units={"document": disp},
extra_units={"document": disp}, )
) cost_cny = float(cost or 0)
cost_cny = float(cost)
except Exception as e:
print(f"[read_document] record_vision_usage failed: {type(e).__name__}: {e}", flush=True)
banner = ( banner = (
f"[read_document] model={model_id} · document={disp}" f"[read_document] model={model_id} · document={disp}"
@ -251,25 +239,3 @@ class ReadDocumentTool(Tool):
except OSError as e: except OSError as e:
return "", f"[Error] 写入 {rel} 失败: {type(e).__name__}: {e}" return "", f"[Error] 写入 {rel} 失败: {type(e).__name__}: {e}"
return self._display(target), "" 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

View File

@ -18,7 +18,6 @@ W×H 由 resolution + ratio 推算(横版 height=resolution_num,竖版 width=res
from __future__ import annotations from __future__ import annotations
import json import json
import secrets
import time import time
from datetime import datetime from datetime import datetime
from pathlib import Path from pathlib import Path
@ -26,9 +25,16 @@ from typing import Any, Callable, Optional
from uuid import UUID from uuid import UUID
from core.ark_client import ArkClient, ArkConfig, ArkError 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 .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) # resolution → 短边像素;W/H 实际由 ratio 决定(横版短边=H,竖版短边=W)
@ -159,18 +165,14 @@ class SeedanceTool(Tool):
return "[Error] prompt 不能为空" return "[Error] prompt 不能为空"
resume_id = (resume_task_id or "").strip() resume_id = (resume_task_id or "").strip()
# 每账号每日配额(yaml quotas.videos_per_day)。失败 / cancel 不计,因为 # 每账号每日配额(yaml quotas.videos_per_day,细节见 media_common.quota_gate)。
# record_video_usage 只在 succeeded+下载完才落库。tool 返串会进 LLM 上下文
# → 模型据此向用户解释,所以**只暴露用户该看的部分**(已用/上限 + 重置时间),
# 内部 yaml 路径不进对话(管理员要改的地方读代码/yaml 自己找)。
# resume 不过配额闸:原次提交已经占过额度,续查不是新生成。 # resume 不过配额闸:原次提交已经占过额度,续查不是新生成。
if self.daily_limit > 0 and not resume_id: if not resume_id:
used, over = check_daily_quota(user_id=self.user_id, kind="video", limit=self.daily_limit) quota_err = quota_gate(
if over: self.user_id, kind="video", limit=self.daily_limit, what="视频生成", noun="",
return ( )
f"[Error] 已达每日视频生成上限({used}/{self.daily_limit} 个)," if quota_err:
f"次日 00:00 重置。" return quota_err
)
cfg = self.cfg cfg = self.cfg
model_id = cfg["model_id"] 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]}" return f"[Error] seedance succeeded 但响应缺 video url: {json.dumps(final_resp, ensure_ascii=False)[:400]}"
# 3. download # 3. download
ts = datetime.now().strftime("%Y%m%d-%H%M%S") dest_mp4 = stamped_path(self.working_dir / "videos", ".mp4")
short = secrets.token_hex(3)
videos_dir = self.working_dir / "videos"
dest_mp4 = videos_dir / f"{ts}-{short}.mp4"
client.download(video_url, dest_mp4, timeout_s=300.0) client.download(video_url, dest_mp4, timeout_s=300.0)
except ArkError as e: except ArkError as e:
return f"[Error] seedance API: {e}" return f"[Error] seedance API: {e}"
@ -284,32 +283,29 @@ class SeedanceTool(Tool):
"cgt_id": cgt_id, "cgt_id": cgt_id,
"ts": datetime.now().isoformat(timespec="seconds"), "ts": datetime.now().isoformat(timespec="seconds"),
} }
meta_path = dest_mp4.with_suffix(".meta.json") write_meta(dest_mp4, meta)
meta_path.write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")
try: record_usage_safe(
record_video_usage( "seedance", record_video_usage,
task_id=self.task_id, task_id=self.task_id,
user_id=self.user_id, user_id=self.user_id,
model_profile=f"doubao.{self.variant_key}", model_profile=f"doubao.{self.variant_key}",
resolution=chosen_resolution, resolution=chosen_resolution,
ratio=chosen_ratio, ratio=chosen_ratio,
duration_s=chosen_duration, duration_s=chosen_duration,
fps=fps, fps=fps,
width=width, width=width,
height=height, height=height,
tokens=tokens_actual, tokens=tokens_actual,
price_cny_per_mtoken=price_t2v, price_cny_per_mtoken=price_t2v,
has_video_input=False, # phase 1 仅 t2v;i2v 接入后这里读 body 判断 has_video_input=False, # phase 1 仅 t2v;i2v 接入后这里读 body 判断
watermark=chosen_watermark, watermark=chosen_watermark,
extra_units={ extra_units={
"cgt_id": cgt_id, "cgt_id": cgt_id,
"elapsed_s": round(elapsed, 1), "elapsed_s": round(elapsed, 1),
"generate_audio": chosen_generate_audio, "generate_audio": chosen_generate_audio,
}, },
) )
except Exception as e:
print(f"[seedance] record_video_usage failed: {type(e).__name__}: {e}", flush=True)
disp = self._display(dest_mp4) disp = self._display(dest_mp4)
# banner 协议与 seedream 一致:首行 `[tool] key=value · key=value ...` # banner 协议与 seedream 一致:首行 `[tool] key=value · key=value ...`
@ -358,26 +354,12 @@ class SeedanceTool(Tool):
if isinstance(v, str) and v.startswith("http"): if isinstance(v, str) and v.startswith("http"):
return v return v
# output[0].content[0].text → json parse(代理路径,极少用到) # output[0].content[0].text → json parse(代理路径,极少用到)
# 递归兜底 # 递归兜底(media_common.find_first_url):video_url 直收;裸 url 须是视频扩展名
def _find_url(o: Any) -> Optional[str]: return find_first_url(
if isinstance(o, dict): resp,
for k, v in o.items(): keys=("video_url", "url"),
if k in ("video_url", "url") and isinstance(v, str) and v.startswith("http"): accept=lambda k, v: k == "video_url" or v.lower().endswith((".mp4", ".webm", ".mov")),
# 过滤明显非视频的(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 ""
@staticmethod @staticmethod
def _extract_tokens(resp: dict) -> Optional[int]: def _extract_tokens(resp: dict) -> Optional[int]:

View File

@ -9,7 +9,6 @@
from __future__ import annotations from __future__ import annotations
import json import json
import secrets
import time import time
from datetime import datetime from datetime import datetime
from pathlib import Path from pathlib import Path
@ -17,10 +16,17 @@ from typing import Any, Optional
from uuid import UUID from uuid import UUID
from core.ark_client import ArkClient, ArkConfig, ArkError 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 .base import Tool
from .image_ref import load_image_as_data_url 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): class SeedreamTool(Tool):
@ -124,17 +130,12 @@ class SeedreamTool(Tool):
ref_disp.append(disp) ref_disp.append(disp)
is_i2i = bool(ref_data_urls) is_i2i = bool(ref_data_urls)
# 每账号每日配额(yaml quotas.images_per_day)。失败 retry 不计,因为 # 每账号每日配额(yaml quotas.images_per_day),细节见 media_common.quota_gate
# record_image_usage 只在成功+下载完才落库。tool 返串会进 LLM 上下文, quota_err = quota_gate(
# 模型据此向用户解释,所以**只暴露用户该看的部分**(已用/上限 + 重置时间), self.user_id, kind="image", limit=self.daily_limit, what="图片生成", noun="",
# 内部 yaml 路径不进对话(管理员要改的地方读代码/yaml 自己找)。 )
if self.daily_limit > 0: if quota_err:
used, over = check_daily_quota(user_id=self.user_id, kind="image", limit=self.daily_limit) return quota_err
if over:
return (
f"[Error] 已达每日图片生成上限({used}/{self.daily_limit} 张),"
f"次日 00:00 重置。"
)
cfg = self.cfg cfg = self.cfg
model_id = cfg["model_id"] 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]}" return f"[Error] seedream response 缺 image url: {json.dumps(resp, ensure_ascii=False)[:300]}"
# 落盘 figures/<ts>-<rand>.png + .meta.json # 落盘 figures/<ts>-<rand>.png + .meta.json
ts = datetime.now().strftime("%Y%m%d-%H%M%S") dest_png = stamped_path(self.working_dir / "figures", ".png")
short = secrets.token_hex(3)
figures_dir = self.working_dir / "figures"
dest_png = figures_dir / f"{ts}-{short}.png"
client.download(image_url, dest_png, timeout_s=120.0) client.download(image_url, dest_png, timeout_s=120.0)
except ArkError as e: except ArkError as e:
return f"[Error] seedream API: {e}" return f"[Error] seedream API: {e}"
@ -203,24 +201,19 @@ class SeedreamTool(Tool):
"response_id": response_id, "response_id": response_id,
"ts": datetime.now().isoformat(timespec="seconds"), "ts": datetime.now().isoformat(timespec="seconds"),
} }
meta_path = dest_png.with_suffix(".meta.json") write_meta(dest_png, meta)
meta_path.write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")
# usage_events 记账;失败不阻塞 tool 返回,但 emit 一条 warn 给 sink 的事走不到这里 record_usage_safe(
# (tool 层没 sink 引用),先 print 兜底;后续可改成 sink 注入。 "seedream", record_image_usage,
try: task_id=self.task_id,
record_image_usage( user_id=self.user_id,
task_id=self.task_id, model_profile=f"doubao.{self.variant_key}",
user_id=self.user_id, n_images=1,
model_profile=f"doubao.{self.variant_key}", size=chosen_size,
n_images=1, price_cny_per_image=float(price),
size=chosen_size, search=chosen_search,
price_cny_per_image=float(price), extra_units={"search_extra_cny": extra_cny} if chosen_search else None,
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)
disp = self._display(dest_png) disp = self._display(dest_png)
# 第一行 banner:前端 SPA 把这行(name===seedream 时)单独提到 details summary # 第一行 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 u = imgs[0].get("url") if isinstance(imgs[0], dict) else None
if isinstance(u, str): if isinstance(u, str):
return u, rid return u, rid
# 兜底:递归搜 # 兜底:递归搜(media_common.find_first_url)
def _find_url(o: Any) -> Optional[str]: return find_first_url(resp, keys=("url", "image_url")), rid
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