mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-05-07 05:43:55 +08:00
0277288a41
Track per-session model and token usage so users can inspect context pressure and cumulative usage with /session_status.
107 lines
3.8 KiB
Python
107 lines
3.8 KiB
Python
import asyncio
|
|
import unittest
|
|
from datetime import datetime
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
from langchain.agents.middleware.types import ModelRequest, ModelResponse
|
|
from langchain_core.messages import AIMessage
|
|
|
|
from app.agent.middleware.usage import UsageMiddleware
|
|
from app.chain.message import MessageChain
|
|
from app.schemas.types import MessageChannel
|
|
|
|
|
|
class TestAgentSessionStatus(unittest.TestCase):
|
|
def test_usage_middleware_records_usage_metadata(self):
|
|
snapshots = []
|
|
middleware = UsageMiddleware(on_usage=snapshots.append)
|
|
request = ModelRequest(
|
|
model=SimpleNamespace(
|
|
model="gpt-4o-mini", profile={"max_input_tokens": 128000}
|
|
),
|
|
messages=[],
|
|
state={},
|
|
runtime=None,
|
|
)
|
|
response = ModelResponse(
|
|
result=[
|
|
AIMessage(
|
|
content="ok",
|
|
usage_metadata={
|
|
"input_tokens": 1200,
|
|
"output_tokens": 300,
|
|
"total_tokens": 1500,
|
|
},
|
|
)
|
|
]
|
|
)
|
|
|
|
async def handler(_: ModelRequest):
|
|
return response
|
|
|
|
result = asyncio.run(middleware.awrap_model_call(request, handler))
|
|
|
|
self.assertIs(result, response)
|
|
self.assertEqual(len(snapshots), 1)
|
|
self.assertEqual(snapshots[0]["model"], "gpt-4o-mini")
|
|
self.assertEqual(snapshots[0]["context_window_tokens"], 128000)
|
|
self.assertEqual(snapshots[0]["input_tokens"], 1200)
|
|
self.assertEqual(snapshots[0]["output_tokens"], 300)
|
|
self.assertEqual(snapshots[0]["total_tokens"], 1500)
|
|
self.assertAlmostEqual(snapshots[0]["context_usage_ratio"], 1200 / 128000)
|
|
|
|
def test_remote_session_status_sends_usage_summary(self):
|
|
chain = MessageChain()
|
|
chain._user_sessions["10001"] = ("session-1", datetime.now())
|
|
status = {
|
|
"session_id": "session-1",
|
|
"model": "gpt-4o-mini",
|
|
"context_window_tokens": 128000,
|
|
"last_input_tokens": 1200,
|
|
"last_output_tokens": 300,
|
|
"last_total_tokens": 1500,
|
|
"last_context_usage_ratio": 1200 / 128000,
|
|
"total_input_tokens": 4500,
|
|
"total_output_tokens": 1500,
|
|
"total_tokens": 6000,
|
|
"model_call_count": 4,
|
|
"last_updated_at": "2026-04-26 12:34:56",
|
|
"is_processing": True,
|
|
"pending_messages": 2,
|
|
}
|
|
|
|
with (
|
|
patch(
|
|
"app.chain.message.agent_manager.get_session_status",
|
|
return_value=status,
|
|
),
|
|
patch.object(chain, "post_message") as post_message,
|
|
):
|
|
chain.remote_session_status(
|
|
channel=MessageChannel.Telegram,
|
|
userid="10001",
|
|
source="telegram-test",
|
|
)
|
|
|
|
notification = post_message.call_args.args[0]
|
|
self.assertEqual(notification.title, "当前智能体会话状态")
|
|
self.assertIn("session-1", notification.text)
|
|
self.assertIn("gpt-4o-mini", notification.text)
|
|
self.assertIn("1,200 / 128,000 (0.94%)", notification.text)
|
|
self.assertIn("输入 4,500 / 输出 1,500 / 总计 6,000", notification.text)
|
|
self.assertIn("运行中", notification.text)
|
|
|
|
def test_remote_session_status_handles_missing_session(self):
|
|
chain = MessageChain()
|
|
|
|
with patch.object(chain, "post_message") as post_message:
|
|
chain.remote_session_status(
|
|
channel=MessageChannel.Telegram,
|
|
userid="10001",
|
|
source="telegram-test",
|
|
)
|
|
|
|
notification = post_message.call_args.args[0]
|
|
self.assertEqual(notification.title, "您当前没有活跃的智能体会话")
|