source_summary.py

import os
from collections import OrderedDict


def should_show_sources():
    return os.getenv("RAG_SHOW_SOURCES", "1") == "1"


def _get_display_source(meta):
    return (
        meta.get("source_name")
        or meta.get("source")
        or meta.get("source_path")
        or "unknown"
    )


def collect_source_summary(docs):
    summary = OrderedDict()
    for doc in docs or []:
        meta = doc.metadata or {}
        source_key = meta.get("source_path") or meta.get("source") or "unknown"
        item = summary.setdefault(
            source_key,
            {
                "source": meta.get("source", "unknown"),
                "source_name": _get_display_source(meta),
                "source_path": meta.get("source_path", meta.get("source", "unknown")),
                "source_type": meta.get("source_type", "unknown"),
                "hits": 0,
                "pages": [],
                "sheets": [],
            },
        )
        item["hits"] += 1

        page = meta.get("page")
        if page is not None and page not in item["pages"]:
            item["pages"].append(page)

        sheet = meta.get("sheet")
        if sheet and sheet not in item["sheets"]:
            item["sheets"].append(sheet)

    return list(summary.values())


def format_source_summary(summary):
    mode = os.getenv("RAG_SOURCE_SUMMARY_MODE", "compact").strip().lower()
    parts = []

    for item in summary or []:
        label = item.get("source_name", "unknown")
        if mode == "path":
            label = item.get("source_path", label)

        extras = []
        if mode in {"verbose", "debug"}:
            extras.append(item.get("source_type", "unknown"))
            extras.append(f"hits={item.get('hits', 0)}")
            if item.get("pages"):
                extras.append(f"pages={item['pages']}")
            if item.get("sheets"):
                extras.append(f"sheets={item['sheets']}")
        elif mode == "compact+hits":
            extras.append(f"x{item.get('hits', 0)}")

        if extras:
            label = f"{label} ({', '.join(extras)})"
        parts.append(label)

    return ", ".join(parts)