mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-16 09:40:30 +08:00
feat: context token counting support for multimodal content (images, audio, and chain-of-thought) (#6596)
EstimateTokenCounter 之前只计算 TextPart,完全忽略 ImageURLPart、 AudioURLPart 和 ThinkPart。多模态对话中图片占 500-2000 token, 不被计入会导致 context 压缩触发过晚,API 先报 context_length_exceeded。 改动: - ImageURLPart 按 765 token 估算(OpenAI vision 低/高分辨率中位数) - AudioURLPart 按 500 token 估算 - ThinkPart 的文本内容正常计算 - 10 个新测试覆盖各类型单独和混合场景 Co-authored-by: Yufeng He <40085740+universeplayer@users.noreply.github.com>
This commit is contained in:
@@ -1,7 +1,7 @@
|
||||
import json
|
||||
from typing import Protocol, runtime_checkable
|
||||
|
||||
from ..message import Message, TextPart
|
||||
from ..message import AudioURLPart, ImageURLPart, Message, TextPart, ThinkPart
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
@@ -28,9 +28,19 @@ class TokenCounter(Protocol):
|
||||
...
|
||||
|
||||
|
||||
# 图片/音频 token 开销估算值,参考 OpenAI vision pricing:
|
||||
# low-res ~85 tokens, high-res ~170 per 512px tile, 通常几百到上千。
|
||||
# 这里取一个保守中位数,宁可偏高触发压缩也不要偏低导致 API 报错。
|
||||
IMAGE_TOKEN_ESTIMATE = 765
|
||||
AUDIO_TOKEN_ESTIMATE = 500
|
||||
|
||||
|
||||
class EstimateTokenCounter:
|
||||
"""Estimate token counter implementation.
|
||||
Provides a simple estimation of token count based on character types.
|
||||
|
||||
Supports multimodal content: images, audio, and thinking parts
|
||||
are all counted so that the context compressor can trigger in time.
|
||||
"""
|
||||
|
||||
def count_tokens(
|
||||
@@ -45,12 +55,16 @@ class EstimateTokenCounter:
|
||||
if isinstance(content, str):
|
||||
total += self._estimate_tokens(content)
|
||||
elif isinstance(content, list):
|
||||
# 处理多模态内容
|
||||
for part in content:
|
||||
if isinstance(part, TextPart):
|
||||
total += self._estimate_tokens(part.text)
|
||||
elif isinstance(part, ThinkPart):
|
||||
total += self._estimate_tokens(part.think)
|
||||
elif isinstance(part, ImageURLPart):
|
||||
total += IMAGE_TOKEN_ESTIMATE
|
||||
elif isinstance(part, AudioURLPart):
|
||||
total += AUDIO_TOKEN_ESTIMATE
|
||||
|
||||
# 处理 Tool Calls
|
||||
if msg.tool_calls:
|
||||
for tc in msg.tool_calls:
|
||||
tc_str = json.dumps(tc if isinstance(tc, dict) else tc.model_dump())
|
||||
|
||||
103
tests/agent/test_token_counter.py
Normal file
103
tests/agent/test_token_counter.py
Normal file
@@ -0,0 +1,103 @@
|
||||
"""Tests for EstimateTokenCounter multimodal support."""
|
||||
|
||||
from astrbot.core.agent.context.token_counter import (
|
||||
AUDIO_TOKEN_ESTIMATE,
|
||||
IMAGE_TOKEN_ESTIMATE,
|
||||
EstimateTokenCounter,
|
||||
)
|
||||
from astrbot.core.agent.message import (
|
||||
AudioURLPart,
|
||||
ImageURLPart,
|
||||
Message,
|
||||
TextPart,
|
||||
ThinkPart,
|
||||
)
|
||||
|
||||
|
||||
counter = EstimateTokenCounter()
|
||||
|
||||
|
||||
def _msg(role: str, content) -> Message:
|
||||
return Message(role=role, content=content)
|
||||
|
||||
|
||||
class TestTextCounting:
|
||||
def test_plain_string(self):
|
||||
tokens = counter.count_tokens([_msg("user", "hello world")])
|
||||
assert tokens > 0
|
||||
|
||||
def test_chinese(self):
|
||||
# 中文字符权重更高
|
||||
en = counter.count_tokens([_msg("user", "abc")])
|
||||
zh = counter.count_tokens([_msg("user", "你好啊")])
|
||||
assert zh > en
|
||||
|
||||
def test_text_part(self):
|
||||
msg = _msg("user", [TextPart(text="hello")])
|
||||
assert counter.count_tokens([msg]) > 0
|
||||
|
||||
|
||||
class TestMultimodalCounting:
|
||||
def test_image_counted(self):
|
||||
msg = _msg("user", [
|
||||
ImageURLPart(image_url=ImageURLPart.ImageURL(url="data:image/png;base64,abc")),
|
||||
])
|
||||
tokens = counter.count_tokens([msg])
|
||||
assert tokens == IMAGE_TOKEN_ESTIMATE
|
||||
|
||||
def test_audio_counted(self):
|
||||
msg = _msg("user", [
|
||||
AudioURLPart(audio_url=AudioURLPart.AudioURL(url="https://x.com/a.mp3")),
|
||||
])
|
||||
tokens = counter.count_tokens([msg])
|
||||
assert tokens == AUDIO_TOKEN_ESTIMATE
|
||||
|
||||
def test_think_counted(self):
|
||||
msg = _msg("assistant", [ThinkPart(think="let me think about this")])
|
||||
tokens = counter.count_tokens([msg])
|
||||
assert tokens > 0
|
||||
|
||||
def test_mixed_content(self):
|
||||
"""文本 + 图片的多模态消息,token 数 = 文本 token + 图片估算。"""
|
||||
text_only = _msg("user", [TextPart(text="describe this image")])
|
||||
mixed = _msg("user", [
|
||||
TextPart(text="describe this image"),
|
||||
ImageURLPart(image_url=ImageURLPart.ImageURL(url="data:image/png;base64,x")),
|
||||
])
|
||||
text_tokens = counter.count_tokens([text_only])
|
||||
mixed_tokens = counter.count_tokens([mixed])
|
||||
assert mixed_tokens == text_tokens + IMAGE_TOKEN_ESTIMATE
|
||||
|
||||
def test_multiple_images(self):
|
||||
"""多张图片应该各自计算。"""
|
||||
msg = _msg("user", [
|
||||
ImageURLPart(image_url=ImageURLPart.ImageURL(url="data:image/png;base64,a")),
|
||||
ImageURLPart(image_url=ImageURLPart.ImageURL(url="data:image/png;base64,b")),
|
||||
ImageURLPart(image_url=ImageURLPart.ImageURL(url="data:image/png;base64,c")),
|
||||
])
|
||||
tokens = counter.count_tokens([msg])
|
||||
assert tokens == IMAGE_TOKEN_ESTIMATE * 3
|
||||
|
||||
|
||||
class TestTrustedUsage:
|
||||
def test_trusted_overrides(self):
|
||||
"""如果 API 返回了 token 数,直接用它不做估算。"""
|
||||
msg = _msg("user", [
|
||||
TextPart(text="hello"),
|
||||
ImageURLPart(image_url=ImageURLPart.ImageURL(url="data:image/png;base64,x")),
|
||||
])
|
||||
tokens = counter.count_tokens([msg], trusted_token_usage=42)
|
||||
assert tokens == 42
|
||||
|
||||
|
||||
class TestToolCalls:
|
||||
def test_tool_calls_counted(self):
|
||||
msg = Message(
|
||||
role="assistant",
|
||||
content="calling tool",
|
||||
tool_calls=[{"type": "function", "id": "1", "function": {"name": "get_weather", "arguments": '{"city": "Beijing"}'}}],
|
||||
)
|
||||
tokens = counter.count_tokens([msg])
|
||||
# 文本 + tool call JSON 都应被计算
|
||||
text_only = counter.count_tokens([_msg("assistant", "calling tool")])
|
||||
assert tokens > text_only
|
||||
Reference in New Issue
Block a user