497 lines
20 KiB
Python
497 lines
20 KiB
Python
"""会话引擎：身份 → 记忆 → LLM → TTS 的核心回路。
|
||
|
||
这一层不依赖 Pipecat、不依赖任何传输方式，只依赖四个接口
|
||
（LLMProvider / MemoryStore / IdentityProvider / LocalTTS）。
|
||
好处有两个：
|
||
|
||
1. **可以用文件音频端到端测试**，不需要麦克风、音箱、WebRTC。
|
||
2. 接 Pipecat 时只需写薄适配层；换传输方式（WebRTC / 本地声卡）不动这里。
|
||
|
||
一轮的完整流程：
|
||
|
||
用户音频
|
||
└→ STT 转写
|
||
└→ 身份裁决（决定用谁的记忆）
|
||
└→ 记忆检索（带超时，失败不阻塞）
|
||
└→ 拼 system（每轮重建，单条）
|
||
└→ LLM 流式
|
||
└→ 按句切分 → 逐句 TTS → 边合成边播
|
||
└→ 回复结束，异步写记忆（绑定创建时的身份）
|
||
|
||
延迟的关键在最后两步：**必须按句切分后立刻合成**。
|
||
VITS 是非自回归的，整段合成完才出声；只有切成短句才能提前开口。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
import re
|
||
import time
|
||
import uuid
|
||
from dataclasses import dataclass, field
|
||
from enum import Enum
|
||
from typing import AsyncIterator, Sequence
|
||
|
||
import numpy as np
|
||
|
||
from .config import AppConfig
|
||
from .identity.base import IdentityDecision, IdentityProvider
|
||
from .llm.base import ChunkKind, LLMProvider, ToolCall
|
||
from .memory.base import (
|
||
ConfirmStatus, Memory, MemoryKind, MemoryStore, Scope, Turn,
|
||
)
|
||
from .prompts import build_system_prompt, trim_history
|
||
from .search import (
|
||
WEB_SEARCH_PROMPT,
|
||
WEB_SEARCH_TOOL,
|
||
SearchError,
|
||
SearXNGSearchClient,
|
||
)
|
||
from .voice.stt import StreamingSTT
|
||
from .voice.tts import LocalTTS, split_sentences, text_for_speech
|
||
|
||
log = logging.getLogger(__name__)
|
||
|
||
_SOURCE_SECTION_RE = re.compile(
|
||
r"(?:^|\n)\s*(?:参考来源|资料来源|来源)\s*[：:]",
|
||
re.IGNORECASE,
|
||
)
|
||
|
||
|
||
class EventKind(str, Enum):
|
||
IDENTITY = "identity" # 本轮的身份裁决结果
|
||
MEMORY_USED = "memory_used" # 检索到的记忆
|
||
TRANSCRIPT = "transcript" # STT 结果
|
||
RESPONSE_DELTA = "response_delta" # LLM 增量文本（未播，仅供显示）
|
||
SPEAK = "speak" # 一段可播放的音频
|
||
TOOL_CALL = "tool_call" # 模型请求调用工具
|
||
METRICS = "metrics" # 本轮计时
|
||
ERROR = "error"
|
||
|
||
|
||
@dataclass
|
||
class Event:
|
||
kind: EventKind
|
||
text: str = ""
|
||
audio: np.ndarray | None = None
|
||
sample_rate: int = 16000
|
||
identity: IdentityDecision | None = None
|
||
memories: list[Memory] = field(default_factory=list)
|
||
tool_call: ToolCall | None = None
|
||
metrics: "TurnMetrics | None" = None
|
||
error: str = ""
|
||
|
||
|
||
@dataclass
|
||
class TurnMetrics:
|
||
"""按方案 §10 要求记录耗时，便于定位瓶颈。"""
|
||
stt_ms: float = 0.0
|
||
identity_ms: float = 0.0
|
||
retrieve_ms: float = 0.0
|
||
llm_first_content_ms: float = 0.0
|
||
first_audio_ms: float = 0.0 # 从开始处理到第一段音频就绪
|
||
total_ms: float = 0.0
|
||
llm_backend: str = ""
|
||
memories_used: int = 0
|
||
reply_chars: int = 0
|
||
speaks: int = 0
|
||
search_ms: float = 0.0
|
||
searches: int = 0
|
||
|
||
def summary(self) -> str:
|
||
return (f"STT {self.stt_ms:.0f}ms | 身份 {self.identity_ms:.0f}ms | "
|
||
f"检索 {self.retrieve_ms:.0f}ms | LLM首token {self.llm_first_content_ms:.0f}ms | "
|
||
f"首音 {self.first_audio_ms:.0f}ms | 总 {self.total_ms:.0f}ms | "
|
||
f"{self.llm_backend} | 记忆{self.memories_used}条 | "
|
||
f"搜索{self.searches}次/{self.search_ms:.0f}ms | {self.speaks}段音频")
|
||
|
||
|
||
class ConversationEngine:
|
||
def __init__(
|
||
self,
|
||
cfg: AppConfig,
|
||
llm: LLMProvider,
|
||
memory: MemoryStore,
|
||
identity: IdentityProvider,
|
||
tts: LocalTTS,
|
||
stt: StreamingSTT | None = None,
|
||
search: SearXNGSearchClient | None = None,
|
||
):
|
||
self.cfg = cfg
|
||
self.llm = llm
|
||
self.memory = memory
|
||
self.identity = identity
|
||
self.tts = tts
|
||
self.stt = stt
|
||
self.search = search or (
|
||
SearXNGSearchClient(cfg.search) if cfg.search.enabled else None
|
||
)
|
||
#: 会话历史。切身份时会清空——方案 §8 要求切换身份
|
||
#: 清理上一人的私有检索上下文。
|
||
self._history: list[dict] = []
|
||
self._history_identity_version: int = -1
|
||
self._session_id: str = ""
|
||
self._current_turn: asyncio.Task | None = None
|
||
|
||
# ------------------------------------------------------------ 会话
|
||
|
||
async def start_session(self, session_id: str | None = None,
|
||
user_id: str | None = None) -> str:
|
||
sid = session_id or f"s_{uuid.uuid4().hex[:12]}"
|
||
if user_id:
|
||
self.identity.switch(user_id, reason="会话指定")
|
||
d = self.identity.current()
|
||
await self.memory.start_session(sid, d.user_id)
|
||
self._session_id = sid
|
||
self._history_identity_version = self._identity_version()
|
||
log.info("会话开始 %s，身份 %s", sid, d.user_id)
|
||
return sid
|
||
|
||
async def end_session(self) -> None:
|
||
if self._session_id:
|
||
await self.memory.end_session(self._session_id)
|
||
|
||
def _identity_version(self) -> int:
|
||
# ManualIdentityProvider 有 version；其他实现没有时退化为 0
|
||
return getattr(self.identity, "version", 0)
|
||
|
||
def _sync_identity_scope(self) -> None:
|
||
"""身份变了就丢弃历史。
|
||
|
||
不做这件事的后果：切到新用户后，模型仍看得到上一人的对话，
|
||
第一条回复就可能带出别人的信息。
|
||
"""
|
||
v = self._identity_version()
|
||
if v != self._history_identity_version:
|
||
if self._history:
|
||
log.info("身份变更，丢弃 %d 条历史消息", len(self._history))
|
||
self._history = []
|
||
self._history_identity_version = v
|
||
|
||
# ------------------------------------------------------------ 主回路
|
||
|
||
async def respond_to_text(self, user_text: str, *,
|
||
session_id: str | None = None) -> AsyncIterator[Event]:
|
||
"""处理一句用户输入，产出事件流。"""
|
||
t_start = time.perf_counter()
|
||
metrics = TurnMetrics()
|
||
sid = session_id or self._session_id or await self.start_session()
|
||
|
||
self._sync_identity_scope()
|
||
identity = self.identity.current()
|
||
metrics.identity_ms = (time.perf_counter() - t_start) * 1000
|
||
yield Event(kind=EventKind.IDENTITY, identity=identity)
|
||
|
||
if not user_text.strip():
|
||
yield Event(kind=EventKind.METRICS, metrics=metrics)
|
||
return
|
||
|
||
# ---- 记忆检索：带超时，失败不阻塞（方案 §8）
|
||
t0 = time.perf_counter()
|
||
memories = await self._retrieve_memories(identity, user_text)
|
||
metrics.retrieve_ms = (time.perf_counter() - t0) * 1000
|
||
metrics.memories_used = len(memories)
|
||
if memories:
|
||
yield Event(kind=EventKind.MEMORY_USED, memories=memories)
|
||
|
||
# ---- 拼 system（每轮重建，单条）
|
||
system = build_system_prompt(
|
||
identity,
|
||
memories,
|
||
extra=WEB_SEARCH_PROMPT if self.search is not None else "",
|
||
)
|
||
history = trim_history(self._history,
|
||
max_turns=self.cfg.context.max_turns)
|
||
messages = [{"role": "system", "content": system}] + history + \
|
||
[{"role": "user", "content": user_text}]
|
||
|
||
# ---- LLM 流式 + 逐句 TTS
|
||
reply_parts: list[str] = []
|
||
buffer = ""
|
||
speech_done = False
|
||
speaks = 0
|
||
async for ev in self._stream_and_speak(messages, metrics):
|
||
if ev.kind == EventKind.RESPONSE_DELTA:
|
||
reply_parts.append(ev.text)
|
||
# 增量文本必须**原样透传**给调用方（显示、字幕、日志都要用）。
|
||
# 这里漏掉 yield 的后果是：音频正常出来了，但上层拿到的是空回复——
|
||
# 因为文本只进了内部累加器，从未离开这个生成器。
|
||
yield ev
|
||
if speech_done:
|
||
continue
|
||
buffer += ev.text
|
||
source_start = _SOURCE_SECTION_RE.search(buffer)
|
||
if source_start:
|
||
buffer = buffer[:source_start.start()]
|
||
speech_done = True
|
||
# 攒到句末就送去合成——这是提前出声的唯一手段
|
||
segs = split_sentences(buffer)
|
||
if len(segs) > 1:
|
||
*ready, buffer = segs
|
||
for seg in ready:
|
||
async for speak_ev in self._speak(seg, metrics):
|
||
speaks += 1
|
||
yield speak_ev
|
||
else:
|
||
yield ev
|
||
|
||
# 收尾：把剩下的半句也合成掉
|
||
if buffer.strip():
|
||
async for speak_ev in self._speak(buffer, metrics):
|
||
speaks += 1
|
||
yield speak_ev
|
||
|
||
reply = "".join(reply_parts).strip()
|
||
metrics.reply_chars = len(reply)
|
||
metrics.speaks = speaks
|
||
metrics.total_ms = (time.perf_counter() - t_start) * 1000
|
||
|
||
# ---- 落历史
|
||
self._history.append({"role": "user", "content": user_text})
|
||
if reply:
|
||
self._history.append({"role": "assistant", "content": reply})
|
||
|
||
# ---- 异步写记忆：**绑定创建时的身份**，不读"当前用户"
|
||
if self.cfg.memory.auto_extract_enabled and reply:
|
||
self._spawn_memory_write(
|
||
user_text=user_text, reply=reply,
|
||
user_id=identity.user_id,
|
||
session_id=sid,
|
||
turn_id=f"t_{uuid.uuid4().hex[:10]}")
|
||
|
||
await self.memory.add_turn(Turn(
|
||
turn_id=f"t_{uuid.uuid4().hex[:10]}", session_id=sid,
|
||
user_id=identity.user_id, role="user", text=user_text))
|
||
|
||
yield Event(kind=EventKind.METRICS, metrics=metrics)
|
||
|
||
async def respond_to_audio(self, samples: np.ndarray, *,
|
||
session_id: str | None = None) -> AsyncIterator[Event]:
|
||
"""从音频走完整条链路（需要 STT）。"""
|
||
if self.stt is None:
|
||
yield Event(kind=EventKind.ERROR, error="未配置 STT，无法处理音频")
|
||
return
|
||
t0 = time.perf_counter()
|
||
stream = self.stt.new_stream()
|
||
stream.accept(samples)
|
||
text = stream.finish()
|
||
stt_ms = (time.perf_counter() - t0) * 1000
|
||
yield Event(kind=EventKind.TRANSCRIPT, text=text)
|
||
async for ev in self.respond_to_text(text, session_id=session_id):
|
||
if ev.kind == EventKind.METRICS and ev.metrics:
|
||
ev.metrics.stt_ms = stt_ms
|
||
yield ev
|
||
|
||
# ------------------------------------------------------------ 内部
|
||
|
||
async def _retrieve_memories(self, identity: IdentityDecision,
|
||
query: str) -> list[Memory]:
|
||
"""带超时的检索。超时/出错一律返回空，绝不让对话卡住。"""
|
||
if not identity.confirmed:
|
||
# 访客不做个性化检索（方案 §6："未知身份只用访客临时上下文"）
|
||
return []
|
||
try:
|
||
return await asyncio.wait_for(
|
||
self.memory.retrieve(
|
||
identity.user_id, query,
|
||
limit=self.cfg.context.max_memories,
|
||
scopes=[Scope.PRIVATE, Scope.SHARED]),
|
||
timeout=self.cfg.context.memory_timeout_s)
|
||
except asyncio.TimeoutError:
|
||
log.warning("记忆检索超时 %.2fs，本轮不带记忆继续",
|
||
self.cfg.context.memory_timeout_s)
|
||
return []
|
||
except Exception as e: # noqa: BLE001
|
||
log.warning("记忆检索失败（%s），本轮不带记忆继续", e)
|
||
return []
|
||
|
||
async def _stream_and_speak(self, messages: list[dict],
|
||
metrics: TurnMetrics) -> AsyncIterator[Event]:
|
||
"""执行 LLM/搜索工具闭环；最终文本仍然流式交给调用方做 TTS。"""
|
||
t0 = time.perf_counter()
|
||
try:
|
||
# 最多允许连续两轮搜索；第三次请求不再暴露工具，强制形成回答，
|
||
# 防止模型在搜索词之间无限循环。
|
||
for round_index in range(3):
|
||
tool_defs = (
|
||
[WEB_SEARCH_TOOL]
|
||
if self.search is not None and round_index < 2
|
||
else None
|
||
)
|
||
tool_acc: dict[int, dict[str, str]] = {}
|
||
assistant_parts: list[str] = []
|
||
|
||
async for chunk in self.llm.stream(messages, tools=tool_defs):
|
||
if chunk.kind == ChunkKind.CONTENT and chunk.text:
|
||
assistant_parts.append(chunk.text)
|
||
if not metrics.llm_first_content_ms:
|
||
metrics.llm_first_content_ms = (
|
||
time.perf_counter() - t0
|
||
) * 1000
|
||
yield Event(kind=EventKind.RESPONSE_DELTA, text=chunk.text)
|
||
elif chunk.kind == ChunkKind.TOOL_CALL:
|
||
slot = tool_acc.setdefault(
|
||
chunk.tool_index,
|
||
{"id": "", "name": "", "arguments": ""},
|
||
)
|
||
# base.py 传来的是截至当前分片的完整累计值，直接覆盖，
|
||
# 不能再次拼接，否则参数会重复。
|
||
slot["id"] = chunk.tool_id or slot["id"]
|
||
slot["name"] = chunk.tool_name or slot["name"]
|
||
slot["arguments"] = chunk.tool_arguments
|
||
|
||
metrics.llm_backend = getattr(
|
||
self.llm, "last_backend", self.llm.name)
|
||
if not tool_acc:
|
||
break
|
||
|
||
calls: list[ToolCall] = []
|
||
api_calls: list[dict] = []
|
||
for index, raw in sorted(tool_acc.items()):
|
||
call = ToolCall(
|
||
id=raw["id"] or f"call_{uuid.uuid4().hex[:12]}",
|
||
name=raw["name"],
|
||
arguments=raw["arguments"] or "{}",
|
||
)
|
||
calls.append(call)
|
||
api_calls.append({
|
||
"id": call.id,
|
||
"type": "function",
|
||
"function": {
|
||
"name": call.name,
|
||
"arguments": call.arguments,
|
||
},
|
||
})
|
||
yield Event(kind=EventKind.TOOL_CALL, tool_call=call)
|
||
|
||
messages.append({
|
||
"role": "assistant",
|
||
"content": "".join(assistant_parts) or None,
|
||
"tool_calls": api_calls,
|
||
})
|
||
for call in calls:
|
||
result = await self._execute_tool(call, metrics)
|
||
messages.append({
|
||
"role": "tool",
|
||
"tool_call_id": call.id,
|
||
"content": result,
|
||
})
|
||
except asyncio.CancelledError:
|
||
log.info("本轮被取消（打断）")
|
||
raise
|
||
except Exception as e: # noqa: BLE001
|
||
log.error("LLM 流失败：%s", e)
|
||
yield Event(kind=EventKind.ERROR, error=f"{type(e).__name__}: {e}")
|
||
metrics.llm_backend = getattr(self.llm, "last_backend", self.llm.name)
|
||
|
||
async def _execute_tool(self, call: ToolCall, metrics: TurnMetrics) -> str:
|
||
"""执行只读搜索工具；错误作为工具结果回灌，不让整轮对话中断。"""
|
||
if call.name != "web_search" or self.search is None:
|
||
return "工具不可用。请不要声称已经执行。"
|
||
query = call.parsed_arguments().get("query")
|
||
if not isinstance(query, str) or not query.strip():
|
||
return "搜索失败：缺少有效的 query 参数。"
|
||
|
||
started = time.perf_counter()
|
||
metrics.searches += 1
|
||
try:
|
||
results = await self.search.search(query)
|
||
log.info("联网搜索 %r：返回 %d 条", query[:80], len(results))
|
||
return self.search.for_model(query, results)
|
||
except SearchError as exc:
|
||
log.warning("联网搜索失败：%s", exc)
|
||
return f"搜索失败：{exc}。请明确告知用户暂时无法取得实时资料。"
|
||
except Exception as exc: # noqa: BLE001 — 工具失败不能拖垮对话
|
||
log.exception("联网搜索异常：%s", exc)
|
||
return "搜索服务发生异常。请明确告知用户暂时无法取得实时资料。"
|
||
finally:
|
||
metrics.search_ms += (time.perf_counter() - started) * 1000
|
||
|
||
async def _speak(self, text: str, metrics: TurnMetrics) -> AsyncIterator[Event]:
|
||
"""合成一段并产出可播放音频。"""
|
||
text = text_for_speech(text)
|
||
if not text:
|
||
return
|
||
t0 = time.perf_counter()
|
||
try:
|
||
# TTS 是 CPU 密集的同步调用，放线程里跑，别阻塞事件循环
|
||
audio = await asyncio.to_thread(self.tts.synthesize, text)
|
||
except Exception as e: # noqa: BLE001
|
||
log.error("TTS 失败：%s", e)
|
||
yield Event(kind=EventKind.ERROR, error=f"TTS 失败：{e}")
|
||
return
|
||
if audio is None or len(audio) == 0:
|
||
return
|
||
if not metrics.first_audio_ms:
|
||
metrics.first_audio_ms = (time.perf_counter() - t0) * 1000
|
||
yield Event(kind=EventKind.SPEAK, text=text, audio=audio,
|
||
sample_rate=self.tts.sample_rate)
|
||
|
||
def _spawn_memory_write(self, *, user_text: str, reply: str,
|
||
user_id: str, session_id: str, turn_id: str) -> None:
|
||
"""起一个后台任务写记忆。
|
||
|
||
**身份是参数传进来的，不在这里读 self.identity.current()。**
|
||
方案 §8 明确要求："异步任务绑定创建时的 user_id/session_id，
|
||
避免迟到结果写入新用户。" 如果这里去读"当前身份"，
|
||
用户在这几百毫秒内切了身份，这条记忆就会写到别人名下。
|
||
"""
|
||
async def _write():
|
||
try:
|
||
for mem in extract_memories(user_text, reply):
|
||
mem.owner_user_id = user_id # 绑定创建时的身份
|
||
mem.source_session_id = session_id
|
||
mem.source_turn_id = turn_id
|
||
mem.confirm_status = ConfirmStatus(
|
||
self.cfg.memory.auto_extract_confirm_status)
|
||
await self.memory.remember(mem)
|
||
except Exception as e: # noqa: BLE001
|
||
log.warning("异步写记忆失败（不影响对话）：%s", e)
|
||
|
||
task = asyncio.create_task(_write())
|
||
# 持有引用，避免被 GC 提前回收
|
||
if not hasattr(self, "_bg_tasks"):
|
||
self._bg_tasks: set[asyncio.Task] = set()
|
||
self._bg_tasks.add(task)
|
||
task.add_done_callback(self._bg_tasks.discard)
|
||
|
||
|
||
# ---------------------------------------------------------------- 记忆抽取
|
||
|
||
#: 显式"记住"的触发词。首版用规则，不额外调一次 LLM——
|
||
#: 那会多一次往返，而语音场景对延迟很敏感。
|
||
_REMEMBER_PATTERNS = (
|
||
("记住", ""), ("记一下", ""), ("别忘了", ""), ("帮我记", ""),
|
||
("我喜欢", "偏好"), ("我不喜欢", "偏好"), ("我习惯", "偏好"),
|
||
)
|
||
|
||
|
||
def extract_memories(user_text: str, reply: str) -> list[Memory]:
|
||
"""从一轮对话里抽取候选记忆。
|
||
|
||
首版只认显式表达。**抽取结果默认是 candidate，不参与检索**——
|
||
方案 §8："模型推测不当作确认事实"。要变成 confirmed 需要用户确认
|
||
或走工具调用。
|
||
"""
|
||
out: list[Memory] = []
|
||
text = (user_text or "").strip()
|
||
for trigger, kind_hint in _REMEMBER_PATTERNS:
|
||
if trigger in text:
|
||
# 取触发词之后的正文
|
||
body = text.split(trigger, 1)[1].strip(" ：:，,。.！!？?")
|
||
if len(body) < 2:
|
||
continue
|
||
out.append(Memory(
|
||
kind=MemoryKind.PREFERENCE if kind_hint else MemoryKind.FACT,
|
||
key=kind_hint or "",
|
||
value=body[:200],
|
||
))
|
||
break
|
||
return out
|
||
|
||
|
||
__all__ = [
|
||
"ConversationEngine", "Event", "EventKind", "TurnMetrics", "extract_memories",
|
||
]
|