From ffdabb6ae6fb9c9f28549304a0f907ee294a9bc9 Mon Sep 17 00:00:00 2001 From: admin Date: Sat, 20 Jun 2026 16:58:55 +0000 Subject: [PATCH] Refactor speaker ID into helpers; run before summarization; fix _chat_completion call --- scraibe/tasks.py | 256 +++++++++++++++++++++++++---------------------- 1 file changed, 139 insertions(+), 117 deletions(-) diff --git a/scraibe/tasks.py b/scraibe/tasks.py index e0cefad..1030d8c 100644 --- a/scraibe/tasks.py +++ b/scraibe/tasks.py @@ -6,6 +6,7 @@ import os import json import logging import tempfile +import re from datetime import datetime from .celery_app import celery_app @@ -247,6 +248,111 @@ def send_error_email(to: str, error_message: str, task_id: str): logger.error("Failed to send error email to %s for job %s: %s", to, task_id, e) +def _accent_color() -> str: + """ + Return the accent color used in emails/templates. + """ + return (os.getenv("EMAIL_ACCENT_COLOR") or "#7C6DA0").strip() + + +def _identify_speakers_via_llm(scraibe, transcript_text: str) -> dict: + """ + Use the summarizer LLM to identify speakers in the transcript. + + Returns: dict mapping normalized speaker labels to name/role, e.g. + {"SPEAKER 1": "JUDGE MARTINEZ", "SPEAKER 2": "DEFENSE COUNSEL"} + """ + try: + scraibe._ensure_summarizer() + summarizer = scraibe._summarizer + except Exception as e: + logger.warning("Failed to initialize summarizer for speaker identification: %s", e) + return {} + + prompt = ( + "Below is a transcript with speaker labels like 'SPEAKER 1', 'SPEAKER 2', etc. " + "Based on the context and how each speaker talks, identify each speaker as:\n" + "- Their real name, if it is clearly mentioned or strongly implied, OR\n" + "- A concise role/position (e.g., Judge, Doctor, Manager, Interviewer, Client, Witness), " + "if their identity is not clear.\n" + "Do not invent random personal names. " + "Do not add extra commentary. Output ONLY a mapping in this exact format, one per line:\n" + "SPEAKER 1: Name or Role\n" + "SPEAKER 2: Name or Role\n" + "SPEAKER 3: Name or Role\n" + "\n" + "Transcript:\n" + + transcript_text + ) + + try: + reply = summarizer._chat_completion( + system_prompt="You are an expert legal and business meeting analyst.", + user_prompt=prompt, + ) + except Exception as e: + logger.warning("LLM call failed during speaker identification: %s", e) + return {} + + speaker_map = {} + for m in re.finditer(r"SPEAKER\s+(\d+)\s*:\s*(.+)", reply, re.IGNORECASE): + spk = f"SPEAKER {m.group(1).strip()}" + name = m.group(2).strip().rstrip(".").upper() + if name: + speaker_map[spk] = name + + logger.info("Speaker identification mapping: %s", speaker_map) + return speaker_map + + +def _apply_speaker_map(transcript_text: str, segments: list, speaker_map: dict) -> tuple: + """ + Apply speaker_map to: + - transcript_text lines like "[00:12] SPEAKER 1: ..." + - segment speaker fields. + + Returns (updated_transcript_text, updated_segments). + """ + if not speaker_map: + return transcript_text, segments + + # Normalize function: e.g. "SPEAKER 1" -> "SPEAKER 1" + def normalize_label(label: str) -> str: + return re.sub(r"\s+", " ", re.sub(r"[^A-Z0-9\s]", "", label.upper())).strip() + + # Replace in transcript lines + def replace_in_line(line: str) -> str: + # Pattern: timestamp bracket + speaker label + colon + # More flexible: allow any non-empty label before the colon + def repl(m): + prefix = m.group(1) # e.g. "[00:12] " + label = m.group(2).strip() + normalized = normalize_label(label) + new_label = speaker_map.get(normalized, label) + return f"{prefix}{new_label}: " + return re.sub( + r"(\[\d+:\d+(?::\d+)?\]\s*)([^\]]+?):\s*", + repl, + line, + ) + + updated_transcript = "\n".join( + replace_in_line(line) for line in transcript_text.splitlines() + ) + + # Update segments + updated_segments = [] + for seg in segments: + sp = (seg.get("speaker") or "").strip() + sp_norm = normalize_label(sp) + sp_new = speaker_map.get(sp_norm, sp) + seg = dict(seg) + seg["speaker"] = sp_new + updated_segments.append(seg) + + return updated_transcript, updated_segments + + @celery_app.task( name="scraibe.tasks.process_transcription_task", bind=True, @@ -266,7 +372,7 @@ def process_transcription_task( identify_speakers: bool = False, ): """ - Async task: transcribe audio, optionally summarize, then email results. + Async task: transcribe audio, optionally identify speakers, optionally summarize, then email results. Cleans up temporary files after completion. """ task_id = self.request.id @@ -294,121 +400,42 @@ def process_transcription_task( ) raise - # 3) Transcription - if task_type == "transcript_and_summarize": - result = scraibe.transcript_and_summarize( - audio_file=audio_path, - language=language or None, - num_speakers=int(num_speakers) if num_speakers else None, - verbose=True, - for_export=True, - ) - transcript_text = result.get("transcript", "") - summary_text = result.get("summary", "") - segments = result.get("segments", []) - raw_result = result.get("raw_result") - else: - result = scraibe.transcribe( - audio_file=audio_path, - language=language or None, - num_speakers=int(num_speakers) if num_speakers else None, - verbose=True, - for_export=True, - ) - transcript_text = result.get("transcript", "") - summary_text = "" - segments = result.get("segments", []) - raw_result = result.get("raw_result") + # 3) Transcription (always first, without summarization) + result = scraibe.transcribe( + audio_file=audio_path, + language=language or None, + num_speakers=int(num_speakers) if num_speakers else None, + verbose=True, + for_export=True, + ) + transcript_text = result.get("transcript", "") + segments = result.get("segments", []) + raw_result = result.get("raw_result") + summary_text = "" - # 3b) Optional speaker identification - speaker_map = {} # e.g. {"SPEAKER 1": "John", "SPEAKER 2": "Maria"} + # 4) Optional speaker identification (before summarization) if identify_speakers: try: - # Use the same summarizer client as transcript_and_summarize - scraibe._ensure_summarizer() - summarizer = scraibe._summarizer - - prompt = ( - "Below is a transcript with speaker labels like 'SPEAKER 1', 'SPEAKER 2', etc. " - "Based on the context and how each speaker talks, identify each speaker as:\n" - "- Their real name, if it is clearly mentioned or strongly implied, OR\n" - "- A concise role/position (e.g., Judge, Doctor, Manager, Interviewer, Client, Witness), " - "if their identity is not clear.\n" - "Do not invent random personal names. " - "Do not add extra commentary. Output ONLY a mapping in this exact format, one per line:\n" - "SPEAKER 1: Name or Role\n" - "SPEAKER 2: Name or Role\n" - "SPEAKER 3: Name or Role\n" - "\n" - "Transcript:\n" - + transcript_text - ) - - response = summarizer._chat_completion( - messages=[{"role": "user", "content": prompt}], - temperature=0.3, - max_tokens=300, - ) - reply = (response or {}).get("choices", [{}])[0].get("message", {}).get("content", "") - - # Parse mapping - import re - for m in re.finditer( - r"SPEAKER\s+(\d+)\s*:\s*(.+)", - reply, - re.IGNORECASE, - ): - spk = f"SPEAKER {m.group(1).strip()}" - name = m.group(2).strip().rstrip(".").upper() - if name: - speaker_map[spk] = name - - logger.info("Speaker identification mapping: %s", speaker_map) - - # Apply mapping to transcript text + speaker_map = _identify_speakers_via_llm(scraibe, transcript_text) if speaker_map: - def replace_speaker(m): - label = m.group(0).strip() - # normalize to "SPEAKER N" - normalized = re.sub( - r"\s+", - " ", - re.sub(r"[^A-Z0-9\s]", "", label.upper()), - ).strip() - return speaker_map.get(normalized, label) - - # Replace in lines like "[00:12] SPEAKER 1:" but preserve timestamp and colon - def replace_in_line(line: str) -> str: - # match after timestamp bracket and space: "SPEAKER N:" - return re.sub( - r"(\[\d+:\d+(?::\d+)?\]\s*)([A-Z\s]+?):\s*", - lambda m: m.group(1) + (speaker_map.get(m.group(2).strip(), m.group(2)) + ": "), - line, - ) - - transcript_lines = transcript_text.splitlines() - transcript_text = "\n".join( - replace_in_line(line) for line in transcript_lines + transcript_text, segments = _apply_speaker_map( + transcript_text, segments, speaker_map ) - - # Also update segments for JSON export - updated_segments = [] - for seg in segments: - sp = (seg.get("speaker") or "").strip() - sp_norm = re.sub(r"[^A-Z0-9\s]", "", sp.upper()).strip() - sp_new = speaker_map.get(sp_norm, sp) - seg = dict(seg) - seg["speaker"] = sp_new - updated_segments.append(seg) - segments = updated_segments - - except (SummarizerError, Exception) as e: + logger.info("Applied speaker identification to transcript and segments.") + except Exception as e: logger.warning( - "Speaker identification failed; falling back to Speaker IDs: %s", e + "Speaker identification failed; continuing with original labels: %s", e ) - speaker_map = {} - # 4) Prepare files + # 5) Summarization (if requested) using the (possibly identified) transcript + if include_summary: + try: + summary_text = scraibe._summarizer.summarize_transcript(transcript_text) + except Exception as e: + logger.warning("Summarization failed; continuing without summary: %s", e) + summary_text = "" + + # 6) Prepare files # Transcript .md md_transcript_path = _safe_filename("TRANSCRIPT", local, date_tag, ".md") @@ -417,7 +444,7 @@ def process_transcription_task( f.write(transcript_text) temp_files.append(md_transcript_path) - # Transcript .docx (with summary appended if transcript_and_summarize) + # Transcript .docx (with summary appended if present) docx_transcript_path = _safe_filename("TRANSCRIPT", local, date_tag, ".docx") if summary_text: create_transcript_docx( @@ -452,19 +479,14 @@ def process_transcription_task( json.dump(json_data, f, indent=2, ensure_ascii=False) temp_files.append(json_path) - # No separate summary DOCX/MD when using transcript_and_summarize - # (summary is now embedded in the transcript DOCX) - - # 5) Build attachments list - - # Always: JSON, transcript MD, transcript DOCX + # 7) Build attachments list attachments = [ md_transcript_path, docx_transcript_path, json_path, ] - # 6) Send success email + # 8) Send success email send_success_email( to=email_to, transcript_text=transcript_text, @@ -484,7 +506,7 @@ def process_transcription_task( ) raise e finally: - # 7) Cleanup + # 9) Cleanup for path in temp_files: _remove_file(path) if audio_path: