vellum-ai 0.5.0__py3-none-any.whl → 0.5.1__py3-none-any.whl

vellum/core/client_wrapper.py

@@ -18,7 +18,7 @@ class BaseClientWrapper:
         headers: typing.Dict[str, str] = {
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "vellum-ai",
-            "X-Fern-SDK-Version": "0.5.0",
+            "X-Fern-SDK-Version": "0.5.1",
         }
         headers["X_API_KEY"] = self.api_key
         return headers

vellum/lib/test_suites/resources.py

@@ -2,8 +2,10 @@ from __future__ import annotations
 
 import logging
 import time
-from typing import Callable, Generator, List
+from functools import cached_property
+from typing import Callable, Generator, List, Any
 
+from vellum import TestSuiteRunRead, TestSuiteRunMetricOutput_Number
 from vellum.client import Vellum
 from vellum.lib.test_suites.constants import (
     DEFAULT_MAX_POLLING_DURATION_MS,
@@ -11,6 +13,7 @@ from vellum.lib.test_suites.constants import (
 )
 from vellum.lib.test_suites.exceptions import TestSuiteRunResultsException
 from vellum.lib.utils.env import get_api_key
+from vellum.lib.utils.paginator import PaginatedResults, get_all_results
 from vellum.types import (
     ExternalTestCaseExecutionRequest,
     NamedTestCaseVariableValueRequest,
@@ -22,9 +25,6 @@ from vellum.types import (
     TestSuiteRunState,
 )
 
-from vellum.lib.utils.paginator import PaginatedResults, get_all_results
-
-
 logger = logging.getLogger(__name__)
 
 
@@ -116,44 +116,141 @@ class VellumTestSuiteRunResults:
 
     def __init__(
         self,
-        test_suite_run_id: str,
+        test_suite_run: TestSuiteRunRead,
         *,
         client: Vellum | None = None,
         polling_interval: int = DEFAULT_POLLING_INTERVAL_MS,
         max_polling_duration: int = DEFAULT_MAX_POLLING_DURATION_MS,
     ) -> None:
-        self._test_suite_run_id = test_suite_run_id
+        self._test_suite_run = test_suite_run
         self._client = client or Vellum(
             api_key=get_api_key(),
         )
-        self._state = "QUEUED"
         self._executions: Generator[VellumTestSuiteRunExecution, None, None] | None = (
             None
         )
         self._polling_interval = polling_interval
         self._max_polling_duration = max_polling_duration
 
+    @property
+    def state(self) -> TestSuiteRunState:
+        return self._test_suite_run.state
+
+    @cached_property
+    def all_executions(self) -> list[VellumTestSuiteRunExecution]:
+        return list(self._get_test_suite_run_executions())
+
     def get_metric_outputs(
         self, metric_identifier: str | None = None, output_identifier: str | None = None
-    ) -> Generator[TestSuiteRunMetricOutput, None, None]:
+    ) -> List[TestSuiteRunMetricOutput]:
         """Retrieve a metric's output across all executions by providing the info needed to uniquely identify it."""
 
-        executions = self._get_test_suite_run_executions()
-
-        for execution in executions:
-            yield execution.get_metric_output(
+        return [
+            execution.get_metric_output(
                 metric_identifier=metric_identifier, output_identifier=output_identifier
             )
+            for execution in self.all_executions
+        ]
+
+    def get_count_metric_outputs(
+        self,
+        metric_identifier: str | None = None,
+        output_identifier: str | None = None,
+        *,
+        predicate: Callable[[TestSuiteRunMetricOutput], bool] | None = None,
+    ) -> int:
+        """Returns the count of all metric outputs that match the given criteria."""
+
+        metric_outputs = self.get_metric_outputs(
+            metric_identifier=metric_identifier, output_identifier=output_identifier
+        )
+
+        if predicate is None:
+            return len(metric_outputs)
+
+        return len([output for output in metric_outputs if predicate(output)])
+
+    def get_numeric_metric_output_values(
+        self,
+        metric_identifier: str | None = None,
+        output_identifier: str | None = None,
+    ) -> List[float]:
+        """Returns the values of a numeric metric output that match the given criteria."""
+
+        metric_outputs: list[TestSuiteRunMetricOutput_Number] = []
+
+        for output in self.get_metric_outputs(
+            metric_identifier=metric_identifier, output_identifier=output_identifier
+        ):
+            if output.type != "NUMBER":
+                raise TestSuiteRunResultsException(
+                    f"Expected a numeric metric output, but got a {output.type} output instead."
+                )
+
+            metric_outputs.append(output)
+
+        return [output.value for output in metric_outputs]
+
+    def get_mean_metric_output(
+        self, metric_identifier: str | None = None, output_identifier: str | None = None
+    ) -> float:
+        """Returns the mean of all metric outputs that match the given criteria."""
+        output_values = self.get_numeric_metric_output_values(
+            metric_identifier=metric_identifier, output_identifier=output_identifier
+        )
+        return sum(output_values) / len(output_values)
+
+    def get_min_metric_output(
+        self, metric_identifier: str | None = None, output_identifier: str | None = None
+    ) -> float:
+        """Returns the min value across= all metric outputs that match the given criteria."""
+        output_values = self.get_numeric_metric_output_values(
+            metric_identifier=metric_identifier, output_identifier=output_identifier
+        )
+        return min(output_values)
+
+    def get_max_metric_output(
+        self, metric_identifier: str | None = None, output_identifier: str | None = None
+    ) -> float:
+        """Returns the max value across all metric outputs that match the given criteria."""
+        output_values = self.get_numeric_metric_output_values(
+            metric_identifier=metric_identifier, output_identifier=output_identifier
+        )
+        return max(output_values)
+
+    def wait_until_complete(self) -> None:
+        """Wait until the Test Suite Run is no longer in a QUEUED or RUNNING state."""
+
+        start_time = time.time_ns()
+        while True:
+            logger.debug("Polling for latest test suite run state...")
+            self._refresh_test_suite_run()
+            if self.state not in {"QUEUED", "RUNNING"}:
+                break
+
+            current_time = time.time_ns()
+            if ((current_time - start_time) / 1e6) > self._max_polling_duration:
+                raise TestSuiteRunResultsException(
+                    "Test suite run timed out polling for executions"
+                )
+
+            time.sleep(self._polling_interval / 1000.0)
+
+        if self.state == "FAILED":
+            raise TestSuiteRunResultsException("Test suite run failed")
 
-    def _refresh_test_suite_run_state(self):
-        test_suite_run = self._client.test_suite_runs.retrieve(self._test_suite_run_id)
-        self._state = test_suite_run.state
+        if self.state == "CANCELLED":
+            raise TestSuiteRunResultsException("Test suite run was cancelled")
+
+    def _refresh_test_suite_run(self):
+        test_suite_run = self._client.test_suite_runs.retrieve(self._test_suite_run.id)
+        self._test_suite_run = test_suite_run
 
     def _list_paginated_executions(
         self, offset: int | None, limit: int | None
     ) -> PaginatedResults[TestSuiteRunExecution]:
         response = self._client.test_suite_runs.list_executions(
-            self._test_suite_run_id,
+            self._test_suite_run.id,
             offset=offset,
             limit=limit,
             expand=[
@@ -175,26 +272,7 @@ class VellumTestSuiteRunResults:
         if self._executions is not None:
             return self._executions
 
-        start_time = time.time_ns()
-        while True:
-            logger.debug("Polling for latest test suite run state...")
-            self._refresh_test_suite_run_state()
-            if self._state not in {"QUEUED", "RUNNING"}:
-                break
-
-            current_time = time.time_ns()
-            if ((current_time - start_time) / 1e6) > self._max_polling_duration:
-                raise TestSuiteRunResultsException(
-                    "Test suite run timed out polling for executions"
-                )
-
-            time.sleep(self._polling_interval / 1000.0)
-
-        if self._state == "FAILED":
-            raise TestSuiteRunResultsException("Test suite run failed")
-
-        if self._state == "CANCELLED":
-            raise TestSuiteRunResultsException("Test suite run was cancelled")
+        self.wait_until_complete()
 
         raw_api_executions = get_all_results(self._list_paginated_executions)
         self._executions = self._wrap_api_executions(raw_api_executions)
@@ -250,4 +328,4 @@ class VellumTestSuite:
                 ),
             ),
         )
-        return VellumTestSuiteRunResults(test_suite_run.id, client=self.client)
+        return VellumTestSuiteRunResults(test_suite_run, client=self.client)

vellum/lib/utils/env.py

@@ -6,6 +6,8 @@ from .exceptions import VellumClientException
 def get_api_key() -> str:
     api_key = os.environ.get("VELLUM_API_KEY")
     if api_key is None:
-        raise VellumClientException("`VELLUM_API_KEY` environment variable id required to be set.")
-    
+        raise VellumClientException(
+            "`VELLUM_API_KEY` environment variable is required to be set."
+        )
+
     return api_key

vellum/lib/utils/exceptions.py

@@ -1,2 +1,2 @@
 class VellumClientException(Exception):
-    pass
+    pass

vellum/lib/utils/paginator.py

@@ -12,7 +12,8 @@ class PaginatedResults(Generic[Result]):
 
 
 def get_all_results(
-    paginated_api: Callable[[int, Union[int, None]], PaginatedResults[Result]], page_size: Union[int, None] = None
+    paginated_api: Callable[[int, Union[int, None]], PaginatedResults[Result]],
+    page_size: Union[int, None] = None,
 ) -> Generator[Result, None, None]:
     offset = 0
     count = 0

{vellum_ai-0.5.0.dist-info → vellum_ai-0.5.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 0.5.0
+Version: 0.5.1
 Summary: 
 Requires-Python: >=3.8,<4.0
 Classifier: Programming Language :: Python :: 3

{vellum_ai-0.5.0.dist-info → vellum_ai-0.5.1.dist-info}/RECORD

@@ -2,7 +2,7 @@ vellum/__init__.py,sha256=RpP5FLUDUph2qBCQ-TlpodjRFc1PDGDZCk-JBGj0UBM,35591
 vellum/client.py,sha256=7JaU104s0u_WhB8QAqIZcMv9IyvU-a0nKVZhTPKiEpw,97089
 vellum/core/__init__.py,sha256=1pNSKkwyQvMl_F0wohBqmoQAITptg3zlvCwsoSSzy7c,853
 vellum/core/api_error.py,sha256=RE8LELok2QCjABadECTvtDp7qejA1VmINCh6TbqPwSE,426
-vellum/core/client_wrapper.py,sha256=YUdARR7B9QqhNV9JOZYVoNyzoORUHfpBag8gvwu-BfA,1697
+vellum/core/client_wrapper.py,sha256=PSXJj42WsaIoL7ae07LoF7941fD683wvW7fzz7T2V1o,1697
 vellum/core/datetime_utils.py,sha256=nBys2IsYrhPdszxGKCNRPSOCwa-5DWOHG95FB8G9PKo,1047
 vellum/core/file.py,sha256=sy1RUGZ3aJYuw998bZytxxo6QdgKmlnlgBaMvwEKCGg,1480
 vellum/core/http_client.py,sha256=5ok6hqgZDJhg57EHvMnr0BBaHdG50QxFPKaCZ9aVWTc,5059
@@ -20,11 +20,11 @@ vellum/lib/__init__.py,sha256=KTSY0V59WEOr5uNyAei1dDfaAatyXw_Aca5kNjo5mY0,79
 vellum/lib/test_suites/__init__.py,sha256=hNsLoHSykqXDJP-MwFvu2lExImxo9KEyEJjt_fdAzpE,77
 vellum/lib/test_suites/constants.py,sha256=Vteml4_csZsMgo_q3-71E3JRCAoN6308TXLu5nfLhmU,116
 vellum/lib/test_suites/exceptions.py,sha256=6Xacoyv43fJvVf6Dt6Io5a-f9vF12Tx51jzsQRNSqhY,56
-vellum/lib/test_suites/resources.py,sha256=rjgPFktL37zNyB0WWErLqjDR1OzmBfjf6Ry6pb97r2A,9197
+vellum/lib/test_suites/resources.py,sha256=hokRS0_wT6IdA_6HkWbrh7iFzFxCtiy8JXbUiGtlwRk,12323
 vellum/lib/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum/lib/utils/env.py,sha256=__k8PagSUxW09x2ZMmwFrM_mwy-ky68aqml-e6jaYys,280
-vellum/lib/utils/exceptions.py,sha256=h9s9PnHqrTX5ohmZyCXovpWoTB7f3tAd5z_5nP0drCM,48
-vellum/lib/utils/paginator.py,sha256=mQwHZEkZHmCNBCctp8zkyEXmcfZtuU5gOL5gC3vNgUA,693
+vellum/lib/utils/env.py,sha256=ySl859lYBfls8hmlaU_RFdquHa_A_7SzaC6KEdFqh1Y,298
+vellum/lib/utils/exceptions.py,sha256=dXMAkzqbHV_AP5FjjbegPlfUE0zQDlpA3qOsoOJUxfg,49
+vellum/lib/utils/paginator.py,sha256=yDvgehocYBDclLt5SewZH4hCIyq0yLHdBzkyPCoYPjs,698
 vellum/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/resources/__init__.py,sha256=pqoVsVVIrUG-v6yt4AMtc7F5O-K7wKlvqhQeht9-Ax4,730
 vellum/resources/deployments/__init__.py,sha256=AE0TcFwLrLBljM0ZDX-pPw4Kqt-1f5JDpIok2HS80QI,157
@@ -376,7 +376,7 @@ vellum/types/workflow_result_event_output_data_search_results.py,sha256=gazaUrC5
 vellum/types/workflow_result_event_output_data_string.py,sha256=aVWIIGbLj4TJJhTTj6WzhbYXQkcZatKuhhNy8UYwXbw,1482
 vellum/types/workflow_stream_event.py,sha256=KA6Bkk_XA6AIPWR-1vKnwF1A8l_Bm5y0arQCWWWRpsk,911
 vellum/version.py,sha256=neLt8HBHHUtDF9M5fsyUzHT-pKooEPvceaLDqqIGb0s,77
-vellum_ai-0.5.0.dist-info/LICENSE,sha256=CcaljEIoOBaU-wItPH4PmM_mDCGpyuUY0Er1BGu5Ti8,1073
-vellum_ai-0.5.0.dist-info/METADATA,sha256=NMcjVbiC1qiFwXV12zgJEaDVbp_Ty7SK85_L1BGTc68,3549
-vellum_ai-0.5.0.dist-info/WHEEL,sha256=Zb28QaM1gQi8f4VCBhsUklF61CTlNYfs9YAZn-TOGFk,88
-vellum_ai-0.5.0.dist-info/RECORD,,
+vellum_ai-0.5.1.dist-info/LICENSE,sha256=CcaljEIoOBaU-wItPH4PmM_mDCGpyuUY0Er1BGu5Ti8,1073
+vellum_ai-0.5.1.dist-info/METADATA,sha256=TtLXbIJmAEV0EkDcuRYX2-SsLPvhP8L6GSE-WpEJqV4,3549
+vellum_ai-0.5.1.dist-info/WHEEL,sha256=Zb28QaM1gQi8f4VCBhsUklF61CTlNYfs9YAZn-TOGFk,88
+vellum_ai-0.5.1.dist-info/RECORD,,