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)

    _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

    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)
    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)


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

    Looks for a `submit_world_definition` tool call first. Falls back to
    JSON-in-text parse for older models that don't honor the tool.
    """
    proposed: Optional[WorldDefinition] = None
    is_final = False
    ai_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", "{}")
                try:
                    data = json.loads(args_str) if args_str else {}
                except json.JSONDecodeError:
                    data = {}
                proposed = _try_build_definition(data)
                is_final = bool(data.get("is_final", False))
                break

    # 2) Fallback: parse a JSON block from the text (older models).
    if proposed is None:
        json_str = _extract_json_block(text)
        if json_str:
            try:
                data = json.loads(json_str)
                if isinstance(data, dict):
                    if "proposed_definition" in data and isinstance(data["proposed_definition"], dict):
                        proposed = _try_build_definition(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):
                        ai_text = data["ai_message"]
            except json.JSONDecodeError:
                pass

    # 3) Heuristic: if a definition was proposed and the text mentions "ready",
    # mark as final.
    if proposed is not None and not is_final:
        low = ai_text.lower()
        if any(kw in low for kw in ["world is ready", "world_ready", "ready to commit", "мир готов"]):
            is_final = True

    return ai_text, proposed, is_final


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))
        return None


def _extract_json_block(text: str) -> Optional[str]:
    """Find the first JSON object/array block in text (fallback path only)."""
    if not text:
        return None
    import re
    m = re.search(r"```(?:json)?\s*(\{[\s\S]*?\})\s*```", text)
    if m:
        return m.group(1)
    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