Mutual_Fund/service/client_agent/memory_extractor.py

"""从客服对话中提取长期记忆候选。"""

from __future__ import annotations

import json
import re
from typing import Any

from service.memory.schemas import MemorySource, MemoryType


class DialogueMemoryExtractor:
    """使用项目统一 LLM 提取结构化客户记忆候选。"""

    SYSTEM_PROMPT = """
你是客服记忆提取器，只提取用户明确表达或稳定陈述的客户信息。
只返回 JSON 数组，不要输出 Markdown 或解释文字。
每项必须包含：tag、content、memory_type、source。
对于用户反复询问的产品或投资主题，也可以提取兴趣信号，并额外返回 signal_type=interest_query。
兴趣信号必须使用 CUSTOMER_PREFERENCE 和 dialogue_inferred，tag 使用稳定、可归一化的英文主题名。
兴趣主题只允许以下五类：conservative_interest、steady_interest、balanced_interest、enterprising_interest、aggressive_interest。
memory_type 只能是 PROFILE_FACT、PROFILE_CANDIDATE、CUSTOMER_PREFERENCE、INVESTMENT_GOAL、SERVICE_FACT。
source 只能是 dialogue_confirmed、dialogue_stated、dialogue_inferred。
单次客服知识问题不要作为长期记忆保存；如果问题反映出客户对某个产品或投资主题的关注，使用 signal_type=interest_query 表示兴趣信号。
不确定的信息使用 dialogue_inferred，无法形成客户画像的信息不要提取。
""".strip()

    def __init__(self, llm_client):
        self.llm_client = llm_client

    async def extract(self, query: str, *, context: dict[str, Any] | None = None) -> list[dict[str, str]]:
        """提取并校验本轮用户消息中的客户记忆候选。"""
        prompt = [{"role": "system", "content": self.SYSTEM_PROMPT}]
        prompt.append(
            {
                "role": "user",
                "content": json.dumps(
                    {"query": query, "existing_memory": context or {}},
                    ensure_ascii=False,
                ),
            }
        )
        interest_signal = self._interest_fallback(query)
        try:
            response = await self.llm_client.chat(prompt, temperature=0)
            result = self._parse(response)
        except Exception:
            if interest_signal:
                return [interest_signal]
            raise
        if interest_signal and not any(
            item.get("signal_type") == "interest_query" for item in result
        ):
            result.append(interest_signal)
        return result

    @staticmethod
    def _interest_fallback(query: str) -> dict[str, str] | None:
        """为产品兴趣问题提供确定性兜底，避免依赖 LLM 输出可选字段。"""
        text = query.strip().lower()
        if not text or not any(
            marker in text
            for marker in ("基金", "理财", "投资", "产品", "fund", "investment")
        ):
            return None
        if not any(
            marker in text
            for marker in ("哪些", "什么", "怎么选", "推荐", "适合", "比较", "了解", "有哪些", "what", "which", "how")
        ):
            return None

        risk_topics = (
            (("保守", "conservative"), "conservative_interest", "用户关注保守型投资产品"),
            (("稳健", "steady", "moderate"), "steady_interest", "用户关注稳健型投资产品"),
            (("平衡", "balanced"), "balanced_interest", "用户关注平衡型投资产品"),
            (("进取", "enterprising"), "enterprising_interest", "用户关注进取型投资产品"),
            (("激进", "aggressive", "高风险", "high risk"), "aggressive_interest", "用户关注激进型投资产品"),
        )
        for markers, topic_tag, topic_content in risk_topics:
            if any(marker in text for marker in markers):
                tag = topic_tag
                content = topic_content
                break
        else:
            return None
        return {
            "tag": tag,
            "content": content,
            "memory_type": "CUSTOMER_PREFERENCE",
            "source": "dialogue_inferred",
            "signal_type": "interest_query",
        }

    @staticmethod
    def _parse(response: str) -> list[dict[str, str]]:
        """解析 LLM JSON，并过滤不符合记忆契约的内容。"""
        text = response.strip()
        fenced = re.search(r"```(?:json)?\s*(.*?)\s*```", text, re.S | re.I)
        if fenced:
            text = fenced.group(1)
        data = json.loads(text)
        if not isinstance(data, list):
            raise ValueError("记忆提取结果必须是数组")
        valid_types = {item.value for item in MemoryType}
        valid_sources = {item.value for item in MemorySource}
        result = []
        for item in data:
            if not isinstance(item, dict):
                continue
            if not all(item.get(key) for key in ("tag", "content", "memory_type", "source")):
                continue
            if item["memory_type"] not in valid_types or item["source"] not in valid_sources:
                continue
            candidate = {
                    "tag": str(item["tag"])[:64],
                    "content": str(item["content"])[:512],
                    "memory_type": item["memory_type"],
                    "source": item["source"],
                }
            if item.get("signal_type") == "interest_query":
                candidate["signal_type"] = "interest_query"
            result.append(candidate)
        return result


__all__ = ["DialogueMemoryExtractor"]