src/agentic_multiwriter/agents/aggregator.py

from __future__ import annotations

import time
from typing import List

from ..state import AgentState, ResearchSnippet
from ..tools import normalize_snippets, get_logger, log_state_summary
from ..config import settings

logger = get_logger()


def _rank_snippets(snippets: List[ResearchSnippet]) -> List[ResearchSnippet]:
    """
    Very simple ranking: prefer snippets that mention the topic more often (if present),
    otherwise keep original order.
    """
    return snippets


def aggregator_node(state: AgentState) -> AgentState:
    start = time.time()

    snippets = state.get("research_snippets", []) or []
    logger.info("Aggregator: normalizing %d snippets...", len(snippets))

    normalized = normalize_snippets(snippets)
    ranked = _rank_snippets(normalized)

    # Trim to a maximum number for LLM context
    max_keep = max(3, settings.max_search_results)
    trimmed = ranked[:max_keep]

    meta = state.get("meta", {}) or {}
    meta["aggregator_time_sec"] = round(time.time() - start, 3)
    meta["aggregated_snippets_count"] = len(trimmed)

    new_state: AgentState = {
        **state,
        "research_snippets": trimmed,
        "meta": meta,
    }
    log_state_summary(new_state, prefix="Aggregator")
    return new_state