Coverage for open_webui/models/chat_messages.py: 64%

1import time 

2import uuid 

3from collections import Counter 

4from datetime import datetime, timedelta 

5from typing import Any, Optional 

6from zoneinfo import ZoneInfo, ZoneInfoNotFoundError 

7 

8from sqlalchemy import select, delete, func, cast, Integer, distinct 

9from sqlalchemy.ext.asyncio import AsyncSession 

10from open_webui.internal.db import Base, get_async_db_context 

11from open_webui.utils.response import merge_usage, normalize_usage 

12from pydantic import BaseModel, ConfigDict 

13from sqlalchemy import ( 

14 JSON, 

15 BigInteger, 

16 Boolean, 

17 Column, 

18 ForeignKey, 

19 Index, 

20 Integer, 

21 Text, 

22 cast, 

23 delete, 

24 func, 

25 select, 

26) 

27from sqlalchemy.ext.asyncio import AsyncSession 

28 

29#################### 

30# Helpers 

31#################### 

32 

33 

34def _normalize_timestamp(timestamp: int) -> float: 

35 """Normalize and validate timestamp. Returns current time if invalid.""" 

36 now = time.time() 

37 

38 # Convert milliseconds to seconds if needed 

39 if timestamp > 10_000_000_000: 

40 timestamp = timestamp / 1000 

41 

42 # Validate: must be after 2020 and not in the future (with 1 day tolerance) 

43 min_valid = 1577836800 # 2020-01-01 00:00:00 UTC 

44 max_valid = now + 86400 # 1 day in the future (clock skew tolerance) 

45 

46 if timestamp < min_valid or timestamp > max_valid: 

47 return now 

48 

49 return timestamp 

50 

51 

52def _timezone(tz: Optional[str]) -> ZoneInfo: 

53 try: 

54 return ZoneInfo(tz or 'UTC') 

55 except ZoneInfoNotFoundError: 

56 return ZoneInfo('UTC') 

57 

58 

59def _date_key(timestamp: int, tz: ZoneInfo) -> str: 

60 return datetime.fromtimestamp(_normalize_timestamp(timestamp), tz=tz).strftime('%Y-%m-%d') 

61 

62 

63def get_usage(data: dict) -> Optional[dict]: 

64 """Extract and normalize usage from message data.""" 

65 usage = data.get('usage') or (data.get('info') or {}).get('usage') 

66 return normalize_usage(usage) if usage else None 

67 

68 

69def _token_columns(dialect: str): 

70 """Return (input_tokens, output_tokens) SQL column expressions. 

71 

72 Falls back to OpenAI-style keys (prompt_tokens / completion_tokens) 

73 when the normalized keys are absent. 

74 """ 

75 if dialect == 'sqlite': 75 ↛ 77line 75 didn't jump to line 77 because the condition on line 75 was always true

76 extract = lambda key: cast(func.json_extract(ChatMessage.usage, f'$.{key}'), Integer) 

77 elif dialect == 'postgresql': 

78 extract = lambda key: cast(func.json_extract_path_text(ChatMessage.usage, key), Integer) 

79 else: 

80 raise NotImplementedError(f'Unsupported dialect: {dialect}') 

81 

82 return ( 

83 func.coalesce(extract('input_tokens'), extract('prompt_tokens')), 

84 func.coalesce(extract('output_tokens'), extract('completion_tokens')), 

85 ) 

86 

87 

88def _extract_tool_names(value: Any) -> list[str]: 

89 names: list[str] = [] 

90 

91 def add(name: Any): 

92 if isinstance(name, str): 

93 cleaned = name.strip() 

94 if cleaned and len(cleaned) <= 128: 

95 names.append(cleaned) 

96 

97 def walk(item: Any): 

98 if isinstance(item, list): 98 ↛ 99line 98 didn't jump to line 99 because the condition on line 98 was never true

99 for child in item: 

100 walk(child) 

101 return 

102 

103 if not isinstance(item, dict): 103 ↛ 106line 103 didn't jump to line 106 because the condition on line 103 was always true

104 return 

105 

106 item_type = str(item.get('type') or '') 

107 looks_like_tool = 'tool' in item_type or item_type in {'function_call', 'function_call_output'} 

108 if looks_like_tool: 

109 add(item.get('name') or item.get('tool_name')) 

110 function = item.get('function') 

111 if isinstance(function, dict): 

112 add(function.get('name')) 

113 

114 for key in ('tool_calls', 'tools', 'output', 'meta'): 

115 if key in item: 

116 walk(item.get(key)) 

117 

118 walk(value) 

119 return names 

120 

121 

122#################### 

123# ChatMessage DB Schema 

124#################### 

125 

126 

127class ChatMessage(Base): 

128 __tablename__ = 'chat_message' 

129 

130 # Identity 

131 id = Column(Text, primary_key=True) 

132 chat_id = Column(Text, ForeignKey('chat.id', ondelete='CASCADE'), nullable=False, index=True) 

133 user_id = Column(Text, index=True) 

134 

135 # Structure 

136 role = Column(Text, nullable=False) # user, assistant, system 

137 parent_id = Column(Text, nullable=True) 

138 

139 # Content 

140 content = Column(JSON, nullable=True) # Can be str or list of blocks 

141 output = Column(JSON, nullable=True) 

142 

143 # Model (for assistant messages) 

144 model_id = Column(Text, nullable=True, index=True) 

145 

146 # Attachments 

147 files = Column(JSON, nullable=True) 

148 sources = Column(JSON, nullable=True) 

149 embeds = Column(JSON, nullable=True) 

150 meta = Column(JSON, nullable=True) 

151 

152 # Status 

153 done = Column(Boolean, default=True) 

154 status_history = Column(JSON, nullable=True) 

155 error = Column(JSON, nullable=True) 

156 

157 # Usage (tokens, timing, etc.) 

158 usage = Column(JSON, nullable=True) 

159 

160 # Context compaction checkpoint 

161 context_summary = Column(Text, nullable=True) 

162 

163 # Timestamps 

164 created_at = Column(BigInteger, index=True) 

165 updated_at = Column(BigInteger) 

166 

167 __table_args__ = ( 

168 Index('chat_message_chat_parent_idx', 'chat_id', 'parent_id'), 

169 Index('chat_message_model_created_idx', 'model_id', 'created_at'), 

170 Index('chat_message_user_created_idx', 'user_id', 'created_at'), 

171 Index('chat_message_chat_role_done_idx', 'chat_id', 'role', 'done'), # unfinished-assistant probe 

172 ) 

173 

174 

175#################### 

176# Pydantic Models 

177#################### 

178 

179 

180class ChatMessageModel(BaseModel): 

181 model_config = ConfigDict(from_attributes=True) 

182 

183 id: str 

184 chat_id: str 

185 user_id: str 

186 role: str 

187 parent_id: Optional[str] = None 

188 content: Optional[Any] = None # str or list of blocks 

189 output: Optional[list] = None 

190 model_id: Optional[str] = None 

191 files: Optional[list] = None 

192 sources: Optional[list] = None 

193 embeds: Optional[list] = None 

194 meta: Optional[dict] = None 

195 done: bool = True 

196 status_history: Optional[list] = None 

197 error: Optional[dict | str] = None 

198 usage: Optional[dict] = None 

199 context_summary: Optional[str] = None 

200 created_at: int 

201 updated_at: int 

202 

203 

204#################### 

205# Table Operations 

206#################### 

207 

208 

209class ChatMessageTable: 

210 @staticmethod 

211 def _apply_message_data(message: ChatMessage, data: dict, now: int) -> None: 

212 """Overwrite only the fields the payload carries.""" 

213 if 'role' in data: 213 ↛ 215line 213 didn't jump to line 215 because the condition on line 213 was always true

214 message.role = data['role'] 

215 if 'parent_id' in data or 'parentId' in data: 215 ↛ 217line 215 didn't jump to line 217 because the condition on line 215 was always true

216 message.parent_id = data.get('parent_id') or data.get('parentId') 

217 if 'content' in data: 217 ↛ 219line 217 didn't jump to line 219 because the condition on line 217 was always true

218 message.content = data.get('content') 

219 if 'output' in data: 219 ↛ 220line 219 didn't jump to line 220 because the condition on line 219 was never true

220 message.output = data.get('output') 

221 if 'model_id' in data or 'model' in data: 221 ↛ 222line 221 didn't jump to line 222 because the condition on line 221 was never true

222 message.model_id = data.get('model_id') or data.get('model') 

223 if 'files' in data: 223 ↛ 224line 223 didn't jump to line 224 because the condition on line 223 was never true

224 message.files = data.get('files') 

225 if 'sources' in data: 225 ↛ 226line 225 didn't jump to line 226 because the condition on line 225 was never true

226 message.sources = data.get('sources') 

227 if 'embeds' in data: 227 ↛ 228line 227 didn't jump to line 228 because the condition on line 227 was never true

228 message.embeds = data.get('embeds') 

229 if 'meta' in data: 229 ↛ 230line 229 didn't jump to line 230 because the condition on line 229 was never true

230 message.meta = data.get('meta') 

231 if 'done' in data: 

232 message.done = data['done'] 

233 if 'status_history' in data or 'statusHistory' in data: 233 ↛ 234line 233 didn't jump to line 234 because the condition on line 233 was never true

234 message.status_history = data.get('status_history') or data.get('statusHistory') 

235 if 'error' in data: 235 ↛ 236line 235 didn't jump to line 236 because the condition on line 235 was never true

236 message.error = data.get('error') 

237 if 'context_summary' in data or 'contextSummary' in data: 237 ↛ 238line 237 didn't jump to line 238 because the condition on line 237 was never true

238 message.context_summary = data.get('context_summary') or data.get('contextSummary') 

239 

240 usage = get_usage(data) 

241 if usage: 241 ↛ 242line 241 didn't jump to line 242 because the condition on line 241 was never true

242 existing_usage = normalize_usage(message.usage) 

243 message.usage = existing_usage if usage == existing_usage else merge_usage(existing_usage, usage) 

244 message.updated_at = now 

245 

246 @staticmethod 

247 def _build_message(composite_id: str, chat_id: str, user_id: str, data: dict, now: int) -> ChatMessage: 

248 return ChatMessage( 

249 id=composite_id, 

250 chat_id=chat_id, 

251 user_id=user_id, 

252 role=data.get('role', 'user'), 

253 parent_id=data.get('parent_id') or data.get('parentId'), 

254 content=data.get('content'), 

255 output=data.get('output'), 

256 model_id=data.get('model_id') or data.get('model'), 

257 files=data.get('files'), 

258 sources=data.get('sources'), 

259 embeds=data.get('embeds'), 

260 meta=data.get('meta'), 

261 done=data.get('done', True), 

262 status_history=data.get('status_history') or data.get('statusHistory'), 

263 error=data.get('error'), 

264 usage=get_usage(data), 

265 context_summary=data.get('context_summary') or data.get('contextSummary'), 

266 created_at=data.get('timestamp', now), 

267 updated_at=now, 

268 ) 

269 

270 async def upsert_message( 

271 self, 

272 message_id: str, 

273 chat_id: str, 

274 user_id: str, 

275 data: dict, 

276 db: Optional[AsyncSession] = None, 

277 ) -> Optional[ChatMessageModel]: 

278 """Insert or update a chat message.""" 

279 async with get_async_db_context(db) as db: 

280 now = int(time.time()) 

281 # Use composite ID: {chat_id}-{message_id} 

282 composite_id = f'{chat_id}-{message_id}' 

283 

284 message = await db.get(ChatMessage, composite_id) 

285 if message: 

286 self._apply_message_data(message, data, now) 

287 else: 

288 message = self._build_message(composite_id, chat_id, user_id, data, now) 

289 db.add(message) 

290 

291 await db.commit() 

292 return ChatMessageModel.model_validate(message) 

293 

294 async def upsert_messages( 

295 self, 

296 chat_id: str, 

297 user_id: str, 

298 messages: dict[str, dict], 

299 db: AsyncSession | None = None, 

300 ) -> None: 

301 """Insert or update the given messages of one chat.""" 

302 if not messages: 

303 return 

304 

305 async with get_async_db_context(db) as db: 

306 now = int(time.time()) 

307 result = await db.execute( 

308 select(ChatMessage).filter(ChatMessage.id.in_([f'{chat_id}-{message_id}' for message_id in messages])) 

309 ) 

310 existing_by_id = {row.id: row for row in result.scalars().all()} 

311 

312 for message_id, data in messages.items(): 

313 composite_id = f'{chat_id}-{message_id}' 

314 message = existing_by_id.get(composite_id) 

315 if message: 

316 self._apply_message_data(message, data, now) 

317 else: 

318 db.add(self._build_message(composite_id, chat_id, user_id, data, now)) 

319 

320 await db.commit() 

321 

322 async def get_message_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[ChatMessageModel]: 

323 async with get_async_db_context(db) as db: 

324 message = await db.get(ChatMessage, id) 

325 return ChatMessageModel.model_validate(message) if message else None 

326 

327 async def has_unfinished_assistant_by_chat_id( 

328 self, 

329 chat_id: str, 

330 db: Optional[AsyncSession] = None, 

331 ) -> bool: 

332 async with get_async_db_context(db) as db: 

333 result = await db.execute( 

334 select(ChatMessage.id) 

335 .where(ChatMessage.chat_id == chat_id) 

336 .where(ChatMessage.role == 'assistant') 

337 .where(ChatMessage.done.is_(False)) 

338 .limit(1) 

339 ) 

340 return result.scalar_one_or_none() is not None 

341 

342 async def get_messages_by_chat_id(self, chat_id: str, db: Optional[AsyncSession] = None) -> list[ChatMessageModel]: 

343 async with get_async_db_context(db) as db: 

344 result = await db.execute( 

345 select(ChatMessage).filter_by(chat_id=chat_id).order_by(ChatMessage.created_at.asc()) 

346 ) 

347 messages = result.scalars().all() 

348 return [ChatMessageModel.model_validate(message) for message in messages] 

349 

350 # DB column names that differ from the JSON message keys. 

351 DB_TO_JSON_KEY_MAP = { 

352 'parent_id': 'parentId', 

353 'model_id': 'model', 

354 'status_history': 'statusHistory', 

355 'context_summary': 'contextSummary', 

356 'created_at': 'timestamp', 

357 } 

358 # DB-internal columns excluded from the reconstructed message dict. 

359 EXCLUDED_COLUMNS = frozenset({'id', 'chat_id', 'user_id', 'updated_at'}) 

360 

361 async def get_messages_map_by_chat_id(self, chat_id: str, db: Optional[AsyncSession] = None) -> Optional[dict]: 

362 """Build a {message_id: message_dict} map from chat_message rows. 

363 

364 Returns the same shape as chat.history.messages so callers 

365 (get_message_list, middleware) work unchanged. Returns None if 

366 no rows exist for the chat (caller should fall back to the 

367 embedded JSON blob for legacy chats). 

368 """ 

369 async with get_async_db_context(db) as db: 

370 result = await db.execute(select(ChatMessage).filter_by(chat_id=chat_id)) 

371 rows = result.scalars().all() 

372 

373 if not rows: 

374 return None 

375 

376 # Strip the composite-id prefix ("{chat_id}-") to recover the 

377 # original message_id used as map key. 

378 prefix = f'{chat_id}-' 

379 prefix_len = len(prefix) 

380 col_keys = [c.key for c in ChatMessage.__table__.columns] 

381 

382 messages_map: dict[str, dict] = {} 

383 for row in rows: 

384 msg_id = row.id[prefix_len:] if row.id.startswith(prefix) else row.id 

385 

386 msg: dict = {'id': msg_id} 

387 for key in col_keys: 

388 if key in self.EXCLUDED_COLUMNS: 

389 continue 

390 val = getattr(row, key) 

391 if val is None: 

392 continue 

393 json_key = self.DB_TO_JSON_KEY_MAP.get(key, key) 

394 msg[json_key] = val 

395 

396 # Ensure content always has a value 

397 msg.setdefault('content', '') 

398 

399 # Mirror usage into info.usage for callers that read it there 

400 if 'usage' in msg: 400 ↛ 401line 400 didn't jump to line 401 because the condition on line 400 was never true

401 msg['info'] = {'usage': msg['usage']} 

402 

403 messages_map[msg_id] = msg 

404 

405 # Reconstruct childrenIds from parentId links so that the map 

406 # is fully navigable (callers like the frontend rely on this). 

407 for msg_id, msg in messages_map.items(): 

408 parent_id = msg.get('parentId') 

409 if parent_id and parent_id in messages_map: 409 ↛ 410line 409 didn't jump to line 410 because the condition on line 409 was never true

410 parent = messages_map[parent_id] 

411 children = parent.get('childrenIds') 

412 if children is None: 

413 parent['childrenIds'] = [msg_id] 

414 elif msg_id not in children: 

415 children.append(msg_id) 

416 

417 # Ensure every message has a childrenIds list (leaf nodes get []) 

418 for msg in messages_map.values(): 

419 if 'childrenIds' not in msg: 419 ↛ 418line 419 didn't jump to line 418 because the condition on line 419 was always true

420 msg['childrenIds'] = [] 

421 

422 return messages_map 

423 

424 async def get_messages_by_user_id( 

425 self, 

426 user_id: str, 

427 skip: int = 0, 

428 limit: int = 50, 

429 db: Optional[AsyncSession] = None, 

430 ) -> list[ChatMessageModel]: 

431 async with get_async_db_context(db) as db: 

432 result = await db.execute( 

433 select(ChatMessage) 

434 .filter_by(user_id=user_id) 

435 .order_by(ChatMessage.created_at.desc()) 

436 .offset(skip) 

437 .limit(limit) 

438 ) 

439 messages = result.scalars().all() 

440 return [ChatMessageModel.model_validate(message) for message in messages] 

441 

442 async def get_messages_by_model_id( 

443 self, 

444 model_id: str, 

445 start_date: Optional[int] = None, 

446 end_date: Optional[int] = None, 

447 skip: int = 0, 

448 limit: int = 100, 

449 db: Optional[AsyncSession] = None, 

450 ) -> list[ChatMessageModel]: 

451 async with get_async_db_context(db) as db: 

452 stmt = select(ChatMessage).filter_by(model_id=model_id) 

453 if start_date: 

454 stmt = stmt.filter(ChatMessage.created_at >= start_date) 

455 if end_date: 455 ↛ 456line 455 didn't jump to line 456 because the condition on line 455 was never true

456 stmt = stmt.filter(ChatMessage.created_at <= end_date) 

457 stmt = stmt.order_by(ChatMessage.created_at.desc()).offset(skip).limit(limit) 

458 result = await db.execute(stmt) 

459 messages = result.scalars().all() 

460 return [ChatMessageModel.model_validate(message) for message in messages] 

461 

462 async def get_chat_ids_by_model_id( 

463 self, 

464 model_id: str, 

465 start_date: Optional[int] = None, 

466 end_date: Optional[int] = None, 

467 skip: int = 0, 

468 limit: int = 50, 

469 db: Optional[AsyncSession] = None, 

470 ) -> list[str]: 

471 """Get distinct chat_ids that used a specific model.""" 

472 

473 async with get_async_db_context(db) as db: 

474 stmt = select( 

475 ChatMessage.chat_id, 

476 func.max(ChatMessage.created_at).label('last_message_at'), 

477 ).filter(ChatMessage.model_id == model_id) 

478 if start_date: 478 ↛ 479line 478 didn't jump to line 479 because the condition on line 478 was never true

479 stmt = stmt.filter(ChatMessage.created_at >= start_date) 

480 if end_date: 480 ↛ 481line 480 didn't jump to line 481 because the condition on line 480 was never true

481 stmt = stmt.filter(ChatMessage.created_at <= end_date) 

482 

483 # Group by chat_id and order by most recent message in each chat 

484 # Secondary sort on chat_id ensures deterministic pagination 

485 stmt = ( 

486 stmt.group_by(ChatMessage.chat_id) 

487 .order_by(func.max(ChatMessage.created_at).desc(), ChatMessage.chat_id) 

488 .offset(skip) 

489 .limit(limit) 

490 ) 

491 result = await db.execute(stmt) 

492 chat_ids = result.all() 

493 return [chat_id for chat_id, _ in chat_ids] 

494 

495 async def delete_messages_by_chat_id(self, chat_id: str, db: Optional[AsyncSession] = None) -> bool: 

496 async with get_async_db_context(db) as db: 

497 await db.execute(delete(ChatMessage).filter_by(chat_id=chat_id)) 

498 await db.commit() 

499 return True 

500 

501 async def delete_message_ids_by_chat_id( 

502 self, 

503 chat_id: str, 

504 message_ids: set[str], 

505 db: Optional[AsyncSession] = None, 

506 ) -> bool: 

507 """Delete specific ``chat_message`` rows by their original message IDs.""" 

508 if not message_ids: 508 ↛ 509line 508 didn't jump to line 509 because the condition on line 508 was never true

509 return True 

510 async with get_async_db_context(db) as db: 

511 await db.execute( 

512 delete(ChatMessage) 

513 .where(ChatMessage.chat_id == chat_id) 

514 .where(ChatMessage.id.in_({f'{chat_id}-{mid}' for mid in message_ids})) 

515 ) 

516 await db.commit() 

517 return True 

518 

519 # Analytics methods 

520 async def get_message_count_by_model( 

521 self, 

522 start_date: Optional[int] = None, 

523 end_date: Optional[int] = None, 

524 group_id: Optional[str] = None, 

525 db: Optional[AsyncSession] = None, 

526 ) -> dict[str, int]: 

527 async with get_async_db_context(db) as db: 

528 from open_webui.models.groups import GroupMember 

529 

530 stmt = select(ChatMessage.model_id, func.count(ChatMessage.id).label('count')).filter( 

531 ChatMessage.role == 'assistant', 

532 ChatMessage.model_id.isnot(None), 

533 ) 

534 

535 if start_date: 

536 stmt = stmt.filter(ChatMessage.created_at >= start_date) 

537 if end_date: 

538 stmt = stmt.filter(ChatMessage.created_at <= end_date) 

539 if group_id: 

540 group_users = select(GroupMember.user_id).filter(GroupMember.group_id == group_id).scalar_subquery() 

541 stmt = stmt.filter(ChatMessage.user_id.in_(group_users)) 

542 

543 stmt = stmt.group_by(ChatMessage.model_id) 

544 result = await db.execute(stmt) 

545 return {row.model_id: row.count for row in result.all()} 

546 

547 async def get_unique_counts_by_model( 

548 self, 

549 start_date: Optional[int] = None, 

550 end_date: Optional[int] = None, 

551 group_id: Optional[str] = None, 

552 db: Optional[AsyncSession] = None, 

553 ) -> dict[str, dict]: 

554 """Count distinct users and chats per model.""" 

555 async with get_async_db_context(db) as db: 

556 from open_webui.models.groups import GroupMember 

557 

558 stmt = select( 

559 ChatMessage.model_id, 

560 func.count(distinct(ChatMessage.user_id)).label('unique_users'), 

561 func.count(distinct(ChatMessage.chat_id)).label('unique_chats'), 

562 ).filter( 

563 ChatMessage.role == 'assistant', 

564 ChatMessage.model_id.isnot(None), 

565 ) 

566 

567 if start_date: 

568 stmt = stmt.filter(ChatMessage.created_at >= start_date) 

569 if end_date: 

570 stmt = stmt.filter(ChatMessage.created_at <= end_date) 

571 if group_id: 

572 group_users = select(GroupMember.user_id).filter(GroupMember.group_id == group_id).scalar_subquery() 

573 stmt = stmt.filter(ChatMessage.user_id.in_(group_users)) 

574 

575 stmt = stmt.group_by(ChatMessage.model_id) 

576 result = await db.execute(stmt) 

577 return { 

578 row.model_id: { 

579 'unique_users': row.unique_users, 

580 'unique_chats': row.unique_chats, 

581 } 

582 for row in result.all() 

583 } 

584 

585 async def get_token_usage_by_model( 

586 self, 

587 start_date: Optional[int] = None, 

588 end_date: Optional[int] = None, 

589 group_id: Optional[str] = None, 

590 db: Optional[AsyncSession] = None, 

591 ) -> dict[str, dict]: 

592 """Aggregate token usage by model using database-level aggregation.""" 

593 async with get_async_db_context(db) as db: 

594 from open_webui.models.groups import GroupMember 

595 

596 # We need the dialect to determine JSON extraction syntax 

597 # For async sessions, access via get_bind() 

598 bind = await db.connection() 

599 dialect = bind.dialect.name 

600 

601 input_tokens, output_tokens = _token_columns(dialect) 

602 

603 stmt = select( 

604 ChatMessage.model_id, 

605 func.coalesce(func.sum(input_tokens), 0).label('input_tokens'), 

606 func.coalesce(func.sum(output_tokens), 0).label('output_tokens'), 

607 func.count(ChatMessage.id).label('message_count'), 

608 ).filter( 

609 ChatMessage.role == 'assistant', 

610 ChatMessage.model_id.isnot(None), 

611 ChatMessage.usage.isnot(None), 

612 ) 

613 

614 if start_date: 

615 stmt = stmt.filter(ChatMessage.created_at >= start_date) 

616 if end_date: 

617 stmt = stmt.filter(ChatMessage.created_at <= end_date) 

618 if group_id: 

619 group_users = select(GroupMember.user_id).filter(GroupMember.group_id == group_id).scalar_subquery() 

620 stmt = stmt.filter(ChatMessage.user_id.in_(group_users)) 

621 

622 stmt = stmt.group_by(ChatMessage.model_id) 

623 result = await db.execute(stmt) 

624 

625 return { 

626 row.model_id: { 

627 'input_tokens': row.input_tokens, 

628 'output_tokens': row.output_tokens, 

629 'total_tokens': row.input_tokens + row.output_tokens, 

630 'message_count': row.message_count, 

631 } 

632 for row in result.all() 

633 } 

634 

635 async def get_token_usage_by_user( 

636 self, 

637 start_date: Optional[int] = None, 

638 end_date: Optional[int] = None, 

639 group_id: Optional[str] = None, 

640 db: Optional[AsyncSession] = None, 

641 ) -> dict[str, dict]: 

642 """Aggregate token usage by user using database-level aggregation.""" 

643 async with get_async_db_context(db) as db: 

644 from open_webui.models.groups import GroupMember 

645 

646 bind = await db.connection() 

647 dialect = bind.dialect.name 

648 

649 input_tokens, output_tokens = _token_columns(dialect) 

650 

651 stmt = select( 

652 ChatMessage.user_id, 

653 func.coalesce(func.sum(input_tokens), 0).label('input_tokens'), 

654 func.coalesce(func.sum(output_tokens), 0).label('output_tokens'), 

655 func.count(ChatMessage.id).label('message_count'), 

656 ).filter( 

657 ChatMessage.role == 'assistant', 

658 ChatMessage.user_id.isnot(None), 

659 ChatMessage.usage.isnot(None), 

660 ) 

661 

662 if start_date: 

663 stmt = stmt.filter(ChatMessage.created_at >= start_date) 

664 if end_date: 

665 stmt = stmt.filter(ChatMessage.created_at <= end_date) 

666 if group_id: 

667 group_users = select(GroupMember.user_id).filter(GroupMember.group_id == group_id).scalar_subquery() 

668 stmt = stmt.filter(ChatMessage.user_id.in_(group_users)) 

669 

670 stmt = stmt.group_by(ChatMessage.user_id) 

671 result = await db.execute(stmt) 

672 

673 return { 

674 row.user_id: { 

675 'input_tokens': row.input_tokens, 

676 'output_tokens': row.output_tokens, 

677 'total_tokens': row.input_tokens + row.output_tokens, 

678 'message_count': row.message_count, 

679 } 

680 for row in result.all() 

681 } 

682 

683 async def get_user_usage_summary( 

684 self, 

685 user_id: str, 

686 start_date: Optional[int] = None, 

687 end_date: Optional[int] = None, 

688 include_active_days: bool = True, 

689 timezone: Optional[str] = None, 

690 db: Optional[AsyncSession] = None, 

691 ) -> dict: 

692 async with get_async_db_context(db) as db: 

693 bind = await db.connection() 

694 dialect = bind.dialect.name 

695 input_tokens, output_tokens = _token_columns(dialect) 

696 

697 messages_stmt = select(ChatMessage.role, func.count(ChatMessage.id).label('count')).filter( 

698 ChatMessage.user_id == user_id, 

699 ) 

700 token_stmt = select( 

701 func.coalesce(func.sum(input_tokens), 0).label('input_tokens'), 

702 func.coalesce(func.sum(output_tokens), 0).label('output_tokens'), 

703 ).filter( 

704 ChatMessage.user_id == user_id, 

705 ChatMessage.role == 'assistant', 

706 ChatMessage.usage.isnot(None), 

707 ) 

708 models_stmt = select(func.count(distinct(ChatMessage.model_id)).label('models_used')).filter( 

709 ChatMessage.user_id == user_id, 

710 ChatMessage.role == 'assistant', 

711 ChatMessage.model_id.isnot(None), 

712 ) 

713 if start_date: 

714 messages_stmt = messages_stmt.filter(ChatMessage.created_at >= start_date) 

715 token_stmt = token_stmt.filter(ChatMessage.created_at >= start_date) 

716 models_stmt = models_stmt.filter(ChatMessage.created_at >= start_date)