sidecar/stream_monitor.py

"""
sidecar/stream_monitor.py — Async sentence-level post-inference monitor.

Wraps the synchronous TextMonitor in an async ThreadPoolExecutor so that
per-sentence analysis runs concurrently in the background while the LLM
stream continues flowing to the client.

When a sentence crosses the threat threshold, a BlockSignal is returned
so the SSE layer can push a block event to the client immediately.
"""

import asyncio
import logging
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from typing import List, Optional

log = logging.getLogger(__name__)


# ── Event types ───────────────────────────────────────────────────────────────

@dataclass
class BlockSignal:
    """Emitted when the post-monitor flags a sentence as harmful."""
    sentence_id:  int
    reason:       str
    threat_score: int
    flags:        List[str]
    type:         str = field(default="block_signal", init=False)


# ── Stream Monitor ─────────────────────────────────────────────────────────────

class StreamMonitor:
    """
    Async wrapper around TextMonitor for sentence-level background screening.

    Usage
    -----
    monitor = StreamMonitor(text_monitor_instance, block_threshold=40)

    # Submit sentences as they complete (non-blocking):
    task = await monitor.submit(sentence_id=1, sentence="...", prompt="...")

    # At stream end — collect any pending block signals:
    signals = await monitor.collect(timeout=1.0)

    Args:
        text_monitor:    Existing TextMonitor instance.
        block_threshold: threat_score at or above which a BlockSignal is raised.
        max_workers:     Thread pool size for concurrent sentence analysis.
    """

    def __init__(
        self,
        text_monitor,
        block_threshold: int = 40,
        max_workers:     int = 4,
    ):
        self._monitor   = text_monitor
        self._threshold = block_threshold
        self._executor  = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="shield-monitor")
        self._tasks:    List[asyncio.Future] = []

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    async def submit(self, sentence_id: int, sentence: str, prompt: str) -> None:
        """
        Non-blocking: submit a sentence for background analysis.
        The result is stored internally; call collect() to retrieve signals.
        """
        loop = asyncio.get_event_loop()
        future = loop.run_in_executor(
            self._executor,
            self._analyze_sync,
            sentence_id, sentence, prompt,
        )
        self._tasks.append(future)

    async def collect(self, timeout: float = 1.5) -> List[BlockSignal]:
        """
        Wait for all pending monitor tasks (up to timeout) and return
        any BlockSignals found.

        Called at stream end to finalise the threat assessment.
        """
        if not self._tasks:
            return []

        signals: List[BlockSignal] = []
        try:
            results = await asyncio.wait_for(
                asyncio.gather(*self._tasks, return_exceptions=True),
                timeout=timeout,
            )
            for result in results:
                if isinstance(result, BlockSignal):
                    signals.append(result)
                elif isinstance(result, Exception):
                    log.warning("[StreamMonitor] Task error: %s", result)
        except asyncio.TimeoutError:
            log.warning("[StreamMonitor] collect() timed out after %.1fs — some sentences unscreened", timeout)

        self._tasks.clear()
        return signals

    def reset(self) -> None:
        """Clear pending tasks (e.g. between requests)."""
        self._tasks.clear()

    def shutdown(self) -> None:
        """Gracefully shut down the thread pool."""
        self._executor.shutdown(wait=False)

    # ------------------------------------------------------------------
    # Internal — runs in thread pool
    # ------------------------------------------------------------------

    def _analyze_sync(self, sentence_id: int, sentence: str, prompt: str) -> Optional[BlockSignal]:
        """
        Synchronous analysis — called inside ThreadPoolExecutor.
        Returns a BlockSignal if the sentence is flagged, else None.
        """
        try:
            result = self._monitor.analyze(prompt=prompt, response=sentence)
            if result["threat_score"] >= self._threshold:
                log.warning(
                    "[StreamMonitor] Sentence %d flagged: score=%d flags=%s",
                    sentence_id, result["threat_score"], result["flags"],
                )
                return BlockSignal(
                    sentence_id  = sentence_id,
                    reason       = result["reason"],
                    threat_score = result["threat_score"],
                    flags        = result["flags"],
                )
        except Exception as exc:
            log.error("[StreamMonitor] Analysis error on sentence %d: %s", sentence_id, exc)
        return None