198 lines
6.6 KiB
Python
198 lines
6.6 KiB
Python
|
|
from sqlalchemy.orm import Session
|
|
from sqlalchemy.sql.expression import select
|
|
|
|
import database as db
|
|
import temporary as tmp
|
|
|
|
MESSAGE_PAGE_SIZE = 256
|
|
|
|
class Message:
|
|
def __init__(self, id, content, sender, time=None, kind='message', subkind=None, format=None, thread_id=None):
|
|
self.id = id
|
|
self.content = content
|
|
self.sender = sender
|
|
self.time = time
|
|
self.kind = kind
|
|
self.subkind = subkind
|
|
self.format = format
|
|
self.thread_id = thread_id
|
|
|
|
def from_db_record(record):
|
|
return Message(record.id, record.content, record.sender, record.time)
|
|
|
|
def is_from_user(self):
|
|
return self.sender == "user"
|
|
|
|
def is_from_agent(self, agent_id):
|
|
return self.sender == str(agent_id)
|
|
|
|
def to_json(self):
|
|
return {
|
|
"id": self.id,
|
|
"time": self.time,
|
|
"sender": self.sender,
|
|
"kind": self.kind,
|
|
"subkind": self.subkind,
|
|
"format": self.format,
|
|
"content": self.content,
|
|
"threadId": self.thread_id,
|
|
}
|
|
|
|
def __str__(self):
|
|
time = self.time if self.time != None else 'no-time'
|
|
content = self.content.replace("\n", "\\n")
|
|
return f"Message(#{self.id} ({time}) {self.sender}: {content})"
|
|
|
|
class Thread:
|
|
def __init__(self, id, chat_id, name=None, description=None):
|
|
self.id = id
|
|
self.chat_id = chat_id
|
|
self.name = name
|
|
self.description = description
|
|
|
|
self.messages = []
|
|
self.notes = []
|
|
self.resources = []
|
|
|
|
class ChatView:
|
|
def __init__(self, app, chat_id, participant_type=None, participant_id=None):
|
|
self.app = app
|
|
self.chat_id = chat_id
|
|
|
|
self.name = None
|
|
self.description = None
|
|
|
|
self.participant_type = participant_type
|
|
self.participant_id = participant_id
|
|
|
|
# Thread holding messages without a set thread id
|
|
self.default_thread = Thread(None, chat_id)
|
|
# Most recent thread objects, earliest to latest
|
|
self.threads = []
|
|
# Map of threads by their id
|
|
self.thread_map = {}
|
|
# Long term summary of this chat
|
|
self.summary = None
|
|
# Separate list for note messages from any thread
|
|
# Is it better to exclude them or include them from threads? atm they're included
|
|
self.notes = []
|
|
|
|
# Id of the thread currently in focus, by default the last thread appearing in the conversation
|
|
self.current_thread_id = None
|
|
|
|
self._earliest_message_id = None
|
|
self._earliest_message_time = None
|
|
|
|
async def init_load(self):
|
|
await self._load_basic_data()
|
|
await self._load_next_messages(MESSAGE_PAGE_SIZE)
|
|
|
|
async def _load_basic_data(self):
|
|
if tmp.is_tmp_id(self.chat_id):
|
|
record = self.app.tmp.chats.get(self.chat_id)
|
|
else:
|
|
# Using get_or_none for cleaner existence checks
|
|
record = await db.Chat.get_or_none(id=self.chat_id)
|
|
|
|
if not record:
|
|
raise Exception(f"Chat not found with id {self.chat_id}")
|
|
|
|
self.name = record.name
|
|
self.description = record.description
|
|
|
|
async def _load_next_messages(self, load_count):
|
|
if tmp.is_tmp_id(self.chat_id):
|
|
chat = self.app.tmp.chats[self.chat_id]
|
|
record_batch = [
|
|
msg for msg in chat.messages
|
|
if not self._earliest_message_time or msg.id < self._earliest_message_time
|
|
][:load_count]
|
|
record_batch.reverse()
|
|
records_rev = record_batch
|
|
else:
|
|
# Start building the filter
|
|
filters = {"chat_id": self.chat_id}
|
|
|
|
# Add the 'earliest message' condition if it exists
|
|
if self._earliest_message_id is not None:
|
|
filters["id__lte"] = self._earliest_message_id
|
|
|
|
# Tortoise query: filter, order, and limit
|
|
# ordering by id is more robust than by time
|
|
records_rev = await db.Message.filter(**filters).order_by("-id").limit(load_count)
|
|
|
|
# Update the earliest id for the next load operation
|
|
if records_rev:
|
|
# Note: Since we ordered by -time (descending), index 0 is the newest in this batch (highest id/time)
|
|
# The "earliest" or oldest is the one with the smallest ID/Time in this set
|
|
self._earliest_message_id = records_rev[-1].id
|
|
self._earliest_message_time = records_rev[-1].time
|
|
|
|
if not self.threads:
|
|
# current thread is the latest known from the newest record in batch
|
|
self.current_thread_id = records_rev[0].thread_id
|
|
|
|
# Go over messages from newest to oldest
|
|
for record in records_rev:
|
|
if record.thread_id is None:
|
|
thread = self.default_thread
|
|
else:
|
|
thread = self.thread_map.get(record.thread_id)
|
|
if not thread:
|
|
# Assuming Thread is a local class, not the DB model
|
|
thread = Thread(record.thread_id, self.chat_id)
|
|
self.thread_map[thread.id] = thread
|
|
self.threads.insert(0, thread)
|
|
|
|
# Create local Message object
|
|
message = Message(
|
|
record.id,
|
|
record.content,
|
|
record.sender,
|
|
record.time,
|
|
record.kind,
|
|
record.subkind
|
|
)
|
|
|
|
if message.kind == 'note':
|
|
thread.notes.insert(0, message)
|
|
elif message.kind == 'resource':
|
|
thread.resources.insert(0, message)
|
|
else:
|
|
thread.messages.insert(0, message)
|
|
|
|
def get_thread(self, thread_id):
|
|
if thread_id == None:
|
|
return self.default_thread
|
|
return self.thread_map.get(thread_id)
|
|
|
|
def get_current_thread(self):
|
|
# TODO: if !thread
|
|
return self.get_thread(self.current_thread_id)
|
|
|
|
def debug_string(self):
|
|
result_lines = []
|
|
def _print(line):
|
|
result_lines.append(line)
|
|
|
|
def print_thread(thread, title=None):
|
|
if title:
|
|
_print(title + ":")
|
|
elif thread.name:
|
|
_print(thread.name + ":")
|
|
else:
|
|
_print(f"Thread {thread.id}:")
|
|
for message in thread.messages:
|
|
_print(" " + str(message))
|
|
|
|
print_thread(self.default_thread, "Default Thread")
|
|
for thread in self.threads:
|
|
print_thread(thread)
|
|
|
|
_print("Current thread id: " + str(self.current_thread_id))
|
|
|
|
return "\n".join(result_lines)
|
|
|
|
|