dnnai/asr_agent/voiceprint.py

"""
Voiceprint (声纹) Recognition — Zvec-backed vector store
- Multi-user registration, identification, management
- PyTorch-based speaker embedding (torchaudio ECAPA-TDNN or fallback)
- Zvec embedded vector DB for storage & similarity search
- Persistent storage via mounted /data/voiceprints volume
"""

from __future__ import annotations

import logging
import os
import time
from typing import Optional

import numpy as np

logger = logging.getLogger("voiceprint")

VP_DB_PATH = os.environ.get("VP_DB_PATH", "/data/voiceprints")
VP_SIMILARITY_THRESHOLD = float(os.environ.get("VP_SIMILARITY_THRESHOLD", "0.65"))

_EMBEDDING_DIM = 192  # ECAPA-TDNN output dim

# ── Lazy-loaded encoder ──
_encoder = None


def _get_encoder():
    global _encoder, _EMBEDDING_DIM
    if _encoder is not None:
        return _encoder
    try:
        from scipy.fft import dct

        class _MFCCEncoder:
            def embed_utterance(self, audio: np.ndarray) -> np.ndarray | None:
                """audio: float32 16kHz mono → MFCC statistics (128-dim)."""
                if len(audio) < 8000:
                    return None
                # Pre-emphasis
                audio = np.append(audio[0], audio[1:] - 0.97 * audio[:-1])
                # STFT
                frame_len = int(0.025 * 16000)
                frame_step = int(0.010 * 16000)
                frames = []
                for start in range(0, len(audio) - frame_len, frame_step):
                    frame = audio[start:start + frame_len]
                    frames.append(frame * np.hamming(len(frame)))
                if len(frames) < 3:
                    return None
                spec = np.abs(np.fft.rfft(np.array(frames), n=512))
                # Mel filterbank
                mel = _mel_filterbank(40, 512, 16000)
                mel_energy = np.dot(spec[:, :257], mel.T) + 1e-10
                mel_db = 20 * np.log10(mel_energy)
                # DCT for MFCC
                mfcc = dct(mel_db, type=2, axis=1, norm="ortho")[:, :20]
                # Delta + delta-delta
                d1 = np.vstack([np.zeros((2, 20)), mfcc[2:] - mfcc[:-2], np.zeros((2, 20))])
                d2 = np.vstack([np.zeros((2, 20)), d1[2:] - d1[:-2], np.zeros((2, 20))])
                feats = np.concatenate([
                    mfcc.mean(axis=0), mfcc.std(axis=0),
                    d1.mean(axis=0), d1.std(axis=0),
                    d2.mean(axis=0), d2.std(axis=0),
                ])
                feats = feats / (np.linalg.norm(feats) + 1e-8)
                return feats.astype(np.float32)

        _EMBEDDING_DIM = 120  # 20 mfcc * 2 stats * 3 (mfcc+delta+delta2)
        _encoder = _MFCCEncoder()
        logger.info("Voiceprint: using scipy MFCC (%d dim)", _EMBEDDING_DIM)
        return _encoder
    except Exception as e:
        logger.warning("Voiceprint encoder init failed: %s", e)
        return None


def _mel_filterbank(n_filters: int, n_fft: int, sr: int) -> np.ndarray:
    """Create mel filterbank matrix."""
    low_mel = 0
    high_mel = 2595 * np.log10(1 + (sr / 2) / 700)
    mel_points = np.linspace(low_mel, high_mel, n_filters + 2)
    hz_points = 700 * (10 ** (mel_points / 2595) - 1)
    bins = np.floor((n_fft + 1) * hz_points / sr).astype(int)
    filters = np.zeros((n_filters, n_fft // 2 + 1))
    for i in range(1, n_filters + 1):
        for j in range(bins[i - 1], bins[i]):
            filters[i - 1, j] = (j - bins[i - 1]) / (bins[i] - bins[i - 1] + 1e-8)
        for j in range(bins[i], bins[i + 1]):
            filters[i - 1, j] = (bins[i + 1] - j) / (bins[i + 1] - bins[i] + 1e-8)
    return filters


# ── Zvec store ──
_collection = None


def _get_collection():
    global _collection, _EMBEDDING_DIM
    if _collection is not None:
        return _collection
    try:
        import zvec

        schema = zvec.CollectionSchema(
            name="voiceprints",
            fields=[
                zvec.FieldSchema(name="display_name", data_type=zvec.DataType.STRING),
                zvec.FieldSchema(name="registered_at", data_type=zvec.DataType.FLOAT64),
                zvec.FieldSchema(name="last_matched_at", data_type=zvec.DataType.FLOAT64),
            ],
            vectors=[
                zvec.VectorSchema(
                    name="embedding",
                    data_type=zvec.DataType.VECTOR_FP32,
                    dimension=_EMBEDDING_DIM,
                    index_param=zvec.HnswIndexParam(metric_type=zvec.MetricType.COSINE),
                ),
            ],
        )

        os.makedirs(VP_DB_PATH, exist_ok=True)
        _collection = zvec.create_and_open(path=VP_DB_PATH, schema=schema)
        _collection.optimize()
        logger.info("Zvec voiceprint DB at %s (dim=%d, rows=%d)",
                    VP_DB_PATH, _EMBEDDING_DIM, _collection.stats.row_count)
    except ImportError:
        logger.warning("zvec not installed — voiceprint disabled")
        return None
    except Exception:
        logger.exception("Failed to open Zvec voiceprint DB")
        return None
    return _collection


class VoiceprintStore:
    """Zvec-backed voiceprint database with persistent storage."""

    def __init__(self):
        self._col = _get_collection()

    @property
    def enabled(self) -> bool:
        return self._col is not None and _get_encoder() is not None

    def register(self, user_id: str, display_name: str, audio: np.ndarray) -> bool:
        enc = _get_encoder()
        col = self._col
        if enc is None or col is None:
            return False
        if len(audio) < 8000:
            logger.warning("Audio too short for voiceprint (%d samples)", len(audio))
            return False
        emb = enc.embed_utterance(audio)
        if emb is None:
            return False

        import zvec
        try:
            col.delete(ids=user_id)
        except Exception:
            pass

        col.insert(zvec.Doc(
            id=user_id,
            vectors={"embedding": emb.tolist()},
            fields={
                "display_name": display_name,
                "registered_at": time.time(),
                "last_matched_at": 0.0,
            },
        ))
        col.optimize()
        logger.info("Voiceprint registered: %s (%s)", display_name, user_id)
        return True

    def identify(self, audio: np.ndarray) -> tuple[Optional[str], Optional[str], float]:
        enc = _get_encoder()
        col = self._col
        if enc is None or col is None:
            return None, None, 0.0
        if len(audio) < 8000:
            return None, None, 0.0
        emb = enc.embed_utterance(audio)
        if emb is None:
            return None, None, 0.0

        import zvec
        result = col.query(
            queries=zvec.Query(field_name="embedding", vector=emb.tolist()),
            topk=1,
        )
        if not result:
            return None, None, 0.0

        doc = result[0]
        sim = float(doc.score)
        if sim >= VP_SIMILARITY_THRESHOLD:
            uid = doc.id
            name = doc.fields.get("display_name", uid)
            return uid, name, sim
        return None, None, sim

    def get_speaker_context(self, speaker_name: str | None) -> str:
        if speaker_name:
            return f"当前说话人是 {speaker_name}，请用对方习惯的方式回应。"
        return ""


_store: VoiceprintStore | None = None


def get_store() -> VoiceprintStore:
    global _store
    if _store is None:
        _store = VoiceprintStore()
    return _store


# ═══════════════════════════════════════════════════════════════
#  Conversation Memory — speaker-aware chat history (JSON)
# ═══════════════════════════════════════════════════════════════

import json as _json
import threading as _threading

_MEM_LOCK = _threading.Lock()
_MEMORY: dict[str, list[dict]] = {}
_MEMORY_LOADED = False


def _load_memories():
    global _MEMORY, _MEMORY_LOADED
    path = os.path.join(VP_DB_PATH, "conversations.json")
    try:
        if os.path.exists(path):
            with open(path) as f:
                _MEMORY = _json.load(f)
        logger.info("Loaded conversation memory: %d speakers", len(_MEMORY))
    except Exception:
        pass
    _MEMORY_LOADED = True


def _save_memories():
    path = os.path.join(VP_DB_PATH, "conversations.json")
    try:
        os.makedirs(VP_DB_PATH, exist_ok=True)
        with open(path, "w") as f:
            _json.dump(_MEMORY, f, ensure_ascii=False, indent=2)
    except Exception:
        logger.exception("Failed to save conversation memory")


def save_conversation(speaker_id: str, speaker_name: str, user_query: str, reply: str):
    """Save a conversation turn to persistent memory."""
    with _MEM_LOCK:
        if not _MEMORY_LOADED:
            _load_memories()
        if speaker_id not in _MEMORY:
            _MEMORY[speaker_id] = []
        entry = {
            "name": speaker_name,
            "query": user_query,
            "reply": reply,
            "time": time.time(),
        }
        _MEMORY[speaker_id].append(entry)
        _MEMORY[speaker_id] = _MEMORY[speaker_id][-20:]  # keep last 20
        _save_memories()
    logger.info("Memory saved: %s (%d turns)", speaker_id, len(_MEMORY[speaker_id]))


def load_conversation_context(speaker_id: str | None, speaker_name: str | None, limit: int = 5) -> str:
    """Retrieve recent conversation history as formatted text for system prompt."""
    if not speaker_id:
        return ""
    with _MEM_LOCK:
        if not _MEMORY_LOADED:
            _load_memories()
        entries = _MEMORY.get(speaker_id, [])
        if not entries:
            return ""
        history = []
        for e in entries[-limit:]:
            q, r = e.get("query", ""), e.get("reply", "")
            if q and r:
                history.append(f"用户: {q}\n助手: {r}")
        if history:
            return "以下是之前和该用户的对话记录，请参考上下文回应:\n" + "\n".join(history)
    return ""


def load_conversation_hist(speaker_id: str, limit: int = 6) -> list[dict]:
    """Return conversation history as OpenAI message list for self.hist."""
    with _MEM_LOCK:
        if not _MEMORY_LOADED:
            _load_memories()
        entries = _MEMORY.get(speaker_id, [])
        hist = []
        for e in entries[-limit:]:
            q, r = e.get("query", ""), e.get("reply", "")
            if q and r:
                hist.append({"role": "user", "content": q})
                hist.append({"role": "assistant", "content": r})
    return hist