transcript.py

import json
from dataclasses import dataclass
from typing import Dict, List, Union

import requests
from bs4 import BeautifulSoup
from openai import OpenAI


@dataclass
class TranscriptSegment:
    speaker_id: str
    start_time: float
    end_time: float
    text: str
    speaker_name: str = ""


@dataclass
class AudioSegment:
    id: int
    transcript: str
    start_time: float
    end_time: float
    speaker_label: str
    original_file: str
    items: List[int]


class TranscriptProcessor:
    def __init__(
        self,
        transcript_file: str = None,
        transcript_data: Union[dict, list] = None,
        call_type: str = "le",
        person_names: list = None,
    ):
        self.transcript_file = transcript_file
        self.transcript_data = transcript_data
        self.formatted_transcript = None
        self.segments = []
        self.speaker_mapping = {}
        self.person_names = person_names

        if self.transcript_file:
            self._load_transcript()
        elif self.transcript_data:
            if call_type == "rp":
                self.merge_transcripts(transcript_data, person_names)
        else:
            raise ValueError(
                "Either transcript_file or transcript_data must be provided."
            )

        self._process_transcript()
        self._create_formatted_transcript()  # Create initial formatted transcript
        if call_type != "si" and call_type != "rp":
            self.map_speaker_ids_to_names()

    def _load_transcript(self) -> None:
        """Load the transcript JSON file."""
        with open(self.transcript_file, "r") as f:
            self.transcript_data = json.load(f)

    def _format_time(self, seconds: float) -> str:
        """Convert seconds to formatted time string (MM:SS)."""
        minutes = int(seconds // 60)
        seconds = int(seconds % 60)
        return f"{minutes:02d}:{seconds:02d}"

    def _process_transcript(self) -> None:
        results = self.transcript_data["results"]
        current_words = []
        current_speaker = None
        current_start = None
        current_items = []

        for item in results["items"]:
            if item["type"] == "pronunciation":
                if not self.person_names:
                    speaker = (
                        item.get("speaker_label", "")
                        .replace("spk_", "")
                        .replace("spk", "")
                    )
                else:
                    speaker = item.get("speaker_label", "")
                    print("ITEM", item)

                # Initialize on first pronunciation item
                if current_speaker is None:
                    current_speaker = speaker
                    current_start = float(item["start_time"])

                # Check for speaker change
                if speaker != current_speaker:
                    if current_items:
                        self._create_segment(
                            current_speaker,
                            current_start,
                            float(item["start_time"]),
                            current_items,
                        )
                    current_items = []
                    current_words = []
                    current_speaker = speaker
                    current_start = float(item["start_time"])

                current_items.append(item)
                current_words.append(item["alternatives"][0]["content"])
            elif item["type"] == "punctuation":
                current_items.append(item)
                # Only check for segment break if we're over 20 words
                if len(current_words) >= 20:
                    # Break on this punctuation
                    next_item = next(
                        (
                            it
                            for it in results["items"][
                                results["items"].index(item) + 1 :
                            ]
                            if it["type"] == "pronunciation"
                        ),
                        None,
                    )
                    if next_item:
                        self._create_segment(
                            current_speaker,
                            current_start,
                            float(next_item["start_time"]),
                            current_items,
                        )
                        current_items = []
                        current_words = []
                        current_start = float(next_item["start_time"])

        # Don't forget the last segment
        if current_items:
            last_time = max(
                float(item["end_time"])
                for item in current_items
                if item["type"] == "pronunciation"
            )
            self._create_segment(
                current_speaker, current_start, last_time, current_items
            )

    def _create_segment(
        self, speaker_id: str, start: float, end: float, items: list
    ) -> None:
        segment_content = []
        for item in items:
            if item["type"] == "pronunciation":
                segment_content.append(item["alternatives"][0]["content"])
            elif item["type"] == "punctuation":
                # Append punctuation to the last word without a space
                if segment_content:
                    segment_content[-1] += item["alternatives"][0]["content"]

        if segment_content:
            self.segments.append(
                TranscriptSegment(
                    speaker_id=speaker_id,
                    start_time=start,
                    end_time=end,
                    text=" ".join(segment_content),
                )
            )

    def correct_speaker_mapping_with_agenda(self, url: str) -> None:
        """Fetch agenda from a URL and correct the speaker mapping using OpenAI."""
        try:
            if not url.startswith("http"):
                # add https to the url
                url = "https://" + url

            response = requests.get(url)
            response.raise_for_status()
            html_content = response.text

            # Parse the HTML to find the desired description
            soup = BeautifulSoup(html_content, "html.parser")
            description_tag = soup.find(
                "script", {"type": "application/ld+json"}
            )  # Find the ld+json metadata block
            agenda = ""

            if description_tag:
                # Extract the JSON content
                json_data = json.loads(description_tag.string)
                if "description" in json_data:
                    agenda = json_data["description"]
                else:
                    print("Agenda description not found in the JSON metadata.")
            else:
                print("No structured data (ld+json) found.")

            if not agenda:
                print("No agenda found in the structured metadata. Trying meta tags.")

                # Fallback: Use meta description if ld+json doesn't have it
                meta_description = soup.find("meta", {"name": "description"})
                agenda = meta_description["content"] if meta_description else ""

            if not agenda:
                print("No agenda found in any description tags.")
                return

            prompt = (
                f"Given the original speaker mapping {self.speaker_mapping}, agenda:\n{agenda}, and the transcript: {self.formatted_transcript}\n\n"
                "Some speaker names in the mapping might have spelling errors or be incomplete."
                "Remember that the content in agenda is accurate and transcript can have errors so prioritize the spellings and names in the agenda content."
                "If the speaker name and introduction is similar to the agenda, update the speaker name in the mapping."
                "Please correct the names based on the agenda. Return the corrected mapping in JSON format as "
                "{'spk_0': 'Correct Name', 'spk_1': 'Correct Name', ...}."
                "You should only update the name if the name sounds very similar, or there is a good spelling overlap/ The Speaker Introduction matches the description of the Talk from Agends. If the name is totally unrelated, keep the original name."
                "You should always include all the speakers in the mapping from the original mapping, even if you don't update their names. i.e if there are 4 speakers in original mapping, new mapping should have 4 speakers always, ignore all the other spekaers in the agenda. I REPEAT DO NOT ADD OTHER NEW SPEAKERS IN THE MAPPING."
            )

            client = OpenAI()

            completion = client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {
                        "role": "system",
                        "content": "You are a helpful assistant. Who analyzes the given transcript, original speaker mapping and agenda. From the Agenda, you fix the spelling mistakes in the speaker names or update the names if they are similar to the agenda. You should only update the name if the name sounds very similar, or there is a good spelling overlap/ The Speaker Introduction matches the description of the Talk from Agends. If the name is totally unrelated, keep the original name.",
                    },
                    {"role": "user", "content": prompt},
                ],
                temperature=0,
            )

            response_text = completion.choices[0].message.content.strip()
            try:
                corrected_mapping = json.loads(response_text)
            except Exception:
                response_text = response_text[
                    response_text.find("{") : response_text.rfind("}") + 1
                ]
                try:
                    corrected_mapping = json.loads(response_text)
                except json.JSONDecodeError:
                    print(
                        "Error parsing corrected speaker mapping JSON, keeping the original mapping."
                    )
                    corrected_mapping = self.speaker_mapping
            # Update the speaker mapping with corrected names
            self.speaker_mapping = corrected_mapping

            # Update the transcript segments with corrected names
            for segment in self.segments:
                spk_id = f"spk_{segment.speaker_id}"
                segment.speaker_name = self.speaker_mapping.get(spk_id, spk_id)

            # Recreate the formatted transcript with corrected names
            formatted_segments = []
            for seg in self.segments:
                start_time_str = self._format_time(seg.start_time)
                end_time_str = self._format_time(seg.end_time)
                formatted_segments.append(
                    f"time_stamp: {start_time_str}-{end_time_str}\n"
                    f"{seg.speaker_name}: {seg.text}\n"
                )
            self.formatted_transcript = "\n".join(formatted_segments)

        except requests.exceptions.RequestException as e:
            print(f"  ching agenda from URL: {str(e)}")
        except Exception as e:
            print(f"Error correcting speaker mapping: {str(e)}")

    def _create_formatted_transcript(self) -> None:
        """Create formatted transcript with default speaker labels."""
        formatted_segments = []
        for seg in self.segments:
            start_time_str = self._format_time(seg.start_time)
            end_time_str = self._format_time(seg.end_time)
            # Use default speaker label (spk_X) if no mapping exists
            if not self.person_names:
                speaker_label = f"spk_{seg.speaker_id}"
            else:
                speaker_label = f"{seg.speaker_id}"
            formatted_segments.append(
                f"time_stamp: {start_time_str}-{end_time_str}\n"
                f"{speaker_label}: {seg.text}\n"
            )
        self.formatted_transcript = "\n".join(formatted_segments)

    def map_speaker_ids_to_names(self) -> None:
        """Map speaker IDs to names based on introductions in the transcript."""
        try:
            transcript = self.formatted_transcript

            prompt = (
                "Given the following transcript where speakers are identified as spk 0, spk 1, spk 2, etc., please map each spk ID to the speaker's name based on their introduction in the transcript. If no name is introduced for a speaker, keep it as spk_id. Return the mapping as a JSON object in the format {'spk_0': 'Speaker Name', 'spk_1': 'Speaker Name', ...}\n\n"
                f"Transcript:\n{transcript}"
            )

            client = OpenAI()

            completion = client.chat.completions.create(
                model="gpt-4o",
                messages=[
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": prompt},
                ],
                temperature=0,
            )

            response_text = completion.choices[0].message.content.strip()
            try:
                self.speaker_mapping = json.loads(response_text)
            except json.JSONDecodeError:
                response_text = response_text[
                    response_text.find("{") : response_text.rfind("}") + 1
                ]
                try:
                    self.speaker_mapping = json.loads(response_text)
                except json.JSONDecodeError:
                    print("Error parsing speaker mapping JSON.")
                    self.speaker_mapping = {}

            # Update segments with speaker names and recreate formatted transcript
            for segment in self.segments:
                spk_id = f"spk_{segment.speaker_id}"
                speaker_name = self.speaker_mapping.get(spk_id, spk_id)
                segment.speaker_name = speaker_name

            self._create_formatted_transcript_with_names()

        except Exception as e:
            print(f"Error mapping speaker IDs to names: {str(e)}")
            self.speaker_mapping = {}

    def _create_formatted_transcript_with_names(self) -> None:
        """Create formatted transcript with mapped speaker names."""
        formatted_segments = []
        for seg in self.segments:
            start_time_str = self._format_time(seg.start_time)
            end_time_str = self._format_time(seg.end_time)
            speaker_name = getattr(seg, "speaker_name", f"spk_{seg.speaker_id}")
            formatted_segments.append(
                f"Start Time: {start_time_str} - End Time: {end_time_str}\n"
                f"{speaker_name}: {seg.text}\n"
            )
        self.formatted_transcript = "\n".join(formatted_segments)

    def get_transcript(self) -> str:
        """Return the formatted transcript with speaker names."""
        return self.formatted_transcript

    def get_transcript_data(self) -> Dict:
        """Return the raw transcript data."""
        return self.transcript_data

    def merge_transcripts(
        self, transcript_files: List[Dict], person_names: List[str]
    ) -> None:
        """
        Merge multiple AWS diarized transcripts while maintaining correct time ordering.
        Each transcript is assumed to have one speaker (spk_0) and person_names list index
        corresponds to transcript file index.
        """
        print(person_names)
        if len(transcript_files) != len(person_names):
            raise ValueError("Number of transcripts must match number of speaker names")

        # Initialize merged structure
        merged_transcript = {
            "jobName": "merged_transcript",
            "status": "COMPLETED",
            "results": {
                "audio_segments": [],
                "items": [],
                "speaker_labels": {"segments": [], "speakers": len(transcript_files)},
            },
        }

        # First collect all items with their original data and file index
        all_items = []
        for file_idx, transcript in enumerate(transcript_files):
            items = transcript["results"].get("items", [])
            speaker_name = person_names[file_idx]

            for item in items:
                # Store original item data along with file index and original ID
                item_data = dict(item)
                # if "speaker_label" in item_data:
                item_data["speaker_label"] = speaker_name
                item_data["file_idx"] = file_idx
                item_data["original_id"] = item["id"]
                item_data["start_time"] = float(item.get("start_time", 0))
                item_data["end_time"] = float(item.get("end_time", 0))
                all_items.append(item_data)

        # Sort items by start time
        all_items.sort(key=lambda x: (x["start_time"], x["end_time"]))

        # Create mapping from (file_idx, original_id) to new sequential ID
        item_id_mapping = {}

        # Assign new sequential IDs and add to merged transcript
        for new_id, item in enumerate(all_items):
            file_idx = item.pop("file_idx")
            original_id = item.pop("original_id")
            item_id_mapping[(file_idx, original_id)] = new_id

            # Update item ID and convert times back to strings
            item["id"] = new_id
            item["start_time"] = str(item["start_time"])
            item["end_time"] = str(item["end_time"])

            merged_transcript["results"]["items"].append(item)

        # Process audio segments
        all_segments = []
        for file_idx, transcript in enumerate(transcript_files):
            file_segments = transcript["results"].get("audio_segments", [])
            speaker_name = person_names[file_idx]

            for segment in file_segments:
                # Map original item IDs to new sequential IDs
                new_items = [
                    item_id_mapping[(file_idx, item_id)]
                    for item_id in segment.get("items", [])
                ]

                all_segments.append(
                    AudioSegment(
                        id=len(all_segments),
                        transcript=segment["transcript"],
                        start_time=float(segment["start_time"]),
                        end_time=float(segment["end_time"]),
                        speaker_label=speaker_name,
                        original_file=f"file_{file_idx}",
                        items=new_items,
                    )
                )

        # Sort segments by start time
        sorted_segments = sorted(all_segments, key=lambda x: x.start_time)

        # Convert segments back to dictionary format
        for idx, segment in enumerate(sorted_segments):
            merged_segment = {
                "id": idx,
                "transcript": segment.transcript,
                "start_time": str(segment.start_time),
                "end_time": str(segment.end_time),
                "speaker_label": segment.speaker_label,
                "source_file": segment.original_file,
                "items": sorted(segment.items),
            }
            merged_transcript["results"]["audio_segments"].append(merged_segment)

            # Add to speaker labels segments
            speaker_segment = {
                "start_time": str(segment.start_time),
                "end_time": str(segment.end_time),
                "speaker_label": segment.speaker_label,
            }
            merged_transcript["results"]["speaker_labels"]["segments"].append(
                speaker_segment
            )

        # Update the instance transcript data with merged result
        self.transcript_data = merged_transcript