Nat20-Notes/backend/app/pipeline/summarize.py

"""
Map-reduce summarization: chunk transcript -> per-chunk DM + player summaries
-> combine into final notes. Backend-agnostic: works against a local Ollama
instance or any OpenAI-compatible hosted API, selected via settings.
"""
from pathlib import Path

import requests

from .turns import chunk_turns, turns_to_text, fmt_time
from ..errors import SummarizationError
from ..logging_config import get_logger

log = get_logger(__name__)

DM_CHUNK_PROMPT = """You are summarizing a chunk of a tabletop RPG session transcript for the Game Master's private notes.
{world_context_block}
Extract concise bullet points:
- Major plot events and decisions made
- NPC interactions (names, what was said/promised/revealed)
- Combat outcomes (who fought what, notable rolls, deaths/near-deaths)
- Loot, rewards, or resources gained/lost
- Any GM secrets, foreshadowing, or plot threads revealed at the table
- Open questions or hooks left dangling

Be concise and factual. Use the speaker names given. Do not invent anything not in the transcript.

Transcript chunk (time range {time_range}):
{chunk_text}
"""

PLAYER_CHUNK_PROMPT = """You are summarizing a chunk of a tabletop RPG session transcript for a player-facing recap.
{world_context_block}
Extract concise bullet points of what happened IN THE STORY from the players' in-character perspective:
- Where the party went, who they met, what was said in-character
- Combat encounters and outcomes
- Loot/items the party found
- Decisions the party made and their in-fiction consequences

Exclude out-of-character rules discussion, GM asides, and anything that reads as a secret not yet revealed to characters.
If unsure whether something is a spoiler, leave it out.

Transcript chunk (time range {time_range}):
{chunk_text}
"""

DM_FINAL_PROMPT = """Combine these chunk summaries from a single session into one cohesive GM session log.
Organize into sections: Recap, NPC Interactions, Combat & Encounters, Loot & Rewards, Decisions & Consequences, Open Threads / Hooks for Next Session.
Write in clear prose/bullets, no fluff, no repetition across chunks.

Chunk summaries:
{summaries}
"""

PLAYER_FINAL_PROMPTS = {
    "story": """Combine these chunk summaries into a single "previously on" style recap for the players, in-character perspective only.
Write it like a story recap they could reread before the next session. No GM secrets, no OOC content. Keep it engaging but concise.

Chunk summaries:
{summaries}
""",
    "diary": """Combine these chunk summaries into a "dear diary" style entry written from the collective party's first-person ("we") perspective.
Focus on the party's emotions, reactions, and personal reflections alongside the events.
No GM secrets, no OOC content.

Chunk summaries:
{summaries}
""",
    "bullets": """Combine these chunk summaries into clean, scannable bullet points for the players.
Group related items under headers (Exploration, Combat, NPCs, Loot, Decisions).
Each bullet is one concise fact — no narrative fluff or transitions. No GM secrets, no OOC content.

Chunk summaries:
{summaries}
""",
}



def _call_ollama(host: str, model: str, prompt: str) -> str:
    url = f"{host.rstrip('/')}/api/generate"
    try:
        resp = requests.post(
            url,
            json={"model": model, "prompt": prompt, "stream": False, "options": {"num_ctx": 8192}},
            timeout=600,
        )
    except requests.exceptions.ConnectionError as e:
        raise SummarizationError(
            f"Could not connect to Ollama at {host}. Check the host/port in Settings and that Ollama is running.", cause=e
        )
    except requests.exceptions.Timeout as e:
        raise SummarizationError(f"Ollama at {host} didn't respond within 10 minutes - it may be overloaded or stuck.", cause=e)

    if resp.status_code == 404:
        raise SummarizationError(
            f"Ollama returned 404 for model '{model}'. This almost always means the model isn't pulled on that host. "
            f"Run: ollama pull {model}", cause=None
        )
    try:
        resp.raise_for_status()
    except requests.exceptions.HTTPError as e:
        raise SummarizationError(f"Ollama returned an error ({resp.status_code}): {resp.text[:300]}", cause=e)

    return resp.json()["response"].strip()


def _call_api(base_url: str, api_key: str, model: str, prompt: str) -> str:
    url = f"{base_url.rstrip('/')}/chat/completions"
    try:
        resp = requests.post(
            url,
            headers={"Authorization": f"Bearer {api_key}"},
            json={"model": model, "messages": [{"role": "user", "content": prompt}]},
            timeout=600,
        )
    except requests.exceptions.ConnectionError as e:
        raise SummarizationError(f"Could not connect to API at {base_url}. Check the base URL in Settings.", cause=e)
    except requests.exceptions.Timeout as e:
        raise SummarizationError(f"API at {base_url} didn't respond within 10 minutes.", cause=e)

    if resp.status_code == 401:
        raise SummarizationError("API rejected the request as unauthorized (401) - check your API key in Settings.", cause=None)
    if resp.status_code == 404:
        raise SummarizationError(f"API returned 404 - check the base URL and that model '{model}' exists for this provider.", cause=None)
    try:
        resp.raise_for_status()
    except requests.exceptions.HTTPError as e:
        raise SummarizationError(f"API returned an error ({resp.status_code}): {resp.text[:300]}", cause=e)

    try:
        return resp.json()["choices"][0]["message"]["content"].strip()
    except (KeyError, IndexError) as e:
        raise SummarizationError(f"API response didn't have the expected shape: {resp.text[:300]}", cause=e)


def make_llm_caller(settings: dict):
    if settings.get("llm_mode") == "api":
        base_url, api_key, model = settings["api_base_url"], settings["api_key"], settings["api_model"]
        return lambda prompt: _call_api(base_url, api_key, model, prompt)
    host, model = settings["ollama_host"], settings["ollama_model"]
    return lambda prompt: _call_ollama(host, model, prompt)


def summarize_session(turns: list[dict], settings: dict, progress_cb=None,
                      player_recap_style: str | None = None,
                      player_recap_custom_prompt: str | None = None) -> tuple[str, str]:
    call = make_llm_caller(settings)
    ctx_path = settings.get("world_context_path") or ""
    if ctx_path and Path(ctx_path).exists():
        world_context = Path(ctx_path).read_text(encoding="utf-8").strip()
    else:
        world_context = (settings.get("world_context") or "").strip()
    world_block = f"\nCampaign context (use this to recognize names/places correctly):\n{world_context}\n" if world_context else ""

    chunks = chunk_turns(turns, int(settings.get("chunk_word_target", 2500)))
    dm_summaries, player_summaries = [], []

    for i, chunk in enumerate(chunks):
        if progress_cb:
            progress_cb(f"Summarizing chunk {i+1}/{len(chunks)}")
        time_range = f"{fmt_time(chunk[0]['start'])}-{fmt_time(chunk[-1]['end'])}"
        chunk_text = turns_to_text(chunk)
        dm_summaries.append(call(DM_CHUNK_PROMPT.format(world_context_block=world_block, time_range=time_range, chunk_text=chunk_text)))
        player_summaries.append(call(PLAYER_CHUNK_PROMPT.format(world_context_block=world_block, time_range=time_range, chunk_text=chunk_text)))

    if progress_cb:
        progress_cb("Writing final notes...")
    dm_final = call(DM_FINAL_PROMPT.format(summaries="\n\n".join(dm_summaries)))

    # Select player recap prompt based on style setting
    style = (player_recap_style or settings.get("player_recap_style") or "story")
    if style == "custom":
        raw = (player_recap_custom_prompt or settings.get("player_recap_custom_prompt") or "").strip()
        player_template = raw if raw else PLAYER_FINAL_PROMPTS["story"]
    else:
        player_template = PLAYER_FINAL_PROMPTS.get(style, PLAYER_FINAL_PROMPTS["story"])
    player_final = call(player_template.format(summaries="\n\n".join(player_summaries)))
    return dm_final, player_final