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

fi/evals/framework/evaluators/non_blocking.py

@@ -0,0 +1,577 @@
+"""
+Non-blocking evaluation executor.
+
+Provides async/background evaluation with zero latency impact on the main thread.
+Uses thread pools for local execution with context propagation.
+"""
+
+from typing import Dict, Any, List, Optional, Callable
+from concurrent.futures import ThreadPoolExecutor, Future
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+import threading
+import time
+
+from ..types import FrameworkEvalResult as EvalResult, EvalStatus, BatchEvalResult
+from ..context import EvalContext
+from ..protocols import BaseEvaluation
+from ..registry import register_current_span
+from ..propagation import ContextCarrier
+
+
+@dataclass
+class EvalFuture:
+    """
+    Future representing a pending evaluation.
+
+    Wraps the underlying executor future and provides convenience methods.
+
+    Example:
+        future = evaluator.evaluate(inputs)
+        # Do other work...
+        result = future.result()  # Block until complete
+    """
+
+    future: Future
+    eval_name: str
+    eval_version: str
+    context: Optional[EvalContext] = None
+    submitted_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+
+    def done(self) -> bool:
+        """Check if evaluation is complete."""
+        return self.future.done()
+
+    def result(self, timeout: Optional[float] = None) -> EvalResult:
+        """
+        Get the evaluation result, blocking if necessary.
+
+        Args:
+            timeout: Maximum seconds to wait (None = wait forever)
+
+        Returns:
+            EvalResult from the evaluation
+
+        Raises:
+            TimeoutError: If timeout exceeded
+            Exception: If evaluation raised an exception
+        """
+        return self.future.result(timeout=timeout)
+
+    def cancel(self) -> bool:
+        """
+        Attempt to cancel the evaluation.
+
+        Returns:
+            True if cancelled, False if already running/complete
+        """
+        return self.future.cancel()
+
+    def cancelled(self) -> bool:
+        """Check if evaluation was cancelled."""
+        return self.future.cancelled()
+
+    def add_done_callback(self, fn: Callable[["EvalFuture"], None]) -> None:
+        """
+        Add callback to run when evaluation completes.
+
+        Args:
+            fn: Callback function taking this EvalFuture
+        """
+
+        def wrapper(f):
+            fn(self)
+
+        self.future.add_done_callback(wrapper)
+
+
+@dataclass
+class BatchEvalFuture:
+    """Future representing a batch of evaluations."""
+
+    futures: List[EvalFuture]
+    submitted_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+
+    def done(self) -> bool:
+        """Check if all evaluations are complete."""
+        return all(f.done() for f in self.futures)
+
+    def results(self, timeout: Optional[float] = None) -> BatchEvalResult:
+        """
+        Get all results, blocking if necessary.
+
+        Args:
+            timeout: Maximum seconds to wait for ALL results
+
+        Returns:
+            BatchEvalResult with all evaluation results
+        """
+        results = []
+        for f in self.futures:
+            try:
+                results.append(f.result(timeout=timeout))
+            except Exception as e:
+                # Create failure result for exceptions
+                results.append(
+                    EvalResult.failure(
+                        eval_name=f.eval_name,
+                        eval_version=f.eval_version,
+                        error=str(e),
+                    )
+                )
+        return BatchEvalResult.from_results(results)
+
+    def cancel_all(self) -> int:
+        """
+        Cancel all pending evaluations.
+
+        Returns:
+            Number of successfully cancelled evaluations
+        """
+        return sum(1 for f in self.futures if f.cancel())
+
+
+class NonBlockingEvaluator:
+    """
+    Non-blocking evaluation executor using thread pool.
+
+    Runs evaluations in background threads with zero latency impact.
+    Automatically enriches spans with results when complete.
+
+    Example:
+        evaluator = NonBlockingEvaluator(
+            evaluations=[ToxicityEval(), BiasEval()],
+            max_workers=4,
+        )
+
+        # Returns immediately
+        future = evaluator.evaluate({"response": "..."})
+
+        # Later, check results
+        if future.done():
+            result = future.result()
+
+        # Or wait for result
+        result = future.result(timeout=5.0)
+
+    Thread Safety:
+        This class is thread-safe. Multiple threads can call evaluate()
+        concurrently.
+    """
+
+    def __init__(
+        self,
+        evaluations: Optional[List[BaseEvaluation]] = None,
+        max_workers: int = 4,
+        auto_enrich_span: bool = True,
+        fail_fast: bool = False,
+        validate_inputs: bool = True,
+        executor: Optional[ThreadPoolExecutor] = None,
+    ):
+        """
+        Initialize the non-blocking evaluator.
+
+        Args:
+            evaluations: List of evaluations to run
+            max_workers: Maximum concurrent evaluations
+            auto_enrich_span: Whether to auto-enrich spans with results
+            fail_fast: Stop on first failure (for batch operations)
+            validate_inputs: Whether to validate inputs before evaluation
+            executor: Custom executor (uses internal pool if None)
+        """
+        self.evaluations = list(evaluations) if evaluations else []
+        self.max_workers = max_workers
+        self.auto_enrich_span = auto_enrich_span
+        self.fail_fast = fail_fast
+        self.validate_inputs = validate_inputs
+
+        # Use provided executor or create our own
+        self._executor = executor
+        self._owns_executor = executor is None
+        self._lock = threading.Lock()
+
+    @property
+    def executor(self) -> ThreadPoolExecutor:
+        """Get or create the thread pool executor."""
+        if self._executor is None:
+            with self._lock:
+                if self._executor is None:
+                    self._executor = ThreadPoolExecutor(
+                        max_workers=self.max_workers,
+                        thread_name_prefix="eval_worker_",
+                    )
+        return self._executor
+
+    def add_evaluation(self, evaluation: BaseEvaluation) -> "NonBlockingEvaluator":
+        """
+        Add an evaluation to run.
+
+        Args:
+            evaluation: The evaluation to add
+
+        Returns:
+            Self for chaining
+        """
+        self.evaluations.append(evaluation)
+        return self
+
+    def evaluate(
+        self,
+        inputs: Dict[str, Any],
+        evaluations: Optional[List[BaseEvaluation]] = None,
+        context: Optional[EvalContext] = None,
+        callback: Optional[Callable[[EvalResult], None]] = None,
+    ) -> BatchEvalFuture:
+        """
+        Run evaluations in background, returning immediately.
+
+        Args:
+            inputs: Input data for evaluations
+            evaluations: Override evaluations (uses instance evals if None)
+            context: Trace context for span enrichment
+            callback: Optional callback for each result
+
+        Returns:
+            BatchEvalFuture to track/retrieve results
+
+        Raises:
+            ValueError: If no evaluations configured
+        """
+        evals_to_run = evaluations if evaluations is not None else self.evaluations
+        if not evals_to_run:
+            raise ValueError("No evaluations to run")
+
+        # Capture context for span enrichment
+        carrier = ContextCarrier.capture() if context is None else ContextCarrier(context)
+
+        # Register current span if enrichment enabled
+        if self.auto_enrich_span:
+            register_current_span()
+
+        # Submit all evaluations
+        futures = []
+        for evaluation in evals_to_run:
+            future = self._submit_evaluation(
+                evaluation=evaluation,
+                inputs=inputs,
+                carrier=carrier,
+                callback=callback,
+            )
+            futures.append(future)
+
+        return BatchEvalFuture(futures=futures)
+
+    def evaluate_single(
+        self,
+        evaluation: BaseEvaluation,
+        inputs: Dict[str, Any],
+        context: Optional[EvalContext] = None,
+        callback: Optional[Callable[[EvalResult], None]] = None,
+    ) -> EvalFuture:
+        """
+        Run a single evaluation in background.
+
+        Args:
+            evaluation: The evaluation to run
+            inputs: Input data
+            context: Trace context
+            callback: Optional result callback
+
+        Returns:
+            EvalFuture to track/retrieve result
+        """
+        carrier = ContextCarrier.capture() if context is None else ContextCarrier(context)
+
+        if self.auto_enrich_span:
+            register_current_span()
+
+        return self._submit_evaluation(
+            evaluation=evaluation,
+            inputs=inputs,
+            carrier=carrier,
+            callback=callback,
+        )
+
+    def _submit_evaluation(
+        self,
+        evaluation: BaseEvaluation,
+        inputs: Dict[str, Any],
+        carrier: ContextCarrier,
+        callback: Optional[Callable[[EvalResult], None]] = None,
+    ) -> EvalFuture:
+        """Submit a single evaluation to the thread pool."""
+        eval_name = getattr(evaluation, "name", evaluation.__class__.__name__)
+        eval_version = getattr(evaluation, "version", "1.0.0")
+
+        future = self.executor.submit(
+            self._run_evaluation,
+            evaluation=evaluation,
+            inputs=inputs,
+            carrier=carrier,
+            callback=callback,
+        )
+
+        return EvalFuture(
+            future=future,
+            eval_name=eval_name,
+            eval_version=eval_version,
+            context=carrier.context,
+        )
+
+    def _run_evaluation(
+        self,
+        evaluation: BaseEvaluation,
+        inputs: Dict[str, Any],
+        carrier: ContextCarrier,
+        callback: Optional[Callable[[EvalResult], None]] = None,
+    ) -> EvalResult:
+        """
+        Run evaluation in background thread.
+
+        Handles:
+        - Input validation
+        - Timing
+        - Error handling
+        - Span enrichment
+        - Callback invocation
+        """
+        eval_name = getattr(evaluation, "name", evaluation.__class__.__name__)
+        eval_version = getattr(evaluation, "version", "1.0.0")
+
+        start_time = time.perf_counter()
+
+        try:
+            # Validate inputs if enabled
+            if self.validate_inputs:
+                validate = getattr(evaluation, "validate_inputs", None)
+                if validate:
+                    errors = validate(inputs)
+                    if errors:
+                        return EvalResult.failure(
+                            eval_name=eval_name,
+                            eval_version=eval_version,
+                            error=f"Validation errors: {errors}",
+                        )
+
+            # Run the evaluation
+            value = evaluation.evaluate(inputs)
+
+            latency_ms = (time.perf_counter() - start_time) * 1000
+
+            result = EvalResult(
+                value=value,
+                eval_name=eval_name,
+                eval_version=eval_version,
+                latency_ms=latency_ms,
+                status=EvalStatus.COMPLETED,
+            )
+
+            # Enrich span with results
+            if self.auto_enrich_span:
+                self._enrich_span(evaluation, result, carrier)
+
+            # Invoke callback
+            if callback:
+                try:
+                    callback(result)
+                except Exception:
+                    pass  # Don't fail evaluation for callback errors
+
+            return result
+
+        except Exception as e:
+            latency_ms = (time.perf_counter() - start_time) * 1000
+
+            result = EvalResult.failure(
+                eval_name=eval_name,
+                eval_version=eval_version,
+                error=str(e),
+            )
+            result.latency_ms = latency_ms
+
+            # Enrich span with failure
+            if self.auto_enrich_span:
+                self._enrich_span_failure(eval_name, result, carrier)
+
+            # Invoke callback even on failure
+            if callback:
+                try:
+                    callback(result)
+                except Exception:
+                    pass
+
+            return result
+
+    def _enrich_span(
+        self,
+        evaluation: BaseEvaluation,
+        result: EvalResult,
+        carrier: ContextCarrier,
+    ) -> bool:
+        """Enrich the original span with evaluation results."""
+        eval_name = result.eval_name.replace("-", "_").replace(" ", "_")
+        prefix = f"eval.{eval_name}"
+
+        # Get span attributes from evaluation
+        get_attrs = getattr(evaluation, "get_span_attributes", None)
+        if get_attrs:
+            try:
+                eval_attrs = get_attrs(result.value)
+            except Exception:
+                eval_attrs = {}
+        else:
+            eval_attrs = {}
+
+        # Build attributes
+        attributes = {
+            f"{prefix}.status": result.status.value,
+            f"{prefix}.latency_ms": result.latency_ms,
+            f"{prefix}.version": result.eval_version,
+        }
+
+        # Add eval-specific attributes
+        for key, value in eval_attrs.items():
+            if isinstance(value, (str, int, float, bool)):
+                attributes[f"{prefix}.{key}"] = value
+
+        return carrier.enrich_span(attributes)
+
+    def _enrich_span_failure(
+        self,
+        eval_name: str,
+        result: EvalResult,
+        carrier: ContextCarrier,
+    ) -> bool:
+        """Enrich span with failure information."""
+        safe_name = eval_name.replace("-", "_").replace(" ", "_")
+        prefix = f"eval.{safe_name}"
+
+        attributes = {
+            f"{prefix}.status": result.status.value,
+            f"{prefix}.latency_ms": result.latency_ms,
+            f"{prefix}.error": result.error or "Unknown error",
+        }
+
+        return carrier.enrich_span(attributes)
+
+    def shutdown(self, wait: bool = True) -> None:
+        """
+        Shutdown the executor.
+
+        Args:
+            wait: Whether to wait for pending tasks to complete
+        """
+        if self._executor and self._owns_executor:
+            self._executor.shutdown(wait=wait)
+            self._executor = None
+
+    def __enter__(self) -> "NonBlockingEvaluator":
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self.shutdown(wait=True)
+
+
+def non_blocking_evaluate(
+    inputs: Dict[str, Any],
+    *evaluations: BaseEvaluation,
+    max_workers: int = 4,
+    auto_enrich_span: bool = True,
+    callback: Optional[Callable[[EvalResult], None]] = None,
+) -> BatchEvalFuture:
+    """
+    Convenience function for non-blocking evaluation.
+
+    Runs evaluations in background and returns immediately.
+    The internal executor is automatically shut down when all futures complete.
+
+    Args:
+        inputs: Input data for evaluations
+        *evaluations: Evaluations to run
+        max_workers: Maximum concurrent evaluations
+        auto_enrich_span: Whether to enrich spans with results
+        callback: Optional callback for each result
+
+    Returns:
+        BatchEvalFuture to track/retrieve results
+
+    Example:
+        future = non_blocking_evaluate(
+            {"response": "..."},
+            ToxicityEval(),
+            BiasEval(),
+        )
+        # Do other work...
+        results = future.results()
+    """
+    evaluator = NonBlockingEvaluator(
+        evaluations=list(evaluations),
+        max_workers=max_workers,
+        auto_enrich_span=auto_enrich_span,
+    )
+    batch_future = evaluator.evaluate(inputs, callback=callback)
+
+    # Auto-shutdown executor when all futures complete
+    pending = [f.future for f in batch_future.futures]
+    remaining = {id(f) for f in pending}
+
+    def _on_done(f):
+        remaining.discard(id(f))
+        if not remaining:
+            evaluator.shutdown(wait=False)
+
+    for f in pending:
+        f.add_done_callback(_on_done)
+
+    return batch_future
+
+
+class EvalResultAggregator:
+    """
+    Aggregates results from multiple async evaluations.
+
+    Useful for collecting results from different evaluation runs.
+
+    Example:
+        aggregator = EvalResultAggregator()
+
+        # Add results as they come in
+        aggregator.add(future1.result())
+        aggregator.add(future2.result())
+
+        # Get aggregated results
+        batch = aggregator.to_batch()
+        print(f"Success rate: {batch.success_rate}")
+    """
+
+    def __init__(self):
+        self._results: List[EvalResult] = []
+        self._lock = threading.Lock()
+
+    def add(self, result: EvalResult) -> None:
+        """Add a result to the aggregator."""
+        with self._lock:
+            self._results.append(result)
+
+    def add_all(self, results: List[EvalResult]) -> None:
+        """Add multiple results."""
+        with self._lock:
+            self._results.extend(results)
+
+    def to_batch(self) -> BatchEvalResult:
+        """Get aggregated results as a batch."""
+        with self._lock:
+            return BatchEvalResult.from_results(list(self._results))
+
+    def clear(self) -> int:
+        """Clear all results and return count cleared."""
+        with self._lock:
+            count = len(self._results)
+            self._results.clear()
+            return count
+
+    @property
+    def count(self) -> int:
+        """Get current result count."""
+        with self._lock:
+            return len(self._results)