main.py

"""astrbot_plugin_readingsteiner —— 基于 EverOS 的 AstrBot 长期记忆插件。

main 只做钩子/命令的绑定与异常兜底，不写业务逻辑。
业务在 core/*；EverOS 故障一律 try/except 降级，绝不阻断对话。
"""

from __future__ import annotations

import asyncio
import contextlib

from astrbot.api import AstrBotConfig, logger
from astrbot.api.event import AstrMessageEvent, filter
from astrbot.api.provider import LLMResponse, ProviderRequest
from astrbot.api.star import Context, Star, StarTools, register

from .commands import handlers
from .core import archiving, forget, injection, relationship, synthesis, visibility
from .core import identity as identity_mod
from .core.constants import (
    ARCHIVE_AUTO,
    DEFAULT_BASE_URL,
    DEFAULT_PROJECT_ID,
    DEFAULT_RETRY_RETRYABLE,
    INJECT_TARGET_USER,
    LOG_PREFIX,
)
from .core.everos_client import EverOSClient, EverOSUnavailable


@register(
    "astrbot_plugin_readingsteiner",
    "Sethyrial",
    "基于 EverOS 自进化记忆引擎的长期记忆插件（持久画像 + 按身份硬隔离）",
    "Crepuscule 0.7.2",
    "https://github.com/LoveS0ph1e/astrbot_plugin_readingsteiner",
)
class ReadingSteinerPlugin(Star):
    def __init__(self, context: Context, config: AstrBotConfig):
        super().__init__(context)
        self.context = context
        self.config = config
        # client/flush_policy 在 __init__ 同步构造：实例存在即记忆能力就绪，不依赖
        # on_astrbot_loaded 触发。该钩子在插件热重载竞态下可能不触发，旧版把构造放在
        # 那里曾致 client 永为 None、注入与归档双双静默失能。二者构造均不需事件循环。
        self.client: EverOSClient | None = EverOSClient(
            config.get("everos_base_url", DEFAULT_BASE_URL),
            float(config.get("request_timeout", 30)),
            retry_retryable=int(config.get("everos_retry_retryable", DEFAULT_RETRY_RETRYABLE)),
        )
        self._flush_policy: archiving.FlushPolicy | None = archiving.FlushPolicy(
            config.get("archive_strategy", ARCHIVE_AUTO),
            int(config.get("flush_idle_seconds", 1800)),
            int(config.get("flush_every_n_turns", 0)),
        )
        self._bg_task: asyncio.Task | None = None
        self._everos_up: bool | None = None  # 健康状态;None=首次未探测，用于跳变告警节流
        # 关系层存储（整体印象缓存）：data_dir 取插件专属数据目录；初始化失败则禁用该
        # 功能、不影响注入/归档主链路。_synth_inflight 去重，防同用户并发重复合成。
        self._rel_store: relationship.RelationshipStore | None = None
        try:
            self._rel_store = relationship.RelationshipStore(str(StarTools.get_data_dir()))
        except Exception as e:
            logger.warning(f"{LOG_PREFIX} 关系层存储初始化失败，整体印象功能禁用: {e}")
        self._synth_inflight: set[str] = set()
        # 遗忘(抑制)层存储：每用户 forget/<id>.md，读路径据此过滤；初始化失败则禁用该功能。
        self._forget_store: forget.ForgetStore | None = None
        try:
            self._forget_store = forget.ForgetStore(str(StarTools.get_data_dir()))
        except Exception as e:
            logger.warning(f"{LOG_PREFIX} 遗忘层存储初始化失败，遗忘功能禁用: {e}")

    # ────────────────── 生命周期 ──────────────────

    @filter.on_astrbot_loaded()
    async def on_loaded(self):
        """AstrBot 就绪后：健康探测打日志 + 启动 auto-flush 后台任务。

        client/flush_policy 已在 __init__ 构造（不依赖本钩子触发）；本钩子只做需要
        事件循环的事——健康探测与后台任务启动。重载竞态下本钩子若未触发，注入/归档/
        轮数 flush 仍正常（靠 __init__ 的 client），仅静默兜底 flush 缺席（消息已落
        buffer 不丢，下次轮数到或手动 flush 照常提取）。
        """
        try:
            ok = await self._healthy()
            state = "已连接" if ok else "暂不可达（将惰性重试）"
            logger.info(f"{LOG_PREFIX} 初始化完成，EverOS {state}")
            if (
                self._flush_policy is not None
                and self._flush_policy.strategy == ARCHIVE_AUTO
                and self._bg_task is None  # 幂等：重载竞态下重复触发不再叠起后台任务
            ):
                self._bg_task = asyncio.create_task(self._auto_flush_loop())
        except Exception as e:
            logger.error(f"{LOG_PREFIX} 初始化失败: {e}", exc_info=True)

    async def terminate(self):
        """卸载/停用时关 client、取消后台任务。"""
        if self._bg_task:
            self._bg_task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await self._bg_task
        if self.client:
            await self.client.close()
        logger.info(f"{LOG_PREFIX} 已停止，资源已释放")

    async def _healthy(self) -> bool:
        """惰性健康探测：任何异常视为不健康（容器晚起也能自愈）。

        健康状态发生跳变时各打一次日志（持续同态不刷屏）：
        可达→不可达打 WARNING（长期记忆降级），不可达→可达打 INFO（恢复）。
        """
        if self.client is None:
            return False
        try:
            ok = await self.client.health()
        except EverOSUnavailable:
            ok = False
        self._note_health(ok)
        return ok

    def _note_health(self, ok: bool) -> None:
        """记录健康跳变并节流告警；首次（None）静默，由 on_loaded 负责首条日志。"""
        if self._everos_up is None:
            self._everos_up = ok
            return
        if ok == self._everos_up:
            return
        self._everos_up = ok
        if ok:
            logger.info(f"{LOG_PREFIX} EverOS 已恢复，长期记忆功能重新启用")
        else:
            logger.warning(
                f"{LOG_PREFIX} EverOS 不可达，本轮起跳过记忆注入与归档（对话不受影响，将惰性重试）"
            )

    def _log_everos_unavailable(self, op: str, e: EverOSUnavailable) -> None:
        """EverOS 报错的分级日志：区分可重试/永久，带 error.code 或 HTTP 状态，不刷 ERROR 栈。

        仅记操作名 + code/status + EverOS 侧文案（无对话内容，无 PII）。
        """
        if e.code:
            tag = e.code
        elif e.status:
            tag = f"HTTP{e.status}"
        else:
            tag = "transport"
        kind = "可重试" if e.retryable else "永久"
        logger.warning(f"{LOG_PREFIX} {op}跳过：EverOS {kind}错误 [{tag}] {e}")

    async def _auto_flush_loop(self):
        """auto 策略后台任务：周期性 flush 静默超时的会话。"""
        interval = max(30, int(self.config.get("flush_idle_seconds", 300)) // 2)
        while True:
            await asyncio.sleep(interval)
            try:
                if self._flush_policy is None or self.client is None:
                    continue
                for sid in self._flush_policy.idle_sessions():
                    with contextlib.suppress(EverOSUnavailable):
                        await self.client.flush(
                            sid,
                            self.config.get("app_id", "astrbot"),
                            self.config.get("project_id", DEFAULT_PROJECT_ID),
                        )
            except asyncio.CancelledError:
                raise
            except Exception as e:
                logger.warning(f"{LOG_PREFIX} auto-flush 循环异常: {e}")

    # ────────────────── 整体印象（关系层合成） ──────────────────

    def _resolve_overall_impression(self, user_id: str, profile: dict) -> str:
        """读缓存的整体印象；画像已变则后台异步重算（不阻塞本轮回复，用旧值/空值过渡）。"""
        profile_data = profile.get("profile_data") or {}
        try:
            rel = self._rel_store.get(user_id)
            if self._rel_store.is_stale(user_id, profile_data):
                self._schedule_synth(user_id, profile_data)
            return rel.overall_impression if rel else ""
        except Exception as e:
            logger.warning(f"{LOG_PREFIX} 读取整体印象失败: {e}")
            return ""

    def _schedule_synth(self, user_id: str, profile_data: dict) -> None:
        """后台合成整体印象并落盘；同用户在途则跳过（去重，防重复 LLM 调用）。"""
        if user_id in self._synth_inflight:
            return
        self._synth_inflight.add(user_id)
        asyncio.create_task(self._synth_and_store(user_id, profile_data))

    async def _synth_and_store(self, user_id: str, profile_data: dict) -> None:
        """后台任务：合成整体印象 → 落 relationship.md（带 source_hash）。失败仅告警。"""
        try:
            text = await synthesis.synthesize_impression(
                profile_data.get("explicit_info"),
                profile_data.get("implicit_traits"),
                chat_fn=self._synth_chat,
                persona_name=self.config.get("persona_display_name", ""),
                max_chars=int(self.config.get("overall_impression_max_chars", 80)),
            )
            if text and self._rel_store is not None:
                self._rel_store.put(user_id, text, relationship.profile_source_hash(profile_data))
                logger.info(f"{LOG_PREFIX} 已更新用户 {user_id} 的整体印象")
        except Exception as e:
            logger.warning(f"{LOG_PREFIX} 整体印象合成失败: {e}")
        finally:
            self._synth_inflight.discard(user_id)

    async def _synth_chat(self, prompt: str) -> str:
        """chat_fn：用机器人当前 provider 跑一次合成；无 provider 抛错→上层吞成空印象。"""
        provider = self.context.get_using_provider()
        if provider is None:
            raise RuntimeError("无可用 LLM provider")
        resp = await provider.text_chat(prompt=prompt)
        return resp.completion_text or ""

    # ────────────────── 开关/会话白名单 ──────────────────

    def _session_enabled(self, event: AstrMessageEvent) -> bool:
        """会话白名单：留空=全部启用；否则按 unified_msg_origin 子串匹配。"""
        raw = self.config.get("enabled_sessions", "")
        if not raw:
            return True
        umo = str(event.unified_msg_origin)
        return any(s.strip() and s.strip() in umo for s in str(raw).split(","))

    # ────────────────── LLM 钩子 ──────────────────

    @filter.on_llm_request()
    async def on_llm_request(self, event: AstrMessageEvent, req: ProviderRequest):
        """LLM 请求前：检索当前用户记忆并注入。"""
        try:
            if not self.config.get("enable_injection", True):
                return
            if not self._session_enabled(event):
                return
            ident = identity_mod.resolve(event, self.config)
            if ident is None:  # 铁律 2：无身份则跳过
                return
            if not await self._healthy():  # EverOS 不可达则跳过
                return
            data = await self.client.search(
                query=req.prompt if isinstance(req.prompt, str) else "",
                user_id=ident.user_id,  # 铁律 1/3：单一真实 QQ 号
                app_id=ident.app_id,
                project_id=ident.project_id,
                method=self.config.get("search_method", "hybrid"),
                top_k=int(self.config.get("search_top_k", 5)),
                include_profile=self.config.get("include_profile", True),
            )
            profiles = data.get("profiles", []) or []
            episodes = data.get("episodes", []) or []
            # 群聊私密层抑制：丢情景细节，且不注入永恒铭契（含山盟海誓/学历=私密）
            suppress_private = self.config.get("group_public_only", True) and ident.is_group
            if suppress_private:
                profiles, episodes = visibility.filter_public(profiles, episodes)
                covenant = ""
            else:
                covenant = injection.resolve_covenant(self.config, ident.user_id)
            # 记忆遗忘(抑制)：必须在 filter_public 之后、overall_impression 之前过滤——
            # profile_source_hash 只 hash explicit+implicit，先过滤才能让整体印象既不重复合成、
            # 又不含已遗忘事实。forget_all 用户直接早退，本轮不注入任何记忆(连铭契一并抑制)。
            if self.config.get("enable_forget", True) and self._forget_store is not None:
                fstate = self._forget_store.get(ident.user_id)
                if fstate and fstate.forget_all:
                    return
                profiles, episodes = forget.apply_forget(profiles, episodes, fstate)
            # 整体印象：仅非群聊（群聊连同私密一并抑制）且开启时注入；画像变化时后台刷新。
            overall_impression = ""
            if (
                not suppress_private
                and profiles
                and self._rel_store is not None
                and self.config.get("enable_overall_impression", False)
            ):
                overall_impression = self._resolve_overall_impression(ident.user_id, profiles[0])
            text = injection.build_text(
                profiles,
                episodes,
                self.config,
                covenant=covenant,
                overall_impression=overall_impression,
                user_id=ident.user_id if self.config.get("inject_user_attribution", True) else "",
            )
            injection.inject(
                req,
                text,
                self.config.get("injection_target", INJECT_TARGET_USER),
                self.config.get("injection_position", "prepend"),
            )
        except EverOSUnavailable as e:
            self._log_everos_unavailable("注入", e)
        except Exception as e:
            logger.error(f"{LOG_PREFIX} on_llm_request 失败: {e}", exc_info=True)

    @filter.on_llm_response()
    async def on_llm_response(self, event: AstrMessageEvent, resp: LLMResponse):
        """LLM 响应后：归档本轮对话。"""
        try:
            if not self.config.get("enable_archiving", True):
                return
            if not self._session_enabled(event):
                return
            ident = identity_mod.resolve(event, self.config)
            if ident is None:
                return
            if not await self._healthy():
                return
            # forget_all 用户：opt-out，不归档(可逆)。
            if self.config.get("enable_forget", True) and self._forget_store is not None:
                fstate = self._forget_store.get(ident.user_id)
                if fstate and fstate.forget_all:
                    return
            user_text = event.message_str or ""
            assistant_text = resp.completion_text or ""
            if not user_text and not assistant_text:
                return
            log_ok = self.config.get("log_archive_success", True)
            msgs = archiving.build_messages(user_text, assistant_text, ident)
            add_data = await self.client.add(ident.session_id, msgs, ident.app_id, ident.project_id)
            self._flush_policy.mark_active(ident.session_id)
            if log_ok:
                logger.info(
                    f"{LOG_PREFIX} 已归档 {len(msgs)} 条消息 "
                    f"(session={ident.session_id}, status={add_data.get('status', '?')})"
                )
            if self._flush_policy.should_flush_now(ident.session_id):
                flush_data = await self.client.flush(
                    ident.session_id, ident.app_id, ident.project_id
                )
                if log_ok:
                    logger.info(
                        f"{LOG_PREFIX} flush 完成 "
                        f"(session={ident.session_id}, status={flush_data.get('status', '?')})"
                    )
        except EverOSUnavailable as e:
            self._log_everos_unavailable("归档", e)
        except Exception as e:
            logger.error(f"{LOG_PREFIX} on_llm_response 失败: {e}", exc_info=True)

    # ────────────────── LLM 工具（function-calling） ──────────────────
    # 让模型在对话中主动查/写记忆。身份只从 event 取（铁律 1/3），
    # 模型只能传 query/content，绝不传 user_id —— 从源头杜绝跨用户串线。

    @filter.llm_tool(name="epk_recall")
    async def epk_recall(self, event: AstrMessageEvent, query: str):
        """Recall the current user's long-term memory relevant to a query.

        Covers the persistent profile and past episodes. Use when you need to
        remember who the user is or what they told you before.

        Args:
            query(string): What to look up in the user's memory, in natural language.
        """
        return await handlers.recall_impl(self, event, query)

    @filter.llm_tool(name="epk_remember")
    async def epk_remember(self, event: AstrMessageEvent, content: str):
        """Save a fact about the current user into long-term memory.

        Use when the user shares a durable preference or fact worth remembering.

        Args:
            content(string): The fact to remember, as a complete self-contained sentence.
        """
        return await handlers.remember_impl(self, event, content)

    @filter.llm_tool(name="epk_diary")
    async def epk_diary(self, event: AstrMessageEvent, query: str):
        """Look up your own past diary entries (only your private diary, if enabled).

        Use to recall what you wrote before—to keep continuity with earlier moods
        and moments. Pass a date (YYYY-MM-DD) for that day, a keyword to search
        across entries, or leave empty for the most recent days.

        Args:
            query(string): A date (YYYY-MM-DD), a keyword, or empty for recent entries.
        """
        return await handlers.diary_impl(self, event, query)

    @filter.llm_tool(name="epk_diary_write")
    async def epk_diary_write(self, event: AstrMessageEvent, content: str, date: str = ""):
        """Write into your own private diary (only if enabled; append-only).

        Use to record something you were moved to write down—in your own voice,
        when a moment feels worth keeping. Entries append to that day's file; they
        never overwrite what you wrote before. You cannot choose the file or path.

        Args:
            content(string): What to write, in your own words. Free-form, no format required.
            date(string): Optional target day as YYYY-MM-DD; leave empty for today.
        """
        result = await handlers.diary_write_impl(self, event, content, date)
        # 审计一行：只记正文长度与结果，绝不记正文本身（日记极私密）。
        logger.info(
            f"{LOG_PREFIX} diary_write: chars={len(content or '')} -> {result}"
        )
        return result

    # ────────────────── /epk 命令组 ──────────────────
    # /epk = El Psy Kongroo 缩写
    # 范式：Mnemosyne/main.py:642-741（command_group + impl 代理 + permission_type + confirm）

    @filter.command_group("epk")
    def epk_group(self):
        """EverOS 长期记忆管理命令组 /epk（El Psy Kongroo）"""
        pass

    @filter.permission_type(filter.PermissionType.ADMIN)
    @epk_group.command("status")  # type: ignore
    async def epk_status(self, event: AstrMessageEvent):
        """[管理员] 查看 EverOS 连接状态、当前身份与记忆计数"""
        async for r in handlers.status_impl(self, event):
            yield r

    @filter.permission_type(filter.PermissionType.ADMIN)
    @epk_group.command("search")  # type: ignore
    async def epk_search(self, event: AstrMessageEvent, query: str = ""):
        """[管理员] 按当前用户检索记忆 /epk search <内容>"""
        async for r in handlers.search_impl(self, event, query):
            yield r

    @epk_group.command("flush")  # type: ignore
    async def epk_flush(self, event: AstrMessageEvent):
        """手动归档当前会话 /epk flush"""
        async for r in handlers.flush_impl(self, event):
            yield r

    @filter.permission_type(filter.PermissionType.ADMIN)
    @epk_group.command("quality")  # type: ignore
    async def epk_quality(self, event: AstrMessageEvent):
        """[管理员] 抽查当前用户画像质量 /epk quality"""
        async for r in handlers.quality_impl(self, event):
            yield r

    @filter.permission_type(filter.PermissionType.ADMIN)
    @epk_group.command("forget")  # type: ignore
    async def epk_forget(self, event: AstrMessageEvent, args: str = ""):
        """[管理员] 遗忘(抑制)当前用户记忆 /epk forget [all|clear|<描述>]"""
        async for r in handlers.forget_impl(self, event, args):
            yield r

    @filter.permission_type(filter.PermissionType.ADMIN)
    @epk_group.command("reflect")  # type: ignore
    async def epk_reflect(self, event: AstrMessageEvent):
        """[管理员] 触发情景合并(Reflection)治碎片/过粒度 /epk reflect（需 EverOS≥1.1.0）"""
        async for r in handlers.reflect_impl(self, event):
            yield r

    @epk_group.command("help")  # type: ignore
    async def epk_help(self, event: AstrMessageEvent):
        """显示帮助 /epk help"""
        async for r in handlers.help_impl(self, event):
            yield r