238 lines
7.8 KiB
Python
238 lines
7.8 KiB
Python
"""
|
||
智能体使用统计工具
|
||
记录和统计智能体的调用次数（PostgreSQL数据库存储）
|
||
"""
|
||
from datetime import datetime
|
||
from typing import Dict, Any
|
||
|
||
from config import POSTGRES_CONNECTION_STRING
|
||
|
||
|
||
ROUTE_FLAG_DISPLAY_NAMES = {
|
||
"问题排查": "问题排查",
|
||
"操作使用": "操作使用",
|
||
"百科问答": "百科问答"
|
||
}
|
||
|
||
ROUTE_FLAG_ICONS = {
|
||
"问题排查": "detection",
|
||
"操作使用": "repair",
|
||
"百科问答": "doc"
|
||
}
|
||
|
||
ROUTE_FLAG_COLORS = {
|
||
"问题排查": "#FFB800",
|
||
"操作使用": "#00C2FF",
|
||
"百科问答": "#9153FF"
|
||
}
|
||
|
||
|
||
async def init_agent_usage_table():
|
||
"""
|
||
初始化 agent_usage_records 表
|
||
"""
|
||
try:
|
||
from psycopg_pool import AsyncConnectionPool
|
||
except ImportError:
|
||
print("[agent_usage_statistics] psycopg_pool 未安装，无法初始化表")
|
||
return
|
||
|
||
async with AsyncConnectionPool(
|
||
POSTGRES_CONNECTION_STRING,
|
||
kwargs={"autocommit": True}
|
||
) as pool:
|
||
async with pool.connection() as conn:
|
||
async with conn.cursor() as cur:
|
||
await cur.execute("""
|
||
CREATE TABLE IF NOT EXISTS agent_usage_records (
|
||
id SERIAL PRIMARY KEY,
|
||
chat_id TEXT NOT NULL DEFAULT '',
|
||
route_flag TEXT NOT NULL DEFAULT '',
|
||
message_id TEXT NOT NULL DEFAULT '',
|
||
created_at TIMESTAMP NOT NULL DEFAULT NOW(),
|
||
UNIQUE(chat_id, route_flag)
|
||
)
|
||
""")
|
||
await cur.execute("""
|
||
CREATE INDEX IF NOT EXISTS agent_usage_records_route_flag_idx ON agent_usage_records(route_flag)
|
||
""")
|
||
await cur.execute("""
|
||
CREATE INDEX IF NOT EXISTS agent_usage_records_created_at_idx ON agent_usage_records(created_at)
|
||
""")
|
||
print("[agent_usage_statistics] agent_usage_records 表初始化完成")
|
||
|
||
|
||
async def record_agent_usage(chat_id: str, route_flag: str, message_id: str = "") -> bool:
|
||
"""
|
||
记录智能体使用情况到数据库
|
||
|
||
Args:
|
||
chat_id: 会话ID
|
||
route_flag: 路由标志（智能体类型）
|
||
message_id: 消息ID
|
||
|
||
Returns:
|
||
是否记录成功
|
||
"""
|
||
if not chat_id or not route_flag:
|
||
return False
|
||
|
||
try:
|
||
from psycopg_pool import AsyncConnectionPool
|
||
except ImportError:
|
||
print("[agent_usage_statistics] psycopg_pool 未安装，无法记录")
|
||
return False
|
||
|
||
try:
|
||
async with AsyncConnectionPool(
|
||
POSTGRES_CONNECTION_STRING,
|
||
kwargs={"autocommit": True}
|
||
) as pool:
|
||
async with pool.connection() as conn:
|
||
async with conn.cursor() as cur:
|
||
await cur.execute(
|
||
"""
|
||
INSERT INTO agent_usage_records (chat_id, route_flag, message_id, created_at)
|
||
VALUES (%s, %s, %s, %s)
|
||
ON CONFLICT (chat_id, route_flag) DO UPDATE SET
|
||
message_id = EXCLUDED.message_id,
|
||
created_at = EXCLUDED.created_at
|
||
""",
|
||
(chat_id, route_flag, message_id or "", datetime.now())
|
||
)
|
||
return True
|
||
except Exception as e:
|
||
print(f"[agent_usage_statistics] 记录智能体使用失败: {e}")
|
||
return False
|
||
|
||
|
||
def _format_count(count: int) -> str:
|
||
if count >= 10000:
|
||
return f"{count / 10000:.1f}万"
|
||
else:
|
||
return f"{count:,}"
|
||
|
||
|
||
async def get_agent_usage_statistics() -> Dict[str, Any]:
|
||
"""
|
||
从数据库查询智能体使用统计
|
||
|
||
统计逻辑：
|
||
- 同一 chat_id + route_flag 只保留一条记录
|
||
- 统计各类型的总使用次数
|
||
- 统计近6个月的月度使用次数
|
||
|
||
Returns:
|
||
统计结果字典
|
||
"""
|
||
try:
|
||
from psycopg_pool import AsyncConnectionPool
|
||
except ImportError:
|
||
print("[agent_usage_statistics] psycopg_pool 未安装，无法查询")
|
||
return {"success": False, "data": None}
|
||
|
||
required_routes = ["操作使用", "问题排查", "百科问答"]
|
||
|
||
now = datetime.now()
|
||
last_6_months = []
|
||
for i in range(5, -1, -1):
|
||
year = now.year
|
||
month = now.month - i
|
||
while month <= 0:
|
||
month += 12
|
||
year -= 1
|
||
last_6_months.append({
|
||
"year": year,
|
||
"month": month,
|
||
"key": f"{year}-{month:02d}",
|
||
"label": f"{month}月"
|
||
})
|
||
|
||
try:
|
||
async with AsyncConnectionPool(
|
||
POSTGRES_CONNECTION_STRING,
|
||
kwargs={"autocommit": True}
|
||
) as pool:
|
||
async with pool.connection() as conn:
|
||
async with conn.cursor() as cur:
|
||
# 查询各类型总使用次数
|
||
route_counts = {}
|
||
for route_flag in required_routes:
|
||
await cur.execute(
|
||
"""
|
||
SELECT COUNT(*) FROM agent_usage_records
|
||
WHERE route_flag = %s
|
||
""",
|
||
(route_flag,)
|
||
)
|
||
row = await cur.fetchone()
|
||
route_counts[route_flag] = row[0] if row else 0
|
||
|
||
# 查询近6个月各类型月度使用次数
|
||
monthly_counts = {route: {m["key"]: 0 for m in last_6_months} for route in required_routes}
|
||
|
||
for route_flag in required_routes:
|
||
for m in last_6_months:
|
||
month_start = datetime(m["year"], m["month"], 1)
|
||
if m["month"] == 12:
|
||
month_end = datetime(m["year"] + 1, 1, 1)
|
||
else:
|
||
month_end = datetime(m["year"], m["month"] + 1, 1)
|
||
|
||
await cur.execute(
|
||
"""
|
||
SELECT COUNT(*) FROM agent_usage_records
|
||
WHERE route_flag = %s AND created_at >= %s AND created_at < %s
|
||
""",
|
||
(route_flag, month_start, month_end)
|
||
)
|
||
row = await cur.fetchone()
|
||
monthly_counts[route_flag][m["key"]] = row[0] if row else 0
|
||
|
||
except Exception as e:
|
||
print(f"[agent_usage_statistics] 查询统计数据失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
return {"success": False, "data": None}
|
||
|
||
total = sum(route_counts.get(route, 0) for route in required_routes)
|
||
|
||
types_data = []
|
||
monthly_stats = []
|
||
for route_flag in required_routes:
|
||
count = route_counts.get(route_flag, 0)
|
||
display_name = ROUTE_FLAG_DISPLAY_NAMES.get(route_flag, route_flag)
|
||
icon = ROUTE_FLAG_ICONS.get(route_flag, "info")
|
||
color = ROUTE_FLAG_COLORS.get(route_flag, "#8C8C8C")
|
||
|
||
types_data.append({
|
||
"name": display_name,
|
||
"count": _format_count(count),
|
||
"icon": icon,
|
||
"color": color
|
||
})
|
||
|
||
monthly_data = []
|
||
for m in last_6_months:
|
||
monthly_data.append({
|
||
"month": m["label"],
|
||
"count": monthly_counts[route_flag][m["key"]]
|
||
})
|
||
|
||
monthly_stats.append({
|
||
"name": display_name,
|
||
"color": color,
|
||
"monthlyCounts": monthly_data
|
||
})
|
||
|
||
return {
|
||
"success": True,
|
||
"data": {
|
||
"total": _format_count(total),
|
||
"types": types_data,
|
||
"monthlyStats": monthly_stats,
|
||
"months": [m["label"] for m in last_6_months]
|
||
}
|
||
}
|
||
|