src/context_assembler.py

from typing import List, Dict, Optional


class ContextAssembler:
    """Assemble richer context from FAISS passage hits and document-level metadata.

    This is intentionally lightweight: it groups passages by document (URN), picks
    top documents, and returns a markdown string containing doc summaries and
    representative passages. The builder should later supply `doc_summary` in
    metadata to improve results.
    """

    def __init__(self, manager=None):
        self.manager = manager

    def assemble_from_faiss(self, hits: List[Dict], top_docs: int = 3, passages_per_doc: int = 2,
                            query: str | None = None, reranker=None) -> str:
        if not hits:
            return ""

        # Group passages by document identifier (prefer URN, fallback to source_title)
        groups: Dict[str, List[Dict]] = {}
        for h in hits:
            urn = h.get("urn") or h.get("doc_urn") or h.get("source_urn") or h.get("source_title") or "unknown"
            groups.setdefault(urn, []).append(h)

        # Optionally rerank raw passages by a query using a provided reranker
        if query and reranker is not None:
            try:
                # reranker.rerank expects (query, candidates) and returns ordered candidates
                ordered = reranker.rerank(query, hits, top_k=min(len(hits), top_docs * passages_per_doc))
                # keep only the ordered subset for grouping
                hits = ordered
            except Exception:
                pass

        # Rank documents by best score (if available) or by number of hits
        def doc_score(items: List[Dict]) -> float:
            scores = [float(i.get("score", 0)) for i in items if i.get("score") is not None]
            if scores:
                return max(scores)
            return float(len(items))

        ranked = sorted(groups.items(), key=lambda kv: doc_score(kv[1]), reverse=True)[:top_docs]

        parts: List[str] = []
        for i, (doc_id, passages) in enumerate(ranked, 1):
            # Use the first passage as representative for metadata lookup
            first = passages[0]
            meta_lines = []
            title = first.get("source_title") or first.get("title") or doc_id
            meta_lines.append(f"Legge: {title}")
            art = first.get("article_number")
            if art:
                meta_lines.append(f"Articolo: {art}")
            urn = doc_id if doc_id != title else first.get("urn", "")
            if urn:
                meta_lines.append(f"URN: {urn}")
            # Include doc-level summary if available
            doc_summary = first.get("doc_summary") or first.get("summary")
            if doc_summary:
                meta_lines.append(f"Sommario: {doc_summary}")

            # Select top passages for this doc and include their text
            selected = passages[:passages_per_doc]
            passage_texts = []
            for p in selected:
                txt = p.get("content") or p.get("chunk_text") or p.get("text") or p.get("snippet") or ""
                if txt:
                    passage_texts.append(txt.strip())

            if not passage_texts:
                # fallback to any stored full text
                full = first.get("content") or first.get("full_text") or ""
                if full:
                    passage_texts.append(full.strip()[:2000])

            parts.append(f"[FONTE {i}]\n" + "\n".join(meta_lines) + "\n\nPassaggi selezionati:\n" + "\n\n".join(passage_texts))

        return "\n---\n".join(parts)