ji

tests/test_speaker_inbox.py

@@ -0,0 +1,335 @@
+"""Unidentified-speaker inbox, sorting clips into profiles, and live samples."""
+from __future__ import annotations
+
+import asyncio
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+import numpy as np
+
+from javis.core.config import Config, PathsConfig
+from javis.core.events import EventBus
+from javis.core.state import AgentState
+from javis.stt.unidentified import UnidentifiedStore
+from javis.stt.voice_profiler import NullProfiler, VoiceProfiler
+
+SR = 16000
+
+
+def _clip(voice: int, seconds: float = 1.0) -> np.ndarray:
+    """Audio whose first sample encodes the (fake) speaker id."""
+    a = np.full(int(SR * seconds), 0.2, dtype=np.float32)
+    a[0] = voice / 100.0
+    return a
+
+
+def _fake_embed(audio: np.ndarray, sr: int) -> np.ndarray:
+    voice = round(float(audio[0]) * 100)
+    v = np.full(192, 0.01, dtype=np.float32)
+    v[voice] = 1.0
+    return v / np.linalg.norm(v)
+
+
+def _profiler(base: Path) -> VoiceProfiler:
+    prof = VoiceProfiler(Config(paths=PathsConfig(base_dir=base)))
+    prof._encoder = object()          # pretend the ECAPA model is loaded
+    prof.embed = _fake_embed          # type: ignore[method-assign]
+    return prof
+
+
+class UnidentifiedStoreTests(unittest.TestCase):
+    def test_add_list_load_remove_and_persist(self) -> None:
+        with tempfile.TemporaryDirectory() as d:
+            store = UnidentifiedStore(Path(d), max_clips=10)
+            emb = _fake_embed(_clip(3), SR)
+            a = store.add(_clip(3), SR, embedding=emb, text=" hallo ", source="local",
+                          best_name="Bob", best_score=0.41)
+            b = store.add(_clip(4, 2.0), SR, text="zwei")
+            entries = store.entries()
+            self.assertEqual([e["id"] for e in entries], [b, a])   # newest first
+            self.assertEqual(entries[1]["text"], "hallo")
+            self.assertEqual(entries[1]["best_name"], "Bob")
+            self.assertAlmostEqual(entries[0]["duration_s"], 2.0)
+
+            audio, sr = store.load_audio(a)
+            self.assertEqual(sr, SR)
+            self.assertEqual(len(audio), SR)
+            np.testing.assert_allclose(store.load_embedding(a), emb, rtol=1e-6)
+            self.assertIsNone(store.load_embedding(b))
+
+            # A fresh instance reads the same index from disk.
+            self.assertEqual(len(UnidentifiedStore(Path(d))), 2)
+
+            self.assertEqual(store.remove([a, "nope"]), 1)
+            self.assertFalse((Path(d) / f"{a}.wav").exists())
+            self.assertFalse((Path(d) / f"{a}.npy").exists())
+            self.assertEqual(store.clear(), 1)
+            self.assertEqual(store.entries(), [])
+            with self.assertRaises(KeyError):
+                store.load_audio(a)
+
+    def test_oldest_clips_are_pruned(self) -> None:
+        with tempfile.TemporaryDirectory() as d:
+            store = UnidentifiedStore(Path(d), max_clips=2)
+            ids = [store.add(_clip(i + 1), SR, text=str(i)) for i in range(3)]
+            self.assertEqual([e["id"] for e in store.entries()], [ids[2], ids[1]])
+            self.assertFalse((Path(d) / f"{ids[0]}.wav").exists())
+
+
+class ProfilerEmbeddingTests(unittest.TestCase):
+    def test_enroll_rank_identify_refine(self) -> None:
+        with tempfile.TemporaryDirectory() as d:
+            prof = _profiler(Path(d))
+            self.assertTrue(prof.enroll_embeddings("Alice", [_fake_embed(_clip(1), SR)] * 2))
+            prof.enroll("Bob", [_clip(2)], SR)
+            ranked = prof.rank(_fake_embed(_clip(1), SR))
+            self.assertEqual(ranked[0][0], "Alice")
+            self.assertGreater(ranked[0][1], 0.99)
+            self.assertEqual(prof.identify_embedding(_fake_embed(_clip(2), SR))[0], "Bob")
+            self.assertIsNone(prof.identify_embedding(_fake_embed(_clip(9), SR)))
+            self.assertEqual(prof.rank(np.zeros(192)), [])
+
+            # Running centroid: a Bob sample pulls Alice's centroid towards Bob.
+            before = dict(prof.rank(_fake_embed(_clip(2), SR)))["Alice"]
+            self.assertTrue(prof.refine_embedding("Alice", _fake_embed(_clip(2), SR)))
+            after = dict(prof.rank(_fake_embed(_clip(2), SR)))["Alice"]
+            self.assertGreater(after, before)
+            alice = next(p for p in prof.list_profiles() if p["name"] == "Alice")
+            self.assertEqual(alice["clip_count"], 3)
+
+            self.assertFalse(prof.refine_embedding("Nobody", _fake_embed(_clip(1), SR)))
+            self.assertFalse(prof.refine_embedding("Alice", np.ones(10)))   # wrong shape
+            self.assertFalse(prof.enroll_embeddings("Zero", [np.zeros(192)]))
+
+    def test_sanitised_name_collision_gets_its_own_file(self) -> None:
+        with tempfile.TemporaryDirectory() as d:
+            prof = _profiler(Path(d))
+            prof.enroll_embeddings("Anna B", [_fake_embed(_clip(1), SR)])
+            prof.enroll_embeddings("Anna_B", [_fake_embed(_clip(2), SR)])
+            files = {p["name"]: p["file"] for p in prof.list_profiles()}
+            self.assertNotEqual(files["Anna B"], files["Anna_B"])
+            prof2 = _profiler(Path(d))
+            prof2._load_index()
+            self.assertEqual(prof2.rank(_fake_embed(_clip(1), SR))[0][0], "Anna B")
+
+    def test_null_profiler_mirrors_api(self) -> None:
+        n = NullProfiler()
+        self.assertFalse(n.is_ready())
+        self.assertFalse(n.has_profile("x"))
+        self.assertEqual(n.rank(np.ones(192)), [])
+        self.assertFalse(n.enroll_embeddings("x", [np.ones(192)]))
+        self.assertFalse(n.refine_embedding("x", np.ones(192)))
+        self.assertIsNone(n.identify_embedding(np.ones(192)))
+
+
+class _FakeMic:
+    chunk_ms = 32
+
+    def __init__(self) -> None:
+        self.drains = 0
+
+    def drain(self) -> None:
+        self.drains += 1
+
+
+def _orchestrator(base: Path):
+    from javis.agent.orchestrator import Orchestrator
+
+    orch = object.__new__(Orchestrator)
+    orch.config = Config(paths=PathsConfig(base_dir=base))
+    orch.bus = EventBus()
+    orch.profiler = _profiler(base)
+    orch.unidentified = UnidentifiedStore(base / "voice_profiles" / "unidentified")
+    orch._enroll_clip_buffer = []
+    orch._enroll_ask_pending = False
+    orch._turn_lock = asyncio.Lock()
+    orch._mic_free = asyncio.Event()
+    orch._mic_free.set()
+    orch.mic = _FakeMic()
+    orch.vad = None
+    orch.state = AgentState.IDLE
+    events: list[tuple[str, object]] = []
+
+    async def _record(name, payload=None):
+        events.append((name, payload))
+
+    for ev in ("unidentified_changed", "voice_profiles_changed", "speaker_enrolled",
+               "speaker_ask_name"):
+        orch.bus.subscribe(ev, lambda p, _n=ev: _record(_n, p))
+    return orch, events
+
+
+class AssignUnidentifiedTests(unittest.IsolatedAsyncioTestCase):
+    async def test_unknown_turn_is_stored_then_sorted_into_new_profile(self) -> None:
+        with tempfile.TemporaryDirectory() as d:
+            orch, events = _orchestrator(Path(d))
+            orch.config.speaker_id.auto_enroll = False
+            await orch._handle_unknown_speaker(_clip(5), SR, "wer bin ich", source="Küche")
+            await orch._handle_unknown_speaker(_clip(5), SR, "   ", source="local")  # no text
+            await asyncio.sleep(0.2)
+            entries = orch.unidentified.entries()
+            self.assertEqual(len(entries), 1)
+            self.assertEqual(entries[0]["source"], "Küche")
+            self.assertIn(("unidentified_changed", {"added": entries[0]["id"]}), events)
+
+            res = await orch.assign_unidentified([entries[0]["id"]], "  Carla ")
+            self.assertEqual(res, {"name": "Carla", "added": 1, "created": True})
+            self.assertTrue(orch.profiler.has_profile("Carla"))
+            self.assertEqual(orch.unidentified.entries(), [])
+            self.assertEqual(orch.profiler.identify_embedding(_fake_embed(_clip(5), SR))[0],
+                             "Carla")
+            names = [e[0] for e in events]
+            self.assertIn("voice_profiles_changed", names)
+
+    async def test_assign_to_existing_profile_refines_it(self) -> None:
+        with tempfile.TemporaryDirectory() as d:
+            orch, _ = _orchestrator(Path(d))
+            orch.profiler.enroll("Dora", [_clip(6)], SR)
+            ids = [orch.unidentified.add(_clip(6), SR, text="a"),
+                   orch.unidentified.add(_clip(6), SR, text="b")]   # no stored embedding
+            res = await orch.assign_unidentified(ids, "Dora")
+            self.assertEqual(res["added"], 2)
+            self.assertFalse(res["created"])
+            dora = next(p for p in orch.profiler.list_profiles() if p["name"] == "Dora")
+            self.assertEqual(dora["clip_count"], 3)
+
+    async def test_assign_validates_input(self) -> None:
+        with tempfile.TemporaryDirectory() as d:
+            orch, _ = _orchestrator(Path(d))
+            with self.assertRaises(ValueError):
+                await orch.assign_unidentified(["x"], " ")
+            with self.assertRaises(RuntimeError):
+                await orch.assign_unidentified(["missing"], "Eve")
+            orch.profiler._encoder = None
+            with self.assertRaises(RuntimeError):
+                await orch.assign_unidentified(["x"], "Eve")
+
+    async def test_name_prompt_confirm_removes_inbox_copies(self) -> None:
+        with tempfile.TemporaryDirectory() as d:
+            orch, events = _orchestrator(Path(d))
+            await orch._handle_unknown_speaker(_clip(7), SR, "hallo", source="local")
+            await orch._handle_unknown_speaker(_clip(7), SR, "noch mal", source="local")
+            self.assertEqual(len(orch.pending_enroll_clips()), 2)
+            self.assertEqual(sum(1 for e in events if e[0] == "speaker_ask_name"), 1)
+            self.assertTrue(await orch.auto_enroll_confirm("k", "Frank"))
+            self.assertEqual(orch.unidentified.entries(), [])
+            self.assertTrue(orch.profiler.has_profile("Frank"))
+            self.assertIn(("speaker_enrolled", {"name": "Frank", "clip_count": 2}), events)
+
+    async def test_dismiss_keeps_inbox_and_allows_new_prompt(self) -> None:
+        with tempfile.TemporaryDirectory() as d:
+            orch, events = _orchestrator(Path(d))
+            await orch._handle_unknown_speaker(_clip(8), SR, "hallo", source="local")
+            await asyncio.sleep(0.2)
+            orch.auto_enroll_dismiss()
+            self.assertEqual(orch.pending_enroll_clips(), [])
+            self.assertEqual(len(orch.unidentified.entries()), 1)
+            await orch._handle_unknown_speaker(_clip(8), SR, "wieder", source="local")
+            self.assertEqual(sum(1 for e in events if e[0] == "speaker_ask_name"), 2)
+
+
+class LiveSampleTests(unittest.IsolatedAsyncioTestCase):
+    async def test_record_and_add_samples(self) -> None:
+        with tempfile.TemporaryDirectory() as d:
+            orch, events = _orchestrator(Path(d))
+            orch.profiler.enroll("Gina", [_clip(9)], SR)
+            seen = {}
+
+            async def fake_record(config, mic, vad, **kw):
+                # The wake loop must be paused and the turn lock held meanwhile.
+                seen["mic_free"] = orch._mic_free.is_set()
+                seen["locked"] = orch._turn_lock.locked()
+                seen["state"] = orch.state
+                return _clip(9)
+
+            async def no_chime():
+                return None
+
+            orch._play_wake_chime = no_chime
+            with patch("javis.agent.orchestrator.record_utterance", fake_record):
+                sample = await orch.record_voice_sample("Gina")
+            self.assertEqual(seen, {"mic_free": False, "locked": True,
+                                    "state": AgentState.ENROLLING})
+            self.assertTrue(orch._mic_free.is_set())
+            self.assertFalse(orch._turn_lock.locked())
+            self.assertEqual(orch.state, AgentState.IDLE)
+            self.assertGreater(sample.score, 0.99)
+            self.assertEqual(sample.best_name, "Gina")
+            self.assertAlmostEqual(sample.duration_s, 1.0)
+
+            used = await orch.add_voice_samples("Gina", [sample, sample])
+            self.assertEqual(used, 2)
+            gina = next(p for p in orch.profiler.list_profiles() if p["name"] == "Gina")
+            self.assertEqual(gina["clip_count"], 3)
+            self.assertIn(("voice_profiles_changed",
+                           {"name": "Gina", "added": 2, "created": False}), events)
+
+    async def test_record_errors_are_user_facing(self) -> None:
+        with tempfile.TemporaryDirectory() as d:
+            orch, _ = _orchestrator(Path(d))
+
+            async def no_speech(*a, **kw):
+                return None
+
+            async def no_chime():
+                return None
+
+            orch._play_wake_chime = no_chime
+            with patch("javis.agent.orchestrator.record_utterance", no_speech), \
+                    self.assertRaisesRegex(RuntimeError, "No speech"):
+                await orch.record_voice_sample("Gina")
+            self.assertTrue(orch._mic_free.is_set())
+            self.assertFalse(orch._turn_lock.locked())
+
+            async with orch._turn_lock:   # a conversation turn is running
+                with self.assertRaisesRegex(RuntimeError, "busy"):
+                    async with orch._exclusive_mic(wait_s=0.05):
+                        pass
+            self.assertTrue(orch._mic_free.is_set())
+
+            orch.mic = None
+            with self.assertRaisesRegex(RuntimeError, "microphone"):
+                await orch.record_voice_sample("Gina")
+
+    async def test_wake_loop_yields_mic_to_exclusive_recording(self) -> None:
+        from javis.agent.orchestrator import Orchestrator
+
+        with tempfile.TemporaryDirectory() as d:
+            orch, _ = _orchestrator(Path(d))
+            orch._stop = asyncio.Event()
+            fed: list[int] = []
+
+            class Wake:
+                def reset(self):
+                    pass
+
+                def feed(self, chunk):
+                    fed.append(len(chunk))
+                    return 0.0
+
+                def triggered(self, score):
+                    return False
+
+            class Mic(_FakeMic):
+                async def chunks(self, stale_timeout=None):
+                    while True:
+                        await asyncio.sleep(0.005)
+                        yield np.zeros(512, dtype=np.float32)
+
+            orch.wake = Wake()
+            orch.mic = Mic()
+            orch._mic_free.clear()
+            task = asyncio.ensure_future(Orchestrator._wait_for_wake(orch))
+            await asyncio.sleep(0.05)
+            self.assertEqual(fed, [])          # waits while the mic is taken
+            orch._mic_free.set()
+            await asyncio.sleep(0.05)
+            self.assertTrue(fed)               # resumed wake detection
+            orch._mic_free.clear()
+            self.assertFalse(await asyncio.wait_for(task, 1.0))
+
+
+if __name__ == "__main__":
+    unittest.main()