agents/router.py

"""
Agent router.

Decides how each incoming query should be processed and gathers the
context needed for the corresponding prompt builder (llm/prompts.py,
built in a later layer). This module does NOT call the LLM — it only
classifies + retrieves.

Routes:
    SINGLE_TOPIC  - e.g. "Explain Merge Sort"          -> retrieve one topic
    COMPARISON    - e.g. "Difference between BFS/DFS"  -> retrieve multiple topics
    FOLLOWUP      - e.g. "Can you explain that again?" -> reuse conversation context
    OUT_OF_SCOPE  - e.g. "Who won the World Cup?"      -> reject politely
"""

from dataclasses import dataclass, field
from enum import Enum

from agents.comparison import gather_comparison_context, has_sufficient_context
from agents.followup import gather_followup_context
from agents.intent_classifier import (
    extract_topics,
    is_comparison_query,
    is_followup_query,
    is_generic_comparison_followup,
    topics_in_recent_history,
)
from agents.topic_categories import get_sibling_topics
from logs.logger import get_logger

logger = get_logger(__name__)


class RouteType(str, Enum):
    SINGLE_TOPIC = "single_topic"
    COMPARISON = "comparison"
    FOLLOWUP = "followup"
    OUT_OF_SCOPE = "out_of_scope"


@dataclass
class RouteDecision:
    route_type: RouteType
    query: str
    topics: list = field(default_factory=list)
    # Populated by gather_context() below, per route type:
    single_chunks: list = field(default_factory=list)
    comparison_context: dict = field(default_factory=dict)
    followup_context: dict = field(default_factory=dict)

    def to_log_dict(self) -> dict:
        return {
            "route_type": self.route_type.value,
            "query": self.query,
            "topics": self.topics,
        }


def classify(query: str, has_conversation_history: bool = False, recent_topics: list = None) -> RouteDecision:
    """
    Pure classification step (no retrieval yet). Order matters:
        1. Follow-up check first — a short pronoun-based query with existing
            history should be treated as a follow-up even if it superficially
            contains comparison words.
        2. Comparison check — requires comparison phrasing AND 2+ topics.
        3. Single topic — exactly one recognized topic mentioned.
        4. Otherwise — out of scope.
    """
    recent_topics = recent_topics or []

    if is_followup_query(query, has_conversation_history):
        decision = RouteDecision(route_type=RouteType.FOLLOWUP, query=query)
        logger.info("Routed to FOLLOWUP: %r", query)
        return decision

    topics = extract_topics(query)

    if is_comparison_query(query) and len(topics) >= 2:
        decision = RouteDecision(route_type=RouteType.COMPARISON, query=query, topics=topics)
        logger.info("Routed to COMPARISON: %r | topics=%s", query, topics)
        return decision

    # NEW: comparison referencing the previous topic implicitly, e.g.
    # "compare with other graph algorithms" after discussing BFS.
    if is_generic_comparison_followup(query) and recent_topics:
        anchor_topic = recent_topics[-1]
        siblings = get_sibling_topics(anchor_topic, exclude=topics, limit=2)
        combined = list(dict.fromkeys(topics + [anchor_topic] + siblings))
        if len(combined) >= 2:
            decision = RouteDecision(route_type=RouteType.COMPARISON, query=query, topics=combined)
            logger.info(
                "Routed to COMPARISON (category follow-up): %r | anchor=%s | topics=%s",
                query, anchor_topic, combined,
            )
            return decision

    if len(topics) >= 1:
        decision = RouteDecision(route_type=RouteType.SINGLE_TOPIC, query=query, topics=topics[:1])
        logger.info("Routed to SINGLE_TOPIC: %r | topic=%s", query, topics[0])
        return decision

    decision = RouteDecision(route_type=RouteType.OUT_OF_SCOPE, query=query)
    logger.info("Routed to OUT_OF_SCOPE: %r", query)
    return decision


def gather_context(
    decision: RouteDecision,
    retriever,
    recent_messages: list = None,
) -> RouteDecision:
    """
    Populates the appropriate context field on `decision` based on its
    route_type. Mutates and returns the same decision object.

    Args:
        decision: output of classify().
        retriever: a rag.retriever.Retriever instance.
        recent_messages: required for FOLLOWUP routes — recent conversation
                          turns, oldest -> newest.
    """
    recent_messages = recent_messages or []

    if decision.route_type == RouteType.SINGLE_TOPIC:
        topic = decision.topics[0] if decision.topics else None
        decision.single_chunks = retriever.retrieve(
            query=decision.query,
            topic_filter=[topic] if topic else None,
            bypass_threshold_if_single_topic=True,  
        )
        if not decision.single_chunks:
            logger.warning(
                "SINGLE_TOPIC query had no chunks pass the similarity "
                "threshold: %r",
                decision.query,
            )

    elif decision.route_type == RouteType.COMPARISON:
        decision.comparison_context = gather_comparison_context(decision.topics, retriever)
        if not has_sufficient_context(decision.comparison_context):
            logger.warning(
                "COMPARISON query missing sufficient context for one or "
                "more topics: %s",
                decision.topics,
            )

    elif decision.route_type == RouteType.FOLLOWUP:
        decision.followup_context = gather_followup_context(
            query=decision.query,
            recent_messages=recent_messages,
            retriever=retriever,
        )

    # OUT_OF_SCOPE needs no retrieval.

    return decision

def route(query: str, retriever, recent_messages: list = None) -> RouteDecision:
    """
    Convenience entrypoint combining classify() + gather_context().

    recent_messages: recent conversation turns (oldest -> newest). Presence
    of any messages here signals "has_conversation_history=True" for
    follow-up detection.
    """
    recent_messages = recent_messages or []
    recent_topics = topics_in_recent_history(recent_messages)
    decision = classify(
        query,
        has_conversation_history=len(recent_messages) > 0,
        recent_topics=recent_topics,
    )
    return gather_context(decision, retriever, recent_messages=recent_messages)