whisper-translation/engine_llm.py

import sys
from unittest.mock import MagicMock

# MANDATORY: Mock lzma BEFORE any other imports
try:
    import lzma
except ImportError:
    mock_lzma = MagicMock()
    mock_lzma.FORMAT_XZ, mock_lzma.FORMAT_ALONE, mock_lzma.FORMAT_RAW = 1, 2, 3
    mock_lzma.CHECK_NONE, mock_lzma.CHECK_CRC32, mock_lzma.CHECK_CRC64, mock_lzma.CHECK_SHA256 = 0, 1, 4, 10
    sys.modules["_lzma"] = MagicMock()
    sys.modules["lzma"] = mock_lzma

import time
import argparse
import multiprocessing
import os
import re
import json
import requests
import queue
from collections import deque
from datetime import datetime, timezone
from freeflow_polish import (
    build_user_prompt,
    deterministic_polish,
    guard_against_truncation,
    is_clean_enough_to_skip_llm,
    match_input_casing,
    normalize_formatting,
    strip_keep_tags,
    system_prompt_for_language,
)


def build_paragraph_messages(previous_refined_context, previous_source_text, new_source_text):
    system_prompt = """You are a careful live transcript editor working from noisy translated source text.

Goal:
Produce the most faithful readable English update.

Decision rules:
1. NEW SOURCE TEXT is the primary evidence and should dominate the output.
2. Use PREVIOUS SOURCE TEXT and PREVIOUS REFINED CONTEXT only when they clearly help resolve or continue the new material.
3. The final answer must be English only.
4. Never copy non-English words into the output unless they are proper names.
5. If a non-English fragment appears as an isolated clause, trailing fragment, or side comment without a clear English anchor, drop it.
6. Only translate a non-English fragment when it is clearly central to the same idea and its meaning is reasonably obvious from context.
7. If the new material starts a fresh thought, output only the new material.
8. Remove exact or near-exact repetition unless the repetition is clearly intentional rhetoric.
9. Do not add explanations, disclaimers, or meta commentary.
10. Prefer conservative wording over guessed meaning.
11. Return only the revised transcript text."""

    prompt = f"""Edit this transcript update.

[PREVIOUS REFINED CONTEXT]
{previous_refined_context}

[PREVIOUS SOURCE TEXT]
{previous_source_text}

[NEW SOURCE TEXT]
{new_source_text}
"""
    return [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": prompt},
    ], prompt


def build_line_messages(prev1, prev2, corrected):
    system_prompt = """You are a real-time English caption corrector for live speech.

Task:
Clean only the current caption line.

Hard rules:
1. Preserve meaning exactly. Never replace current content with prior context.
2. Remove disfluencies and false starts.
3. Fix punctuation, casing, and obvious STT typos.
4. If uncertain, return the original line unchanged.
5. Output only one corrected English line."""

    prompt = (
        "Context (reference only):\n"
        f"Previous line 1: {prev1}\n"
        f"Previous line 2: {prev2}\n"
        "Current line to correct:\n"
        f"{corrected}\n"
        "Corrected:"
    )
    return [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": prompt},
    ], prompt


def build_freeflow_line_messages(text, language=None, model=""):
    system_prompt = system_prompt_for_language(language, model=model)
    prompt = build_user_prompt(text, language=language)
    return [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": prompt},
    ], prompt


def append_llm_request_log(log_path, entry):
    if not log_path:
        return
    log_dir = os.path.dirname(log_path)
    if log_dir:
        os.makedirs(log_dir, exist_ok=True)
    with open(log_path, "a", encoding="utf-8") as handle:
        json.dump(entry, handle, ensure_ascii=False)
        handle.write("\n")


def load_llm_request_log_entry(log_path, entry_index):
    if not log_path or not os.path.exists(log_path):
        raise FileNotFoundError(f"LLM request log not found: {log_path}")

    with open(log_path, "r", encoding="utf-8") as handle:
        entries = [json.loads(line) for line in handle if line.strip()]

    if not entries:
        raise ValueError(f"LLM request log is empty: {log_path}")

    if entry_index is None or entry_index == -1:
        return entries[-1]

    if entry_index < 0:
        entry_index = len(entries) + entry_index

    if entry_index < 0 or entry_index >= len(entries):
        raise IndexError(f"LLM request log index {entry_index} is out of range for {len(entries)} entries")

    return entries[entry_index]


def run_llm_prompt_test(args):
    from dotenv import load_dotenv
    load_dotenv()

    api_key = os.environ.get("OPENAI_API_KEY")
    log_path = getattr(args, "llm_request_log_path", "logs/llm_requests.jsonl")

    def call_openai(model, messages):
        response = requests.post(
            "https://api.openai.com/v1/chat/completions",
            json={"model": model, "messages": messages},
            headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
            timeout=getattr(args, "post_correct_llm_timeout", 8.0),
        )
        if response.status_code >= 400:
            try:
                payload = response.json()
                detail = json.dumps(payload, ensure_ascii=True)
            except Exception:
                detail = response.text.strip()
            raise RuntimeError(f"OpenAI API error {response.status_code}: {detail}")
        return response.json()["choices"][0]["message"]["content"].strip()

    def call_ollama(model, prompt, temperature, system=None, extra_options=None):
        options = {"temperature": temperature}
        if extra_options:
            options.update(extra_options)
        response = requests.post(
            getattr(args, "post_correct_ollama_url", "http://127.0.0.1:11434/api/generate"),
            json={
                "model": model,
                "prompt": prompt,
                "system": system or "",
                "stream": False,
                "options": options,
                "keep_alive": getattr(args, "post_correct_keep_alive", "30m"),
            },
            timeout=getattr(args, "post_correct_llm_timeout", 8.0),
        )
        response.raise_for_status()
        return response.json().get("response", "").strip()

    def send_test_request(mode, messages, prompt, temperature, metadata, model=None, extra_options=None):
        request_model = model or args.post_correct_model
        is_openai = request_model.startswith("gpt-")
        entry = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "mode": mode,
            "provider": "openai" if is_openai else "ollama",
            "model": request_model,
            "temperature": temperature,
            "options": extra_options or {},
            "messages": messages,
            "metadata": metadata,
        }
        started_at = time.time()
        if is_openai:
            if not api_key:
                raise RuntimeError("OPENAI_API_KEY is not set")
            try:
                response_text = call_openai(request_model, messages)
                entry["response"] = {"status": "ok", "content": response_text}
                return response_text
            except Exception as exc:
                entry["response"] = {"status": "error", "error": str(exc)}
                raise
            finally:
                entry["duration_ms"] = round((time.time() - started_at) * 1000, 2)
                append_llm_request_log(log_path, entry)
        system = ""
        for message in messages:
            if message.get("role") == "system":
                system = message.get("content", "")
                break
        try:
            response_text = call_ollama(request_model, prompt, temperature, system=system, extra_options=extra_options)
            entry["response"] = {"status": "ok", "content": response_text}
            return response_text
        except Exception as exc:
            entry["response"] = {"status": "error", "error": str(exc)}
            raise
        finally:
            entry["duration_ms"] = round((time.time() - started_at) * 1000, 2)
            append_llm_request_log(log_path, entry)

    if getattr(args, "llm_test_from_log", None) is not None:
        logged_entry = load_llm_request_log_entry(log_path, args.llm_test_from_log)
        logged_messages = logged_entry.get("messages")
        if not logged_messages:
            raise ValueError("Selected log entry does not contain messages")
        logged_prompt = next((m.get("content", "") for m in logged_messages if m.get("role") == "user"), "")
        logged_model = logged_entry.get("model", args.post_correct_model)
        replay_model = logged_model if getattr(args, "llm_test_use_logged_model", False) else args.post_correct_model
        response = send_test_request(
            logged_entry.get("mode", "replay"),
            logged_messages,
            logged_prompt,
            logged_entry.get("temperature", 0.1),
            {
                "replay_source_log_path": log_path,
                "replay_source_index": args.llm_test_from_log,
                "replay_source_timestamp": logged_entry.get("timestamp"),
                "replay_source_model": logged_model,
            },
            model=replay_model,
        )
        print("[LLM TEST] Replayed request from log:")
        print(f"source_index={args.llm_test_from_log} source_model={logged_model} replay_model={replay_model}")
        print(response)

    if getattr(args, "llm_test_line", None):
        use_freeflow = getattr(args, "post_correct_prompt_style", "freeflow") != "legacy"
        if use_freeflow:
            preprocessed, substituted = deterministic_polish(args.llm_test_line)
            messages, prompt = build_freeflow_line_messages(
                substituted,
                language=getattr(args, "lang", None),
                model=getattr(args, "post_correct_model", ""),
            )
            metadata_input = preprocessed
        else:
            messages, prompt = build_line_messages(
                getattr(args, "llm_test_prev1", ""),
                getattr(args, "llm_test_prev2", ""),
                args.llm_test_line,
            )
            metadata_input = args.llm_test_line
        response = send_test_request(
            "line",
            messages,
            prompt,
            0.1,
            {
                "prev1": getattr(args, "llm_test_prev1", ""),
                "prev2": getattr(args, "llm_test_prev2", ""),
                "input": metadata_input,
                "prompt_style": getattr(args, "post_correct_prompt_style", "freeflow"),
            },
        )
        print("[LLM TEST] Line response:")
        print(response)

    if getattr(args, "llm_test_segments", None):
        messages, prompt = build_paragraph_messages(
            getattr(args, "llm_test_context", ""),
            "",
            args.llm_test_segments,
        )
        response = send_test_request(
            "paragraph",
            messages,
            prompt,
            0.1,
            {
                "context": getattr(args, "llm_test_context", ""),
                "segments": args.llm_test_segments,
            },
        )
        print("[LLM TEST] Paragraph response:")
        print(response)

def run_llm_processor(in_queue, out_queue, args):
    from dotenv import load_dotenv
    load_dotenv()
    
    print("[LLM] Starting Rolling Paragraph Engine...")
    
    api_key = os.environ.get("OPENAI_API_KEY")
    is_openai = args.post_correct_model.startswith("gpt-")
    llm_state = {"line_warned": False, "paragraph_warned": False}
    log_path = getattr(args, "llm_request_log_path", "logs/llm_requests.jsonl")

    def check_ollama_health():
        try:
            response = requests.get("http://127.0.0.1:11434/api/tags", timeout=2.5)
            response.raise_for_status()
            data = response.json()
            models = data.get("models", [])
            available = {model.get("name") for model in models if model.get("name")}
            if args.post_correct_model not in available and f"{args.post_correct_model}:latest" not in available:
                print(f"[LLM] Configured Ollama model '{args.post_correct_model}' is not installed.")
            return True
        except Exception as exc:
            print(f"[LLM] Ollama health check failed ({exc}). Local LLM features may fall back.")
            return False

    def warmup_ollama_model():
        try:
            response = requests.post(
                getattr(args, "post_correct_ollama_url", "http://127.0.0.1:11434/api/generate"),
                json={
                    "model": args.post_correct_model,
                    "prompt": "Return exactly: ok",
                    "system": "You are a concise assistant.",
                    "stream": False,
                    "options": {"temperature": 0.0},
                    "keep_alive": getattr(args, "post_correct_keep_alive", "30m"),
                },
                timeout=getattr(args, "post_correct_warmup_timeout", 20.0),
            )
            response.raise_for_status()
            body = response.json()
            if body.get("response", "").strip():
                total_s = body.get("total_duration", 0) / 1_000_000_000
                load_s = body.get("load_duration", 0) / 1_000_000_000
                print(f"[LLM] Ollama warmup OK for {args.post_correct_model} (load={load_s:.2f}s total={total_s:.2f}s).")
                return True
        except Exception as exc:
            print(f"[LLM] Ollama warmup failed for {args.post_correct_model} ({exc}).")
        return False

    if not is_openai:
        check_ollama_health()
        if getattr(args, "post_correct_llm", False) or getattr(args, "llm_paragraph", False):
            warmup_ollama_model()

    def normalize_english_caption(text):
        normalized = " ".join((text or "").split()).strip()
        if normalized and normalized[0].isalpha():
            normalized = normalized[0].upper() + normalized[1:]
        return normalized

    def apply_rule_based_post_correction(text):
        if getattr(args, "freeflow_polish", True):
            corrected, _ = deterministic_polish(text)
            return corrected
        corrected = normalize_english_caption(text)
        corrected = re.sub(r"\s+", " ", corrected).strip()
        corrected = re.sub(r"\s+([,.;!?])", r"\1", corrected)
        corrected = re.sub(r"([,.;!?]){2,}", r"\1", corrected)
        corrected = re.sub(r"\b(uh+|um+|erm+|ah+|hmm+)\b", "", corrected, flags=re.IGNORECASE)
        corrected = re.sub(r"\s+", " ", corrected).strip()
        corrected = re.sub(r"\b(\w+)\s+\1\b", r"\1", corrected, flags=re.IGNORECASE)
        return normalize_english_caption(corrected)

    def word_overlap_ratio(source_text, candidate_text):
        src_tokens = set(re.findall(r"[a-z0-9']+", source_text.lower()))
        cand_tokens = set(re.findall(r"[a-z0-9']+", candidate_text.lower()))
        if not src_tokens:
            return 1.0
        return len(src_tokens & cand_tokens) / len(src_tokens)

    def call_openai(messages):
        response = requests.post(
            "https://api.openai.com/v1/chat/completions",
            json={"model": args.post_correct_model, "messages": messages},
            headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
            timeout=getattr(args, "post_correct_llm_timeout", 8.0),
        )
        if response.status_code >= 400:
            try:
                payload = response.json()
                detail = json.dumps(payload, ensure_ascii=True)
            except Exception:
                detail = response.text.strip()
            raise RuntimeError(f"OpenAI API error {response.status_code}: {detail}")
        return response.json()["choices"][0]["message"]["content"].strip()

    def call_ollama(prompt, temperature, system=None, extra_options=None):
        options = {"temperature": temperature}
        if extra_options:
            options.update(extra_options)
        response = requests.post(
            getattr(args, "post_correct_ollama_url", "http://127.0.0.1:11434/api/generate"),
            json={
                "model": args.post_correct_model,
                "prompt": prompt,
                "system": system or "",
                "stream": False,
                "options": options,
                "keep_alive": getattr(args, "post_correct_keep_alive", "30m"),
            },
            timeout=getattr(args, "post_correct_llm_timeout", 8.0),
        )
        response.raise_for_status()
        return response.json().get("response", "").strip()

    def call_llm(messages, prompt, temperature, warn_key, extra_options=None):
        entry = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "mode": "paragraph" if warn_key == "paragraph_warned" else "line",
            "provider": "openai" if is_openai else "ollama",
            "model": args.post_correct_model,
            "temperature": temperature,
            "options": extra_options or {},
            "messages": messages,
        }
        started_at = time.time()
        try:
            if is_openai:
                if not api_key:
                    raise RuntimeError("OPENAI_API_KEY is not set")
                response_text = call_openai(messages)
                entry["response"] = {"status": "ok", "content": response_text}
                return response_text
            system = ""
            if messages:
                for message in messages:
                    if message.get("role") == "system":
                        system = message.get("content", "")
                        break
            response_text = call_ollama(prompt, temperature, system=system, extra_options=extra_options)
            entry["response"] = {"status": "ok", "content": response_text}
            return response_text
        except Exception as exc:
            entry["response"] = {"status": "error", "error": str(exc)}
            if not llm_state[warn_key]:
                label = "Paragraph LLM" if warn_key == "paragraph_warned" else "Line LLM"
                print(f"[LLM] {label} unavailable ({exc}). Falling back to deterministic text.")
                llm_state[warn_key] = True
            return ""
        finally:
            entry["duration_ms"] = round((time.time() - started_at) * 1000, 2)
            append_llm_request_log(log_path, entry)

    # State: Keep a short rolling window of refined paragraphs as context.
    refined_context_history = deque(maxlen=2)
    previous_source_window = ""
    recent_lines = deque(maxlen=2)
    
    def call_llm_rolling_refine(new_segments_str, previous_refined_context, previous_source_text):
        messages, prompt = build_paragraph_messages(previous_refined_context, previous_source_text, new_segments_str)
        paragraph_options = {
            "top_p": getattr(args, "llm_paragraph_top_p", 0.8),
            "repeat_penalty": getattr(args, "llm_paragraph_repeat_penalty", 1.0),
        }
        return call_llm(
            messages,
            prompt,
            getattr(args, "llm_paragraph_temperature", 0.2),
            "paragraph_warned",
            extra_options=paragraph_options,
        )

    def post_correct_line(text):
        if not getattr(args, "post_correct", False):
            return normalize_english_caption(text)

        corrected, substituted = deterministic_polish(text) if getattr(args, "freeflow_polish", True) else (apply_rule_based_post_correction(text), text)
        if not getattr(args, "post_correct_llm", False):
            return corrected
        if getattr(args, "post_correct_skip_clean", True) and is_clean_enough_to_skip_llm(corrected):
            if getattr(args, "post_correct_debug", False):
                print("[LLM] Skipping line LLM polish; deterministic output looks clean.")
            return corrected

        prompt_style = getattr(args, "post_correct_prompt_style", "freeflow")
        if prompt_style == "legacy":
            prev2 = recent_lines[-2] if len(recent_lines) >= 2 else ""
            prev1 = recent_lines[-1] if len(recent_lines) >= 1 else ""
            messages, prompt = build_line_messages(prev1, prev2, corrected)
        else:
            language = getattr(args, "lang", None)
            messages, prompt = build_freeflow_line_messages(
                substituted,
                language=language,
                model=args.post_correct_model if prompt_style == "qwen" else "",
            )
        candidate = call_llm(messages, prompt, 0.1, "line_warned").strip()
        if prompt_style != "legacy":
            fallback = guard_against_truncation(candidate, corrected)
            if fallback:
                return fallback
            candidate = strip_keep_tags(candidate)
            candidate = normalize_formatting(candidate)
            candidate = match_input_casing(candidate, substituted)
        else:
            candidate = normalize_english_caption(candidate)
        if candidate and word_overlap_ratio(corrected, candidate) >= getattr(args, "post_correct_min_overlap", 0.45):
            return candidate
        return corrected

    pending_buffer = []

    def flush_pending_outputs(processed_items):
        nonlocal pending_buffer, previous_source_window
        if not processed_items:
            return

        structured_paragraph = None
        paragraph_from_llm = False
        if args.llm_paragraph and pending_buffer:
            new_batch = " ".join(pending_buffer)
            previous_refined_context = "\n\n".join(refined_context_history)
            result = call_llm_rolling_refine(new_batch, previous_refined_context, previous_source_window)

            if result:
                if "\n\n" in result:
                    parts = result.split("\n\n")
                    refined_context_history.clear()
                    for part in parts:
                        cleaned_part = part.strip()
                        if cleaned_part:
                            refined_context_history.append(cleaned_part)
                    structured_paragraph = result
                else:
                    cleaned_result = result.strip()
                    if cleaned_result:
                        refined_context_history.append(cleaned_result)
                    structured_paragraph = result
                paragraph_from_llm = True
            else:
                structured_paragraph = new_batch
                cleaned_batch = new_batch.strip()
                if cleaned_batch:
                    refined_context_history.append(cleaned_batch)

            previous_source_window = new_batch

        last_index = len(processed_items) - 1
        for index, processed_item in enumerate(processed_items):
            out_queue.put({
                "raw": processed_item["raw"],
                "corrected": processed_item["corrected"],
                "paragraph": structured_paragraph if index == last_index else None,
                "en_bridge": processed_item["en_bridge"],
                "used_llm_line": bool(getattr(args, "post_correct_llm", False)),
                "used_llm_paragraph": paragraph_from_llm if index == last_index else False,
                "detected_lang": processed_item.get("detected_lang"),
                "speaker": processed_item.get("speaker"),
                "ts": processed_item.get("ts", time.time())
            })

        pending_buffer = []
    
    while True:
        try:
            item = in_queue.get()
            if item is None: break
            if "draft" in item:
                out_queue.put(item)
                continue
            processed_items = []
            pending_item = item

            while pending_item is not None:
                raw_text = pending_item.get("original", "").strip()
                bridge_text = pending_item.get("en_bridge", raw_text).strip()
                if raw_text:
                    corrected_text = post_correct_line(bridge_text)

                    # Hallucination check
                    words = corrected_text.lower().split()
                    if not (len(words) > 10 and len(set(words)) < 3):
                        pending_buffer.append(corrected_text)
                        recent_lines.append(corrected_text)
                        processed_items.append({
                            "raw": raw_text,
                            "corrected": corrected_text,
                            "en_bridge": bridge_text,
                            "detected_lang": pending_item.get("detected_lang"),
                            "speaker": pending_item.get("speaker"),
                            "ts": pending_item.get("ts", time.time()),
                        })

                try:
                    pending_item = in_queue.get_nowait()
                    if pending_item is None:
                        flush_pending_outputs(processed_items)
                        return
                    if "draft" in pending_item:
                        out_queue.put(pending_item)
                        pending_item = None
                except queue.Empty:
                    pending_item = None

            flush_pending_outputs(processed_items)
            
        except Exception as e:
            print(f"[LLM] Error: {e}")

if __name__ == "__main__":
    run_llm_processor(multiprocessing.Queue(), multiprocessing.Queue(), argparse.Namespace())