Spaces:
Runtime error
Runtime error
File size: 5,409 Bytes
02422e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | """
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
|