forked from zovos/vk_hackathon
107 lines
2.9 KiB
Python
107 lines
2.9 KiB
Python
"""Message-based chunking with window by count, length, and time gap."""
|
|
|
|
from cleaning import CleanedMessage, clean_message
|
|
from rendering import render_dense_content, render_page_content, render_sparse_content
|
|
from index_schemas import IndexAPIItem, Message
|
|
|
|
WINDOW_MAX_MESSAGES = 5
|
|
WINDOW_MAX_CHARS = 512
|
|
TIME_GAP_SECONDS = 3600
|
|
OVERLAP_MESSAGES = 2
|
|
|
|
|
|
def _clean_all(messages: list[Message]) -> list[CleanedMessage]:
|
|
cleaned = [clean_message(m) for m in messages]
|
|
return [c for c in cleaned if not c.is_empty]
|
|
|
|
|
|
def _render_chunk(
|
|
overlap: list[CleanedMessage],
|
|
window: list[CleanedMessage],
|
|
) -> IndexAPIItem:
|
|
page_lines: list[str] = []
|
|
dense_lines: list[str] = []
|
|
sparse_tokens: list[str] = []
|
|
|
|
for msg in overlap + window:
|
|
page = render_page_content(msg)
|
|
dense = render_dense_content(msg)
|
|
sparse = render_sparse_content(msg)
|
|
if page:
|
|
page_lines.append(page)
|
|
if dense:
|
|
dense_lines.append(dense)
|
|
if sparse:
|
|
sparse_tokens.append(sparse)
|
|
|
|
return IndexAPIItem(
|
|
page_content="\n".join(page_lines),
|
|
dense_content="\n".join(dense_lines),
|
|
sparse_content=" ".join(sparse_tokens),
|
|
message_ids=[msg.id for msg in window],
|
|
)
|
|
|
|
|
|
def _split_windows(messages: list[CleanedMessage]) -> list[list[CleanedMessage]]:
|
|
"""Split cleaned messages into windows respecting count, length, and time gap."""
|
|
if not messages:
|
|
return []
|
|
|
|
windows: list[list[CleanedMessage]] = []
|
|
current: list[CleanedMessage] = []
|
|
current_chars = 0
|
|
|
|
for msg in messages:
|
|
msg_text = render_page_content(msg)
|
|
msg_chars = len(msg_text)
|
|
|
|
time_break = (
|
|
current
|
|
and (msg.time - current[-1].time) > TIME_GAP_SECONDS
|
|
)
|
|
size_break = (
|
|
current
|
|
and (
|
|
len(current) >= WINDOW_MAX_MESSAGES
|
|
or current_chars + msg_chars > WINDOW_MAX_CHARS
|
|
)
|
|
)
|
|
|
|
if time_break or size_break:
|
|
if current:
|
|
windows.append(current)
|
|
current = [msg]
|
|
current_chars = msg_chars
|
|
else:
|
|
current.append(msg)
|
|
current_chars += msg_chars
|
|
|
|
if current:
|
|
windows.append(current)
|
|
|
|
return windows
|
|
|
|
|
|
def build_chunks(
|
|
overlap_messages: list[Message],
|
|
new_messages: list[Message],
|
|
) -> list[IndexAPIItem]:
|
|
clean_overlap = _clean_all(overlap_messages)
|
|
clean_new = _clean_all(new_messages)
|
|
|
|
if not clean_new:
|
|
return []
|
|
|
|
overlap_tail = clean_overlap[-OVERLAP_MESSAGES:] if clean_overlap else []
|
|
windows = _split_windows(clean_new)
|
|
|
|
result: list[IndexAPIItem] = []
|
|
prev_window_tail: list[CleanedMessage] = overlap_tail
|
|
|
|
for window in windows:
|
|
chunk = _render_chunk(prev_window_tail, window)
|
|
if chunk.message_ids:
|
|
result.append(chunk)
|
|
prev_window_tail = window[-OVERLAP_MESSAGES:]
|
|
|
|
return result
|