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

"""World builder: multi-turn dialogue to produce a finalized WorldDefinition.

Design (v2 — tool-calling-first):
  The world-builder LLM is given a single tool, `submit_world_definition`,
  which it calls when it has enough information to propose a world. The LLM's
  text response is the conversational reply shown to the player (in
  world.language). This replaces the old "return JSON in your text response"
  pattern which conflicted with tool use and caused raw JSON to leak into
  the chat.
"""
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 Preset, User, World
from app.prompts.templates import get_prompt
from app.schemas import WorldBuilderReply, WorldDefinition
from app.engine.tools.tools import WORLD_BUILDER_TOOL_SCHEMAS

log = get_logger("world_builder")


# In-memory store of world-builder dialogues (session_id -> dialogue state).
# For production scale, move this to Redis. For MVP single-instance it's fine.
_DIALOGUES: Dict[uuid.UUID, Dict[str, Any]] = {}


async def start_world_builder(
    db: AsyncSession,
    user: User,
    world_name: str,
    language: str,
    preset_id: Optional[uuid.UUID],
    setting_brief: str,
    character_brief: str,
    rules_brief: str,
    notes: str,
) -> WorldBuilderReply:
    """Kick off a new world-builder dialogue. Returns the first AI reply."""
    session_id = uuid.uuid4()
    settings_map = await get_all_settings(db)
    llm = LlmClient(settings_map)

    preset_payload: Optional[Dict[str, Any]] = None
    if preset_id:
        result = await db.execute(select(Preset).where(Preset.id == preset_id))
        preset = result.scalars().first()
        if preset:
            preset_payload = preset.payload

    user_brief = _build_user_brief(
        world_name=world_name,
        setting_brief=setting_brief,
        character_brief=character_brief,
        rules_brief=rules_brief,
        notes=notes,
        preset_payload=preset_payload,
        language=language,
    )

    # System prompt is English (system content convention). The LLM is told to
    # produce player-facing text in `language` (interpolated as world_language).
    system_prompt = get_prompt("world_builder", language).format(world_language=language)
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_brief},
    ]

    response = await llm.chat(
        messages=messages,
        tools=WORLD_BUILDER_TOOL_SCHEMAS,
        temperature=float(cast_setting("llm.temperature", settings_map.get("llm.temperature", 0.7))),
        purpose="world_builder",
        user_id=user.id,
        db=db,
    )

    ai_text, proposed, is_final = _extract_world_definition(
        response.text, response.tool_calls, user_confirmation=False,
    )

    _DIALOGUES[session_id] = {
        "user_id": user.id,
        "world_name": world_name,
        "language": language,
        "preset_id": preset_id,
        "messages": messages + [
            {
                "role": "assistant",
                "content": response.text or "",
                "tool_calls": response.tool_calls or None,
            },
        ],
        "turn": 1,
        "last_proposed": proposed.model_dump() if proposed else None,
    }

    return WorldBuilderReply(
        session_id=session_id,
        turn=1,
        ai_message=ai_text,
        proposed_definition=proposed,
        is_final=is_final,
        followup_questions=[],
    )


async def continue_world_builder(
    db: AsyncSession,
    user: User,
    session_id: uuid.UUID,
    user_message: str,
) -> WorldBuilderReply:
    """Continue an existing world-builder dialogue."""
    dialogue = _DIALOGUES.get(session_id)
    if not dialogue:
        raise ValueError("dialogue_not_found")
    if dialogue["user_id"] != user.id:
        raise ValueError("forbidden")

    settings_map = await get_all_settings(db)
    llm = LlmClient(settings_map)
    dialogue["messages"].append({"role": "user", "content": user_message})
    dialogue["turn"] += 1

    # Detect if the player is explicitly confirming the world is ready.
    # If so, we'll force is_final=True on the extracted definition (the model
    # often forgets to set is_final even when the player clearly accepted).
    user_confirms = _is_user_confirmation(user_message)

    response = await llm.chat(
        messages=dialogue["messages"],
        tools=WORLD_BUILDER_TOOL_SCHEMAS,
        temperature=float(cast_setting("llm.temperature", settings_map.get("llm.temperature", 0.7))),
        purpose="world_builder",
        user_id=user.id,
        db=db,
    )
    dialogue["messages"].append({
        "role": "assistant",
        "content": response.text or "",
        "tool_calls": response.tool_calls or None,
    })

    ai_text, proposed, is_final = _extract_world_definition(
        response.text, response.tool_calls, user_confirmation=user_confirms,
    )

    # If the model didn't propose a new definition this turn but we already had
    # one stored and the player just confirmed, reuse the stored definition
    # and mark it final.
    if proposed is None and user_confirms and dialogue.get("last_proposed"):
        proposed = WorldDefinition.model_validate(dialogue["last_proposed"])
        is_final = True

    if proposed:
        dialogue["last_proposed"] = proposed.model_dump()

    return WorldBuilderReply(
        session_id=session_id,
        turn=dialogue["turn"],
        ai_message=ai_text,
        proposed_definition=proposed,
        is_final=is_final,
        followup_questions=[],
    )


async def commit_world_builder(
    db: AsyncSession,
    user: User,
    session_id: uuid.UUID,
    name: Optional[str] = None,
) -> World:
    """Commit the proposed world definition into a real World row."""
    dialogue = _DIALOGUES.get(session_id)
    if not dialogue:
        raise ValueError("dialogue_not_found")
    if dialogue["user_id"] != user.id:
        raise ValueError("forbidden")
    proposed = dialogue.get("last_proposed")
    if not proposed:
        raise ValueError("no_proposed_definition")

    definition = WorldDefinition.model_validate(proposed)
    world = World(
        owner_id=user.id,
        name=name or dialogue.get("world_name") or "New World",
        language=dialogue.get("language", "en"),
        definition=definition.model_dump(),
        state=definition.initial_state or {},
        current_time=definition.initial_time,
        status="ready",
        preset_id=dialogue.get("preset_id"),
    )
    db.add(world)
    await db.commit()
    await db.refresh(world)

    # Clean up dialogue
    _DIALOGUES.pop(session_id, None)
    return world


def _build_user_brief(
    world_name: str,
    setting_brief: str,
    character_brief: str,
    rules_brief: str,
    notes: str,
    preset_payload: Optional[Dict[str, Any]],
    language: str,
) -> str:
    parts = [f"=== WORLD BRIEF ==="]
    parts.append(f"Player-facing language: {language}")
    parts.append(f"Name: {world_name}")
    if preset_payload:
        parts.append(f"Preset seed: {preset_payload.get('world_seed_prompt', '')}")
        parts.append(f"Suggested rules: {json.dumps(preset_payload.get('rules', {}), ensure_ascii=False)[:400]}")
    if setting_brief:
        parts.append(f"Setting: {setting_brief}")
    if character_brief:
        parts.append(f"Character: {character_brief}")
    if rules_brief:
        parts.append(f"Rules: {rules_brief}")
    if notes:
        parts.append(f"Notes: {notes}")
    parts.append("\nAsk 2-4 clarifying questions OR call submit_world_definition with a proposed world.")
    return "\n".join(parts)


# Phrases in EN/RU that the player might type to confirm a proposed world is
# ready to commit. Matched case-insensitively against the player's message.
# Keep this list SHORT and specific — false positives would auto-finalize a
# world the player didn't intend to accept.
_CONFIRMATION_PHRASES = (
    # English
    "ready", "looks good", "looks great", "perfect", "ok", "okay", "fine",
    "yes", "yep", "yeah", "sure", "go ahead", "confirm", "confirmed",
    "approve", "approved", "let's go", "lets go", "do it", "lgtm",
    "i'm happy", "im happy", "ship it", "all good",
    # Russian
    "готово", "готов", "супер", "ок", "окей", "хорошо", "отлично",
    "да", "согласен", "согласна", "подтверждаю", "одобряю", "норм",
    "нормально", "поехали", "создай", "сохраняй", "принимаю",
)


def _is_user_confirmation(message: str) -> bool:
    """Return True if the player's message looks like explicit confirmation.

    Heuristic: the message is short (under 60 chars) AND contains one of the
    known confirmation phrases. Longer messages are treated as edits/feedback,
    not confirmation, even if they contain a "yes".
    """
    if not message:
        return False
    msg = message.strip().lower()
    if not msg or len(msg) > 60:
        return False
    # Exact match or substring — both work. The phrase list is short enough
    # that false positives are rare in normal player edits.
    return any(phrase in msg for phrase in _CONFIRMATION_PHRASES)


def _extract_world_definition(
    text: str,
    tool_calls: Optional[List[Dict[str, Any]]],
    user_confirmation: bool = False,
) -> tuple[str, Optional[WorldDefinition], bool]:
    """Extract AI message text, proposed definition (if any), and is_final flag.

    Strategy (in order):
    1. If the model returned a `submit_world_definition` tool_call — use its
       arguments directly. This is the preferred path for models that support
       OpenAI-style function calling.
    2. Otherwise, scan the text for a JSON object (in a ```json fenced block
       or a bare `{...}` block). Many smaller local models emit the structured
       payload as text instead of using tool_calls; we still want to honor it.
    3. If `user_confirmation=True` (the player just said something like
       "ok" / "ready" / "go") and we have a proposed definition, force
       `is_final=True` even if the model forgot to set it.
    """
    proposed: Optional[WorldDefinition] = None
    is_final = False
    ai_text = _sanitize_model_text(text or "")

    # 1) Prefer the submit_world_definition tool call (the proper way).
    if tool_calls:
        for tc in tool_calls:
            if tc.get("function", {}).get("name") == "submit_world_definition":
                args_str = tc.get("function", {}).get("arguments", "{}")
                data = _safe_json_loads(args_str, {})
                if data:
                    proposed = _try_build_definition(data)
                    is_final = bool(data.get("is_final", False))
                    break

    # 2) Fallback: parse a JSON block from the text (models without tool_calls).
    if proposed is None:
        json_str = _extract_json_block(ai_text)
        if json_str:
            data = _safe_json_loads(json_str, None)
            if isinstance(data, dict):
                # Accept either a flat world-definition object or a wrapper
                # like {"proposed_definition": {...}, "is_final": bool, "ai_message": "..."}.
                target = data
                if "proposed_definition" in data and isinstance(data["proposed_definition"], dict):
                    target = data["proposed_definition"]
                    if "is_final" in data:
                        is_final = bool(data["is_final"])
                    if "ai_message" in data and isinstance(data["ai_message"], str):
                        # Strip the JSON wrapper from the visible text so the
                        # player doesn't see raw JSON in chat.
                        ai_text = data["ai_message"]
                else:
                    if "is_final" in data:
                        is_final = bool(data["is_final"])
                proposed = _try_build_definition(target)

    # 3) If the player just confirmed and we have a definition, force is_final.
    if proposed is not None and user_confirmation and not is_final:
        is_final = True

    # 4) Strip any leaked JSON block from the visible AI text so the player
    # never sees raw JSON in chat. Keep the prose portion only.
    if proposed is not None:
        ai_text = _strip_json_blocks(ai_text).strip()
        if not ai_text:
            # Model returned only JSON with no prose — synthesize a short
            # confirmation message in the player's language.
            ai_text = "(definition ready)"

    return ai_text, proposed, is_final


def _sanitize_model_text(text: str) -> str:
    """Remove common model-output artifacts that would break JSON parsing.

    Strips:
    - End-of-sequence tokens (model-specific control tokens that occasionally
      leak into the decoded text).
    - Leading/trailing whitespace per line.
    - Empty leading lines.
    """
    if not text:
        return ""
    # Remove EOS-style control tokens (single token on its own line, or
    # repeated tokens, or trailing tokens).
    import re as _re
    # Collapse repeated EOS tokens anywhere in the text.
    cleaned = _re.sub(r"<\s*/?\s*[a-zA-Z]+\s*>", "", text)
    # Collapse 3+ blank lines into 1.
    cleaned = _re.sub(r"\n{3,}", "\n\n", cleaned)
    return cleaned.strip()


def _strip_json_blocks(text: str) -> str:
    """Remove fenced ```json ... ``` blocks and bare trailing {...} blocks
    from `text`. Used to clean the player-visible AI message after we've
    already extracted the structured data.
    """
    if not text:
        return ""
    import re as _re
    # Strip fenced code blocks (any language).
    cleaned = _re.sub(r"```[a-zA-Z]*\s*[\s\S]*?```", "", text)
    # Strip trailing bare JSON object (last {...} block in the text).
    # We only remove it if it's the LAST substantial thing in the text,
    # to avoid mangling prose that legitimately contains braces.
    m = _re.search(r"\n\{[\s\S]*\}\s*$", cleaned)
    if m:
        cleaned = cleaned[: m.start()] + cleaned[m.end():]
    return cleaned


def _safe_json_loads(s: str, default: Any) -> Any:
    """Tolerant JSON loader. Returns `default` on failure.

    Attempts standard json.loads first. If that fails, tries to repair common
    model-output mistakes:
    - Trailing commas before } or ].
    - Single quotes instead of double quotes.
    - Unescaped newlines inside string values.
    """
    if not s:
        return default
    try:
        return json.loads(s)
    except json.JSONDecodeError:
        pass
    # Repair attempt 1: remove trailing commas.
    import re as _re
    repaired = _re.sub(r",\s*([}\]])", r"\1", s)
    try:
        return json.loads(repaired)
    except json.JSONDecodeError:
        pass
    # Repair attempt 2: replace single quotes with double quotes (naive, but
    # catches many cases where the model emits pseudo-JSON).
    try:
        repaired2 = repaired.replace("'", '"')
        return json.loads(repaired2)
    except json.JSONDecodeError:
        return default


def _try_build_definition(data: Dict[str, Any]) -> Optional[WorldDefinition]:
    try:
        return WorldDefinition.model_validate(data)
    except Exception as e:
        log.warning("world_definition_invalid", error=str(e), keys=list(data.keys()))
        return None


def _extract_json_block(text: str) -> Optional[str]:
    """Find the first JSON object block in text.

    Prefers a fenced ```json ... ``` block. Falls back to the largest balanced
    {...} block in the text.
    """
    if not text:
        return None
    import re as _re
    # Fenced block (```json ... ``` or ``` ... ```).
    m = _re.search(r"```(?:json)?\s*(\{[\s\S]*?\})\s*```", text)
    if m:
        return m.group(1)
    # Bare block: find the first `{` and balance braces, respecting strings
    # and escape sequences.
    start = text.find("{")
    if start == -1:
        return None
    depth = 0
    in_str = False
    esc = False
    for i in range(start, len(text)):
        c = text[i]
        if in_str:
            if esc:
                esc = False
            elif c == "\\":
                esc = True
            elif c == '"':
                in_str = False
        else:
            if c == '"':
                in_str = True
            elif c == "{":
                depth += 1
            elif c == "}":
                depth -= 1
                if depth == 0:
                    return text[start:i + 1]
    return None