Fix unknown speaker filtering and discovery

Author: AnkushMalaker

backends/advanced/src/advanced_omi_backend/controllers/data_audit_controller.py

@@ -21,7 +21,10 @@
 
 from advanced_omi_backend.client_manager import synthetic_client_id
 from advanced_omi_backend.config import get_diarization_settings
-from advanced_omi_backend.constants import is_non_enrollable_speaker
+from advanced_omi_backend.constants import (
+    is_non_enrollable_speaker,
+    is_unknown_speaker_label,
+)
 from advanced_omi_backend.controllers.conversation_controller import (
     archive_conversation_audio,
 )
@@ -309,6 +312,20 @@ def _speakers_for_doc(doc: dict) -> List[str]:
     return sorted(speakers)
 
 
+def _speaker_facets_for_doc(doc: dict) -> tuple[List[str], bool]:
+    """Return global identities separately from conversation-local unknowns.
+
+    ``Unknown Speaker N`` is numbered independently inside each conversation.
+    It is a display description, not a corpus-wide identity, so it must never
+    enter a global list of named speakers.
+    """
+    labels = _speakers_for_doc(doc)
+    return (
+        [label for label in labels if not is_unknown_speaker_label(label)],
+        any(is_unknown_speaker_label(label) for label in labels),
+    )
+
+
 def _unknown_speech_count(doc: dict) -> int:
     """Speech segments in the audited version not matched to an enrolled speaker.
 
@@ -401,6 +418,7 @@ async def list_for_audit(
     created_before: Optional[datetime] = None,
     include_speakers: Optional[List[str]] = None,
     exclude_speakers: Optional[List[str]] = None,
+    unknown_speakers: Optional[str] = None,
     dataset_id: Optional[str] = None,
     exported: Optional[str] = None,
     archived_only: bool = False,
@@ -508,12 +526,15 @@ async def list_for_audit(
         # compound predicate), so the filter UI offers exactly the labels that
         # exist in this view — and isn't narrowed by its own selection.
         available_speakers: set = set()
+        has_unknown_speakers = False
 
         for doc in raw_docs:
             duration = doc.get("audio_total_duration") or 0.0
             va = doc.get("vad_analysis")
-            doc_speakers = _speakers_for_doc(doc)
+            all_doc_speakers = _speakers_for_doc(doc)
+            doc_speakers, doc_has_unknown = _speaker_facets_for_doc(doc)
             available_speakers.update(doc_speakers)
+            has_unknown_speakers = has_unknown_speakers or doc_has_unknown
             unknown_count = _unknown_speech_count(doc)
 
             speech_fraction = None
@@ -539,10 +560,16 @@ async def list_for_audit(
                         continue
                     if speech_fraction < min_speech_fraction:
                         continue
-                if include_set and not include_set.intersection(doc_speakers):
+                include_matches = bool(include_set.intersection(doc_speakers))
+                if (include_set or unknown_speakers == "include") and not (
+                    include_matches
+                    or (unknown_speakers == "include" and doc_has_unknown)
+                ):
                     continue
                 if exclude_set and exclude_set.intersection(doc_speakers):
                     continue
+                if unknown_speakers == "exclude" and doc_has_unknown:
+                    continue
                 # Opt-in: drop conversations with nothing left to triage (every
                 # speech segment already has an identified_as).
                 if hide_reviewed and unknown_count == 0:
@@ -563,7 +590,7 @@ async def list_for_audit(
                     "client_id": doc.get("client_id"),
                     "created_at": created_at.isoformat() if created_at else None,
                     "duration_seconds": duration,
-                    "speakers": doc_speakers,
+                    "speakers": all_doc_speakers,
                     "unknown_speech_segments": unknown_count,
                     "marginal_identified_segments": _marginal_identified_count(
                         doc, similarity_threshold, marginal_margin
@@ -635,6 +662,7 @@ async def list_for_audit(
             "marginal_margin": marginal_margin,
             "unanalyzed_count": unanalyzed_count,
             "speakers": sorted(available_speakers),
+            "has_unknown_speakers": has_unknown_speakers,
             "datasets": available_datasets,
         }
 

backends/advanced/src/advanced_omi_backend/controllers/guided_enrollment_controller.py

@@ -30,10 +30,12 @@
 
 from advanced_omi_backend.config import get_diarization_settings
 from advanced_omi_backend.constants import is_non_enrollable_speaker
+from advanced_omi_backend.controllers import data_audit_controller
 from advanced_omi_backend.controllers.queue_controller import (
     JOB_RESULT_TTL,
     default_queue,
 )
+from advanced_omi_backend.models.annotation import Annotation, AnnotationType
 from advanced_omi_backend.models.conversation import Conversation
 from advanced_omi_backend.speaker_recognition_client import SpeakerRecognitionClient
 from advanced_omi_backend.users import User
@@ -44,6 +46,9 @@
 from advanced_omi_backend.workers.speaker_discovery_jobs import (
     discover_speaker_candidates_job,
 )
+from advanced_omi_backend.workers.unknown_speaker_jobs import (
+    discover_unknown_speakers_job,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -74,6 +79,196 @@ def _discovery_runs_collection():
     return Conversation.get_pymongo_collection().database["speaker_discovery_runs"]
 
 
+def _unknown_clusters_collection():
+    return Conversation.get_pymongo_collection().database["unknown_speaker_clusters"]
+
+
+def _unknown_runs_collection():
+    return Conversation.get_pymongo_collection().database["unknown_speaker_runs"]
+
+
+async def enqueue_unknown_discovery(user: User):
+    """Start or reattach to corpus-wide unknown-identity clustering."""
+    user_id = str(user.user_id)
+    existing = await _unknown_runs_collection().find_one({"requested_by": user_id})
+    status = _job_status(existing.get("job_id") if existing else None)
+    if status in {"queued", "started", "deferred", "scheduled"}:
+        return {"job_id": existing["job_id"], "status": status, "reused": True}
+    job = default_queue.enqueue(
+        discover_unknown_speakers_job,
+        requested_by=user_id,
+        job_timeout=14400,
+        result_ttl=JOB_RESULT_TTL,
+        description="Discover unknown speakers across corpus",
+    )
+    await _unknown_runs_collection().update_one(
+        {"requested_by": user_id},
+        {"$set": {"job_id": job.id, "queued_at": datetime.now(timezone.utc)}},
+        upsert=True,
+    )
+    return {"job_id": job.id, "status": "queued", "reused": False}
+
+
+async def list_unknown_clusters(user: User, limit: int = 50):
+    cursor = (
+        _unknown_clusters_collection()
+        .find({"requested_by": str(user.user_id), "status": "pending"}, {"_id": 0})
+        .sort("segment_count", -1)
+        .limit(limit)
+    )
+    return {"clusters": await cursor.to_list(length=limit)}
+
+
+async def decide_unknown_cluster(
+    user: User,
+    cluster_id: str,
+    run_fingerprint: str,
+    action: str,
+    speaker_name: Optional[str],
+    accepted_identity_keys: list[str],
+    enrollment_clips: list[dict],
+):
+    """Relabel accepted members and enroll only explicitly selected clips."""
+    collection = _unknown_clusters_collection()
+    cluster = await collection.find_one(
+        {
+            "requested_by": str(user.user_id),
+            "cluster_id": cluster_id,
+            "run_fingerprint": run_fingerprint,
+            "status": "pending",
+        }
+    )
+    if not cluster:
+        return JSONResponse(status_code=409, content={"error": "Cluster is stale"})
+    if action == "dismiss":
+        await collection.update_one(
+            {"_id": cluster["_id"]},
+            {"$set": {"status": "dismissed", "decided_at": datetime.now(timezone.utc)}},
+        )
+        return {"status": "dismissed"}
+    if not speaker_name or is_non_enrollable_speaker(speaker_name):
+        return JSONResponse(
+            status_code=422, content={"error": "A real speaker name is required"}
+        )
+
+    accepted = {
+        member["identity_key"]: member
+        for member in cluster["members"]
+        if member["identity_key"] in set(accepted_identity_keys)
+    }
+    if not accepted:
+        return JSONResponse(
+            status_code=422, content={"error": "Select at least one identity"}
+        )
+    if not enrollment_clips:
+        return JSONResponse(
+            status_code=422, content={"error": "Select at least one enrollment clip"}
+        )
+
+    validated_segments = []
+    for member in accepted.values():
+        conversation = await Conversation.find_one(
+            Conversation.conversation_id == member["conversation_id"]
+        )
+        if not conversation or conversation.user_id != str(user.user_id):
+            return JSONResponse(
+                status_code=409,
+                content={"error": "Conversation is no longer available"},
+            )
+        segments = _active_segments(conversation.model_dump())
+        for ref in member["segments"]:
+            index = ref["segment_index"]
+            if (
+                index >= len(segments)
+                or abs(float(segments[index]["start"]) - ref["start"]) > 0.01
+                or not is_non_enrollable_speaker(
+                    segments[index].get("identified_as")
+                    or segments[index].get("speaker")
+                )
+                or (
+                    segments[index].get("identified_as")
+                    or segments[index].get("speaker")
+                )
+                != member["local_label"]
+            ):
+                return JSONResponse(
+                    status_code=409, content={"error": "Transcript changed since scan"}
+                )
+            validated_segments.append((member, ref, index, segments[index]))
+
+    annotations = 0
+    for member, ref, index, segment in validated_segments:
+        await Annotation(
+            annotation_type=AnnotationType.DIARIZATION,
+            user_id=str(user.user_id),
+            conversation_id=member["conversation_id"],
+            segment_index=index,
+            original_speaker=segment.get("speaker") or "",
+            corrected_speaker=speaker_name,
+            segment_start_time=ref["start"],
+            processed=False,
+        ).insert()
+        annotations += 1
+
+    client = SpeakerRecognitionClient()
+    existing_speaker = await client.get_speaker_by_name(
+        speaker_name, user_id=str(user.user_id)
+    )
+    enrolled = appended = 0
+    for clip in enrollment_clips:
+        member = accepted.get(clip.get("identity_key"))
+        if not member:
+            continue
+        ref = next(
+            (
+                segment
+                for segment in member["segments"]
+                if segment["segment_index"] == clip.get("segment_index")
+            ),
+            None,
+        )
+        if not ref or ref["duration"] < MIN_CLIP_SECONDS:
+            continue
+        wav = await reconstruct_audio_segment(
+            member["conversation_id"], ref["start"], ref["end"]
+        )
+        if existing_speaker:
+            result = await client.append_to_speaker(
+                existing_speaker["id"], wav, user_id=str(user.user_id)
+            )
+            if not result.get("error"):
+                appended += 1
+        else:
+            result = await client.enroll_new_speaker(
+                speaker_name, wav, user_id=str(user.user_id)
+            )
+            if not result.get("error"):
+                enrolled += 1
+                existing_speaker = await client.get_speaker_by_name(
+                    speaker_name, user_id=str(user.user_id)
+                )
+
+    await collection.update_one(
+        {"_id": cluster["_id"]},
+        {
+            "$set": {
+                "status": "confirmed",
+                "speaker_name": speaker_name,
+                "accepted_identity_keys": list(accepted),
+                "decided_at": datetime.now(timezone.utc),
+            }
+        },
+    )
+    apply_result = await data_audit_controller.apply_triage(user)
+    return {
+        "status": "confirmed",
+        "annotations_created": annotations,
+        "enrolled_new": enrolled,
+        "appended": appended,
+        "corrections_applied": getattr(apply_result, "body", apply_result),
+    }
+
+
 def _job_status(job_id: Optional[str]) -> Optional[str]:
     if not job_id:
         return None

backends/advanced/src/advanced_omi_backend/routers/modules/data_audit_routes.py

@@ -68,6 +68,19 @@ class MergeRequest(BaseModel):
     )
 
 
+class UnknownClusterClip(BaseModel):
+    identity_key: str
+    segment_index: int
+
+
+class UnknownClusterDecision(BaseModel):
+    run_fingerprint: str
+    action: Literal["confirm", "dismiss"]
+    speaker_name: Optional[str] = None
+    accepted_identity_keys: List[str] = Field(default_factory=list)
+    enrollment_clips: List[UnknownClusterClip] = Field(default_factory=list)
+
+
 class ScreenRequest(BaseModel):
     conversation_ids: List[str] = Field(..., min_length=1, max_length=200)
     policy: Optional[str] = Field(
@@ -177,6 +190,11 @@ async def list_conversations(
     exclude_speakers: Optional[str] = Query(
         None, description="Comma-separated speakers a conversation must contain none of"
     ),
+    unknown_speakers: Optional[Literal["include", "exclude"]] = Query(
+        None,
+        description="Include conversations containing any conversation-local unknown "
+        "speaker, or exclude every such conversation",
+    ),
     dataset_id: Optional[str] = Query(
         None,
         max_length=200,
@@ -220,6 +238,7 @@ def _csv(v: Optional[str]) -> Optional[list]:
         created_before=created_before,
         include_speakers=_csv(include_speakers),
         exclude_speakers=_csv(exclude_speakers),
+        unknown_speakers=unknown_speakers,
         dataset_id=dataset_id,
         exported=exported,
         archived_only=archived_only,
@@ -446,6 +465,38 @@ async def guided_enrollment_discover(
     )
 
 
+@router.post("/enrollment/unknown/discover")
+async def unknown_speaker_discover(
+    current_user: User = Depends(current_active_user),
+):
+    return await guided_enrollment_controller.enqueue_unknown_discovery(current_user)
+
+
+@router.get("/enrollment/unknown/clusters")
+async def unknown_speaker_clusters(
+    limit: int = Query(50, ge=1, le=200),
+    current_user: User = Depends(current_active_user),
+):
+    return await guided_enrollment_controller.list_unknown_clusters(current_user, limit)
+
+
+@router.post("/enrollment/unknown/clusters/{cluster_id}/decide")
+async def unknown_speaker_cluster_decide(
+    cluster_id: str,
+    body: UnknownClusterDecision,
+    current_user: User = Depends(current_active_user),
+):
+    return await guided_enrollment_controller.decide_unknown_cluster(
+        current_user,
+        cluster_id,
+        body.run_fingerprint,
+        body.action,
+        body.speaker_name,
+        body.accepted_identity_keys,
+        [clip.model_dump() for clip in body.enrollment_clips],
+    )
+
+
 @router.post("/enrollment/guided/mine")
 async def guided_enrollment_mine(
     speaker_name: str = Form(..., description="Enrolled speaker to mine for"),