scriptoriumadmin/ReMe
1
0
Fork 0
mirror of https://github.com/agentscope-ai/ReMe.git synced 2026-11-06 07:51:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 11

support agent scope 2.0.5

Browse source
This commit is contained in:
酱牛肉 2026-07-30 18:14:58 +08:00
parent 41d947bf28
commit b5471e668d
11 changed files with 137 additions and 156 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
benchmark/beam/run.py
View file

@ -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",
)

3
benchmark/longmemeval/run.py
View file

@ -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",
)

2
pyproject.toml
View file

@ -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",

43
reme/components/agent_wrapper/as_agent_wrapper.py
View file

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

9
reme/components/agent_wrapper/base_agent_wrapper.py
View file

@ -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:

10
reme/components/agent_wrapper/cc_agent_wrapper.py
View file

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

6
reme/components/agent_wrapper/codex_agent_wrapper.py
View file

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

54
reme/schema/token_usage.py
View file

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

22
reme/utils/evaluation_interface.py
View file

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

24
tests/unit/test_evaluation_interface.py
View file

@ -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,
},
}

117
tests/unit/test_token_usage.py
View file

@ -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."""