102 lines
3.2 KiB
Python
102 lines
3.2 KiB
Python
from dataclasses import dataclass
|
|
from typing import List, Callable
|
|
from enum import Enum
|
|
import datetime
|
|
|
|
from .io_buffer import IOBuffer
|
|
from .util import format_timestamp
|
|
|
|
class MessageType(Enum):
|
|
USER = "user"
|
|
ASSISTANT = "assistant"
|
|
|
|
@dataclass
|
|
class ChatMessage:
|
|
id: str # Formatted timestamp
|
|
content: str
|
|
message_type: MessageType
|
|
|
|
class ChatIOBuffer(IOBuffer):
|
|
def __init__(self):
|
|
self._messages: List[ChatMessage] = []
|
|
self._last_read_timestamp: str = format_timestamp(datetime.datetime.now())
|
|
self._change_handlers: List[Callable] = []
|
|
|
|
@property
|
|
def messages(self) -> List[ChatMessage]:
|
|
return self._messages.copy()
|
|
|
|
@property
|
|
def last_read_timestamp(self) -> str:
|
|
return self._last_read_timestamp
|
|
|
|
def read(self) -> str:
|
|
unread_messages = self._unread_messages()
|
|
self._last_read_timestamp = self._format_now()
|
|
self._notify_handlers([], [])
|
|
return unread_messages
|
|
|
|
def write(self, content: str) -> None:
|
|
message = ChatMessage(
|
|
id=self._format_now(),
|
|
content=content,
|
|
message_type=MessageType.ASSISTANT
|
|
)
|
|
self._messages.append(message)
|
|
self._notify_handlers(
|
|
new_messages=[message],
|
|
deleted_message_ids=[]
|
|
)
|
|
|
|
def buffer_length(self) -> int:
|
|
return len(self._unread_messages())
|
|
|
|
def add_user_message(self, content: str):
|
|
message = ChatMessage(
|
|
id=self._format_now(),
|
|
content=content,
|
|
message_type=MessageType.USER
|
|
)
|
|
self._messages.append(message)
|
|
self._notify_handlers(
|
|
new_messages=[message],
|
|
deleted_message_ids=[]
|
|
)
|
|
|
|
def delete_message(self, message_id: str):
|
|
for i, msg in enumerate(self._messages):
|
|
if msg.id == message_id:
|
|
self._messages.pop(i)
|
|
self._notify_handlers(
|
|
new_messages=[],
|
|
deleted_message_ids=[message_id]
|
|
)
|
|
break
|
|
|
|
def clear(self) -> None:
|
|
deleted_ids = [msg.id for msg in self._messages]
|
|
self._messages.clear()
|
|
self._last_read_timestamp = self._format_now()
|
|
self._notify_handlers(
|
|
new_messages=[],
|
|
deleted_message_ids=deleted_ids
|
|
)
|
|
|
|
def add_change_handler(self, handler: Callable[[
|
|
List[ChatMessage], # new messages
|
|
List[str], # deleted message ids
|
|
str # last read timestamp
|
|
], None]) -> None:
|
|
self._change_handlers.append(handler)
|
|
|
|
def _notify_handlers(self, new_messages: List[ChatMessage],
|
|
deleted_message_ids: List[str]) -> None:
|
|
for handler in self._change_handlers:
|
|
handler(new_messages, deleted_message_ids, self._last_read_timestamp)
|
|
|
|
def _unread_messages(self) -> List[ChatMessage]:
|
|
return "\n".join(msg.content for msg in self._messages
|
|
if msg.message_type == MessageType.USER and msg.id > self._last_read_timestamp)
|
|
|
|
def _format_now(self) -> str:
|
|
return format_timestamp(datetime.datetime.now()) |