diff --git a/benchmark/beam/run.py b/benchmark/beam/run.py index cc278fd0..351658f2 100644 --- a/benchmark/beam/run.py +++ b/benchmark/beam/run.py @@ -787,9 +787,6 @@ def main( # pylint: disable=too-many-statements _TOKEN_USAGE_METRICS = ( "input_tokens", "output_tokens", - "cache_read_tokens", - "cache_write_tokens", - "reasoning_tokens", "total_tokens", ) diff --git a/benchmark/longmemeval/run.py b/benchmark/longmemeval/run.py index b7047595..08ad8cde 100644 --- a/benchmark/longmemeval/run.py +++ b/benchmark/longmemeval/run.py @@ -749,9 +749,6 @@ def _print_summary(results: list[dict], start_time: float) -> None: _TOKEN_USAGE_METRICS = ( "input_tokens", "output_tokens", - "cache_read_tokens", - "cache_write_tokens", - "reasoning_tokens", "total_tokens", ) diff --git a/pyproject.toml b/pyproject.toml index 2e1efc6f..cb0077db 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,7 +42,7 @@ dependencies = [ [project.optional-dependencies] core = [ - "agentscope==2.0.4.post1", + "agentscope==2.0.5", "claude-agent-sdk>=0.2.126", "dingtalk-stream>=0.24.3", "openai-codex>=0.144.4", diff --git a/reme/components/agent_wrapper/as_agent_wrapper.py b/reme/components/agent_wrapper/as_agent_wrapper.py index d919f02f..41886a05 100644 --- a/reme/components/agent_wrapper/as_agent_wrapper.py +++ b/reme/components/agent_wrapper/as_agent_wrapper.py @@ -149,15 +149,10 @@ class AsAgentWrapper(BaseAgentWrapper): SDK_PACKAGE = "agentscope" - def _agentscope_usage(self, usage: Any) -> TokenUsage: - """Normalize AgentScope usage while preserving provider cache semantics.""" - model = self.as_llm.model if self.as_llm is not None else None - module = type(model).__module__ if model is not None else "" - # Anthropic reports normal, cache-read, and cache-write input tokens - # separately. AgentScope normalizes every provider's usage into the - # same ChatUsage type, so the model implementation identifies the - # provider instead of the usage object's module. - return TokenUsage.from_provider(usage, input_includes_cache="_anthropic" not in module) + @staticmethod + def _agentscope_usage(usage: Any) -> TokenUsage: + """Normalize AgentScope's portable input/output usage.""" + return TokenUsage.from_provider(usage) def __init__(self, as_llm: str = "default", session_retention_days: int = 10, **kwargs): super().__init__(**kwargs) @@ -353,20 +348,9 @@ class AsAgentWrapper(BaseAgentWrapper): kwargs = self._merged_kwargs(kwargs) agent, inputs = await self._build_agent(inputs, **kwargs) - usages: list[TokenUsage] = [] await agent.observe(inputs) - async for event in agent.reply_stream(): - if isinstance(event, ModelCallEndEvent): - # AgentScope's event intentionally contains only the portable - # input/output pair. Cache dimensions are unavailable here. - usages.append( - TokenUsage( - input_tokens=event.input_tokens, - output_tokens=event.output_tokens, - ), - ) + last_msg = await agent.reply() await self._dump_state(agent.state) - last_msg = agent.state.context[-1] result = { "session_id": agent.state.session_id, @@ -374,6 +358,10 @@ class AsAgentWrapper(BaseAgentWrapper): "result": last_msg.get_text_content(), } + usage = self._agentscope_usage(last_msg.usage) if hasattr(last_msg, "usage") and last_msg.usage else None + if usage is None: + self.logger.error("AgentScope did not return token usage; token accounting is unavailable for this reply.") + output_schema: dict | None = kwargs.get("output_schema") if output_schema is not None: assert self.as_llm is not None, "AsAgentWrapper requires a bound as_llm component with a valid model." @@ -385,12 +373,15 @@ class AsAgentWrapper(BaseAgentWrapper): tool_choice=ToolChoice(mode="auto"), ) result["structured_output"] = res.content - if res.usage is not None: - usages.append(self._agentscope_usage(res.usage)) + if res.usage is None: + usage = None + self.logger.error("AgentScope did not return structured-output token usage; token accounting is unavailable.") + elif usage is not None: + usage = TokenUsage.combine([usage, self._agentscope_usage(res.usage)]) - usage = TokenUsage.combine(usages) - result["usage"] = usage.model_dump() - self._record_token_usage(usage) + result["usage"] = usage.model_dump() if usage is not None else None + if usage is not None: + self._record_token_usage(usage) return result diff --git a/reme/components/agent_wrapper/base_agent_wrapper.py b/reme/components/agent_wrapper/base_agent_wrapper.py index 46006122..078d81bb 100644 --- a/reme/components/agent_wrapper/base_agent_wrapper.py +++ b/reme/components/agent_wrapper/base_agent_wrapper.py @@ -242,15 +242,6 @@ class BaseAgentWrapper(BaseComponent): [self.TOKEN_COUNTER_PREFIX, self.name, field], getattr(usage, field), ) - for field in ("cache_read_tokens", "cache_write_tokens", "reasoning_tokens"): - value = getattr(usage, field) - if value is not None: - global_counter_add(counters, [self.TOKEN_COUNTER_PREFIX, self.name, field], value) - global_counter_add( - counters, - [self.TOKEN_COUNTER_PREFIX, self.name, f"{field}_reported_calls"], - 1, - ) @abstractmethod async def reply(self, inputs: Any, **kwargs) -> dict: diff --git a/reme/components/agent_wrapper/cc_agent_wrapper.py b/reme/components/agent_wrapper/cc_agent_wrapper.py index fa2c95f2..068175df 100644 --- a/reme/components/agent_wrapper/cc_agent_wrapper.py +++ b/reme/components/agent_wrapper/cc_agent_wrapper.py @@ -38,8 +38,8 @@ class CcAgentWrapper(BaseAgentWrapper): @staticmethod def _claude_usage(usage: dict[str, Any] | None) -> TokenUsage: - """Normalize Claude CLI usage, whose cache dimensions are separate.""" - return TokenUsage.from_provider(usage or {}, input_includes_cache=False) + """Normalize Claude CLI's input/output usage.""" + return TokenUsage.from_provider(usage or {}) @property def session_path(self) -> Path: @@ -455,13 +455,15 @@ class CcAgentWrapper(BaseAgentWrapper): if last_msg is None: raise ValueError("No message received from Claude Code.") + usage = self._claude_usage(last_msg.usage) if last_msg.usage is not None else None result = { "session_id": last_msg.session_id or "", "last_message": asdict(last_msg), "result": last_msg.result, - "usage": self._claude_usage(last_msg.usage).model_dump(), + "usage": usage.model_dump() if usage is not None else None, } - self._record_token_usage(TokenUsage.model_validate(result["usage"])) + if usage is not None: + self._record_token_usage(usage) if kwargs.get("output_schema") is not None: result["structured_output"] = last_msg.structured_output return result diff --git a/reme/components/agent_wrapper/codex_agent_wrapper.py b/reme/components/agent_wrapper/codex_agent_wrapper.py index 5bced0b7..882fbfbf 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -81,8 +81,8 @@ class CodexAgentWrapper(BaseAgentWrapper): @staticmethod def _codex_usage(usage: Any) -> TokenUsage: - """Normalize Codex's per-turn token usage snapshot.""" - return TokenUsage.from_provider(usage, input_includes_cache=True) + """Normalize Codex's full-turn input/output usage snapshot.""" + return TokenUsage.from_provider(usage) # pylint: disable=too-many-arguments def __init__( @@ -426,7 +426,7 @@ class CodexAgentWrapper(BaseAgentWrapper): "last_message": final_response, "result": final_response, "turn": self._serialize(result), - "usage": TokenUsage().model_dump(), + "usage": None, } if (raw_usage := getattr(result, "usage", None)) is not None: usage = self._codex_usage(raw_usage.last) diff --git a/reme/schema/token_usage.py b/reme/schema/token_usage.py index 76013d94..f0ba3587 100644 --- a/reme/schema/token_usage.py +++ b/reme/schema/token_usage.py @@ -8,18 +8,12 @@ from pydantic import BaseModel, Field, model_validator class TokenUsage(BaseModel): """Token usage for one completed agent invocation. - ``input_tokens`` is the complete prompt size. Cache-read and cache-write - tokens are included in it, but are also exposed separately when the - backend reports them. ``reasoning_tokens`` is a subset of output tokens. - A ``None`` cache field means that the backend did not report it; it must - not be interpreted as zero. + ``input_tokens`` and ``output_tokens`` are the complete usage reported by + one backend invocation. ``total_tokens`` is their derived sum. """ input_tokens: int = Field(default=0, ge=0) output_tokens: int = Field(default=0, ge=0) - cache_read_tokens: int | None = Field(default=None, ge=0) - cache_write_tokens: int | None = Field(default=None, ge=0) - reasoning_tokens: int | None = Field(default=None, ge=0) total_tokens: int = Field(default=0, ge=0) @model_validator(mode="after") @@ -31,15 +25,8 @@ class TokenUsage(BaseModel): def from_provider( cls, usage: Any, - *, - input_includes_cache: bool, ) -> "TokenUsage": - """Normalize a provider usage object or mapping. - - Providers that report cache tokens separately from normal input (for - example Claude) pass ``False``. Providers whose input count already - includes cached input (for example Codex/OpenAI) pass ``True``. - """ + """Normalize a provider usage object or mapping to input/output totals.""" def get(*names: str) -> int | None: for name in names: @@ -48,38 +35,15 @@ class TokenUsage(BaseModel): return int(value) return None - reported_input = get("input_tokens", "prompt_tokens") or 0 - cache_read = get( - "cache_read_input_tokens", - "cache_input_tokens", - "cached_input_tokens", - ) - cache_write = get("cache_creation_input_tokens", "cache_write_input_tokens") - input_tokens = reported_input - if not input_includes_cache: - input_tokens += (cache_read or 0) + (cache_write or 0) return cls( - input_tokens=input_tokens, + input_tokens=get("input_tokens", "prompt_tokens") or 0, output_tokens=get("output_tokens", "completion_tokens") or 0, - cache_read_tokens=cache_read, - cache_write_tokens=cache_write, - reasoning_tokens=get("reasoning_output_tokens", "reasoning_tokens"), ) @classmethod def combine(cls, usages: list["TokenUsage"]) -> "TokenUsage": - """Combine completed model calls without turning partial data into a total. - - A cache or reasoning metric is reported only when every model call - supplied that metric. A sum over only the reporting calls would look - complete while silently undercounting the invocation. - """ - optional = ("cache_read_tokens", "cache_write_tokens", "reasoning_tokens") - values: dict[str, int | None] = { - "input_tokens": sum(item.input_tokens for item in usages), - "output_tokens": sum(item.output_tokens for item in usages), - } - for field in optional: - reported = [getattr(item, field) for item in usages if getattr(item, field) is not None] - values[field] = sum(reported) if reported and len(reported) == len(usages) else None - return cls(**values) + """Combine completed model calls into one full-invocation usage.""" + return cls( + input_tokens=sum(item.input_tokens for item in usages), + output_tokens=sum(item.output_tokens for item in usages), + ) diff --git a/reme/utils/evaluation_interface.py b/reme/utils/evaluation_interface.py index c8e57fc0..49e226e9 100644 --- a/reme/utils/evaluation_interface.py +++ b/reme/utils/evaluation_interface.py @@ -20,9 +20,6 @@ _JOB_ENTRY_CLASSES = {CronJob, StreamJob, BackgroundJob, BaseJob} _TOKEN_METRICS = ( "input_tokens", "output_tokens", - "cache_read_tokens", - "cache_write_tokens", - "reasoning_tokens", "total_tokens", ) @@ -102,8 +99,8 @@ def check_agent_token_count( def check_agent_token_usage(agent_name: str, app_context: "ApplicationContext") -> dict[str, int | None]: """Return all token metrics currently accumulated for one agent wrapper. - Optional cache and reasoning metrics remain ``None`` until the backend has - reported them at least once. This keeps unknown usage distinct from zero. + A metric remains ``None`` until the backend has reported it. This keeps + unavailable usage distinct from zero. """ tree = global_counter_get_all(app_context.metadata, ["__token_counter", agent_name]) children = tree.get("children", {}) if tree is not None else {} @@ -176,18 +173,10 @@ class AgentTokenUsageTracker: self.agent_names = list(dict.fromkeys(agent_names)) self.app_context = app_context self._start_usage: dict[str, dict[str, int | None]] = {} - self._start_report_counts: dict[str, dict[str, int]] = {} self.usages: dict[str, dict[str, int | None]] = {} def __enter__(self) -> dict[str, dict[str, int | None]]: self._start_usage = {name: check_agent_token_usage(name, self.app_context) for name in self.agent_names} - self._start_report_counts = { - name: { - metric: check_agent_token_count(name, self.app_context, f"{metric}_reported_calls") - for metric in ("cache_read_tokens", "cache_write_tokens", "reasoning_tokens") - } - for name in self.agent_names - } return self.usages def __exit__(self, exc_type, exc_value, traceback) -> bool: @@ -197,12 +186,7 @@ class AgentTokenUsageTracker: for metric in _TOKEN_METRICS: current = end_usage[metric] start = self._start_usage[name][metric] - if metric in self._start_report_counts[name]: - end_reports = check_agent_token_count(name, self.app_context, f"{metric}_reported_calls") - if end_reports == self._start_report_counts[name][metric]: - delta[metric] = None - continue - delta[metric] = (current or 0) - (start or 0) + delta[metric] = None if current is None else current - (start or 0) self.usages[name] = delta return False diff --git a/tests/unit/test_evaluation_interface.py b/tests/unit/test_evaluation_interface.py index 8e486781..a5904866 100644 --- a/tests/unit/test_evaluation_interface.py +++ b/tests/unit/test_evaluation_interface.py @@ -108,8 +108,8 @@ def test_track_agent_token_counts_returns_delta_for_one_agent(): assert check_agent_token_count("bench", app_context) == 35 -def test_track_agent_token_usage_preserves_unreported_cache_as_none(): - """Detailed usage tracking does not turn an unknown cache value into zero.""" +def test_track_agent_token_usage_reports_only_supported_metrics(): + """Detailed usage tracking reports the shared input/output contract.""" app_context = SimpleNamespace(metadata={}, jobs={}) for metric, value in (("input_tokens", 10), ("output_tokens", 5), ("total_tokens", 15)): global_counter_add(app_context.metadata, ["__token_counter", "bench", metric], value) @@ -122,10 +122,22 @@ def test_track_agent_token_usage_preserves_unreported_cache_as_none(): "bench": { "input_tokens": 20, "output_tokens": 7, - "cache_read_tokens": None, - "cache_write_tokens": None, - "reasoning_tokens": None, "total_tokens": 27, }, } - assert check_agent_token_usage("bench", app_context)["cache_read_tokens"] is None + + +def test_track_agent_token_usage_keeps_unavailable_usage_as_none(): + """A backend that reports no usage remains unavailable to benchmarks.""" + app_context = SimpleNamespace(metadata={}, jobs={}) + + with track_agent_token_usage(["bench"], app_context) as usages: + pass + + assert usages == { + "bench": { + "input_tokens": None, + "output_tokens": None, + "total_tokens": None, + }, + } diff --git a/tests/unit/test_token_usage.py b/tests/unit/test_token_usage.py index 6d52efa4..81b595b1 100644 --- a/tests/unit/test_token_usage.py +++ b/tests/unit/test_token_usage.py @@ -1,5 +1,7 @@ """Tests for unified agent token accounting.""" +from types import SimpleNamespace + from agentscope.model._model_usage import ChatUsage import pytest @@ -14,55 +16,44 @@ class _UsageWrapper(BaseAgentWrapper): raise NotImplementedError -def test_claude_style_usage_normalizes_cache_into_complete_input(): - """Cache-excluded provider input is normalized into the complete input total.""" +def test_provider_usage_keeps_only_input_and_output_tokens(): + """Provider-specific cache and reasoning details are not persisted.""" usage = TokenUsage.from_provider( { "input_tokens": 10, "output_tokens": 4, "cache_read_input_tokens": 20, "cache_creation_input_tokens": 30, - }, - input_includes_cache=False, + "reasoning_output_tokens": 2, + } ) assert usage.model_dump() == { - "input_tokens": 60, + "input_tokens": 10, "output_tokens": 4, - "cache_read_tokens": 20, - "cache_write_tokens": 30, - "reasoning_tokens": None, - "total_tokens": 64, + "total_tokens": 14, } -def test_codex_style_usage_does_not_double_count_cached_input(): - """Cache-inclusive provider input is preserved without adding cached tokens twice.""" +def test_provider_usage_uses_reported_input_without_cache_adjustment(): + """Reported input is kept unchanged across backend-specific usage shapes.""" usage = TokenUsage.from_provider( { "input_tokens": 60, "output_tokens": 4, "cached_input_tokens": 20, "reasoning_output_tokens": 2, - }, - input_includes_cache=True, + } ) assert usage.input_tokens == 60 - assert usage.cache_read_tokens == 20 - assert usage.reasoning_tokens == 2 assert usage.total_tokens == 64 -def test_agentscope_anthropic_usage_includes_cache_tokens(tmp_path): - """AgentScope uses one usage type, so provider identity comes from its model.""" +def test_agentscope_usage_keeps_portable_input_and_output(tmp_path): + """AgentScope usage has the same input/output-only contract.""" context = ApplicationContext(workspace_dir=str(tmp_path)) wrapper = AsAgentWrapper(name="research", as_llm="", app_context=context) - wrapper.as_llm = type( - "AnthropicLLM", - (), - {"model": type("AnthropicModel", (), {"__module__": "agentscope.model._anthropic._model"})()}, - )() usage = ChatUsage( input_tokens=10, output_tokens=4, @@ -72,20 +63,17 @@ def test_agentscope_anthropic_usage_includes_cache_tokens(tmp_path): ) assert wrapper._agentscope_usage(usage).model_dump() == { # pylint: disable=protected-access - "input_tokens": 60, + "input_tokens": 10, "output_tokens": 4, - "cache_read_tokens": 20, - "cache_write_tokens": 30, - "reasoning_tokens": None, - "total_tokens": 64, + "total_tokens": 14, } -def test_combined_usage_marks_partially_reported_metrics_as_unknown(): - """A partial cache/reasoning sum must not be presented as a full total.""" +def test_combined_usage_sums_all_model_calls(): + """One wrapper invocation records aggregate input/output usage.""" usage = TokenUsage.combine( [ - TokenUsage(input_tokens=10, output_tokens=4, cache_read_tokens=6), + TokenUsage(input_tokens=10, output_tokens=4), TokenUsage(input_tokens=5, output_tokens=2), ], ) @@ -93,19 +81,16 @@ def test_combined_usage_marks_partially_reported_metrics_as_unknown(): assert usage.model_dump() == { "input_tokens": 15, "output_tokens": 6, - "cache_read_tokens": None, - "cache_write_tokens": None, - "reasoning_tokens": None, "total_tokens": 21, } def test_token_counter_is_a_per_agent_metric_tree(tmp_path): - """Recorded usage accumulates per agent, and optional metrics track reported calls.""" + """Recorded usage accumulates input, output, and total tokens per agent.""" context = ApplicationContext(workspace_dir=str(tmp_path)) wrapper = _UsageWrapper(name="research", app_context=context) wrapper._record_token_usage( # pylint: disable=protected-access - TokenUsage(input_tokens=10, output_tokens=4, cache_read_tokens=6), + TokenUsage(input_tokens=10, output_tokens=4), ) wrapper._record_token_usage(TokenUsage(input_tokens=3, output_tokens=2)) # pylint: disable=protected-access @@ -115,12 +100,70 @@ def test_token_counter_is_a_per_agent_metric_tree(tmp_path): "input_tokens": {"value": 13, "children": {}}, "output_tokens": {"value": 6, "children": {}}, "total_tokens": {"value": 19, "children": {}}, - "cache_read_tokens": {"value": 6, "children": {}}, - "cache_read_tokens_reported_calls": {"value": 1, "children": {}}, }, } +@pytest.mark.asyncio +async def test_agentscope_reply_records_final_message_usage(tmp_path, monkeypatch): + """AgentScope 2.0.5 exposes full-reply usage on the final message.""" + context = ApplicationContext(workspace_dir=str(tmp_path)) + wrapper = AsAgentWrapper(name="research", as_llm="", app_context=context) + message = SimpleNamespace( + usage=SimpleNamespace(input_tokens=10, output_tokens=4), + model_dump=lambda: {"text": "answer"}, + get_text_content=lambda: "answer", + ) + agent = SimpleNamespace( + state=SimpleNamespace(session_id="session-1", context=[message]), + observe=lambda _inputs: _async_none(), + reply=lambda: _async_value(message), + ) + + async def build_agent(inputs, **_kwargs): + return agent, inputs + + monkeypatch.setattr(wrapper, "_build_agent", build_agent) + monkeypatch.setattr(wrapper, "_dump_state", _async_none) + + result = await wrapper.reply("hello") + + assert result["usage"] == {"input_tokens": 10, "output_tokens": 4, "total_tokens": 14} + assert global_counter_get_all(context.metadata, ["__token_counter", "research"])["children"]["total_tokens"]["value"] == 14 + + +@pytest.mark.asyncio +async def test_agentscope_reply_without_usage_leaves_token_counters_unset(tmp_path, monkeypatch): + """AgentScope 2.0.4 messages without usage remain visibly unavailable.""" + context = ApplicationContext(workspace_dir=str(tmp_path)) + wrapper = AsAgentWrapper(name="research", as_llm="", app_context=context) + message = SimpleNamespace(model_dump=lambda: {"text": "answer"}, get_text_content=lambda: "answer") + agent = SimpleNamespace( + state=SimpleNamespace(session_id="session-1", context=[message]), + observe=lambda _inputs: _async_none(), + reply=lambda: _async_value(message), + ) + + async def build_agent(inputs, **_kwargs): + return agent, inputs + + monkeypatch.setattr(wrapper, "_build_agent", build_agent) + monkeypatch.setattr(wrapper, "_dump_state", _async_none) + + result = await wrapper.reply("hello") + + assert result["usage"] is None + assert "__token_counter" not in context.metadata + + +async def _async_none(*_args, **_kwargs): + return None + + +async def _async_value(value): + return value + + @pytest.mark.asyncio async def test_agentscope_stream_reply_does_not_record_token_usage(tmp_path, monkeypatch): """Only non-streaming AgentScope replies contribute to token accounting."""