momentry_core/scripts/redis_publisher.py

#!/opt/homebrew/bin/python3.11
"""
Redis Progress Publisher
Common module for publishing progress to Redis

Usage:
    from redis_publisher import RedisPublisher

    pub = RedisPublisher("video-uuid-123")
    pub.info("asr", "Starting ASR processing")
    pub.progress("asr", current=50, total=100, message="Processing segment")
    pub.complete("asr", "Transcription complete")
"""

import os
import json
import time
import redis
from typing import Optional, Any, Dict
from dataclasses import dataclass, asdict
from enum import Enum


class MessageType(Enum):
    INFO = "info"
    PROGRESS = "progress"
    COMPLETE = "complete"
    ERROR = "error"
    WARNING = "warning"


@dataclass
class ProgressData:
    message: Optional[str] = None
    current: Optional[int] = None
    total: Optional[int] = None
    output_count: Optional[int] = None
    output_type: Optional[str] = None
    extra: Optional[Dict[str, Any]] = None


@dataclass
class StructuredMessage:
    type: str
    processor: str
    uuid: str
    timestamp: int
    data: ProgressData


class RedisPublisher:
    def __init__(self, uuid: str):
        self.uuid = uuid
        prefix = os.environ.get("MOMENTRY_REDIS_PREFIX", "momentry:")
        self.channel = f"{prefix}progress:{uuid}"
        self._enabled = False
        self._client = None
        self._connect()

    def _connect(self) -> None:
        redis_url = os.environ.get("REDIS_URL")
        if not redis_url:
            password = os.environ.get("REDIS_PASSWORD", "accusys")
            redis_url = f"redis://:{password}@localhost:6379"

        try:
            self._client = redis.from_url(redis_url, decode_responses=True)
            self._client.ping()
            self._enabled = True
        except redis.ConnectionError as e:
            import sys

            print(f"[RedisPublisher] Connection failed: {e}", file=sys.stderr)
        except Exception as e:
            import sys

            print(f"[RedisPublisher] Redis not available: {e}", file=sys.stderr)

    @property
    def enabled(self) -> bool:
        return self._enabled

    def _publish_json(self, msg: StructuredMessage) -> bool:
        if not self._enabled or self._client is None:
            return False

        try:
            client: redis.Redis = self._client
            client.publish(self.channel, json.dumps(asdict(msg)))
            return True
        except Exception as e:
            import sys

            print(f"[RedisPublisher] Publish error: {e}", file=sys.stderr)
            return False

        try:
            self._client.publish(self.channel, json.dumps(asdict(msg)))
            return True
        except Exception as e:
            import sys

            print(f"[RedisPublisher] Publish error: {e}", file=sys.stderr)
            return False

    def publish(
        self,
        msg_type: MessageType,
        processor: str,
        message: Optional[str] = None,
        current: Optional[int] = None,
        total: Optional[int] = None,
        output_count: Optional[int] = None,
        output_type: Optional[str] = None,
        extra: Optional[Dict[str, Any]] = None,
    ) -> bool:
        if not self._enabled:
            return False

        msg = StructuredMessage(
            type=msg_type.value,
            processor=processor,
            uuid=self.uuid,
            timestamp=int(time.time()),
            data=ProgressData(
                message=message,
                current=current,
                total=total,
                output_count=output_count,
                output_type=output_type,
                extra=extra,
            ),
        )

        return self._publish_json(msg)

    def info(self, processor: str, message: str) -> bool:
        return self.publish(MessageType.INFO, processor, message=message)

    def progress(
        self,
        processor: str,
        current: int,
        total: int,
        message: str = "",
        output_count: Optional[int] = None,
        output_type: Optional[str] = None,
    ) -> bool:
        return self.publish(
            MessageType.PROGRESS,
            processor,
            message=message,
            current=current,
            total=total,
            output_count=output_count,
            output_type=output_type,
        )

    def complete(self, processor: str, message: str = "") -> bool:
        return self.publish(MessageType.COMPLETE, processor, message=message)

    def error(self, processor: str, message: str) -> bool:
        return self.publish(MessageType.ERROR, processor, message=message)

    def warning(self, processor: str, message: str) -> bool:
        return self.publish(MessageType.WARNING, processor, message=message)

    def percentage(self, processor: str, percent: float, message: str = "") -> bool:
        return self.publish(
            MessageType.PROGRESS,
            processor,
            message=message,
            current=int(percent),
            total=100,
            extra={"percentage": percent},
        )


class ProgressContext:
    """Context manager for tracking processor progress"""

    def __init__(self, publisher: RedisPublisher, processor: str):
        self.publisher = publisher
        self.processor = processor

    def __enter__(self):
        self.publisher.info(self.processor, f"{self.processor.upper()} started")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            self.publisher.error(self.processor, str(exc_val))
        else:
            self.publisher.complete(self.processor)
        return False