feat: OpenMesh 基础平台与 MD/PDF 转换技能
- 后端: coworker 智能体框架, WS API, 文件上传, 附件处理 - 前端: Open WebUI, 文件全量走 upload API (含 MD/TXT/JSON 等文本类) - 技能: md-to-office (pandoc + wkhtmltopdf) - 修复: 上传文件路径丢失, Agent 搜索浪费, 输出文件跑到 uploads/ - 打包: PyInstaller one-dir, 预打包 pandoc/wkhtmltopdf/chromium
This commit is contained in:
362
tests/test_token_usage.py
Normal file
362
tests/test_token_usage.py
Normal file
@@ -0,0 +1,362 @@
|
||||
"""Token-usage metering — provider capture, normalization, engine plumbing.
|
||||
|
||||
Fakes follow the provider test convention: SimpleNamespace objects mimicking each
|
||||
SDK's response surface, dict events for Bedrock's Converse stream.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace
|
||||
|
||||
import aisuite as ai
|
||||
from coworker.engine import TurnEngine
|
||||
from coworker.events import EventType
|
||||
from coworker.permissions import PermissionEngine
|
||||
from coworker.providers import (
|
||||
AssistantTurn,
|
||||
ModelCapabilities,
|
||||
ProviderClient,
|
||||
)
|
||||
from coworker.providers.anthropic_provider import AnthropicProvider
|
||||
from coworker.providers.base import TokenUsage
|
||||
from coworker.providers.bedrock_provider import _BedrockConverseClient
|
||||
from coworker.providers.gemini_provider import GeminiProvider
|
||||
from coworker.providers.matrix import model_context_windows
|
||||
from coworker.providers.openai_provider import OpenAIProvider
|
||||
from coworker.tools import ToolRegistry
|
||||
|
||||
|
||||
def _final_turn(chunks):
|
||||
return chunks[-1].turn
|
||||
|
||||
|
||||
# -- TokenUsage ---------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_context_tokens_is_prompt_side_total():
|
||||
usage = TokenUsage(input=100, output=50, cache_read=300, cache_write=20)
|
||||
assert usage.context_tokens == 420
|
||||
assert usage.as_dict() == {
|
||||
"input": 100,
|
||||
"output": 50,
|
||||
"cache_read": 300,
|
||||
"cache_write": 20,
|
||||
}
|
||||
|
||||
|
||||
# -- Anthropic ----------------------------------------------------------------------
|
||||
|
||||
|
||||
class _FakeAnthropicClient:
|
||||
def __init__(self, events):
|
||||
def create(**kwargs):
|
||||
self.kwargs = kwargs
|
||||
return events
|
||||
|
||||
@contextmanager
|
||||
def stream(**kwargs):
|
||||
self.kwargs = kwargs
|
||||
yield SimpleNamespace(get_final_message=lambda: events)
|
||||
|
||||
self.messages = SimpleNamespace(create=create, stream=stream)
|
||||
self.beta = SimpleNamespace(messages=SimpleNamespace(create=create, stream=stream))
|
||||
|
||||
|
||||
def test_anthropic_stream_captures_usage():
|
||||
events = [
|
||||
SimpleNamespace(
|
||||
type="message_start",
|
||||
message=SimpleNamespace(
|
||||
usage=SimpleNamespace(
|
||||
input_tokens=7,
|
||||
output_tokens=1,
|
||||
cache_read_input_tokens=100,
|
||||
cache_creation_input_tokens=25,
|
||||
)
|
||||
),
|
||||
),
|
||||
SimpleNamespace(
|
||||
type="content_block_start",
|
||||
index=0,
|
||||
content_block=SimpleNamespace(type="text"),
|
||||
),
|
||||
SimpleNamespace(
|
||||
type="content_block_delta",
|
||||
index=0,
|
||||
delta=SimpleNamespace(type="text_delta", text="hi"),
|
||||
),
|
||||
SimpleNamespace(
|
||||
type="message_delta",
|
||||
delta=SimpleNamespace(stop_reason="end_turn"),
|
||||
usage=SimpleNamespace(output_tokens=42),
|
||||
),
|
||||
SimpleNamespace(type="message_stop"),
|
||||
]
|
||||
provider = AnthropicProvider(client=_FakeAnthropicClient(events))
|
||||
turn = _final_turn(
|
||||
list(provider.stream(model="m", messages=[{"role": "user", "content": "x"}]))
|
||||
)
|
||||
assert turn.usage == TokenUsage(input=7, output=42, cache_read=100, cache_write=25)
|
||||
|
||||
|
||||
def test_anthropic_complete_captures_usage():
|
||||
response = SimpleNamespace(
|
||||
content=[SimpleNamespace(type="text", text="hi")],
|
||||
stop_reason="end_turn",
|
||||
usage=SimpleNamespace(
|
||||
input_tokens=10,
|
||||
output_tokens=5,
|
||||
cache_read_input_tokens=0,
|
||||
cache_creation_input_tokens=0,
|
||||
),
|
||||
)
|
||||
provider = AnthropicProvider(client=_FakeAnthropicClient(response))
|
||||
turn = provider.complete(model="m", messages=[{"role": "user", "content": "x"}])
|
||||
assert turn.usage == TokenUsage(input=10, output=5)
|
||||
|
||||
|
||||
def test_anthropic_stream_without_usage_leaves_none():
|
||||
events = [
|
||||
SimpleNamespace(type="message_start"), # no message/usage attrs
|
||||
SimpleNamespace(
|
||||
type="message_delta", delta=SimpleNamespace(stop_reason="end_turn")
|
||||
),
|
||||
]
|
||||
provider = AnthropicProvider(client=_FakeAnthropicClient(events))
|
||||
turn = _final_turn(
|
||||
list(provider.stream(model="m", messages=[{"role": "user", "content": "x"}]))
|
||||
)
|
||||
assert turn.usage is None
|
||||
|
||||
|
||||
# -- OpenAI-compat ------------------------------------------------------------------
|
||||
|
||||
|
||||
class _FakeOpenAIClient:
|
||||
def __init__(self, chunks, *, reject_stream_options=False):
|
||||
self.calls = []
|
||||
|
||||
def create(**kwargs):
|
||||
self.calls.append(kwargs)
|
||||
if reject_stream_options and "stream_options" in kwargs:
|
||||
raise RuntimeError("unknown parameter: 'stream_options'")
|
||||
return chunks
|
||||
|
||||
self.chat = SimpleNamespace(completions=SimpleNamespace(create=create))
|
||||
|
||||
|
||||
def _openai_chunks():
|
||||
return [
|
||||
SimpleNamespace(
|
||||
choices=[
|
||||
SimpleNamespace(
|
||||
delta=SimpleNamespace(content="hi", tool_calls=None),
|
||||
finish_reason="stop",
|
||||
)
|
||||
],
|
||||
usage=None,
|
||||
),
|
||||
# Usage arrives on a final empty-choices chunk (include_usage contract).
|
||||
SimpleNamespace(
|
||||
choices=[],
|
||||
usage=SimpleNamespace(
|
||||
prompt_tokens=140,
|
||||
completion_tokens=9,
|
||||
prompt_tokens_details=SimpleNamespace(cached_tokens=40),
|
||||
),
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_openai_stream_requests_and_captures_usage():
|
||||
fake = _FakeOpenAIClient(_openai_chunks())
|
||||
provider = OpenAIProvider(client=fake)
|
||||
turn = _final_turn(
|
||||
list(provider.stream(model="m", messages=[{"role": "user", "content": "x"}]))
|
||||
)
|
||||
assert fake.calls[0]["stream_options"] == {"include_usage": True}
|
||||
# Cached share is carved out of prompt_tokens into cache_read.
|
||||
assert turn.usage == TokenUsage(input=100, output=9, cache_read=40)
|
||||
|
||||
|
||||
def test_openai_stream_retries_without_stream_options_when_rejected():
|
||||
fake = _FakeOpenAIClient(_openai_chunks(), reject_stream_options=True)
|
||||
provider = OpenAIProvider(client=fake)
|
||||
turn = _final_turn(
|
||||
list(provider.stream(model="m", messages=[{"role": "user", "content": "x"}]))
|
||||
)
|
||||
assert "stream_options" not in fake.calls[-1]
|
||||
assert turn.text == "hi" # the turn still completes; only metering is lost
|
||||
|
||||
|
||||
def test_openai_complete_captures_usage_without_cache_details():
|
||||
response = SimpleNamespace(
|
||||
choices=[
|
||||
SimpleNamespace(
|
||||
message=SimpleNamespace(content="hi", tool_calls=None),
|
||||
finish_reason="stop",
|
||||
)
|
||||
],
|
||||
usage=SimpleNamespace(
|
||||
prompt_tokens=30, completion_tokens=4, prompt_tokens_details=None
|
||||
),
|
||||
)
|
||||
fake = _FakeOpenAIClient(response)
|
||||
provider = OpenAIProvider(client=fake)
|
||||
turn = provider.complete(model="m", messages=[{"role": "user", "content": "x"}])
|
||||
assert turn.usage == TokenUsage(input=30, output=4)
|
||||
|
||||
|
||||
# -- Gemini -------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _FakeGeminiClient:
|
||||
def __init__(self, responses):
|
||||
def generate_content_stream(**kwargs):
|
||||
self.kwargs = kwargs
|
||||
return iter(responses)
|
||||
|
||||
def generate_content(**kwargs):
|
||||
self.kwargs = kwargs
|
||||
return responses[0]
|
||||
|
||||
self.models = SimpleNamespace(
|
||||
generate_content=generate_content,
|
||||
generate_content_stream=generate_content_stream,
|
||||
)
|
||||
|
||||
|
||||
def _gemini_response(text, usage_metadata=None):
|
||||
return SimpleNamespace(
|
||||
candidates=[
|
||||
SimpleNamespace(
|
||||
content=SimpleNamespace(
|
||||
parts=[SimpleNamespace(text=text, function_call=None)]
|
||||
),
|
||||
finish_reason=SimpleNamespace(name="STOP"),
|
||||
)
|
||||
],
|
||||
usage_metadata=usage_metadata,
|
||||
)
|
||||
|
||||
|
||||
def test_gemini_stream_keeps_last_usage_metadata():
|
||||
responses = [
|
||||
_gemini_response(
|
||||
"he",
|
||||
SimpleNamespace(
|
||||
prompt_token_count=90,
|
||||
candidates_token_count=1,
|
||||
cached_content_token_count=50,
|
||||
thoughts_token_count=0,
|
||||
),
|
||||
),
|
||||
_gemini_response(
|
||||
"y",
|
||||
# Cumulative — the last chunk carries the final totals.
|
||||
SimpleNamespace(
|
||||
prompt_token_count=90,
|
||||
candidates_token_count=12,
|
||||
cached_content_token_count=50,
|
||||
thoughts_token_count=6,
|
||||
),
|
||||
),
|
||||
]
|
||||
provider = GeminiProvider(client=_FakeGeminiClient(responses))
|
||||
turn = _final_turn(
|
||||
list(provider.stream(model="m", messages=[{"role": "user", "content": "x"}]))
|
||||
)
|
||||
# input = prompt minus cached; thinking tokens fold into output.
|
||||
assert turn.usage == TokenUsage(input=40, output=18, cache_read=50)
|
||||
|
||||
|
||||
# -- Bedrock (Converse) -------------------------------------------------------------
|
||||
|
||||
|
||||
def test_bedrock_converse_stream_captures_metadata_usage():
|
||||
fake = SimpleNamespace(
|
||||
converse_stream=lambda **kwargs: {
|
||||
"stream": [
|
||||
{"contentBlockDelta": {"delta": {"text": "hi"}, "contentBlockIndex": 0}},
|
||||
{"messageStop": {"stopReason": "end_turn"}},
|
||||
{
|
||||
"metadata": {
|
||||
"usage": {
|
||||
"inputTokens": 11,
|
||||
"outputTokens": 3,
|
||||
"cacheReadInputTokens": 8,
|
||||
"cacheWriteInputTokens": 2,
|
||||
}
|
||||
}
|
||||
},
|
||||
]
|
||||
}
|
||||
)
|
||||
client = _BedrockConverseClient(client=fake)
|
||||
turn = _final_turn(
|
||||
list(client.stream(model="m", messages=[{"role": "user", "content": "x"}]))
|
||||
)
|
||||
assert turn.usage == TokenUsage(input=11, output=3, cache_read=8, cache_write=2)
|
||||
|
||||
|
||||
# -- engine plumbing ----------------------------------------------------------------
|
||||
|
||||
|
||||
class _UsageProvider(ProviderClient):
|
||||
def complete(self, *, model, messages, tools=None, **settings):
|
||||
return AssistantTurn(
|
||||
text="done",
|
||||
finish_reason="stop",
|
||||
usage=TokenUsage(input=100, output=20, cache_read=5),
|
||||
)
|
||||
|
||||
def capabilities(self, model):
|
||||
return ModelCapabilities()
|
||||
|
||||
|
||||
def _run_engine(tmp_path):
|
||||
registry = ToolRegistry()
|
||||
registry.register_all(ai.toolkits.files(root=str(tmp_path), allow_write=True))
|
||||
engine = TurnEngine(
|
||||
provider=_UsageProvider(),
|
||||
registry=registry,
|
||||
permissions=PermissionEngine(workspace_root=tmp_path),
|
||||
model="gpt-5.5",
|
||||
)
|
||||
|
||||
async def _collect():
|
||||
return [ev async for ev in engine.run("hello")]
|
||||
|
||||
return engine, asyncio.run(_collect())
|
||||
|
||||
|
||||
def test_engine_attaches_usage_to_event_and_message(tmp_path):
|
||||
engine, events = _run_engine(tmp_path)
|
||||
assistant = next(ev for ev in events if ev.type == EventType.ASSISTANT_MESSAGE)
|
||||
expected = {
|
||||
"model": "gpt-5.5",
|
||||
"input": 100,
|
||||
"output": 20,
|
||||
"cache_read": 5,
|
||||
"cache_write": 0,
|
||||
}
|
||||
assert assistant.data["usage"] == expected
|
||||
persisted = next(m for m in engine.messages if m.get("role") == "assistant")
|
||||
assert persisted["usage"] == expected
|
||||
|
||||
|
||||
def test_outbound_messages_strip_usage_sidecar(tmp_path):
|
||||
engine, _ = _run_engine(tmp_path)
|
||||
assert all("usage" not in m for m in engine._outbound_messages())
|
||||
|
||||
|
||||
# -- matrix -------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_model_context_windows_covers_verified_entries_only():
|
||||
windows = model_context_windows()
|
||||
assert windows["anthropic:claude-fable-5"] == 1_000_000
|
||||
assert "together:thinkingmachines/Inkling" not in windows # unverified stays absent
|
||||
assert all(isinstance(v, int) and v > 0 for v in windows.values())
|
||||
Reference in New Issue
Block a user