chroma/summarize.py

"""Summarize agenda text with the LLM using a map-reduce over chunks.

For short documents the text is summarized in a single pass. For longer ones each
chunk is summarized (the "map" step) and the partial summaries are then combined
into one final summary (the "reduce" step).

The LLM call is injectable via the ``complete`` parameter -- a
``Callable[[str], str]`` (or one also accepting a ``system`` kwarg) -- so callers
can swap in a fake for testing or a different backend. It defaults to
:func:`chroma.llm.chat_complete`.
"""

from __future__ import annotations

import os
from typing import Callable

from .chunking import chunk_text

Completer = Callable[..., str]

# Final summary length cap (env-tunable); raise SUMMARY_REDUCE_TOKENS for longer summaries.
_SUMMARY_REDUCE_TOKENS = int(os.getenv("SUMMARY_REDUCE_TOKENS", "1024"))

_MAP_SYSTEM = (
    "You summarize sections of public-meeting agendas. Be factual and concise; "
    "preserve agenda item numbers, topics, motions, names, dollar amounts, and dates."
)
_MAP_PROMPT = (
    "Summarize this section of a public meeting agenda in a few bullet points, "
    "capturing the concrete agenda items and topics.\n\nSECTION:\n{chunk}"
)
_REDUCE_SYSTEM = (
    "You combine section summaries of a single public-meeting agenda into one "
    "coherent summary. Do not invent content."
)
_REDUCE_PROMPT = (
    "Below are summaries of consecutive sections of ONE meeting agenda. Write a "
    "single clear summary ({style}) covering the main agenda items, decisions, and "
    "topics, in agenda order.\n\nSECTION SUMMARIES:\n{joined}"
)


def _default_completer() -> Completer:
    from .llm import chat_complete

    return chat_complete


def _call(complete: Completer, prompt: str, system: str | None) -> str:
    """Invoke ``complete`` passing ``system`` when its signature accepts it."""
    try:
        return complete(prompt, system=system)
    except TypeError:
        return complete(prompt)


def summarize_text(
    text: str,
    *,
    complete: Completer | None = None,
    chunk_size: int = 4000,
    overlap: int = 300,
    style: str = "of about 5-10 bullet points",
    map_max_tokens: int = 512,
    reduce_max_tokens: int = _SUMMARY_REDUCE_TOKENS,
) -> str:
    """Summarize arbitrary text via single-pass or map-reduce over chunks.

    ``complete`` defaults to :func:`chroma.llm.chat_complete`. To control token
    limits per step, pass a partial/wrapper that fixes ``max_tokens``.
    """
    text = (text or "").strip()
    if not text:
        return ""
    complete = complete or _default_completer()

    chunks = chunk_text(text, chunk_size=chunk_size, overlap=overlap)
    if len(chunks) <= 1:
        prompt = _REDUCE_PROMPT.format(style=style, joined=text)
        return _call(complete, prompt, _REDUCE_SYSTEM)

    # Map: summarize each chunk.
    partials: list[str] = []
    for i, ch in enumerate(chunks, 1):
        part = _call(complete, _MAP_PROMPT.format(chunk=ch), _MAP_SYSTEM)
        if part:
            partials.append(f"[Section {i}/{len(chunks)}]\n{part}")

    # Reduce: combine partial summaries into one.
    joined = "\n\n".join(partials)
    prompt = _REDUCE_PROMPT.format(style=style, joined=joined)
    return _call(complete, prompt, _REDUCE_SYSTEM)


def summarize_document(
    document: dict,
    *,
    complete: Completer | None = None,
    **kwargs,
) -> dict:
    """Summarize one loader document; returns ``{doc_id, metadata, summary, ...}``."""
    summary = summarize_text(document.get("text", ""), complete=complete, **kwargs)
    return {
        "doc_id": document.get("doc_id"),
        "metadata": document.get("metadata", {}),
        "summary": summary,
    }