vk_hackathon/search/aggregation.py

23 lines
607 B
Python

from typing import Any
from .config import TOP_K
from .retrieval import extract_message_ids
def aggregate_message_ids(
reranked_head: list[Any],
retrieval_tail: list[Any],
) -> list[str]:
"""Deduplicate and collect top-K message_ids, reranked head first."""
seen: set[str] = set()
result: list[str] = []
for point in reranked_head + retrieval_tail:
for mid in extract_message_ids(point):
if mid not in seen:
seen.add(mid)
result.append(mid)
if len(result) >= TOP_K:
return result
return result