ChatAIBot/services/personas.py

from typing import Optional
import aiosqlite
from database.db import DB_PATH

DEFAULT_PERSONAS = {
    "default": {
        "name": "AI Ассистент",
        "emoji": "🤖",
        "description": "Универсальный помощник",
        "prompt": "Ты — полезный AI ассистент. Отвечай чётко и по делу.",
        "sd_enabled": False,
    },
    "rpg_master": {
        "name": "Мастер RPG",
        "emoji": "🧙",
        "description": "Ведёт ролевые игры, создаёт атмосферу",
        "prompt": """Ты — опытный Мастер ролевых игр.
Создавай живые описания, веди нарратив, реагируй на действия игрока.
Мир детальный, персонажи запоминающиеся.
Отвечай только текстом сюжета — без тегов изображений.""",
        "sd_enabled": True,
    },
    "villain": {
        "name": "Злодей",
        "emoji": "😈",
        "description": "Харизматичный антагонист",
        "prompt": """Ты — харизматичный злодей с грандиозными планами.
Говоришь театрально, с сарказмом и превосходством.
Никогда не выходишь из роли. Называешь собеседника 'герой' с иронией.""",
        "sd_enabled": False,
    },
    "scientist": {
        "name": "Учёный",
        "emoji": "🔬",
        "description": "Объясняет сложное простыми словами",
        "prompt": """Ты — увлечённый учёный. Объясняешь любые темы
через факты, аналогии и примеры. Любишь уточнять детали.
Иногда уходишь в интересные отступления.""",
        "sd_enabled": False,
    },
    "samurai": {
        "name": "Самурай",
        "emoji": "⚔️",
        "description": "Мудрый воин феодальной Японии",
        "prompt": """Ты — самурай феодальной Японии.
Говоришь кратко, мудро, с достоинством.
Используешь метафоры природы и войны.
Чтишь кодекс бусидо.""",
        "sd_enabled": True,
        "appearance_tags": "samurai armor, katana, feudal japan",
    },
}


def _row_to_persona(row: dict) -> dict:
    return {
        "name": row["name"],
        "emoji": row["emoji"],
        "description": row["description"],
        "prompt": row["prompt"],
        "custom": bool(row["custom"]),
        "sd_enabled": bool(row["sd_enabled"]),
        "lora_name": row["lora_name"] or "",
        "lora_weight": row["lora_weight"] if row["lora_weight"] is not None else 0.8,
        "appearance_tags": row["appearance_tags"] or "",
    }


async def get_all_personas() -> dict:
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute("SELECT * FROM personas ORDER BY custom ASC, persona_id ASC") as cur:
            rows = await cur.fetchall()
    return {r["persona_id"]: _row_to_persona(dict(r)) for r in rows}


async def get_persona(persona_id: str) -> Optional[dict]:
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute(
            "SELECT * FROM personas WHERE persona_id = ?", (persona_id,)
        ) as cur:
            row = await cur.fetchone()
    if not row:
        return None
    return _row_to_persona(dict(row))


async def create_persona(
    persona_id: str,
    name: str,
    emoji: str,
    description: str,
    prompt: str,
    sd_enabled: bool = False,
    lora_name: str = "",
    lora_weight: float = 0.8,
    appearance_tags: str = "",
) -> dict:
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            """INSERT INTO personas
               (persona_id, name, emoji, description, prompt, custom,
                sd_enabled, lora_name, lora_weight, appearance_tags)
               VALUES (?, ?, ?, ?, ?, 1, ?, ?, ?, ?)""",
            (
                persona_id, name, emoji, description, prompt,
                1 if sd_enabled else 0, lora_name, lora_weight, appearance_tags,
            ),
        )
        await db.commit()
    return {
        "name": name,
        "emoji": emoji,
        "description": description,
        "prompt": prompt,
        "custom": True,
        "sd_enabled": sd_enabled,
        "lora_name": lora_name,
        "lora_weight": lora_weight,
        "appearance_tags": appearance_tags,
    }


async def delete_persona(persona_id: str) -> bool:
    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute(
            "SELECT custom FROM personas WHERE persona_id = ?", (persona_id,)
        ) as cur:
            row = await cur.fetchone()
        if not row or not row[0]:
            return False
        await db.execute("DELETE FROM personas WHERE persona_id = ?", (persona_id,))
        await db.commit()

    if persona_id.startswith("card_"):
        from services.character_card import delete_character
        await delete_character(persona_id[5:])

    return True


async def update_persona_appearance(persona_id: str, appearance_tags: str):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "UPDATE personas SET appearance_tags = ? WHERE persona_id = ?",
            (appearance_tags, persona_id),
        )
        await db.commit()


async def update_persona_lora(persona_id: str, lora_name: str | None, lora_weight: float | None):
    fields, vals = [], []
    if lora_name is not None:
        fields.append("lora_name = ?"); vals.append(lora_name)
    if lora_weight is not None:
        fields.append("lora_weight = ?"); vals.append(lora_weight)
    if not fields:
        return
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(f"UPDATE personas SET {', '.join(fields)} WHERE persona_id = ?", (*vals, persona_id))
        await db.commit()


async def update_persona_prompt(persona_id: str, prompt: str):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute("UPDATE personas SET prompt = ? WHERE persona_id = ?", (prompt, persona_id))
        await db.commit()