feat: add /stats command to view conversation token usage

- Add stats() method to ConversationCommands that queries ProviderStat
  records by conversation_id and aggregates token breakdowns
- Register /stats command in main.py
This commit is contained in:
Soulter
2026-04-27 11:52:10 +08:00
parent 415da218f6
commit f02444146d
2 changed files with 60 additions and 0 deletions

View File

@@ -1,3 +1,5 @@
from sqlalchemy import select
from astrbot.api import sp, star
from astrbot.api.event import AstrMessageEvent, MessageEventResult
from astrbot.core import logger
@@ -7,6 +9,7 @@ from astrbot.core.agent.runners.deerflow.constants import (
DEERFLOW_THREAD_ID_KEY,
)
from astrbot.core.agent.runners.deerflow.deerflow_api_client import DeerFlowAPIClient
from astrbot.core.db.po import ProviderStat
from astrbot.core.utils.active_event_registry import active_event_registry
from .utils.rst_scene import RstScene
@@ -246,3 +249,55 @@ class ConversationCommands:
f"✅ Switched to new conversation: {cid[:4]}"
),
)
async def stats(self, message: AstrMessageEvent) -> None:
"""Show token usage statistics for the current conversation."""
umo = message.unified_msg_origin
cid = await self.context.conversation_manager.get_curr_conversation_id(umo)
if not cid:
message.set_result(
MessageEventResult().message(
"❌ You are not in a conversation. Use /new to create one."
),
)
return
db = self.context.get_db()
async with db.get_db() as session:
result = await session.execute(
select(ProviderStat).where(
ProviderStat.agent_type == "internal",
ProviderStat.conversation_id == cid,
)
)
records = result.scalars().all()
if not records:
message.set_result(
MessageEventResult().message(
"📊 No stats available for this conversation yet."
),
)
return
total_calls = len(records)
total_input_other = sum(r.token_input_other for r in records)
total_input_cached = sum(r.token_input_cached for r in records)
total_output = sum(r.token_output for r in records)
total_tokens = total_input_other + total_input_cached + total_output
success_count = sum(1 for r in records if r.status != "error")
ret = (
f"📊 Conversation Stats (ID: {cid[:8]}...)\n"
f"───\n"
f"Total Calls: {total_calls}\n"
f"Successful: {success_count}\n"
f"───\n"
f"Input Tokens (other): {total_input_other:,}\n"
f"Input Tokens (cached): {total_input_cached:,}\n"
f"Output Tokens: {total_output:,}\n"
f"Total Tokens: {total_tokens:,}"
)
message.set_result(MessageEventResult().message(ret))

View File

@@ -47,6 +47,11 @@ class Main(star.Star):
"""Create new conversation"""
await self.conversation_c.new_conv(message)
@filter.command("stats")
async def stats(self, message: AstrMessageEvent) -> None:
"""Show token usage statistics for the current conversation"""
await self.conversation_c.stats(message)
@filter.permission_type(filter.PermissionType.ADMIN)
@filter.command("provider")
async def provider(