core/generator.py

"""Narrative generator: creates locations, NPCs, and quests from parsed entities."""

from core.models import Entity, WorldContext, _new_id
import config


def generate_location(
    entity: Entity,
    context: WorldContext,
    generate_fn: callable,
) -> Entity:
    """Generate a richly described location from a parsed entity.

    Args:
        entity: The parsed entity to expand into a location.
        context: Current world state for context.
        generate_fn: Function that takes a prompt string and returns model output.

    Returns:
        Entity with populated description (location type).
    """
    prompt = config.GENERATE_LOCATION_PROMPT_ZH.format(
        entity_name=entity.name,
        entity_desc=entity.description,
        emotions=", ".join(context.emotions) if context.emotions else "神秘",
        existing_locations=", ".join(context.nearby_locations[:5]) if context.nearby_locations else "无",
    )

    full_prompt = f"{config.SYSTEM_PROMPT_ZH}\n\n{prompt}"
    description = generate_fn(full_prompt).strip()

    return Entity(
        id=entity.id or _new_id(),
        name=entity.name,
        entity_type="location",
        description=description,
        dream_id=entity.dream_id,
        metadata={"emotions": context.emotions},
    )


def generate_npc(
    entity: Entity,
    location: Entity,
    context: WorldContext,
    generate_fn: callable,
) -> Entity:
    """Generate an NPC for a given location.

    Args:
        entity: The parsed entity hint (creature/person from dream).
        location: The location this NPC resides in.
        context: Current world state for context.
        generate_fn: Model generation function.

    Returns:
        Entity with NPC type, including dialog line.
    """
    prompt = config.GENERATE_NPC_PROMPT_ZH.format(
        location_name=location.name,
        location_desc=location.description[:200],
        emotions=", ".join(context.emotions) if context.emotions else "平静",
    )

    full_prompt = f"{config.SYSTEM_PROMPT_ZH}\n\n{prompt}"
    raw_output = generate_fn(full_prompt).strip()

    # Extract dialog line if present (between quotes)
    import re
    dialog_match = re.search(r'["""「](.+?)["""」]', raw_output)
    dialog = dialog_match.group(1) if dialog_match else ""

    # Use entity name if provided, otherwise extract from output
    npc_name = entity.name if entity.entity_type in ("npc", "creature") else ""
    if not npc_name:
        # Try to extract name from the first line
        first_line = raw_output.split("\n")[0]
        name_match = re.search(r'(?:名叫|叫做|名为)\s*[：:]?\s*(\S+)', first_line)
        npc_name = name_match.group(1) if name_match else f"神秘旅人"

    return Entity(
        id=_new_id(),
        name=npc_name,
        entity_type="npc",
        description=raw_output,
        dream_id=entity.dream_id,
        dialog=dialog,
        location_id=location.id,
        metadata={"emotions": context.emotions},
    )


def generate_quest(
    npc: Entity,
    location: Entity,
    context: WorldContext,
    generate_fn: callable,
) -> Entity:
    """Generate a quest given by an NPC at a location.

    Args:
        npc: The quest-giving NPC.
        location: The location where the quest is found.
        context: Current world state.
        generate_fn: Model generation function.

    Returns:
        Entity with quest type.
    """
    prompt = config.GENERATE_QUEST_PROMPT_ZH.format(
        location_name=location.name,
        location_desc=location.description[:200],
        npc_name=npc.name,
        npc_desc=npc.description[:200],
        emotions=", ".join(context.emotions) if context.emotions else "冒险",
    )

    full_prompt = f"{config.SYSTEM_PROMPT_ZH}\n\n{prompt}"
    raw_output = generate_fn(full_prompt).strip()

    # Extract quest name from first line
    first_line = raw_output.split("\n")[0]
    import re
    quest_name_match = re.search(r'(?:任务[：:]?\s*|名称[：:]?\s*)(.+)', first_line)
    quest_name = quest_name_match.group(1).strip() if quest_name_match else f"{npc.name}的委托"

    return Entity(
        id=_new_id(),
        name=quest_name,
        entity_type="quest",
        description=raw_output,
        dream_id=npc.dream_id,
        quest_status="active",
        giver_id=npc.id,
        location_id=location.id,
        metadata={"emotions": context.emotions},
    )