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, "apple_warned": False} log_path = getattr(args, "llm_request_log_path", "logs/llm_requests.jsonl") # --- Apple on-device LLM (ANE) – primary if available --- apple_llm = None apple_provider_enabled = getattr(args, "apple_llm", True) # on by default use_apple_llm = False if apple_provider_enabled and not is_openai: try: from engine_apple_llm import resolve_llm_binary, check_apple_llm_available, AppleLLM bp = resolve_llm_binary() if bp and check_apple_llm_available(bp): print(f"[LLM] Apple on-device LLM available at {bp} — using ANE path (primary), Ollama fallback") try: apple_llm = AppleLLM(bp, verbose=getattr(args, "verbose", False)) use_apple_llm = True print(f"[LLM] Apple LLM ready: {bp} (ANE-backed)") except Exception as e: print(f"[LLM] Apple LLM start failed ({e}) — falling back to Ollama") else: if getattr(args, "verbose", False): print(f"[LLM] Apple LLM not available (bin={bp}) — using Ollama") except Exception as e: if getattr(args, "verbose", False): print(f"[LLM] Apple LLM probe failed ({e})") 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 and not use_apple_llm: 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): # Apple LLM path (ANE) - try first if enabled if use_apple_llm and apple_llm is not None: # Determine mode from warn_key is_paragraph = warn_key == "paragraph_warned" # Extract from prompt/messages best-effort # For line: prompt is the substituted line, we need recent lines for prev1/prev2 # For paragraph: prompt built from paragraph_messages – we need previous_refined_context, previous_source, new_segments # We will keep the signature but also intercept via wrapper functions below that pass raw components via closure. # This generic call_llm fallback will still attempt Ollama if Apple fails inside specific wrapper. pass # actual Apple dispatch is in wrappers below, this is preserved for direct callers entry = { "timestamp": datetime.now(timezone.utc).isoformat(), "mode": "paragraph" if warn_key == "paragraph_warned" else "line", "provider": "openai" if is_openai else ("apple" if use_apple_llm else "ollama"), "model": args.post_correct_model if not use_apple_llm else "apple-fm", "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 # If we are in Apple mode, this generic path should not be hit for line/para polish # because post_correct_line and call_llm_rolling_refine now directly call Apple. # But keep Ollama fallback for any direct call_llm usage. 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): # Apple ANE path first if use_apple_llm and apple_llm is not None: t0 = time.time() try: polished, ok, apple_ms = apple_llm.polish_paragraph( new_text=new_segments_str, context=previous_refined_context, prev_source=previous_source_text, language=getattr(args, "lang", None), timeout=getattr(args, "post_correct_llm_timeout", 8.0), ) entry = { "timestamp": datetime.now(timezone.utc).isoformat(), "mode": "paragraph", "provider": "apple", "model": "apple-fm", "duration_ms": int((time.time()-t0)*1000), "apple_ms": apple_ms, "ok": ok, } append_llm_request_log(log_path, entry) if ok and polished: print(f"[LLM] Apple paragraph polish OK {apple_ms:.0f}ms -> {polished[:80]}") return polished else: print(f"[LLM] Apple paragraph fallback ok={ok} ms={apple_ms} err, trying Ollama") except Exception as e: print(f"[LLM] Apple paragraph error {e}, falling back to Ollama") if not llm_state["apple_warned"]: llm_state["apple_warned"] = True 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 # Try Apple ANE path first (fast, offline) if use_apple_llm and apple_llm is not None: try: t0 = time.time() prev1 = recent_lines[-1] if len(recent_lines) >= 1 else "" prev2 = recent_lines[-2] if len(recent_lines) >= 2 else "" polished, ok, apple_ms = apple_llm.polish_line( text=substituted if getattr(args, "freeflow_polish", True) else corrected, prev1=prev1, prev2=prev2, language=getattr(args, "lang", None), timeout=getattr(args, "post_correct_llm_timeout", 8.0), ) dur = int((time.time()-t0)*1000) # Log Apple call entry = { "timestamp": datetime.now(timezone.utc).isoformat(), "mode": "line", "provider": "apple", "model": "apple-fm", "duration_ms": dur, "apple_ms": apple_ms, "ok": ok, "input": substituted[:200], "response": {"content": polished}, } append_llm_request_log(log_path, entry) if ok and polished: # Apply same post-processing as Ollama path prompt_style = getattr(args, "post_correct_prompt_style", "freeflow") cand = polished if prompt_style != "legacy": fb = guard_against_truncation(cand, corrected) if fb: if getattr(args, "post_correct_debug", False): print(f"[LLM] Apple guard fallback triggered, keeping deterministic: {fb[:60]}") return fb cand = strip_keep_tags(cand) cand = normalize_formatting(cand) cand = match_input_casing(cand, substituted) else: cand = normalize_english_caption(cand) if cand and word_overlap_ratio(corrected, cand) >= getattr(args, "post_correct_min_overlap", 0.45): print(f"[LLM] Apple line OK {apple_ms:.0f}ms ({dur}ms total) -> {cand[:80]}") return cand else: if getattr(args, "verbose", False): print(f"[LLM] Apple line low overlap {word_overlap_ratio(corrected, cand):.2f}, keep deterministic") return corrected else: print(f"[LLM] Apple line not ok ok={ok} ms={apple_ms}, falling back to Ollama") except Exception as e: print(f"[LLM] Apple line error {e}, falling back to Ollama") if not llm_state["apple_warned"]: llm_state["apple_warned"] = True 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())