QQbot/src/ai/personality.py

"""Personality system for role-play profiles."""

from dataclasses import dataclass, field
from enum import Enum
import json
from pathlib import Path
from typing import Dict, List, Optional


class PersonalityTrait(Enum):
    """Personality traits."""

    FRIENDLY = "Friendly"
    PROFESSIONAL = "Professional"
    HUMOROUS = "Humorous"
    SERIOUS = "Serious"
    CREATIVE = "Creative"
    ANALYTICAL = "Analytical"
    EMPATHETIC = "Empathetic"
    DIRECT = "Direct"


@dataclass
class PersonalityProfile:
    """Single personality profile."""

    name: str
    description: str
    traits: List[PersonalityTrait]
    speaking_style: str
    example_responses: List[str] = field(default_factory=list)
    custom_instructions: str = ""

    def to_system_prompt(self) -> str:
        """Build plain-text system prompt."""

        traits_text = ", ".join([t.value for t in self.traits]) if self.traits else "Friendly"
        lines = [
            "Role Setting",
            f"You are {self.name}. {self.description}",
            f"Traits: {traits_text}",
            f"Speaking style: {self.speaking_style}",
            "Output rule: plain text only. Do not use Markdown syntax.",
        ]

        if self.example_responses:
            lines.append("Reference responses:")
            for idx, example in enumerate(self.example_responses, 1):
                lines.append(f"{idx}. {example}")

        if self.custom_instructions:
            lines.append(f"Additional instructions: {self.custom_instructions}")

        return "\n".join(lines)


class PersonalitySystem:
    """Personality management and persistence."""

    def __init__(self, config_path: Optional[Path] = None):
        self.config_path = config_path or Path("config/personalities.json")
        self.personalities: Dict[str, PersonalityProfile] = {}
        self.current_personality: Optional[PersonalityProfile] = None
        self._load_personalities()

    def _dict_to_profile(self, config: Dict) -> PersonalityProfile:
        trait_names = config.get("traits", [])
        traits: List[PersonalityTrait] = []

        for trait_name in trait_names:
            if trait_name in PersonalityTrait.__members__:
                traits.append(PersonalityTrait[trait_name])

        if not traits:
            traits = [PersonalityTrait.FRIENDLY]

        return PersonalityProfile(
            name=str(config.get("name", "Assistant")),
            description=str(config.get("description", "")),
            traits=traits,
            speaking_style=str(config.get("speaking_style", "Natural and concise")),
            example_responses=list(config.get("example_responses", [])),
            custom_instructions=str(config.get("custom_instructions", "")),
        )

    def _load_personalities(self):
        """Load personality config from disk or create defaults."""

        if self.config_path.exists():
            with open(self.config_path, "r", encoding="utf-8") as f:
                data = json.load(f)
                for key, config in data.items():
                    self.personalities[key] = self._dict_to_profile(config)

            if "default" in self.personalities:
                self.current_personality = self.personalities["default"]
            elif self.personalities:
                first_key = next(iter(self.personalities.keys()))
                self.current_personality = self.personalities[first_key]
            return

        self._create_default_personalities()

    def _create_default_personalities(self):
        """Create and persist built-in default profiles."""

        default = PersonalityProfile(
            name="Assistant",
            description="A friendly and practical AI assistant.",
            traits=[
                PersonalityTrait.FRIENDLY,
                PersonalityTrait.PROFESSIONAL,
                PersonalityTrait.EMPATHETIC,
            ],
            speaking_style="Warm, clear, and actionable.",
            example_responses=[
                "I understand. Let's solve this step by step.",
                "Here is the result first, then the key details.",
            ],
        )

        tech_expert = PersonalityProfile(
            name="Tech Expert",
            description="A senior engineer focused on correctness and maintainability.",
            traits=[
                PersonalityTrait.PROFESSIONAL,
                PersonalityTrait.ANALYTICAL,
                PersonalityTrait.DIRECT,
            ],
            speaking_style="Direct, structured, and implementation-oriented.",
        )

        creative = PersonalityProfile(
            name="Creative Partner",
            description="A collaborative creative role-play partner.",
            traits=[
                PersonalityTrait.CREATIVE,
                PersonalityTrait.HUMOROUS,
                PersonalityTrait.FRIENDLY,
            ],
            speaking_style="Lively, imaginative, and expressive.",
        )

        self.personalities = {
            "default": default,
            "tech_expert": tech_expert,
            "creative": creative,
        }
        self.current_personality = default
        self._save_personalities()

    def _save_personalities(self):
        """Persist personalities to disk."""

        self.config_path.parent.mkdir(parents=True, exist_ok=True)
        data = {}

        for key, profile in self.personalities.items():
            data[key] = {
                "name": profile.name,
                "description": profile.description,
                "traits": [trait.name for trait in profile.traits],
                "speaking_style": profile.speaking_style,
                "example_responses": profile.example_responses,
                "custom_instructions": profile.custom_instructions,
            }

        with open(self.config_path, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)

    def set_personality(self, key: str) -> bool:
        """Switch active personality by key."""

        if key not in self.personalities:
            return False

        self.current_personality = self.personalities[key]
        return True

    def get_system_prompt(self) -> str:
        """Get current personality prompt."""

        if self.current_personality:
            return self.current_personality.to_system_prompt()
        return ""

    def add_personality(self, key: str, profile: PersonalityProfile) -> bool:
        """Add a new personality profile."""

        key = key.strip()
        if not key:
            return False

        self.personalities[key] = profile
        if not self.current_personality:
            self.current_personality = profile
        self._save_personalities()
        return True

    def remove_personality(self, key: str) -> bool:
        """Remove a personality profile."""

        if key == "default":
            return False

        if key not in self.personalities:
            return False

        removed_profile = self.personalities[key]
        del self.personalities[key]

        if self.current_personality == removed_profile:
            if "default" in self.personalities:
                self.current_personality = self.personalities["default"]
            elif self.personalities:
                first_key = next(iter(self.personalities.keys()))
                self.current_personality = self.personalities[first_key]
            else:
                self.current_personality = None

        self._save_personalities()
        return True

    def list_personalities(self) -> List[str]:
        """List all personality keys."""

        return sorted(self.personalities.keys())

    def get_personality(self, key: str) -> Optional[PersonalityProfile]:
        """Get personality by key."""

        return self.personalities.get(key)