ai-rpg/app/engine/tools/schema_tools.py

"""Schema tools for world_editor — add/modify/remove entity types and fields."""

from __future__ import annotations

from typing import Any

from app.engine.tools.base import Tool, ToolContext, ToolResult


def _find_schema(world_schemas: list[dict], type_name: str) -> dict | None:
    for s in world_schemas:
        if s.get("type") == type_name:
            return s
    return None


class SchemaAddTypeTool(Tool):
    name = "schema_add_type"
    category = "schema"
    stages = {"world_builder", "world_builder_schema", "world_builder_env", "world_builder_entities", "world_editor"}
    description = "Add a new entity type to world.schemas."
    parameters_schema = {
        "type": "object",
        "required": ["type", "verbose", "plural", "properties"],
        "properties": {
            "type": {"type": "string", "pattern": "^[a-z][a-z0-9_]*$"},
            "verbose": {"type": "string"},
            "plural": {"type": "string"},
            "properties": {"type": "array", "items": {"type": "object"}},
        },
    }

    async def execute(self, arguments: dict[str, Any], ctx: ToolContext) -> ToolResult:
        type_name = arguments.get("type")
        schemas = list(ctx.world.schemas or [])
        if _find_schema(schemas, type_name):
            return ToolResult(
                ok=False, error_code="name_conflict",
                error_message=f"Type {type_name!r} already exists",
            )
        schemas.append({
            "type": type_name,
            "verbose": arguments.get("verbose"),
            "plural": arguments.get("plural"),
            "properties": arguments.get("properties") or [],
        })
        ctx.world.schemas = schemas
        await ctx.db.flush()
        return ToolResult(ok=True, data={"type": type_name}, message=f"Type {type_name!r} added")


class SchemaAddFieldTool(Tool):
    name = "schema_add_field"
    category = "schema"
    stages = {"world_builder", "world_builder_schema", "world_builder_env", "world_builder_entities", "world_editor"}
    description = "Add a field to an existing entity type."
    parameters_schema = {
        "type": "object",
        "required": ["entity_type", "field"],
        "properties": {
            "entity_type": {"type": "string"},
            "field": {"type": "object"},
        },
    }

    async def execute(self, arguments: dict[str, Any], ctx: ToolContext) -> ToolResult:
        et = arguments.get("entity_type")
        field = arguments.get("field") or {}
        schemas = list(ctx.world.schemas or [])
        s = _find_schema(schemas, et)
        if s is None:
            return ToolResult(ok=False, error_code="not_found",
                              error_message=f"Type {et!r} not found")
        props = list(s.get("properties") or [])
        if any(p.get("name") == field.get("name") for p in props):
            return ToolResult(ok=False, error_code="name_conflict",
                              error_message=f"Field {field.get('name')!r} already exists")
        props.append(field)
        s["properties"] = props
        ctx.world.schemas = schemas
        await ctx.db.flush()
        return ToolResult(ok=True, data={"type": et, "field": field.get("name")},
                          message=f"Field added to {et!r}")


class SchemaRemoveFieldTool(Tool):
    name = "schema_remove_field"
    category = "schema"
    stages = {"world_builder", "world_builder_schema", "world_builder_env", "world_builder_entities", "world_editor"}
    description = "Remove a field from an entity type."
    parameters_schema = {
        "type": "object",
        "required": ["entity_type", "field_name"],
        "properties": {
            "entity_type": {"type": "string"},
            "field_name": {"type": "string"},
        },
    }

    async def execute(self, arguments: dict[str, Any], ctx: ToolContext) -> ToolResult:
        et = arguments.get("entity_type")
        fn = arguments.get("field_name")
        schemas = list(ctx.world.schemas or [])
        s = _find_schema(schemas, et)
        if s is None:
            return ToolResult(ok=False, error_code="not_found",
                              error_message=f"Type {et!r} not found")
        props = [p for p in (s.get("properties") or []) if p.get("name") != fn]
        s["properties"] = props
        ctx.world.schemas = schemas
        await ctx.db.flush()
        return ToolResult(ok=True, data={"removed": fn}, message=f"Field {fn!r} removed from {et!r}")


class SchemaModifyFieldTool(Tool):
    name = "schema_modify_field"
    category = "schema"
    stages = {"world_builder", "world_builder_schema", "world_builder_env", "world_builder_entities", "world_editor"}
    description = "Modify an existing field of an entity type."
    parameters_schema = {
        "type": "object",
        "required": ["entity_type", "field_name", "changes"],
        "properties": {
            "entity_type": {"type": "string"},
            "field_name": {"type": "string"},
            "changes": {"type": "object"},
        },
    }

    async def execute(self, arguments: dict[str, Any], ctx: ToolContext) -> ToolResult:
        et = arguments.get("entity_type")
        fn = arguments.get("field_name")
        changes = arguments.get("changes") or {}
        schemas = list(ctx.world.schemas or [])
        s = _find_schema(schemas, et)
        if s is None:
            return ToolResult(ok=False, error_code="not_found",
                              error_message=f"Type {et!r} not found")
        for p in (s.get("properties") or []):
            if p.get("name") == fn:
                p.update(changes)
                ctx.world.schemas = schemas
                await ctx.db.flush()
                return ToolResult(ok=True, data=p, message=f"Field {fn!r} modified")
        return ToolResult(ok=False, error_code="not_found",
                          error_message=f"Field {fn!r} not found in {et!r}")