forked from zovos/vk_hackathon
102 lines
2.7 KiB
Python
102 lines
2.7 KiB
Python
"""Render cleaned messages into page_content, dense_content, sparse_content."""
|
|
|
|
import datetime
|
|
|
|
from .cleaning import CleanedMessage
|
|
|
|
|
|
def _format_time(ts: int) -> str:
|
|
try:
|
|
return datetime.datetime.fromtimestamp(ts, tz=datetime.timezone.utc).strftime("%Y-%m-%d %H:%M")
|
|
except (OSError, OverflowError, ValueError):
|
|
return str(ts)
|
|
|
|
|
|
def render_page_content(msg: CleanedMessage) -> str:
|
|
"""Human-readable text for the chunk payload."""
|
|
lines: list[str] = []
|
|
|
|
prefix = f"{msg.sender_id}:"
|
|
if msg.text:
|
|
lines.append(f"{prefix} {msg.text}")
|
|
elif not msg.parts and not msg.member_event_text:
|
|
lines.append(prefix)
|
|
|
|
for part in msg.parts:
|
|
lines.append(part["text"])
|
|
|
|
if msg.member_event_text:
|
|
lines.append(msg.member_event_text)
|
|
|
|
for fi in msg.file_info:
|
|
name = fi.get("name") or fi.get("url") or "file"
|
|
lines.append(f"[attachment: {name}]")
|
|
|
|
return "\n".join(lines)
|
|
|
|
|
|
def render_dense_content(msg: CleanedMessage) -> str:
|
|
"""Text optimized for semantic dense embedding: role markers + full context."""
|
|
lines: list[str] = []
|
|
ts = _format_time(msg.time)
|
|
|
|
header_parts = [f"[{ts}]", f"sender:{msg.sender_id}"]
|
|
if msg.is_forward:
|
|
header_parts.append("type:forward")
|
|
if msg.is_quote:
|
|
header_parts.append("type:quote")
|
|
if msg.is_system:
|
|
header_parts.append("type:system")
|
|
if msg.thread_sn:
|
|
header_parts.append(f"thread:{msg.thread_sn}")
|
|
lines.append(" ".join(header_parts))
|
|
|
|
if msg.text:
|
|
lines.append(msg.text)
|
|
|
|
for part in msg.parts:
|
|
lines.append(part["text"])
|
|
|
|
if msg.member_event_text:
|
|
lines.append(msg.member_event_text)
|
|
|
|
for fi in msg.file_info:
|
|
fi_parts = []
|
|
if fi.get("name"):
|
|
fi_parts.append(fi["name"])
|
|
if fi.get("mime"):
|
|
fi_parts.append(fi["mime"])
|
|
if fi.get("url"):
|
|
fi_parts.append(fi["url"])
|
|
if fi_parts:
|
|
lines.append(f"[file: {' '.join(fi_parts)}]")
|
|
|
|
if msg.mentions:
|
|
lines.append("mentions: " + ", ".join(msg.mentions))
|
|
|
|
return "\n".join(lines)
|
|
|
|
|
|
def render_sparse_content(msg: CleanedMessage) -> str:
|
|
"""Keyword-heavy text for sparse/BM25 embedding."""
|
|
tokens: list[str] = []
|
|
|
|
tokens.append(msg.sender_id)
|
|
tokens.extend(msg.mentions)
|
|
|
|
if msg.text:
|
|
tokens.append(msg.text)
|
|
|
|
for part in msg.parts:
|
|
tokens.append(part["text"])
|
|
|
|
if msg.member_event_text:
|
|
tokens.append(msg.member_event_text)
|
|
|
|
for fi in msg.file_info:
|
|
for key in ("name", "mime", "url"):
|
|
val = fi.get(key)
|
|
if val:
|
|
tokens.append(val)
|
|
|
|
return " ".join(tokens)
|