whisper-translation/engine_distribute.py

import requests
import time
import argparse
import multiprocessing
import sys
import os

INGEST_URL = "https://emiapi.reynafamily.com/live-captions/ingest"

def post_with_retries(payload, timeout, max_attempts=3):
    delay = 1
    for attempt in range(1, max_attempts + 1):
        try:
            response = requests.post(INGEST_URL, json=payload, timeout=timeout)
            if response.status_code == 200:
                return True
            print(f"[Distribute Error] {response.status_code} on attempt {attempt}/{max_attempts}.")
        except Exception as exc:
            print(f"[Distribute Error] {exc} on attempt {attempt}/{max_attempts}.")
        if attempt < max_attempts:
            time.sleep(delay)
            delay = min(delay * 2, 5)
    return False

def log_debug(message):
    try:
        with open("distribute_debug.log", "a") as f:
            f.write(f"{time.strftime('%H:%M:%S')} {message}\n")
            f.flush()
            os.fsync(f.fileno())
    except:
        pass

def run_distribution(in_queue, args):
    print("[Distribute] Starting distribution service...", flush=True)
    log_debug("Starting distribution service...")
    
    while True:
        try:
            payload = in_queue.get()
            if payload is None: break
            
            # Draft support
            if "draft" in payload:
                draft_text = payload["draft"]
                sys.stdout.write(f"\r\033[K[DRAFT]: {draft_text}")
                sys.stdout.flush()
                if args.ingest:
                    post_with_retries({"draft": draft_text}, timeout=2, max_attempts=2)
                continue

            # Finalized segment
            speaker_label = payload.get('speaker')
            speaker_prefix = f"[{speaker_label}]: " if speaker_label else ""
            
            msg = f"[Final]: {speaker_prefix}{payload.get('original', '')}"
            print(msg, flush=True)
            log_debug(msg)

            bridge_text = payload.get("en_bridge")
            original_text = payload.get("original", "")
            corrected_text = payload.get("corrected")
            paragraph_text = payload.get("paragraph")

            if bridge_text and bridge_text.strip() != original_text.strip():
                bridge_msg = f"[EN-BRIDGE]: {speaker_prefix}{bridge_text}"
                print(bridge_msg, flush=True)
                log_debug(bridge_msg)

            if payload.get("used_llm_line") and corrected_text and corrected_text.strip() != bridge_text.strip():
                llm_line_msg = f"[LLM-LINE]: {speaker_prefix}{corrected_text}"
                print(llm_line_msg, flush=True)
                log_debug(llm_line_msg)

            if payload.get("used_llm_paragraph") and paragraph_text:
                print("", flush=True)
                log_debug("")
                llm_paragraph_msg = f"[LLM-PARAGRAPH]: {speaker_prefix}{paragraph_text}"
                print(llm_paragraph_msg, flush=True)
                log_debug(llm_paragraph_msg)
                print("", flush=True)
                log_debug("")

            if payload.get("paragraph_fallback") and paragraph_text:
                paragraph_msg = f"[PARAGRAPH]: {speaker_prefix}{paragraph_text}"
                print(paragraph_msg, flush=True)
                log_debug(paragraph_msg)
            
            for lang in ["es", "fr", "ar", "en"]:
                if payload.get(lang):
                    lang_msg = f"[{lang.upper()}]: {speaker_prefix}{payload[lang]}"
                    print(lang_msg, flush=True)
                    log_debug(lang_msg)

            if args.ingest:
                # Flat JSON Schema: Keep only what's needed for the server
                ingest_payload = {
                    "original": payload.get("original"),
                    "speaker": payload.get("speaker"),
                    "ts": payload.get("ts")
                }
                english_output = payload.get("english_output") or payload.get("en")
                if english_output:
                    ingest_payload["en"] = english_output
                has_any_translation = False
                for lang in ["es", "fr", "ar"]:
                    if payload.get(lang):
                        ingest_payload[lang] = payload[lang]
                        has_any_translation = True
                if english_output:
                    has_any_translation = True
                
                # If only-translate-llm is on, ONLY send when we have a translation (paragraph-level)
                should_send = True
                if getattr(args, "only_translate_llm", False) and not has_any_translation:
                    should_send = False
                
                if should_send:
                    success = post_with_retries(ingest_payload, timeout=5, max_attempts=3)
                    if not success:
                        log_debug(f"Dropped ingest payload after retries: {ingest_payload}")
                
        except Exception as e:
            print(f"[Distribute] Error: {e}")

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