codekingpro/portable-devtools

1"""Make approximate assertions as "expectations" on test results.2 3This module is designed to be used within test cases decorated with the4`@pytest.mark.decorator` decorator5 6It allows you to log scores about a test case and optionally make assertions that log as7"expectation" feedback to LangSmith.8 9Example:10    ```python11    import pytest12    from langsmith import expect13 14 15    @pytest.mark.langsmith16    def test_output_semantically_close():17        response = oai_client.chat.completions.create(18            model="gpt-3.5-turbo",19            messages=[20                {"role": "system", "content": "You are a helpful assistant."},21                {"role": "user", "content": "Say hello!"},22            ],23        )24        response_txt = response.choices[0].message.content25        # Intended usage26        expect.embedding_distance(27            prediction=response_txt,28            reference="Hello!",29        ).to_be_less_than(0.9)30 31        # Score the test case32        matcher = expect.edit_distance(33            prediction=response_txt,34            reference="Hello!",35        )36        # Apply an assertion and log 'expectation' feedback to LangSmith37        matcher.to_be_less_than(1)38 39        # You can also directly make assertions on values directly40        expect.value(response_txt).to_contain("Hello!")41        # Or using a custom check42        expect.value(response_txt).against(lambda x: "Hello" in x)43 44        # You can even use this for basic metric logging within tests45 46        expect.score(0.8)47        expect.score(0.7, key="similarity").to_be_greater_than(0.7)48    ```49"""  # noqa: E50150 51from __future__ import annotations52 53import atexit54import inspect55from typing import (56    TYPE_CHECKING,57    Any,58    Callable,59    Literal,60    Optional,61    Union,62    overload,63)64 65from langsmith import client as ls_client66from langsmith import run_helpers as rh67from langsmith import run_trees as rt68from langsmith import utils as ls_utils69 70if TYPE_CHECKING:71    from langsmith._internal._edit_distance import EditDistanceConfig72    from langsmith._internal._embedding_distance import EmbeddingConfig73 74 75# Sentinel class used until PEP 0661 is accepted76class _NULL_SENTRY:77    """A sentinel singleton class used to distinguish omitted keyword arguments78    from those passed in with the value None (which may have different behavior).79    """  # noqa: D20580 81    def __bool__(self) -> Literal[False]:82        return False83 84    def __repr__(self) -> str:85        return "NOT_GIVEN"86 87 88NOT_GIVEN = _NULL_SENTRY()89 90 91class _Matcher:92    """A class for making assertions on expectation values."""93 94    def __init__(95        self,96        client: Optional[ls_client.Client],97        key: str,98        value: Any,99        _executor: Optional[ls_utils.ContextThreadPoolExecutor] = None,100        run_id: Optional[str] = None,101    ):102        self._client = client103        self.key = key104        self.value = value105        self._executor = _executor or ls_utils.ContextThreadPoolExecutor(max_workers=3)106        self._rt = rh.get_current_run_tree()107        self._run_id = self._rt.trace_id if self._rt else run_id108 109    def _submit_feedback(self, score: int, message: Optional[str] = None) -> None:110        if not ls_utils.test_tracking_is_disabled():111            if not self._client:112                self._client = rt.get_cached_client()113            self._executor.submit(114                self._client.create_feedback,115                run_id=self._run_id,116                key="expectation",117                score=score,118                comment=message,119                session_id=self._rt.session_id if self._rt else None,120                start_time=self._rt.start_time if self._rt else None,121            )122 123    def _assert(self, condition: bool, message: str, method_name: str) -> None:124        try:125            assert condition, message126            self._submit_feedback(1, message=f"Success: {self.key}.{method_name}")127        except AssertionError as e:128            self._submit_feedback(0, repr(e))129            raise e from None130 131    def to_be_less_than(self, value: float) -> None:132        """Assert that the expectation value is less than the given value.133 134        Args:135            value: The value to compare against.136 137        Raises:138            AssertionError: If the expectation value is not less than the given value.139        """140        self._assert(141            self.value < value,142            f"Expected {self.key} to be less than {value}, but got {self.value}",143            "to_be_less_than",144        )145 146    def to_be_greater_than(self, value: float) -> None:147        """Assert that the expectation value is greater than the given value.148 149        Args:150            value: The value to compare against.151 152        Raises:153            AssertionError: If the expectation value is not154            greater than the given value.155        """156        self._assert(157            self.value > value,158            f"Expected {self.key} to be greater than {value}, but got {self.value}",159            "to_be_greater_than",160        )161 162    def to_be_between(self, min_value: float, max_value: float) -> None:163        """Assert that the expectation value is between the given min and max values.164 165        Args:166            min_value: The minimum value (exclusive).167            max_value: The maximum value (exclusive).168 169        Raises:170            AssertionError: If the expectation value is not between the min and max.171        """172        self._assert(173            min_value < self.value < max_value,174            f"Expected {self.key} to be between {min_value} and {max_value},"175            f" but got {self.value}",176            "to_be_between",177        )178 179    def to_be_approximately(self, value: float, precision: int = 2) -> None:180        """Assert that the expectation value is approximately equal to the given value.181 182        Args:183            value: The value to compare against.184            precision: The number of decimal places to round to for comparison.185 186        Raises:187            AssertionError: If the rounded expectation value188                does not equal the rounded given value.189        """190        self._assert(191            round(self.value, precision) == round(value, precision),192            f"Expected {self.key} to be approximately {value}, but got {self.value}",193            "to_be_approximately",194        )195 196    def to_equal(self, value: float) -> None:197        """Assert that the expectation value equals the given value.198 199        Args:200            value: The value to compare against.201 202        Raises:203            AssertionError: If the expectation value does204                not exactly equal the given value.205        """206        self._assert(207            self.value == value,208            f"Expected {self.key} to be equal to {value}, but got {self.value}",209            "to_equal",210        )211 212    def to_be_none(self) -> None:213        """Assert that the expectation value is `None`.214 215        Raises:216            AssertionError: If the expectation value is not `None`.217        """218        self._assert(219            self.value is None,220            f"Expected {self.key} to be None, but got {self.value}",221            "to_be_none",222        )223 224    def to_contain(self, value: Any) -> None:225        """Assert that the expectation value contains the given value.226 227        Args:228            value: The value to check for containment.229 230        Raises:231            AssertionError: If the expectation value does not contain the given value.232        """233        self._assert(234            value in self.value,235            f"Expected {self.key} to contain {value}, but it does not",236            "to_contain",237        )238 239    # Custom assertions240    def against(self, func: Callable, /) -> None:241        """Assert the expectation value against a custom function.242 243        Args:244            func: A custom function that takes the expectation value as input.245 246        Raises:247            AssertionError: If the custom function returns False.248        """249        func_signature = inspect.signature(func)250        self._assert(251            func(self.value),252            f"Assertion {func_signature} failed for {self.key}",253            "against",254        )255 256 257class _Expect:258    """A class for setting expectations on test results."""259 260    def __init__(self, *, client: Optional[ls_client.Client] = None):261        self._client = client262        self.executor = ls_utils.ContextThreadPoolExecutor(max_workers=3)263        atexit.register(self.executor.shutdown, wait=True)264 265    def embedding_distance(266        self,267        prediction: str,268        reference: str,269        *,270        config: Optional[EmbeddingConfig] = None,271    ) -> _Matcher:272        """Compute the embedding distance between the prediction and reference.273 274        This logs the embedding distance to LangSmith and returns a `_Matcher` instance275        for making assertions on the distance value.276 277        By default, this uses the OpenAI API for computing embeddings.278 279        Args:280            prediction: The predicted string to compare.281            reference: The reference string to compare against.282            config: Optional configuration for the embedding distance evaluator.283 284                Supported options:285 286                - `encoder`: A custom encoder function to encode the list of input287                    strings to embeddings.288 289                    Defaults to the OpenAI API.290                - `metric`: The distance metric to use for comparison.291 292                    Supported values: `'cosine'`, `'euclidean'`, `'manhattan'`,293                    `'chebyshev'`, `'hamming'`.294 295        Returns:296            A `_Matcher` instance for the embedding distance value.297 298 299        Example:300            ```python301            expect.embedding_distance(302                prediction="hello",303                reference="hi",304            ).to_be_less_than(1.0)305            ```306        """  # noqa: E501307        from langsmith._internal._embedding_distance import EmbeddingDistance308 309        config = config or {}310        encoder_func = "custom" if config.get("encoder") else "openai"311        evaluator = EmbeddingDistance(config=config)312        score = evaluator.evaluate(prediction=prediction, reference=reference)313        src_info = {"encoder": encoder_func, "metric": evaluator.distance}314        self._submit_feedback(315            "embedding_distance",316            {317                "score": score,318                "source_info": src_info,319                "comment": f"Using {encoder_func}, Metric: {evaluator.distance}",320            },321        )322        return _Matcher(323            self._client, "embedding_distance", score, _executor=self.executor324        )325 326    def edit_distance(327        self,328        prediction: str,329        reference: str,330        *,331        config: Optional[EditDistanceConfig] = None,332    ) -> _Matcher:333        """Compute the string distance between the prediction and reference.334 335        This logs the string distance (Damerau-Levenshtein) to LangSmith and returns336        a `_Matcher` instance for making assertions on the distance value.337 338        This depends on the `rapidfuzz` package for string distance computation.339 340        Args:341            prediction: The predicted string to compare.342            reference: The reference string to compare against.343            config: Optional configuration for the string distance evaluator.344 345                Supported options:346 347                - `metric`: The distance metric to use for comparison.348 349                    Supported values: `'damerau_levenshtein'`, `'levenshtein'`,350                    `'jaro'`, `'jaro_winkler'`, `'hamming'`, `'indel'`.351                - `normalize_score`: Whether to normalize the score between `0` and `1`.352 353        Returns:354            A `_Matcher` instance for the string distance value.355 356        Examples:357            ```python358            expect.edit_distance("hello", "helo").to_be_less_than(1)359            ```360        """361        from langsmith._internal._edit_distance import EditDistance362 363        config = config or {}364        metric = config.get("metric") or "damerau_levenshtein"365        normalize = config.get("normalize_score", True)366        evaluator = EditDistance(config=config)367        score = evaluator.evaluate(prediction=prediction, reference=reference)368        src_info = {"metric": metric, "normalize": normalize}369        self._submit_feedback(370            "edit_distance",371            {372                "score": score,373                "source_info": src_info,374                "comment": f"Using {metric}, Normalize: {normalize}",375            },376        )377        return _Matcher(378            self._client,379            "edit_distance",380            score,381            _executor=self.executor,382        )383 384    def value(self, value: Any) -> _Matcher:385        """Create a `_Matcher` instance for making assertions on the given value.386 387        Args:388            value: The value to make assertions on.389 390        Returns:391            A `_Matcher` instance for the given value.392 393        Example:394            ```python395            expect.value(10).to_be_less_than(20)396            ```397        """398        return _Matcher(self._client, "value", value, _executor=self.executor)399 400    def score(401        self,402        score: Union[float, int, bool],403        *,404        key: str = "score",405        source_run_id: Optional[ls_client.ID_TYPE] = None,406        comment: Optional[str] = None,407    ) -> _Matcher:408        """Log a numeric score to LangSmith.409 410        Args:411            score: The score value to log.412            key: The key to use for logging the score. Defaults to `'score'`.413 414        Example:415            ```python416            expect.score(0.8)  # doctest: +ELLIPSIS417            <langsmith._expect._Matcher object at ...>418 419            expect.score(0.8, key="similarity").to_be_greater_than(0.7)420            ```421        """422        self._submit_feedback(423            key,424            {425                "score": score,426                "source_info": {"method": "expect.score"},427                "source_run_id": source_run_id,428                "comment": comment,429            },430        )431        return _Matcher(self._client, key, score, _executor=self.executor)432 433    ## Private Methods434 435    @overload436    def __call__(self, value: Any, /) -> _Matcher: ...437 438    @overload439    def __call__(self, /, *, client: ls_client.Client) -> _Expect: ...440 441    def __call__(442        self,443        value: Optional[Any] = NOT_GIVEN,444        /,445        client: Optional[ls_client.Client] = None,446    ) -> Union[_Expect, _Matcher]:447        expected = _Expect(client=client)448        if value is not NOT_GIVEN:449            return expected.value(value)450        return expected451 452    def _submit_feedback(self, key: str, results: dict):453        current_run = rh.get_current_run_tree()454        run_id = current_run.trace_id if current_run else None455        if not ls_utils.test_tracking_is_disabled():456            if not self._client:457                self._client = rt.get_cached_client()458            self.executor.submit(459                self._client.create_feedback, run_id=run_id, key=key, **results460            )461 462 463expect = _Expect()464 465__all__ = ["expect"]466