forked from zovos/vk_hackathon
133 lines
3.8 KiB
Python
133 lines
3.8 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 _append_limited(parts: list[str], piece: str, limit: int, sep: str) -> bool:
|
|
"""Append text piece to parts while respecting the final joined length limit."""
|
|
if not piece or limit <= 0:
|
|
return False
|
|
|
|
current_len = sum(len(p) for p in parts) + max(0, len(parts)) * len(sep)
|
|
extra_sep = len(sep) if parts else 0
|
|
remaining = limit - current_len - extra_sep
|
|
if remaining <= 0:
|
|
return False
|
|
|
|
parts.append(piece[:remaining])
|
|
return len(piece) <= remaining
|
|
|
|
|
|
def _join_limited(pieces: list[str], sep: str, limit: int) -> str:
|
|
if limit <= 0:
|
|
return ""
|
|
result: list[str] = []
|
|
for piece in pieces:
|
|
fully_added = _append_limited(result, piece, limit, sep)
|
|
if not fully_added:
|
|
break
|
|
return sep.join(result)
|
|
|
|
|
|
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=_join_limited(page_lines, "\n", WINDOW_MAX_CHARS),
|
|
dense_content=_join_limited(dense_lines, "\n", WINDOW_MAX_CHARS),
|
|
sparse_content=_join_limited(sparse_tokens, " ", WINDOW_MAX_CHARS),
|
|
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 = min(len(msg_text), WINDOW_MAX_CHARS)
|
|
|
|
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
|