ai-rpg/backend/app/engine/context.py

"""Context manager: builds the LLM prompt context with guaranteed-recent + dynamic summarization."""
from __future__ import annotations

import json
import uuid
from typing import Any, Dict, List, Optional

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.llm import LlmClient
from app.core.settings_service import cast_setting, get_all_settings
from app.logging_setup import get_logger
from app.models import Message, World
from app.prompts.templates import get_prompt

log = get_logger("context")


async def build_orchestrator_messages(
    db: AsyncSession,
    world: World,
    session_id: uuid.UUID,
    action_text: str,
) -> tuple[List[Dict[str, Any]], Dict[str, Any]]:
    """Build the messages list for the orchestrator LLM call.

    Returns (messages, settings_used).
    """
    settings_map = await get_all_settings(db)
    recent_n = int(cast_setting("context.recent_messages", settings_map.get("context.recent_messages", 10)))
    threshold = int(cast_setting("context.compress_threshold", settings_map.get("context.compress_threshold", 20)))
    summary_n = int(cast_setting("context.summary_messages", settings_map.get("context.summary_messages", 10)))

    # Load all messages ordered by seq
    result = await db.execute(
        select(Message).where(Message.session_id == session_id).order_by(Message.seq)
    )
    all_msgs: List[Message] = list(result.scalars().all())

    # Check if we need to compress
    if len(all_msgs) >= threshold:
        await _maybe_compress(db, session_id, all_msgs, summary_n, recent_n, world, settings_map)
        # Reload after compression
        result = await db.execute(
            select(Message).where(Message.session_id == session_id).order_by(Message.seq)
        )
        all_msgs = list(result.scalars().all())

    # Get summary message (the latest summary before the recent window)
    summary_text = ""
    visible_msgs = [m for m in all_msgs if not m.hidden]
    if len(visible_msgs) > recent_n:
        # Look for the latest summary
        summaries = [m for m in all_msgs if m.kind == "summary"]
        if summaries:
            summary_text = summaries[-1].content

    recent = visible_msgs[-recent_n:] if visible_msgs else []

    # Build orchestrator system prompt with current state.
    # NOTE: prompts are always English (system content convention). The LLM
    # produces player-facing text in world.language when relevant (the
    # step-writer prompt interpolates world_language explicitly).
    defn = world.definition or {}
    system_prompt_template = get_prompt("orchestrator", world.language)
    player_state = world.state.get("player", {}) if world.state else {}
    system_prompt = system_prompt_template.format(
        world_name=world.name,
        setting_description=defn.get("setting_description", "")[:800],
        rules=json.dumps(defn.get("rules", {}), ensure_ascii=False)[:600],
        current_time=world.current_time or "",
        player_state=json.dumps(player_state, ensure_ascii=False)[:600],
        plot_rails=json.dumps(defn.get("plot_rails", {}), ensure_ascii=False)[:400],
        summary=summary_text or "(no summary yet)",
    )

    messages: List[Dict[str, Any]] = [{"role": "system", "content": system_prompt}]

    # Add summary as a system note if present
    if summary_text:
        messages.append({"role": "system", "content": f"Past summary:\n{summary_text}"})

    # Add recent visible messages
    for m in recent:
        if m.kind == "player_action":
            messages.append({"role": "user", "content": m.content})
        elif m.kind == "narrative_step":
            messages.append({"role": "assistant", "content": m.content})

    # The current action
    messages.append({"role": "user", "content": f'Player action: "{action_text}"'})

    return messages, settings_map


async def _maybe_compress(
    db: AsyncSession,
    session_id: uuid.UUID,
    all_msgs: List[Message],
    summary_n: int,
    recent_n: int,
    world: World,
    settings_map: Dict[str, Any],
) -> None:
    """If history exceeds threshold, summarize older messages into a single summary message.

    Uses the `submit_summary` tool (tool-calling-first design) to get structured
    output from the summarizer LLM.
    """
    visible = [m for m in all_msgs if not m.hidden]
    if len(visible) <= recent_n + summary_n:
        return

    # Take the messages that will be summarized (everything before the recent window)
    to_summarize = visible[:-recent_n]
    if not to_summarize:
        return

    # Build summarization input (English labels — system content convention)
    summary_input_lines = []
    for m in to_summarize:
        prefix = {
            "player_action": "Player",
            "narrative_step": "Scene",
            "summary": "Summary",
            "orchestrator_plan": "GM",
            "technical_offscreen": "Offscreen",
        }.get(m.kind, m.kind)
        summary_input_lines.append(f"{prefix}: {m.content[:300]}")
    summary_input = "\n\n".join(summary_input_lines)

    llm = LlmClient(settings_map)
    system_prompt = get_prompt("summarizer", world.language)
    from app.engine.tools.tools import SUMMARIZER_TOOL_SCHEMAS
    response = await llm.chat(
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": summary_input[:4000]},
        ],
        tools=SUMMARIZER_TOOL_SCHEMAS,
        temperature=float(cast_setting("llm.summary_temperature", settings_map.get("llm.summary_temperature", 0.3))),
        max_tokens=400,
        purpose="summary",
        session_id=session_id,
        db=db,
    )

    # Extract from submit_summary tool call; fall back to text parse.
    summary_text = response.text or ""
    facts: List[Dict[str, Any]] = []
    extracted = False
    for tc in (response.tool_calls or []):
        if tc.get("function", {}).get("name") == "submit_summary":
            args_str = tc.get("function", {}).get("arguments", "{}")
            try:
                data = json.loads(args_str) if args_str else {}
                summary_text = data.get("summary", response.text or "")
                facts = data.get("facts", []) or []
                extracted = True
            except json.JSONDecodeError:
                pass
            break
    if not extracted:
        # Fallback: extract JSON from text response (older models).
        import re as _re
        json_match = _re.search(r"\{[\s\S]*\}", response.text or "")
        if json_match:
            try:
                data = json.loads(json_match.group(0))
                summary_text = data.get("summary", response.text or "")
                facts = data.get("facts", []) or []
            except json.JSONDecodeError:
                pass

    # Create summary message
    next_seq = (max((m.seq for m in all_msgs), default=0)) + 1
    summary_msg = Message(
        session_id=session_id,
        seq=next_seq,
        role="system",
        kind="summary",
        content=summary_text,
        payload={"summarized_count": len(to_summarize), "facts": facts},
        is_pinned=True,
        hidden=False,
    )
    db.add(summary_msg)

    # Hide the summarized messages (but keep them in DB)
    for m in to_summarize:
        m.hidden = True

    # Index facts into RAG glossary
    if facts:
        from app.core.rag import get_rag
        from app.models import GlossaryEntry
        rag = await get_rag(settings_map)
        for f in facts:
            if not isinstance(f, dict):
                continue
            entry = GlossaryEntry(
                world_id=world.id,
                session_id=session_id,
                kind=f.get("kind", "lore"),
                name=f.get("name", "unknown"),
                description=f.get("description", ""),
                payload={},
            )
            db.add(entry)
            await db.flush()
            await rag.upsert_glossary(
                world_id=world.id,
                entry_id=entry.id,
                kind=entry.kind,
                name=entry.name,
                description=entry.description,
                payload={},
                settings_map=settings_map,
            )

    await db.commit()
    log.info("context_compressed", session_id=str(session_id), summarized=len(to_summarize))


async def build_step_writer_messages(
    db: AsyncSession,
    world: World,
    session_id: uuid.UUID,
    outcome: str,
    narrative_prompt: str,
) -> List[Dict[str, Any]]:
    """Build messages for the step writer LLM call.

    The step-writer prompt is in English (system content convention) but
    instructs the LLM to produce the narrative in world.language.
    """
    defn = world.definition or {}
    player_state = world.state.get("player", {}) if world.state else {}
    system_prompt = get_prompt("step_writer", world.language).format(
        setting_description=defn.get("setting_description", "")[:600],
        current_time=world.current_time or "",
        player_state=json.dumps(player_state, ensure_ascii=False)[:400],
        outcome=outcome,
        narrative_prompt=narrative_prompt[:600],
        world_language=world.language or "en",
    )
    return [{"role": "system", "content": system_prompt}]


async def build_subagent_messages(
    world: World,
    task: str,
    context: str,
) -> List[Dict[str, Any]]:
    """Build messages for a clean-context sub-agent call."""
    system_prompt = get_prompt("subagent", world.language).format(task=task, context=context[:600])
    return [{"role": "system", "content": system_prompt}]