refac
This commit is contained in:
@@ -331,6 +331,63 @@ class ChatMessageTable:
|
||||
for row in results
|
||||
}
|
||||
|
||||
def get_token_usage_by_user(
|
||||
self,
|
||||
start_date: Optional[int] = None,
|
||||
end_date: Optional[int] = None,
|
||||
db: Optional[Session] = None,
|
||||
) -> dict[str, dict]:
|
||||
"""Aggregate token usage by user using database-level aggregation."""
|
||||
with get_db_context(db) as db:
|
||||
from sqlalchemy import func, cast, Integer
|
||||
|
||||
dialect = db.bind.dialect.name
|
||||
|
||||
if dialect == "sqlite":
|
||||
input_tokens = cast(
|
||||
func.json_extract(ChatMessage.usage, "$.input_tokens"), Integer
|
||||
)
|
||||
output_tokens = cast(
|
||||
func.json_extract(ChatMessage.usage, "$.output_tokens"), Integer
|
||||
)
|
||||
elif dialect == "postgresql":
|
||||
input_tokens = cast(
|
||||
ChatMessage.usage["input_tokens"].astext, Integer
|
||||
)
|
||||
output_tokens = cast(
|
||||
ChatMessage.usage["output_tokens"].astext, Integer
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(f"Unsupported dialect: {dialect}")
|
||||
|
||||
query = db.query(
|
||||
ChatMessage.user_id,
|
||||
func.coalesce(func.sum(input_tokens), 0).label("input_tokens"),
|
||||
func.coalesce(func.sum(output_tokens), 0).label("output_tokens"),
|
||||
func.count(ChatMessage.id).label("message_count"),
|
||||
).filter(
|
||||
ChatMessage.role == "assistant",
|
||||
ChatMessage.user_id.isnot(None),
|
||||
ChatMessage.usage.isnot(None),
|
||||
)
|
||||
|
||||
if start_date:
|
||||
query = query.filter(ChatMessage.created_at >= start_date)
|
||||
if end_date:
|
||||
query = query.filter(ChatMessage.created_at <= end_date)
|
||||
|
||||
results = query.group_by(ChatMessage.user_id).all()
|
||||
|
||||
return {
|
||||
row.user_id: {
|
||||
"input_tokens": row.input_tokens,
|
||||
"output_tokens": row.output_tokens,
|
||||
"total_tokens": row.input_tokens + row.output_tokens,
|
||||
"message_count": row.message_count,
|
||||
}
|
||||
for row in results
|
||||
}
|
||||
|
||||
def get_message_count_by_user(
|
||||
self,
|
||||
start_date: Optional[int] = None,
|
||||
|
||||
@@ -33,6 +33,9 @@ class UserAnalyticsEntry(BaseModel):
|
||||
name: Optional[str] = None
|
||||
email: Optional[str] = None
|
||||
count: int
|
||||
input_tokens: int = 0
|
||||
output_tokens: int = 0
|
||||
total_tokens: int = 0
|
||||
|
||||
|
||||
class UserAnalyticsResponse(BaseModel):
|
||||
@@ -70,12 +73,15 @@ async def get_user_analytics(
|
||||
user=Depends(get_admin_user),
|
||||
db: Session = Depends(get_session),
|
||||
):
|
||||
"""Get message counts per user with user info."""
|
||||
"""Get message counts and token usage per user with user info."""
|
||||
from open_webui.models.users import Users
|
||||
|
||||
counts = ChatMessages.get_message_count_by_user(
|
||||
start_date=start_date, end_date=end_date, db=db
|
||||
)
|
||||
token_usage = ChatMessages.get_token_usage_by_user(
|
||||
start_date=start_date, end_date=end_date, db=db
|
||||
)
|
||||
|
||||
# Get user info for top users
|
||||
top_user_ids = [uid for uid, _ in sorted(counts.items(), key=lambda x: -x[1])[:limit]]
|
||||
@@ -84,11 +90,15 @@ async def get_user_analytics(
|
||||
users = []
|
||||
for user_id in top_user_ids:
|
||||
u = user_info.get(user_id)
|
||||
tokens = token_usage.get(user_id, {})
|
||||
users.append(UserAnalyticsEntry(
|
||||
user_id=user_id,
|
||||
name=u.name if u else None,
|
||||
email=u.email if u else None,
|
||||
count=counts[user_id]
|
||||
count=counts[user_id],
|
||||
input_tokens=tokens.get("input_tokens", 0),
|
||||
output_tokens=tokens.get("output_tokens", 0),
|
||||
total_tokens=tokens.get("total_tokens", 0),
|
||||
))
|
||||
|
||||
return UserAnalyticsResponse(users=users)
|
||||
|
||||
@@ -310,6 +310,7 @@
|
||||
{/if}
|
||||
</div>
|
||||
</th>
|
||||
<th scope="col" class="px-2.5 py-2 text-right">{$i18n.t('Tokens')}</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
@@ -327,10 +328,11 @@
|
||||
</div>
|
||||
</td>
|
||||
<td class="px-3 py-1 text-right">{user.count.toLocaleString()}</td>
|
||||
<td class="px-3 py-1 text-right">{formatNumber(user.total_tokens ?? 0)}</td>
|
||||
</tr>
|
||||
{/each}
|
||||
{#if sortedUsers.length === 0}
|
||||
<tr><td colspan="3" class="px-3 py-2 text-center text-gray-400">{$i18n.t('No data')}</td></tr>
|
||||
<tr><td colspan="4" class="px-3 py-2 text-center text-gray-400">{$i18n.t('No data')}</td></tr>
|
||||
{/if}
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
Reference in New Issue
Block a user