book-ingestor/book_ingestor/summarizer.py

"""
summarizer.py — the ONLY module that calls an LLM (Groq/Llama 4).
Generates book-level and chapter-level summaries.
Keeps prompts tight to minimise token spend.
"""

from __future__ import annotations

import logging

from groq import Groq

from .config import cfg

log = logging.getLogger(__name__)

_client: Groq | None = None


def _get_client() -> Groq:
    global _client
    if _client is None:
        _client = Groq(api_key=cfg.groq_api_key)
    return _client


def _call(prompt: str, max_tokens: int = 512) -> str:
    """Single Groq call. Returns text response."""
    response = _get_client().chat.completions.create(
        model=cfg.groq_model,
        messages=[{"role": "user", "content": prompt}],
        max_tokens=max_tokens,
        temperature=0.3,   # low temp = factual, consistent summaries
    )
    return response.choices[0].message.content.strip()


# ── Public API ─────────────────────────────────────────────────────────────────

def summarize_book(title: str, chapter_summaries: list[str]) -> str:
    """
    Generate a high-level book summary from the chapter summaries.
    Input is cheap: we only send summaries, not raw text.
    """
    joined = "\n\n".join(
        f"[Section {i+1}]: {s}" for i, s in enumerate(chapter_summaries)
    )
    prompt = (
        f'You are summarizing the book "{title}".\n'
        f"Below are summaries of each chapter/section.\n"
        f"Write a concise overall summary (4-6 sentences) covering the main thesis, "
        f"key ideas, and conclusions. Be factual and dense — no filler.\n\n"
        f"{joined}"
    )
    log.info("Generating book summary for: %s", title)
    return _call(prompt, max_tokens=400)


def summarize_chapter(
    book_title: str,
    chapter_title: str,
    chapter_text: str,
    max_input_chars: int = 6000,
) -> str:
    """
    Summarize a single chapter. Truncates input to keep token cost low.
    6000 chars ≈ ~1500 tokens — well within Llama 4 context.
    """
    # Truncate raw text to control input tokens
    truncated = chapter_text[:max_input_chars]
    if len(chapter_text) > max_input_chars:
        truncated += "\n[... text truncated for summary ...]"

    prompt = (
        f'From the book "{book_title}", summarize the chapter "{chapter_title}".\n'
        f"Write 3-5 sentences covering the key points, arguments, and conclusions. "
        f"Be specific and factual.\n\n"
        f"{truncated}"
    )
    log.debug("Summarizing chapter: %s", chapter_title)
    return _call(prompt, max_tokens=300)


def summarize_flat_document(title: str, full_text: str, max_input_chars: int = 8000) -> str:
    """
    Summarize a flat (unstructured) document.
    For long docs, summarizes the first portion — sufficient for most reference material.
    """
    truncated = full_text[:max_input_chars]
    if len(full_text) > max_input_chars:
        truncated += "\n[... text truncated for summary ...]"

    prompt = (
        f'Summarize the following document titled "{title}".\n'
        f"Write 4-6 sentences covering the main topic, key points, and conclusions. "
        f"Be specific and factual.\n\n"
        f"{truncated}"
    )
    log.info("Summarizing flat document: %s", title)
    return _call(prompt, max_tokens=400)