mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-15 17:30:13 +08:00
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:
@@ -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))
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user