codekingpro/portable-devtools

1"""A tracer that runs evaluators over completed runs."""2 3from __future__ import annotations4 5import logging6import threading7import weakref8from concurrent.futures import Future, ThreadPoolExecutor, wait9from typing import TYPE_CHECKING, Any, cast10from uuid import UUID11 12import langsmith13from langsmith.evaluation.evaluator import EvaluationResult, EvaluationResults14 15from langchain_core.tracers import langchain as langchain_tracer16from langchain_core.tracers._compat import run_copy17from langchain_core.tracers.base import BaseTracer18from langchain_core.tracers.context import tracing_v2_enabled19from langchain_core.tracers.langchain import _get_executor20 21if TYPE_CHECKING:22    from collections.abc import Sequence23 24    from langchain_core.tracers.schemas import Run25 26logger = logging.getLogger(__name__)27 28_TRACERS: weakref.WeakSet[EvaluatorCallbackHandler] = weakref.WeakSet()29 30 31def wait_for_all_evaluators() -> None:32    """Wait for all tracers to finish."""33    for tracer in list(_TRACERS):34        if tracer is not None:35            tracer.wait_for_futures()36 37 38class EvaluatorCallbackHandler(BaseTracer):39    """Tracer that runs a run evaluator whenever a run is persisted.40 41    Attributes:42        client: The LangSmith client instance used for evaluating the runs.43    """44 45    name: str = "evaluator_callback_handler"46 47    example_id: UUID | None = None48    """The example ID associated with the runs."""49 50    client: langsmith.Client51    """The LangSmith client instance used for evaluating the runs."""52 53    evaluators: Sequence[langsmith.RunEvaluator] = ()54    """The sequence of run evaluators to be executed."""55 56    executor: ThreadPoolExecutor | None = None57    """The thread pool executor used for running the evaluators."""58 59    futures: weakref.WeakSet[Future] = weakref.WeakSet()60    """The set of futures representing the running evaluators."""61 62    skip_unfinished: bool = True63    """Whether to skip runs that are not finished or raised an error."""64 65    project_name: str | None = None66    """The LangSmith project name to be organize eval chain runs under."""67 68    logged_eval_results: dict[tuple[str, str], list[EvaluationResult]]69 70    lock: threading.Lock71 72    def __init__(73        self,74        evaluators: Sequence[langsmith.RunEvaluator],75        client: langsmith.Client | None = None,76        example_id: UUID | str | None = None,77        skip_unfinished: bool = True,  # noqa: FBT001,FBT00278        project_name: str | None = "evaluators",79        max_concurrency: int | None = None,80        **kwargs: Any,81    ) -> None:82        """Create an EvaluatorCallbackHandler.83 84        Args:85            evaluators: The run evaluators to apply to all top level runs.86            client: The LangSmith client instance to use for evaluating the runs.87 88                If not specified, a new instance will be created.89            example_id: The example ID to be associated with the runs.90            skip_unfinished: Whether to skip unfinished runs.91            project_name: The LangSmith project name to be organize eval chain runs92                under.93            max_concurrency: The maximum number of concurrent evaluators to run.94        """95        super().__init__(**kwargs)96        self.example_id = (97            UUID(example_id) if isinstance(example_id, str) else example_id98        )99        self.client = client or langchain_tracer.get_client()100        self.evaluators = evaluators101        if max_concurrency is None:102            self.executor = _get_executor()103        elif max_concurrency > 0:104            self.executor = ThreadPoolExecutor(max_workers=max_concurrency)105            weakref.finalize(106                self,107                lambda: cast("ThreadPoolExecutor", self.executor).shutdown(wait=True),108            )109        else:110            self.executor = None111        self.futures = weakref.WeakSet[Future[None]]()112        self.skip_unfinished = skip_unfinished113        self.project_name = project_name114        self.logged_eval_results = {}115        self.lock = threading.Lock()116        _TRACERS.add(self)117 118    def _evaluate_in_project(self, run: Run, evaluator: langsmith.RunEvaluator) -> None:119        """Evaluate the run in the project.120 121        Args:122            run: The run to be evaluated.123            evaluator: The evaluator to use for evaluating the run.124        """125        try:126            if self.project_name is None:127                eval_result = self.client.evaluate_run(run, evaluator)128                eval_results = [eval_result]129            with tracing_v2_enabled(130                project_name=self.project_name, tags=["eval"], client=self.client131            ) as cb:132                reference_example = (133                    self.client.read_example(run.reference_example_id)134                    if run.reference_example_id135                    else None136                )137                evaluation_result = evaluator.evaluate_run(138                    # This is subclass, but getting errors for some reason139                    run,  # type: ignore[arg-type]140                    example=reference_example,141                )142                eval_results = self._log_evaluation_feedback(143                    evaluation_result,144                    run,145                    source_run_id=cb.latest_run.id if cb.latest_run else None,146                )147        except Exception:148            logger.exception(149                "Error evaluating run %s with %s",150                run.id,151                evaluator.__class__.__name__,152            )153            raise154        example_id = str(run.reference_example_id)155        with self.lock:156            for res in eval_results:157                run_id = str(getattr(res, "target_run_id", run.id))158                self.logged_eval_results.setdefault((run_id, example_id), []).append(159                    res160                )161 162    @staticmethod163    def _select_eval_results(164        results: EvaluationResult | EvaluationResults,165    ) -> list[EvaluationResult]:166        if isinstance(results, EvaluationResult):167            results_ = [results]168        elif isinstance(results, dict) and "results" in results:169            results_ = results["results"]170        else:171            msg = (172                f"Invalid evaluation result type {type(results)}."173                " Expected EvaluationResult or EvaluationResults."174            )175            raise TypeError(msg)176        return results_177 178    def _log_evaluation_feedback(179        self,180        evaluator_response: EvaluationResult | EvaluationResults,181        run: Run,182        source_run_id: UUID | None = None,183    ) -> list[EvaluationResult]:184        results = self._select_eval_results(evaluator_response)185        for res in results:186            source_info_: dict[str, Any] = {}187            if res.evaluator_info:188                source_info_ = {**res.evaluator_info, **source_info_}189            run_id_ = getattr(res, "target_run_id", None)190            if run_id_ is None:191                run_id_ = run.id192            self.client.create_feedback(193                run_id_,194                res.key,195                score=res.score,196                value=res.value,197                comment=res.comment,198                correction=res.correction,199                source_info=source_info_,200                source_run_id=res.source_run_id or source_run_id,201                feedback_source_type=langsmith.schemas.FeedbackSourceType.MODEL,202            )203        return results204 205    def _persist_run(self, run: Run) -> None:206        """Run the evaluator on the run.207 208        Args:209            run: The run to be evaluated.210        """211        if self.skip_unfinished and not run.outputs:212            logger.debug("Skipping unfinished run %s", run.id)213            return214        run_ = run_copy(run)215        run_.reference_example_id = self.example_id216        for evaluator in self.evaluators:217            if self.executor is None:218                self._evaluate_in_project(run_, evaluator)219            else:220                self.futures.add(221                    self.executor.submit(self._evaluate_in_project, run_, evaluator)222                )223 224    def wait_for_futures(self) -> None:225        """Wait for all futures to complete."""226        wait(self.futures)227