Refactor speaker ID into helpers; run before summarization; fix _chat_completion call

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
+        # 3) Transcription (always first, without summarization)
-        if task_type == "transcript_and_summarize":
+        result = scraibe.transcribe(
-            result = scraibe.transcript_and_summarize(
+            audio_file=audio_path,
-                audio_file=audio_path,
+            language=language or None,
-                language=language or None,
+            num_speakers=int(num_speakers) if num_speakers else None,
-                num_speakers=int(num_speakers) if num_speakers else None,
+            verbose=True,
-                verbose=True,
+            for_export=True,
-                for_export=True,
+        )
-            )
+        transcript_text = result.get("transcript", "")
-            transcript_text = result.get("transcript", "")
+        segments = result.get("segments", [])
-            summary_text = result.get("summary", "")
+        raw_result = result.get("raw_result")
-            segments = result.get("segments", [])
+        summary_text = ""
-            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")
 
-        # 3b) Optional speaker identification
+        # 4) Optional speaker identification (before summarization)
-        speaker_map = {}  # e.g. {"SPEAKER 1": "John", "SPEAKER 2": "Maria"}
         if identify_speakers:
             try:
-                # Use the same summarizer client as transcript_and_summarize
+                speaker_map = _identify_speakers_via_llm(scraibe, transcript_text)
-                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
                 if speaker_map:
-                    def replace_speaker(m):
+                    transcript_text, segments = _apply_speaker_map(
-                        label = m.group(0).strip()
+                        transcript_text, segments, speaker_map
-                        # 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
                     )
-
+                    logger.info("Applied speaker identification to transcript and segments.")
-                    # Also update segments for JSON export
+            except Exception as e:
-                    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.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
+        # 7) Build attachments list
-        # (summary is now embedded in the transcript DOCX)
-
-        # 5) Build attachments list
-
-        # Always: JSON, transcript MD, transcript DOCX
         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: