ai-companion/app/agent/flows_runtime.py

"""Pipecat Flows 运行时装配（Phase 3）。

**只接管网页文字路径。** 语音仍走 legacy 的 `EngineProcessor`——交接文档
明确要求 "语音路径暂时保持 legacy"，等 Phase 3、4 稳定后再迁。

管线形态：

    transport.input()
      → EngineProcessor(owns_text=False)   语音仍由它处理；文字帧往下传
      → flows 的 context aggregator.user()
      → DeepSeekLLMService                 Flows 驱动
      → FlowsBridge                        LLM 文本 → RTVI 消息 + 分句 TTS
      → flows 的 context aggregator.assistant()
      → transport.output()

两个关键设计：

  * **只用 `task_messages`，不用 `role_message`。** Phase 1 spike 实测：
    两者会各产生一条 system 消息进请求。只用一种。这里把人格与任务都写进
    `task_messages`，与 legacy"每轮重建单条 system"的做法一致。
  * **工具执行复用 `app/agent/tools.py`**，不另写一套。身份由服务端注入，
    模型不得传。
"""

from __future__ import annotations

import asyncio
import logging
from datetime import UTC, datetime
from typing import Any, AsyncIterator, Callable

import numpy as np

from pipecat.frames.frames import (
    Frame,
    LLMFullResponseEndFrame,
    LLMFullResponseStartFrame,
    TextFrame,
)
from pipecat.processors.aggregators.llm_context import LLMContext
from pipecat.processors.aggregators.llm_response_universal import (
    LLMContextAggregatorPair,
)
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor

from ..config import AppConfig
from ..memory.base import MemoryStore
from ..prompts import PERSONA_RULES
from ..voice.tts import split_sentences
from . import tools as agent_tools

log = logging.getLogger(__name__)


# ---------------------------------------------------------------- 本轮用户原话

#: 本轮用户原话。策略层要它区分"用户明确要求记住"和"模型自己推断"，
#: 而 Flows 的工具 handler 只看得到工具参数，看不到原话。
#: 用单元素列表做可变闭包——handler 是闭包，拿不到外层的非局部变量。
#:
#: 注意：这是**单用户在用的原型**。多用户并发会话时这里会互相覆盖，
#: 要支持并发得改成按 session 分键（Phase 5 迁移语音路径时一并处理）。
_LAST_USER_TEXT: list[str] = [""]


def set_last_user_text(text: str) -> None:
    _LAST_USER_TEXT[0] = text or ""


# ---------------------------------------------------------------- 节点

#: 会话控制状态。**不要把长期记忆复制进来**——Flows 的 state 不是数据库，
#: 长期数据一律走 MemoryStore（交接文档 §七）。
def _conversation_node(*, identity_decision, memories: list,
                       tool_schemas: list[dict],
                       ctx_factory: Callable[[], agent_tools.ToolContext]) -> dict:
    """默认节点：普通陪伴对话 + 可用工具。

    人格与记忆都放进**同一条** system 消息（`task_messages`）——
    **不用 role_message**，避免两条 system 消息（见模块文档）。

    这里复用 legacy 的 `build_system_prompt`，保证两条路径：
      - 人设与规则完全一致
      - 记忆注入格式一致
      - 身份未确认时的"按访客对待"口径一致
    """
    from ..prompts import build_system_prompt
    persona = build_system_prompt(identity_decision, memories)

    from pipecat.flows import FlowsFunctionSchema

    functions = []
    for schema in tool_schemas:
        fn = schema["function"]
        functions.append(FlowsFunctionSchema(
            name=fn["name"],
            description=fn["description"],
            properties=fn["parameters"].get("properties", {}),
            required=fn["parameters"].get("required", []),
            handler=_make_handler(fn["name"], ctx_factory),
        ))

    return {
        "name": "conversation",
        "task_messages": [
            {"role": "system", "content": persona},
        ],
        "functions": functions,
        # 进入节点不自动开口：文字路径要等用户先说话。
        # （spike 实测默认 True 会在 initialize 时立刻触发一轮回复，
        #   那条回复与用户输入无关，会污染归因，也会让网页莫名多一段话。）
        "respond_immediately": False,
    }


def _make_handler(name: str, ctx_factory: Callable[[], agent_tools.ToolContext]):
    """把工具层的执行器包成 Flows 的 handler。

    **用闭包捕获 ctx_factory，不用模块级全局**——全局在并发连接时会串
    （两个会话各自的 owner_user_id 会互相覆盖）。
    """
    async def handler(args, flow_manager=None):
        outcome = await agent_tools.execute(name, dict(args or {}), ctx_factory())
        return outcome.content
    return handler


# ---------------------------------------------------------------- 输入记录

class UserTextRecorder(FrameProcessor):
    """记录本轮用户原话，供策略层判定语义等级。

    策略层要区分"用户明确要求记住"和"模型自己推断"，依据是**用户原话**，
    而 Flows 的工具 handler 只看得到工具参数。所以这里把原话截下来。
    """

    def __init__(self, extractor):
        super().__init__()
        self._extract = extractor

    async def process_frame(self, frame: Frame, direction: FrameDirection):
        from pipecat.frames.frames import LLMMessagesAppendFrame
        await super().process_frame(frame, direction)
        if isinstance(frame, LLMMessagesAppendFrame):
            try:
                text = self._extract(frame.messages)
                if text:
                    set_last_user_text(text)
            except Exception:  # noqa: BLE001 — 记录失败不该影响对话
                log.debug("记录用户原话失败", exc_info=True)
        await self.push_frame(frame, direction)


# ---------------------------------------------------------------- 输出桥

class FlowsBridge(FrameProcessor):
    """把 Flows/LLM 的文本帧转成 RTVI 消息 + 分句 TTS。

    产出**与 legacy 的 `_run_turn` 保持一致**（网页依赖这些消息渲染）：
    bot-llm-started → 逐 token 的 bot-llm-text → bot-llm-stopped
    → bot-transcription / bot-output（完整文本）
    """

    def __init__(self, *, rtvi, tts, cfg: AppConfig, audio_enabled: bool = True):
        super().__init__()
        self.rtvi = rtvi
        self.tts = tts
        self.cfg = cfg
        self.audio_enabled = audio_enabled
        self._buffer = ""
        self._reply: list[str] = []
        self._started = False
        self._speak_tasks: list[asyncio.Task] = []

    async def process_frame(self, frame: Frame, direction: FrameDirection):
        await super().process_frame(frame, direction)

        if isinstance(frame, LLMFullResponseStartFrame):
            self._started = True
            self._buffer = ""
            self._reply = []
            await self.push_frame(frame, direction)
            await self.rtvi.push_transport_message(_rtvi_bot_started())

        elif isinstance(frame, TextFrame) and frame.text:
            self._buffer += frame.text
            self._reply.append(frame.text)
            # 逐 token 给网页（Playground 的 Conversation 视图需要）
            await self.rtvi.push_transport_message(_rtvi_bot_text(frame.text))
            await self.push_frame(frame, direction)
            # 攒到句子边界就送去合成——这是提前出声的唯一手段
            segs = split_sentences(self._buffer)
            if len(segs) > 1:
                *ready, self._buffer = segs
                for seg in ready:
                    await self._speak(seg)

        elif isinstance(frame, LLMFullResponseEndFrame):
            if self._buffer.strip():
                await self._speak(self._buffer)
                self._buffer = ""
            reply = "".join(self._reply).strip()
            if reply:
                # 完整文本单发一次：Playground 1.7 的 Conversation 视图
                # 主要消费完整转写/bot-output，只给 token 级不会稳定渲染。
                await self.rtvi.push_transport_message(_rtvi_bot_transcription(reply))
                await self.rtvi.push_transport_message(
                    _rtvi_bot_output(reply, spoken=self.audio_enabled))
            if self._started:
                await self.rtvi.push_transport_message(_rtvi_bot_stopped())
                self._started = False
            await self.push_frame(frame, direction)

        else:
            await self.push_frame(frame, direction)

    async def _speak(self, text: str) -> None:
        text = (text or "").strip()
        if not text or not self.audio_enabled:
            return
        try:
            audio = await asyncio.to_thread(self.tts.synthesize, text)
        except Exception as e:  # noqa: BLE001 — TTS 失败不该中断文字显示
            log.warning("Flows 路径 TTS 失败（文字仍显示）：%s", e)
            return
        if audio is None or not len(audio):
            return
        await self._push_audio(np.asarray(audio, dtype=np.float32),
                               getattr(self.tts, "sample_rate", 24000))

    async def _push_audio(self, samples: np.ndarray, sample_rate: int) -> None:
        from pipecat.frames.frames import OutputAudioRawFrame
        out_rate = self.cfg.voice.output_sample_rate if hasattr(
            self.cfg.voice, "output_sample_rate") else 24000
        if sample_rate != out_rate:
            n = int(len(samples) * out_rate / sample_rate)
            samples = np.interp(np.linspace(0, len(samples) - 1, n),
                                np.arange(len(samples)), samples).astype(np.float32)
        pcm = (np.clip(samples, -1, 1) * 32767).astype(np.int16)
        step = out_rate // 50          # 20ms 一块
        for i in range(0, len(pcm), step):
            await self.push_frame(OutputAudioRawFrame(
                audio=pcm[i:i + step].tobytes(), sample_rate=out_rate,
                num_channels=1))


# ---------------------------------------------------------------- RTVI 消息

def _rtvi_bot_started():
    import pipecat.processors.frameworks.rtvi.models as R
    return R.BotLLMStartedMessage()


def _rtvi_bot_stopped():
    import pipecat.processors.frameworks.rtvi.models as R
    return R.BotLLMStoppedMessage()


def _rtvi_bot_text(text: str):
    import pipecat.processors.frameworks.rtvi.models as R
    return R.BotLLMTextMessage(data=R.TextMessageData(text=text))


def _rtvi_bot_transcription(text: str):
    import pipecat.processors.frameworks.rtvi.models as R
    return R.BotTranscriptionMessage(data=R.TextMessageData(text=text))


def _rtvi_bot_output(text: str, *, spoken: bool):
    import pipecat.processors.frameworks.rtvi.models as R
    return R.BotOutputMessage(data=R.BotOutputMessageData(
        text=text, aggregated_by="sentence", spoken=spoken))


# ---------------------------------------------------------------- 装配

async def build_flows_parts(
    *,
    cfg: AppConfig,
    memory: MemoryStore,
    identity,
    session_id: str,
    tts,
    rtvi,
    audio_enabled: bool = True,
):
    """造出 Flows 路径需要的**零件**，不含 FlowManager。

    为什么不在这里建 FlowManager：它要求传 `worker`（PipelineWorker），
    而 worker 又必须在管线建好之后才能建，是个先后依赖。
    所以顺序是：零件 → 管线 → worker → FlowManager → initialize。

    返回 (llm, aggregator, bridge, ctx_factory, tool_schemas)
    """
    from .llm_service import create_llm_service

    # 用 agent 的 max_tokens，不是语音的 200——带工具需要余量
    llm = create_llm_service(cfg.llm, max_tokens=cfg.agent.max_tokens)
    if llm is None:
        raise RuntimeError(
            f"Flows 路径不支持后端 {cfg.llm.name!r}（只支持 deepseek）。"
            f"请把 AI_AGENT_RUNTIME 改回 legacy。")

    context = LLMContext()
    aggregator = LLMContextAggregatorPair(context)

    def ctx_factory() -> agent_tools.ToolContext:
        """每次工具调用现算身份——用户可能在会话中途切换。"""
        d = identity.current()
        return agent_tools.ToolContext(
            owner_user_id=d.user_id,          # ← 服务端注入，模型不得传
            session_id=session_id,
            turn_id=f"t_{datetime.now(UTC).strftime('%H%M%S%f')}",
            identity_confirmed=d.confirmed,
            user_text=_LAST_USER_TEXT[0],
            store=memory,
            # Phase 3 先不暴露 web_search：搜索仍走 legacy 路径，
            # 等 flows 文字路径验收通过后再接（交接文档要求分阶段）。
            search=None,
            memory_limit=cfg.context.max_memories,
            memory_timeout_s=cfg.context.memory_timeout_s,
        )

    bridge = FlowsBridge(rtvi=rtvi, tts=tts, cfg=cfg, audio_enabled=audio_enabled)
    tool_schemas = agent_tools.available_schemas(ctx_factory())
    return llm, aggregator, bridge, ctx_factory, tool_schemas


def make_flow_manager(*, worker, llm, aggregator):
    """worker 建好之后再建 FlowManager。"""
    from pipecat.flows import FlowManager
    return FlowManager(worker=worker, llm=llm, context_aggregator=aggregator)


async def load_stable_memories(memory, identity, *, limit: int = 8) -> list:
    """加载"稳定档案"：已确认、未过期的记忆。

    方案 §四要求"稳定档案（姓名、称呼、语言、明确的稳定偏好）
    **连接开始时加载**"，而不是等模型自己想起来查。

    身份未确认 → 返回空。访客不加载任何私人记忆（验收第 6 条）。
    """
    d = identity.current()
    if not d.confirmed or memory is None:
        return []
    try:
        rows = await memory.list_memories(d.user_id, include_candidates=False)
    except Exception as e:  # noqa: BLE001 — 加载失败不阻塞会话
        log.warning("加载稳定记忆失败（继续）：%s", e)
        return []
    # 最近更新的排前面
    rows.sort(key=lambda m: getattr(m, "updated_at", 0), reverse=True)
    return rows[:limit]


async def initialize_conversation_node(
    flow, *, identity_decision, memories: list, tool_schemas: list[dict],
    ctx_factory: Callable[[], agent_tools.ToolContext],
) -> None:
    """设置初始节点。

    **必须 await。** `FlowManager.initialize` 是协程，但签名显示 `-> None`，
    漏 await 不报错，只是人格/任务/工具**全部静默失效**——
    Phase 1 spike 在这里栽过一次，表现极具误导性：
    模型会正常回复（因为 LLM 照样被调用），但没有人格、没有工具。
    """
    node = _conversation_node(identity_decision=identity_decision,
                              memories=memories,
                              tool_schemas=tool_schemas,
                              ctx_factory=ctx_factory)
    await flow.initialize(node)


__all__ = [
    "FlowsBridge", "UserTextRecorder", "build_flows_parts", "make_flow_manager",
    "initialize_conversation_node", "load_stable_memories",
    "set_last_user_text", "PERSONA_RULES",
]