ai-rpg/app/engine/world_editor.py

"""World Editor — chat-based editing of an existing world.

Flow:
1. Player sends an instruction via POST /api/worlds/{id}/edit.
2. SSE stream opens (GET /api/sessions/worlds/{id}/editor/stream?instruction=...).
3. LLM gets the instruction + world state + tools. It can:
   - Call entity_create/update/delete, env_update, schema_* tools directly (applied immediately).
   - Call ask_user to clarify (emits clarification SSE event, stream stays open for 60s
     waiting for POST /api/sessions/worlds/{id}/answer).
   - Call propose_changes to suggest a batch (emits change_proposed SSE event, stream
     stays open waiting for POST /api/sessions/worlds/{id}/apply or /discard).
   - Call comment_to_user to send text.
4. When LLM stops calling tools, stream ends with done event.
"""

from __future__ import annotations

import asyncio
import json
import uuid
from typing import Any

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

from app.core.llm import LlmClient, MockLlmClient
from app.core.logging import get_logger
from app.core.time_utils import summarize_schemas
from app.engine.sse import SseEmitter
from app.engine.tools.base import ToolContext, get_registry
from app.models import Entity, World
from app.prompts.registry import get_prompt

_logger = get_logger(__name__)


# In-memory pending propose_changes diffs, keyed by world_id.
# The frontend accepts/rejects via REST, which resolves the Future.
_pending_changes: dict[uuid.UUID, asyncio.Future[bool]] = {}

# In-memory pending ask_user clarifications, keyed by world_id.
_pending_clarifications: dict[uuid.UUID, asyncio.Future[str]] = {}


async def run_world_editor(
    *,
    db: AsyncSession,
    world: World,
    instruction: str,
    llm: LlmClient | MockLlmClient,
    sse: SseEmitter,
    max_iterations: int = 8,
    answer_timeout: float = 120.0,
) -> None:
    """Run a world_editor iteration."""
    try:
        registry = get_registry()
        ctx = ToolContext(db=db, world=world, stage="world_editor", sse_emitter=sse.emit)

        entities = (
            await db.execute(
                select(Entity).where(
                    Entity.world_id == world.id, Entity.deleted_at.is_(None)
                ).limit(30)
            )
        ).scalars().all()
        entities_summary = "\n".join(
            f"- {e.entity_type}: {e.name}" for e in entities
        ) or "(no entities yet)"
        env_json = json.dumps(world.environment or {}, ensure_ascii=False, indent=2)
        if env_json == "{}":
            env_json = "(empty — world has no environment yet. Use env_update to add player, current_location, plot_rails.)"
        sys_prompt = get_prompt("world_editor", "en").format(
            world_name=world.name,
            world_description=world.description or "",
            language=world.language,
            schemas_summary=summarize_schemas(world.schemas or []) or "(no schemas yet)",
            environment_json=env_json,
            entities_summary=entities_summary,
            instruction=instruction,
        )
        messages: list[dict[str, Any]] = [
            {"role": "system", "content": sys_prompt},
            {"role": "user", "content": instruction},
        ]
        tools = registry.to_openai_format("world_editor")
        for _ in range(max_iterations):
            resp = await llm.complete(
                stage="world_editor",
                messages=messages,
                tools=tools,
                temperature=0.5,
                max_tokens=2048,
                world_id=world.id,
                session=db,
            )
            msg = resp.get("message", {})
            # Normalize tool_calls: handle cases where the model returns them as
            # strings or in non-standard formats.
            raw_tcs = msg.get("tool_calls") or []
            tcs: list[dict] = []
            for tc in raw_tcs:
                if isinstance(tc, str):
                    # Try to parse as JSON
                    try:
                        tc = json.loads(tc)
                    except json.JSONDecodeError:
                        continue
                if not isinstance(tc, dict):
                    continue
                tcs.append(tc)

            if not tcs:
                content = msg.get("content", "")
                if content:
                    await sse.emit("comment", {"text": content})
                break

            messages.append(msg)
            done = False
            for tc in tcs:
                fn = tc.get("function") or {}
                if not isinstance(fn, dict):
                    fn = {"name": tc.get("name", ""), "arguments": tc.get("arguments", "{}")}
                tname = fn.get("name", "")
                args_str = fn.get("arguments", "{}")
                if not isinstance(args_str, str):
                    args_str = json.dumps(args_str)
                try:
                    targs = json.loads(args_str) if args_str else {}
                except json.JSONDecodeError:
                    targs = {}
                if tname == "ask_user":
                    await sse.emit("clarification", {
                        "question": targs.get("question", ""),
                        "options": targs.get("options"),
                    })
                    fut: asyncio.Future[str] = asyncio.get_event_loop().create_future()
                    _pending_clarifications[world.id] = fut
                    try:
                        answer = await asyncio.wait_for(fut, timeout=answer_timeout)
                    except asyncio.TimeoutError:
                        await sse.emit("warning", {"code": "answer_timeout",
                                                   "message": "User did not answer in time"})
                        answer = ""
                    finally:
                        _pending_clarifications.pop(world.id, None)
                    messages.append({
                        "role": "tool",
                        "tool_call_id": tc.get("id", ""),
                        "name": tname,
                        "content": json.dumps({"answer": answer}),
                    })
                    continue
                if tname == "propose_changes":
                    diff = targs.get("diff", [])
                    comment = targs.get("comment", "")
                    await sse.emit("change_proposed", {
                        "diff": diff,
                        "comment": comment,
                    })
                    fut2: asyncio.Future[bool] = asyncio.get_event_loop().create_future()
                    _pending_changes[world.id] = fut2
                    try:
                        accepted = await asyncio.wait_for(fut2, timeout=answer_timeout)
                    except asyncio.TimeoutError:
                        accepted = False
                        await sse.emit("warning", {"code": "accept_timeout",
                                                   "message": "User did not respond in time, discarding changes"})
                    finally:
                        _pending_changes.pop(world.id, None)
                    if accepted:
                        await _apply_diff(world, diff)
                        await db.commit()
                        await sse.emit("apply_changes", {"diff": diff})
                    else:
                        await sse.emit("discard_changes", {})
                    # After propose_changes (accepted or rejected), END the loop.
                    # Don't make another LLM call — the user's decision is final.
                    done = True
                    break
                # Regular tool — execute
                result = await registry.execute(tname, targs, ctx)
                messages.append({
                    "role": "tool",
                    "tool_call_id": tc.get("id", ""),
                    "name": tname,
                    "content": json.dumps(result.to_dict(), ensure_ascii=False),
                })
            if done:
                break
        await sse.done({"status": "completed"})
    except Exception as e:  # noqa: BLE001
        _logger.exception("world_editor_failed", world_id=str(world.id), error=str(e))
        await sse.error("internal_error", str(e))


def submit_propose_changes_answer(world_id: uuid.UUID, accepted: bool) -> bool:
    """Resolve a pending propose_changes Future. Returns True if a pending call existed."""
    fut = _pending_changes.get(world_id)
    if fut is None or fut.done():
        return False
    fut.set_result(accepted)
    return True


def submit_clarification_answer(world_id: uuid.UUID, answer: str) -> bool:
    """Resolve a pending ask_user Future. Returns True if a pending call existed."""
    fut = _pending_clarifications.get(world_id)
    if fut is None or fut.done():
        return False
    fut.set_result(answer)
    return True


async def _apply_diff(world: World, diff: list[dict[str, Any]]) -> None:
    """Apply a propose_changes diff to the world.

    Supported path formats:
    - `world.name` → set world.name
    - `world.description` → set world.description
    - `world.language` → set world.language
    - `schemas` → replace entire world.schemas
    - `schemas.<type>` → replace a single schema type
    - `environment.<field_path>` → apply_patch to world.environment
    - `environment` → replace entire world.environment
    - `player.<field>` → shorthand for environment.player.<field>
    - `current_location` → shorthand for environment.current_location
    - `plot_rails.<field>` → shorthand for environment.plot_rails.<field>
    """
    from app.core.state_validator import apply_patch

    env_patch: dict[str, Any] = {}
    for d in diff:
        if not isinstance(d, dict):
            continue
        path = d.get("path", "") or ""
        op = d.get("op", "replace")
        new = d.get("new")
        old = d.get("old")  # noqa: F841 — not used but part of the schema

        # Skip if new is empty dict {} (model sometimes returns empty)
        if isinstance(new, dict) and not new and op == "replace":
            _logger.warning("apply_diff_skip_empty", path=path)
            continue
        if new is None and op == "replace":
            continue

        # World-level fields
        if path == "world.name" or path == "name":
            if isinstance(new, str) and new.strip():
                world.name = new
        elif path == "world.description" or path == "description":
            if isinstance(new, str):
                world.description = new
        elif path == "world.language" or path == "language":
            if isinstance(new, str):
                world.language = new
        # Schemas
        elif path == "schemas":
            if isinstance(new, list):
                world.schemas = new
        elif path.startswith("schemas."):
            if isinstance(new, list):
                world.schemas = new
        # Environment (full replace)
        elif path == "environment":
            if isinstance(new, dict):
                world.environment = new
        # Environment field paths
        elif path.startswith("environment."):
            field = path[len("environment."):]
            env_patch[field] = new
        # Shorthand: player.xxx → environment.player.xxx
        elif path.startswith("player."):
            field = path[len("player."):]
            env_patch[f"player.{field}"] = new
        elif path == "player":
            if isinstance(new, dict):
                env_patch["player"] = new
        # Shorthand: current_location
        elif path == "current_location":
            if isinstance(new, str):
                env_patch["current_location"] = new
            elif isinstance(new, dict) and "name" in new:
                env_patch["current_location"] = new["name"]
        # Shorthand: plot_rails.xxx
        elif path.startswith("plot_rails."):
            field = path[len("plot_rails."):]
            env_patch[f"plot_rails.{field}"] = new
        elif path == "plot_rails":
            if isinstance(new, dict):
                env_patch["plot_rails"] = new
        else:
            _logger.warning("apply_diff_unknown_path", path=path)

    if env_patch:
        new_env, errors = apply_patch(dict(world.environment or {}), env_patch)
        if not errors:
            world.environment = new_env
            # Sync plot_rails column
            if isinstance(new_env.get("plot_rails"), dict):
                world.plot_rails = new_env["plot_rails"]
        else:
            _logger.warning("apply_diff_errors", errors=errors)