336 lines
15 KiB
Python
336 lines
15 KiB
Python
"""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()
|