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:
parent
7568ceeb48
commit
8038fb7491
|
|
@ -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()
|
||||
|
|
@ -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,12 +130,10 @@ 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(
|
||||
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}",
|
||||
|
|
@ -149,8 +142,6 @@ class GptImageTool(Tool):
|
|||
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)
|
||||
|
||||
disp = self._display(dest_png)
|
||||
# 首行 banner 协议同 seedream(`key=value · ` 分隔,前端 extractMediaBanner 解析);
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
# 透明重试(超时/网络抖动 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",
|
||||
)
|
||||
time.sleep(2 ** attempt)
|
||||
except ArkError as e:
|
||||
return f"[Error] look_at_image API: {e}"
|
||||
if api_err:
|
||||
return api_err
|
||||
|
||||
answer = self._extract_answer(resp)
|
||||
answer, _truncated = extract_chat_answer(resp)
|
||||
if not answer:
|
||||
return (
|
||||
"[Error] vision 响应缺内容(模型未返回文本)。"
|
||||
|
|
@ -149,10 +137,8 @@ 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(
|
||||
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}",
|
||||
|
|
@ -162,9 +148,7 @@ class LookAtImageTool(Tool):
|
|||
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_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 ""
|
||||
|
|
|
|||
|
|
@ -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 ""
|
||||
|
|
@ -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,
|
||||
# 超时透明重试(细节见 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",
|
||||
)
|
||||
time.sleep(2 ** attempt)
|
||||
except ArkError as e:
|
||||
return f"[Error] read_document API: {e}"
|
||||
if api_err:
|
||||
return api_err
|
||||
|
||||
answer, truncated = self._extract_answer(resp)
|
||||
answer, truncated = extract_chat_answer(resp)
|
||||
if not answer:
|
||||
return (
|
||||
"[Error] 文档理解响应缺内容(模型未返回文本)。"
|
||||
|
|
@ -188,9 +179,8 @@ 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(
|
||||
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}",
|
||||
|
|
@ -200,9 +190,7 @@ class ReadDocumentTool(Tool):
|
|||
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_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
|
||||
|
|
|
|||
|
|
@ -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,11 +283,10 @@ 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(
|
||||
record_usage_safe(
|
||||
"seedance", record_video_usage,
|
||||
task_id=self.task_id,
|
||||
user_id=self.user_id,
|
||||
model_profile=f"doubao.{self.variant_key}",
|
||||
|
|
@ -308,8 +306,6 @@ class SeedanceTool(Tool):
|
|||
"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)
|
||||
# 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]:
|
||||
|
|
|
|||
|
|
@ -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/<ts>-<rand>.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,13 +201,10 @@ 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(
|
||||
record_usage_safe(
|
||||
"seedream", record_image_usage,
|
||||
task_id=self.task_id,
|
||||
user_id=self.user_id,
|
||||
model_profile=f"doubao.{self.variant_key}",
|
||||
|
|
@ -219,8 +214,6 @@ class SeedreamTool(Tool):
|
|||
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)
|
||||
# 第一行 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
|
||||
|
|
|
|||
Loading…
Reference in New Issue