wehub-resource-sync
99 lines
3.9 KiB
Python
99 lines
3.9 KiB
Python
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from langchain.agents.middleware.types import ExtendedModelResponse, ModelResponse
|
|
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
|
|
|
|
from yuxi.agents.middlewares.token_usage import TokenUsageMiddleware
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_token_usage_middleware_records_request_and_state_tokens() -> None:
|
|
middleware = TokenUsageMiddleware()
|
|
tool_schema = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "search_docs",
|
|
"description": "Search project documents.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {"query": {"type": "string"}},
|
|
"required": ["query"],
|
|
},
|
|
},
|
|
}
|
|
request = SimpleNamespace(
|
|
model=SimpleNamespace(profile={"max_input_tokens": 2000}),
|
|
state={"messages": [HumanMessage(content="old message")]},
|
|
messages=[
|
|
HumanMessage(content="current message"),
|
|
ToolMessage(content="tool result", tool_call_id="call_1"),
|
|
],
|
|
system_message=SystemMessage(content="system prompt"),
|
|
tools=[tool_schema],
|
|
runtime=SimpleNamespace(context=SimpleNamespace(summary_threshold=2)),
|
|
)
|
|
|
|
async def handler(_request):
|
|
return ModelResponse(
|
|
result=[
|
|
AIMessage(
|
|
content="answer",
|
|
usage_metadata={"input_tokens": 12, "output_tokens": 5, "total_tokens": 17},
|
|
)
|
|
]
|
|
)
|
|
|
|
result = await middleware.awrap_model_call(request, handler)
|
|
|
|
assert isinstance(result, ExtendedModelResponse)
|
|
token_usage = result.command.update["token_usage"]
|
|
assert token_usage["state_message_count"] == 2
|
|
assert token_usage["state_message_count_before_call"] == 1
|
|
assert token_usage["llm_message_count"] == 2
|
|
assert token_usage["llm_content_message_count"] == 1
|
|
assert token_usage["llm_content_message_tokens"] > 0
|
|
assert token_usage["llm_tool_message_count"] == 1
|
|
assert token_usage["llm_tool_message_tokens"] > 0
|
|
assert token_usage["state_messages_tokens"] >= token_usage["state_messages_tokens_before_call"]
|
|
assert token_usage["llm_input_tokens"] >= token_usage["llm_messages_tokens"]
|
|
assert token_usage["system_tokens"] > 0
|
|
assert token_usage["tools_tokens"] > 0
|
|
assert token_usage["tool_count"] == 1
|
|
assert token_usage["context_window"] == 2000
|
|
assert token_usage["remaining_context_tokens"] == 2000 - token_usage["llm_input_tokens"]
|
|
assert token_usage["summary_trigger_tokens"] == 2048
|
|
assert "summary_keep_tokens" not in token_usage
|
|
assert token_usage["model_usage"] == {"input_tokens": 12, "output_tokens": 5, "total_tokens": 17}
|
|
assert token_usage["estimate"] is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_token_usage_middleware_detects_effective_summary_message() -> None:
|
|
middleware = TokenUsageMiddleware()
|
|
summary_message = HumanMessage(
|
|
content="conversation summary",
|
|
additional_kwargs={"lc_source": "summarization"},
|
|
)
|
|
request = SimpleNamespace(
|
|
model=SimpleNamespace(profile={}),
|
|
state={"messages": [HumanMessage(content="raw history")]},
|
|
messages=[summary_message, HumanMessage(content="recent user turn")],
|
|
system_message=None,
|
|
tools=[],
|
|
runtime=SimpleNamespace(context=SimpleNamespace(summary_threshold=100)),
|
|
)
|
|
|
|
async def handler(_request):
|
|
return ModelResponse(result=[AIMessage(content="answer")])
|
|
|
|
result = await middleware.awrap_model_call(request, handler)
|
|
token_usage = result.command.update["token_usage"]
|
|
|
|
assert token_usage["summary_active"] is True
|
|
assert token_usage["summary_message_tokens"] > 0
|
|
assert token_usage["context_window"] is None
|
|
assert token_usage["context_usage_ratio"] is None
|