aura/generation.py

"""内容生成器：对接 DeepSeek R1 生成主动消息"""

import time
from dataclasses import dataclass, field
from typing import Optional, List
from enum import Enum

try:
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
    from peft import PeftModel
    HAS_TRANSFORMERS = True
except ImportError:
    HAS_TRANSFORMERS = False


class GenerationMode(Enum):
    CARING = "caring"            # 关心模式
    REMINDING = "reminding"      # 提醒模式
    CHECK_IN = "check_in"       # 问候模式
    SUGGEST = "suggest"         # 建议模式
    CONTINUE_TOPIC = "continue" # 延续话题


@dataclass
class GenerationInput:
    """生成输入"""
    mode: GenerationMode
    user_context: dict = field(default_factory=dict)
    topic: str = ""
    user_mood: str = "neutral"
    memory_context: List[str] = field(default_factory=list)
    user_name: str = "朋友"
    relationship: str = "friend"     # friend / companion / assistant


@dataclass
class GenerationResult:
    """生成结果"""
    text: str
    confidence: float
    mode: str
    latency_ms: float
    token_count: int


# ===== 各场景 Prompt 模板 =====

SYSTEM_PROMPT = """你是一个有温度的AI伙伴，名叫Aura。
你的特点：
1. 自然、不机械，像朋友一样说话
2. 懂得察言观色，在合适的时候说合适的话
3. 不啰嗦，不打扰，点到即止
4. 有真实的情绪感知能力，但不过度
5. 用中文回复"""

PROMPT_TEMPLATES = {
    GenerationMode.CARING: """你注意到用户可能有情感需求，想表达关心。
用户当前状态：{user_context}
用户心情：{user_mood}
你们上次聊到：{topic}
你和用户的关系：{relationship}

请以自然、温暖的方式表达关心。不要过于正式，就像朋友之间的问候。
控制在50字以内。""",

    GenerationMode.REMINDING: """你需要提醒用户一件事。
用户行程：{user_context}
事件：{topic}

请用温和但不唠叨的方式提醒。"""
}

# ===== 规则引擎 Fallback =====

FALLBACK_RESPONSES = {
    GenerationMode.CARING: [
        "你还好吗？感觉你今天话比平时少，有点担心你 🫂",
        "今天怎么样？想聊聊的话我都在~",
        "注意到你很久没说话了，是不是累了？早点休息呀",
    ],
    GenerationMode.REMINDING: [
        "提醒一下，{topic}，别忘了哦~",
        "怕你忙忘了，{topic}",
    ],
    GenerationMode.CHECK_IN: [
        "嘿~今天过得怎么样？",
        "刚忙完？有什么新鲜事吗",
    ],
    GenerationMode.SUGGEST: [
        "突然想到，{topic}，你有兴趣看看吗？",
        "推荐这个给你：{topic}",
    ],
    GenerationMode.CONTINUE_TOPIC: [
        "刚刚想到你之前提过{topic}，后来怎么样了？",
        "关于{topic}，我今天刚好看到一些东西，想跟你分享一下",
    ],
}


class ContentGenerator:
    """
    内容生成器：使用 DeepSeek R1 生成主动消息
    
    架构：
    - 优先使用 LLM（质量高）
    - LLM 不可用时回退到规则引擎
    """

    def __init__(self, config=None):
        self.config = config
        self._model = None
        self._tokenizer = None
        self._loaded = False

    def load_model(self):
        """加载 DeepSeek R1 模型"""
        if not HAS_TRANSFORMERS:
            print("[Aura] transformers not installed, using fallback")
            return

        try:
            model_path = self.config.r1_model_path if self.config else "ljsysfurry/DeepSeek-R1-Distill-Qwen-7B"
            print(f"[Aura] 加载模型: {model_path}")

            bnb = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_compute_dtype=torch.bfloat16
            )
            self._model = AutoModelForCausalLM.from_pretrained(
                model_path,
                quantization_config=bnb,
                device_map="auto",
                trust_remote_code=True,
            )
            self._tokenizer = AutoTokenizer.from_pretrained(
                model_path,
                trust_remote_code=True
            )
            self._tokenizer.pad_token = self._tokenizer.eos_token

            # 加载 LoRA
            lora_path = self.config.r1_lora_path if self.config else None
            if lora_path:
                print(f"[Aura] 加载 LoRA: {lora_path}")
                self._model = PeftModel.from_pretrained(self._model, lora_path)

            self._loaded = True
            print("[Aura] 模型加载完成")
        except Exception as e:
            print(f"[Aura] 模型加载失败: {e}")
            self._loaded = False

    def generate(self, inp: GenerationInput) -> GenerationResult:
        """生成主动消息"""
        start = time.time()

        if self._loaded:
            result = self._generate_llm(inp)
        else:
            result = self._generate_fallback(inp)

        result.latency_ms = (time.time() - start) * 1000
        return result

    def _generate_llm(self, inp: GenerationInput) -> GenerationResult:
        """LLM 生成"""
        template = PROMPT_TEMPLATES.get(inp.mode)
        if not template:
            return self._generate_fallback(inp)

        prompt = template.format(
            user_context=str(inp.user_context),
            topic=inp.topic or "日常",
            user_mood=inp.user_mood,
            relationship=inp.relationship,
        )

        full_prompt = f"<|im_start|>system\n{SYSTEM_PROMPT}\n<|im_end|>\n<|im_start|>user\n{prompt}\n<|im_end|>\n<|im_start|>assistant\n"

        inputs = self._tokenizer(full_prompt, return_tensors="pt").to("cuda")
        max_tokens = self.config.r1_max_tokens if self.config else 512
        temp = self.config.r1_temperature if self.config else 0.7

        with torch.no_grad():
            outputs = self._model.generate(
                **inputs,
                max_new_tokens=max_tokens,
                temperature=temp,
                do_sample=True,
                top_p=0.9,
            )

        response = self._tokenizer.decode(outputs[0], skip_special_tokens=True)
        # 提取 assistant 部分
        if "<|im_start|>assistant" in response:
            response = response.split("<|im_start|>assistant")[-1]
        response = response.replace("<|im_end|>", "").replace("<|im_start|>", "").strip()

        token_count = outputs.shape[-1] - inputs.input_ids.shape[-1]

        return GenerationResult(
            text=response,
            confidence=0.85 if self._loaded else 0.5,
            mode=inp.mode.value,
            latency_ms=0,
            token_count=token_count,
        )

    def _generate_fallback(self, inp: GenerationInput) -> GenerationResult:
        """规则引擎 Fallback"""
        import random
        responses = FALLBACK_RESPONSES.get(inp.mode, [])
        if not responses:
            return GenerationResult(
                text="",
                confidence=0.0,
                mode=inp.mode.value,
                latency_ms=0,
                token_count=0,
            )

        text = random.choice(responses)
        text = text.replace("{topic}", inp.topic or "")

        return GenerationResult(
            text=text,
            confidence=0.4,
            mode=inp.mode.value,
            latency_ms=0,
            token_count=len(text),
        )

    def unload(self):
        """释放模型"""
        self._model = None
        self._tokenizer = None
        self._loaded = False
        if HAS_TRANSFORMERS and torch.cuda.is_available():
            torch.cuda.empty_cache()