agent-learning-kit 0.1.0__py3-none-any.whl

fi/evals/streaming/evaluator.py

@@ -0,0 +1,551 @@
+"""Streaming Evaluator for real-time LLM output evaluation.
+
+Evaluates LLM outputs in real-time as tokens stream in, with support for
+early stopping based on configurable policies.
+"""
+
+import asyncio
+import time
+from dataclasses import dataclass
+from typing import (
+    AsyncIterator,
+    Callable,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+)
+
+from .types import (
+    ChunkResult,
+    EarlyStopReason,
+    StreamingConfig,
+    StreamingEvalResult,
+    StreamingState,
+)
+from .buffer import ChunkBuffer
+from .policy import EarlyStopPolicy
+
+
+@dataclass
+class EvalSpec:
+    """Specification for a streaming evaluation."""
+
+    name: str
+    eval_fn: Callable[[str, str], float]  # (chunk, cumulative) -> score
+    threshold: float = 0.7
+    weight: float = 1.0
+    pass_above: bool = True  # True if higher scores are better
+
+
+class StreamingEvaluator:
+    """
+    Evaluates LLM outputs in real-time as tokens stream in.
+
+    Supports early stopping based on configurable policies and provides
+    detailed per-chunk and aggregate evaluation results.
+
+    Example:
+        evaluator = StreamingEvaluator(config)
+        evaluator.add_eval("toxicity", toxicity_scorer, threshold=0.7, pass_above=False)
+        evaluator.set_policy(EarlyStopPolicy.default())
+
+        # Synchronous iteration
+        for token in stream:
+            result = evaluator.process_token(token)
+            if result and result.should_stop:
+                break
+
+        final_result = evaluator.finalize()
+
+        # Or async iteration
+        async for token in async_stream:
+            result = await evaluator.process_token_async(token)
+            if result and result.should_stop:
+                break
+    """
+
+    def __init__(
+        self,
+        config: Optional[StreamingConfig] = None,
+        policy: Optional[EarlyStopPolicy] = None,
+    ):
+        """
+        Initialize the streaming evaluator.
+
+        Args:
+            config: Streaming configuration (uses defaults if None)
+            policy: Early stop policy (uses default if None)
+        """
+        self.config = config or StreamingConfig()
+        self._policy = policy or EarlyStopPolicy.default()
+        self._buffer = ChunkBuffer(self.config)
+        self._evals: List[EvalSpec] = []
+        self._chunk_results: List[ChunkResult] = []
+        self._state = StreamingState.IDLE
+        self._start_time: float = 0.0
+        self._stop_reason = EarlyStopReason.NONE
+        self._stopped_at_chunk: Optional[int] = None
+
+    def add_eval(
+        self,
+        name: str,
+        eval_fn: Callable[[str, str], float],
+        threshold: float = 0.7,
+        weight: float = 1.0,
+        pass_above: bool = True,
+    ) -> "StreamingEvaluator":
+        """
+        Add an evaluation function.
+
+        Args:
+            name: Name of the evaluation
+            eval_fn: Function that takes (chunk_text, cumulative_text) and returns score
+            threshold: Passing threshold
+            weight: Weight for final score calculation
+            pass_above: If True, scores above threshold pass; if False, below
+
+        Returns:
+            Self for chaining
+        """
+        self._evals.append(
+            EvalSpec(
+                name=name,
+                eval_fn=eval_fn,
+                threshold=threshold,
+                weight=weight,
+                pass_above=pass_above,
+            )
+        )
+        return self
+
+    def set_policy(self, policy: EarlyStopPolicy) -> "StreamingEvaluator":
+        """
+        Set the early stop policy.
+
+        Args:
+            policy: Policy to use for early stopping
+
+        Returns:
+            Self for chaining
+        """
+        self._policy = policy
+        return self
+
+    def reset(self) -> None:
+        """Reset the evaluator for a new stream."""
+        self._buffer.reset()
+        self._policy.reset()
+        self._chunk_results = []
+        self._state = StreamingState.IDLE
+        self._start_time = 0.0
+        self._stop_reason = EarlyStopReason.NONE
+        self._stopped_at_chunk = None
+
+    def process_token(self, token: str) -> Optional[ChunkResult]:
+        """
+        Process a single token from the stream.
+
+        Args:
+            token: The token text
+
+        Returns:
+            ChunkResult if evaluation was triggered, None otherwise
+        """
+        if self._state == StreamingState.IDLE:
+            self._state = StreamingState.STREAMING
+            self._start_time = time.perf_counter()
+
+        if self._state in (StreamingState.STOPPED, StreamingState.COMPLETED, StreamingState.ERROR):
+            return None
+
+        # Add token to buffer
+        self._buffer.add(token)
+
+        # Check for limits
+        should_stop_limits, limit_reason = self._buffer.should_stop_for_limits()
+        if should_stop_limits:
+            return self._handle_limit_stop(limit_reason)
+
+        # Check if we should evaluate
+        if not self._buffer.should_evaluate():
+            return None
+
+        # Run evaluation
+        return self._evaluate_chunk()
+
+    async def process_token_async(self, token: str) -> Optional[ChunkResult]:
+        """
+        Process a single token asynchronously.
+
+        Args:
+            token: The token text
+
+        Returns:
+            ChunkResult if evaluation was triggered, None otherwise
+        """
+        # For now, wrap sync processing - can be optimized later
+        return await asyncio.to_thread(self.process_token, token)
+
+    def process_chunk(self, chunk: str) -> Optional[ChunkResult]:
+        """
+        Process a larger chunk of text (multiple tokens).
+
+        Args:
+            chunk: The chunk text
+
+        Returns:
+            ChunkResult if evaluation was triggered, None otherwise
+        """
+        if self._state == StreamingState.IDLE:
+            self._state = StreamingState.STREAMING
+            self._start_time = time.perf_counter()
+
+        if self._state in (StreamingState.STOPPED, StreamingState.COMPLETED, StreamingState.ERROR):
+            return None
+
+        # Add chunk to buffer
+        self._buffer.add_chunk(chunk)
+
+        # Check for limits
+        should_stop_limits, limit_reason = self._buffer.should_stop_for_limits()
+        if should_stop_limits:
+            return self._handle_limit_stop(limit_reason)
+
+        # Check if we should evaluate
+        if not self._buffer.should_evaluate():
+            return None
+
+        # Run evaluation
+        return self._evaluate_chunk()
+
+    def _evaluate_chunk(self) -> ChunkResult:
+        """Run evaluations on the current chunk."""
+        chunk_start = time.perf_counter()
+
+        chunk_text = self._buffer.get_chunk()
+        cumulative_text = self._buffer.get_cumulative()
+        chunk_index = self._buffer.get_chunk_index()
+
+        scores: Dict[str, float] = {}
+        flags: Dict[str, bool] = {}
+
+        # Run all evaluations
+        for eval_spec in self._evals:
+            try:
+                score = eval_spec.eval_fn(chunk_text, cumulative_text)
+                scores[eval_spec.name] = score
+
+                # Determine if passed
+                if eval_spec.pass_above:
+                    flags[eval_spec.name] = score >= eval_spec.threshold
+                else:
+                    flags[eval_spec.name] = score <= eval_spec.threshold
+            except Exception:
+                # Handle eval errors gracefully
+                scores[eval_spec.name] = 0.0
+                flags[eval_spec.name] = False
+
+        latency_ms = (time.perf_counter() - chunk_start) * 1000
+
+        # Create chunk result
+        chunk_result = ChunkResult(
+            chunk_index=chunk_index,
+            chunk_text=chunk_text,
+            cumulative_text=cumulative_text,
+            scores=scores,
+            flags=flags,
+            latency_ms=latency_ms,
+        )
+
+        # Check policy for early stop
+        if self.config.enable_early_stop:
+            should_stop, stop_reason = self._policy.check(chunk_result)
+
+            if should_stop:
+                chunk_result.should_stop = True
+                chunk_result.stop_reason = stop_reason
+                self._state = StreamingState.STOPPED
+                self._stop_reason = stop_reason
+                self._stopped_at_chunk = chunk_index
+
+                # Trigger callback if configured
+                if self.config.on_stop_callback:
+                    self.config.on_stop_callback(stop_reason, cumulative_text)
+
+            # Check stop on first failure
+            elif self.config.stop_on_first_failure and not chunk_result.all_passed:
+                chunk_result.should_stop = True
+                chunk_result.stop_reason = EarlyStopReason.THRESHOLD
+                self._state = StreamingState.STOPPED
+                self._stop_reason = EarlyStopReason.THRESHOLD
+                self._stopped_at_chunk = chunk_index
+
+                # Trigger callback if configured
+                if self.config.on_stop_callback:
+                    self.config.on_stop_callback(EarlyStopReason.THRESHOLD, cumulative_text)
+
+        # Store result
+        self._chunk_results.append(chunk_result)
+
+        # Mark as evaluated
+        self._buffer.mark_evaluated()
+
+        # Trigger chunk callback if configured
+        if self.config.on_chunk_callback:
+            self.config.on_chunk_callback(chunk_result)
+
+        return chunk_result
+
+    def _handle_limit_stop(self, reason: str) -> ChunkResult:
+        """Handle stopping due to limits."""
+        # Evaluate any remaining content first
+        if self._buffer.has_pending:
+            chunk_result = self._evaluate_chunk()
+        else:
+            # Create a minimal result for the stop
+            chunk_result = ChunkResult(
+                chunk_index=self._buffer.get_chunk_index(),
+                chunk_text="",
+                cumulative_text=self._buffer.get_cumulative(),
+                scores={},
+                flags={},
+            )
+
+        # Map limit reason to stop reason
+        if reason == "max_tokens":
+            stop_reason = EarlyStopReason.MAX_TOKENS
+        elif reason == "max_chars":
+            stop_reason = EarlyStopReason.MAX_CHARS
+        elif reason == "timeout":
+            stop_reason = EarlyStopReason.TIMEOUT
+        else:
+            stop_reason = EarlyStopReason.ERROR
+
+        chunk_result.should_stop = True
+        chunk_result.stop_reason = stop_reason
+        self._state = StreamingState.STOPPED
+        self._stop_reason = stop_reason
+        self._stopped_at_chunk = chunk_result.chunk_index
+
+        return chunk_result
+
+    def finalize(self) -> StreamingEvalResult:
+        """
+        Finalize evaluation and return results.
+
+        Should be called after stream completes or after early stop.
+
+        Returns:
+            StreamingEvalResult with all evaluation data
+        """
+        # Evaluate any remaining pending content
+        if self._buffer.has_pending and self._state == StreamingState.STREAMING:
+            self._evaluate_chunk()
+
+        # Mark as completed if not already stopped
+        if self._state == StreamingState.STREAMING:
+            self._state = StreamingState.COMPLETED
+
+        # Calculate final scores (weighted average across chunks)
+        final_scores = self._calculate_final_scores()
+
+        # Determine overall pass/fail
+        passed = self._determine_passed(final_scores)
+
+        # Calculate total latency
+        total_latency_ms = (time.perf_counter() - self._start_time) * 1000 if self._start_time else 0.0
+
+        return StreamingEvalResult(
+            passed=passed,
+            final_text=self._buffer.get_cumulative(),
+            total_chunks=len(self._chunk_results),
+            chunk_results=self._chunk_results,
+            final_scores=final_scores,
+            early_stopped=self._state == StreamingState.STOPPED,
+            stop_reason=self._stop_reason,
+            stopped_at_chunk=self._stopped_at_chunk,
+            total_latency_ms=total_latency_ms,
+            state=self._state,
+            metadata={
+                "buffer_stats": self._buffer.get_stats(),
+                "policy_stats": self._policy.get_stats(),
+            },
+        )
+
+    def _calculate_final_scores(self) -> Dict[str, float]:
+        """Calculate final weighted average scores."""
+        if not self._chunk_results:
+            return {}
+
+        # Build weight lookup from eval specs
+        weights: Dict[str, float] = {e.name: e.weight for e in self._evals}
+
+        final_scores: Dict[str, float] = {}
+        weighted_sums: Dict[str, float] = {}
+        weight_totals: Dict[str, float] = {}
+
+        for chunk_result in self._chunk_results:
+            for name, score in chunk_result.scores.items():
+                w = weights.get(name, 1.0)
+                if name not in weighted_sums:
+                    weighted_sums[name] = 0.0
+                    weight_totals[name] = 0.0
+                weighted_sums[name] += score * w
+                weight_totals[name] += w
+
+        for name in weighted_sums:
+            if weight_totals[name] > 0:
+                final_scores[name] = weighted_sums[name] / weight_totals[name]
+            else:
+                final_scores[name] = 0.0
+
+        return final_scores
+
+    def _determine_passed(self, final_scores: Dict[str, float]) -> bool:
+        """Determine if evaluation passed overall."""
+        if self._state == StreamingState.STOPPED:
+            # If stopped early due to safety/toxicity, fail
+            if self._stop_reason in (
+                EarlyStopReason.TOXICITY,
+                EarlyStopReason.SAFETY,
+                EarlyStopReason.PII,
+                EarlyStopReason.JAILBREAK,
+            ):
+                return False
+
+        # Check final scores against thresholds
+        for eval_spec in self._evals:
+            if eval_spec.name in final_scores:
+                score = final_scores[eval_spec.name]
+                if eval_spec.pass_above:
+                    if score < eval_spec.threshold:
+                        return False
+                else:
+                    if score > eval_spec.threshold:
+                        return False
+
+        return True
+
+    @property
+    def state(self) -> StreamingState:
+        """Get current evaluation state."""
+        return self._state
+
+    @property
+    def is_stopped(self) -> bool:
+        """Check if evaluation has stopped."""
+        return self._state in (StreamingState.STOPPED, StreamingState.COMPLETED, StreamingState.ERROR)
+
+    @property
+    def chunk_count(self) -> int:
+        """Get number of chunks evaluated."""
+        return len(self._chunk_results)
+
+    def evaluate_stream(
+        self,
+        stream: Iterator[str],
+    ) -> StreamingEvalResult:
+        """
+        Evaluate an entire stream synchronously.
+
+        Args:
+            stream: Iterator yielding tokens/chunks
+
+        Returns:
+            StreamingEvalResult after processing complete stream
+        """
+        self.reset()
+
+        for token in stream:
+            result = self.process_token(token)
+            if result and result.should_stop:
+                break
+
+        return self.finalize()
+
+    async def evaluate_stream_async(
+        self,
+        stream: AsyncIterator[str],
+    ) -> StreamingEvalResult:
+        """
+        Evaluate an entire stream asynchronously.
+
+        Args:
+            stream: Async iterator yielding tokens/chunks
+
+        Returns:
+            StreamingEvalResult after processing complete stream
+        """
+        self.reset()
+
+        async for token in stream:
+            result = await self.process_token_async(token)
+            if result and result.should_stop:
+                break
+
+        return self.finalize()
+
+    @classmethod
+    def with_defaults(cls) -> "StreamingEvaluator":
+        """
+        Create an evaluator with default configuration.
+
+        Returns:
+            StreamingEvaluator with default settings
+        """
+        return cls(
+            config=StreamingConfig(),
+            policy=EarlyStopPolicy.default(),
+        )
+
+    @classmethod
+    def for_safety(
+        cls,
+        toxicity_threshold: float = 0.5,
+        safety_threshold: float = 0.5,
+    ) -> "StreamingEvaluator":
+        """
+        Create an evaluator optimized for safety monitoring.
+
+        Args:
+            toxicity_threshold: Threshold for toxicity (stop if above)
+            safety_threshold: Threshold for safety (stop if below)
+
+        Returns:
+            StreamingEvaluator configured for safety
+        """
+        config = StreamingConfig(
+            enable_early_stop=True,
+            stop_on_first_failure=True,
+            toxicity_threshold=toxicity_threshold,
+            safety_threshold=safety_threshold,
+        )
+        policy = EarlyStopPolicy.strict()
+        return cls(config=config, policy=policy)
+
+    @classmethod
+    def for_quality(
+        cls,
+        min_chunk_size: int = 50,
+        eval_interval_ms: int = 500,
+    ) -> "StreamingEvaluator":
+        """
+        Create an evaluator optimized for quality assessment.
+
+        Args:
+            min_chunk_size: Minimum characters before evaluation
+            eval_interval_ms: Milliseconds between evaluations
+
+        Returns:
+            StreamingEvaluator configured for quality
+        """
+        config = StreamingConfig(
+            min_chunk_size=min_chunk_size,
+            max_chunk_size=200,
+            eval_interval_ms=eval_interval_ms,
+            enable_early_stop=False,  # Don't stop early for quality
+            eval_on_sentence_end=True,
+        )
+        policy = EarlyStopPolicy.permissive()
+        return cls(config=config, policy=policy)