48 lines
No EOL
1.6 KiB
Python
48 lines
No EOL
1.6 KiB
Python
"""Chat store for MAX surface."""
|
|
from dataclasses import dataclass, field
|
|
from typing import Dict, Optional
|
|
|
|
|
|
@dataclass
|
|
class RoomMeta:
|
|
platform_chat_id: str
|
|
max_chat_id: str
|
|
name: str
|
|
user_id: str
|
|
agent_id: str
|
|
|
|
|
|
@dataclass
|
|
class ChatStore:
|
|
rooms: Dict[str, RoomMeta] = field(default_factory=dict)
|
|
staged_attachments: Dict[str, list] = field(default_factory=dict)
|
|
|
|
def get_room_by_max_chat_id(self, max_chat_id: str) -> Optional[RoomMeta]:
|
|
return self.rooms.get(max_chat_id)
|
|
|
|
def get_room_by_platform_chat_id(self, platform_chat_id: str) -> Optional[RoomMeta]:
|
|
for room in self.rooms.values():
|
|
if room.platform_chat_id == platform_chat_id:
|
|
return room
|
|
return None
|
|
|
|
def add_room(self, room: RoomMeta) -> None:
|
|
self.rooms[room.max_chat_id] = room
|
|
|
|
def remove_room(self, max_chat_id: str) -> None:
|
|
self.rooms.pop(max_chat_id, None)
|
|
self.staged_attachments.pop(max_chat_id, None)
|
|
|
|
def list_rooms_for_user(self, user_id: str) -> list:
|
|
return [r for r in self.rooms.values() if r.user_id == user_id]
|
|
|
|
def stage_attachment(self, max_chat_id: str, attachment: tuple) -> None:
|
|
if max_chat_id not in self.staged_attachments:
|
|
self.staged_attachments[max_chat_id] = []
|
|
self.staged_attachments[max_chat_id].append(attachment)
|
|
|
|
def pop_attachments(self, max_chat_id: str) -> list:
|
|
return self.staged_attachments.pop(max_chat_id, [])
|
|
|
|
def get_attachments(self, max_chat_id: str) -> list:
|
|
return self.staged_attachments.get(max_chat_id, []) |