DungeonsDragons/backend/app/llm/orchestrator.py

import json
import logging
import uuid
from collections.abc import Awaitable, Callable

from sqlalchemy.ext.asyncio import AsyncSession

from app.config import settings
from app.llm.client import get_dm_system_prompt, get_llm_client
from app.llm.context import build_context
from app.llm.json_utils import fix_double_escaped_unicode
from app.llm.tools import character_sheet, dice, end_game, monster
from app.models.game import Game
from app.models.message import Message
from app.rag.retrieval import build_rag_block

logger = logging.getLogger("app.llm.orchestrator")

MAX_TOOL_ROUNDS = 5
MAX_TOKENS = 4096

TOOLS = [
    {"type": "function", "function": dice.TOOL_SCHEMA},
    {"type": "function", "function": character_sheet.TOOL_SCHEMA},
    {"type": "function", "function": end_game.TOOL_SCHEMA},
    {"type": "function", "function": monster.TOOL_SCHEMA},
]


async def _execute_tool_call(
    session: AsyncSession, game_id: uuid.UUID, tool_name: str, tool_input: dict
) -> dict:
    """Runs one tool call. Errors are returned as a payload (not raised), so the DM sees them
    as a tool result and can recover instead of the whole turn crashing."""
    try:
        if tool_name == "roll_dice":
            return dice.roll(tool_input["notation"])
        if tool_name == "upsert_character_sheet":
            return await character_sheet.upsert(session, game_id, tool_input)
        if tool_name == "end_game":
            return await end_game.end(session, game_id, tool_input)
        if tool_name == "update_monster_hp":
            return await monster.update(session, game_id, tool_input)
        return {"error": f"Unknown tool {tool_name!r}"}
    except Exception as exc:  # noqa: BLE001
        logger.warning("Tool call %s failed: %s", tool_name, exc)
        await session.rollback()
        return {"error": str(exc)}


async def _check_hp_game_over(session: AsyncSession, game_id: uuid.UUID, result: dict) -> str | None:
    """After an upsert_character_sheet call, auto-ends the game if the character it touched
    dropped to 0 HP or below. Returns the end reason if the game was just ended, else None."""
    current_hp = result.get("current_hp")
    if current_hp is None or current_hp > 0:
        return None

    game = await session.get(Game, game_id)
    if game is None or game.status == "ended":
        return None

    reason = f"{result.get('name', 'Ein Charakter')} ist bei {current_hp} Trefferpunkten zusammengebrochen."
    game.status = "ended"
    game.ended_reason = reason
    await session.commit()
    return reason


async def run_dm_turn(
    session: AsyncSession,
    game_id: uuid.UUID,
    latest_player_message: str | None = None,
    on_roll: Callable[[str], Awaitable[None]] | None = None,
    on_game_ended: Callable[[str], Awaitable[None]] | None = None,
) -> Message:
    client = get_llm_client()
    system_prompt = get_dm_system_prompt()

    game = await session.get(Game, game_id)
    if game is not None:
        system_prompt = (
            f"{system_prompt}\n\n"
            "Spiel-Rahmendaten (vom Ersteller beim Anlegen bereits festgelegt — NICHT erneut "
            "abfragen, in Session 0 höchstens knapp bestätigen):\n"
            f"Name: {game.name}\nBeschreibung: {game.description}"
        )

    if latest_player_message:
        rag_block = await build_rag_block(session, latest_player_message)
        if rag_block:
            system_prompt = f"{system_prompt}\n\n{rag_block}"

    messages = [{"role": "system", "content": system_prompt}]
    messages.extend(await build_context(session, game_id))

    final_text = ""
    for _ in range(MAX_TOOL_ROUNDS):
        response = await client.chat.completions.create(
            model=settings.dm_model,
            max_tokens=MAX_TOKENS,
            messages=messages,
            tools=TOOLS,
        )
        choice = response.choices[0]
        message = choice.message
        final_text = message.content or ""

        if not message.tool_calls:
            break

        messages.append(message.model_dump(exclude_unset=True))

        for tool_call in message.tool_calls:
            tool_input = fix_double_escaped_unicode(json.loads(tool_call.function.arguments))
            if tool_call.function.name == "roll_dice" and on_roll is not None:
                await on_roll(tool_input.get("notation", ""))
            result = await _execute_tool_call(session, game_id, tool_call.function.name, tool_input)

            end_reason = None
            if tool_call.function.name == "end_game" and "reason" in result:
                end_reason = result["reason"]
            elif tool_call.function.name == "upsert_character_sheet":
                end_reason = await _check_hp_game_over(session, game_id, result)
                if end_reason:
                    result = {**result, "game_ended": True, "end_reason": end_reason}
            if end_reason and on_game_ended is not None:
                await on_game_ended(end_reason)

            messages.append(
                {
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "content": json.dumps(result, ensure_ascii=False),
                }
            )
    else:
        if not final_text:
            final_text = "(Der Dungeon Master braucht einen Moment länger als erwartet — bitte versuche es erneut.)"

    dm_message = Message(game_id=game_id, sender_type="dm", content=final_text)
    session.add(dm_message)
    await session.commit()
    await session.refresh(dm_message)
    return dm_message