arize 8.0.0b2__py3-none-any.whl → 8.0.1__py3-none-any.whl

arize/experiments/functions.py

@@ -7,7 +7,14 @@ import json
 import logging
 import traceback
 from binascii import hexlify
-from collections.abc import Awaitable, Callable, Mapping, Sequence
+from collections.abc import (
+    Awaitable,
+    Callable,
+    Coroutine,
+    Iterable,
+    Mapping,
+    Sequence,
+)
 from contextlib import ExitStack
 from copy import deepcopy
 from datetime import date, datetime, time, timedelta, timezone
@@ -34,7 +41,7 @@ from openinference.semconv.trace import (
 )
 from opentelemetry.context import Context
 from opentelemetry.sdk.resources import Resource
-from opentelemetry.trace import Status, StatusCode, Tracer
+from opentelemetry.trace import NoOpTracer, Status, StatusCode, Tracer
 
 if TYPE_CHECKING:
     from opentelemetry.sdk.trace import Span
@@ -48,6 +55,7 @@ from arize.experiments.evaluators.types import (
     EvaluationResult,
     EvaluationResultFieldNames,
     EvaluatorName,
+    JSONSerializable,
 )
 from arize.experiments.evaluators.utils import create_evaluator
 from arize.experiments.tracing import capture_spans, flatten
@@ -64,6 +72,9 @@ RateLimitErrors: TypeAlias = type[BaseException] | Sequence[type[BaseException]]
 
 logger = logging.getLogger(__name__)
 
+# Module-level singleton for no-op tracing
+_NOOP_TRACER = NoOpTracer()
+
 
 def run_experiment(
     experiment_name: str,
@@ -116,11 +127,12 @@ def run_experiment(
         error: BaseException | None = None
         status = Status(StatusCode.OK)
         with ExitStack() as stack:
-            span: Span = stack.enter_context(
+            # Type ignore: OpenTelemetry interface vs implementation type mismatch
+            span: Span = stack.enter_context(  # type: ignore[assignment]
                 cm=tracer.start_as_current_span(
                     name=root_span_name, context=Context()
                 )
-            )  # type:ignore
+            )
             stack.enter_context(capture_spans(resource))
             span.set_attribute(METADATA, json.dumps(md, ensure_ascii=False))
             try:
@@ -146,9 +158,12 @@ def run_experiment(
                     raise TypeError(sync_error_message)
                 output = _output
 
-            output = jsonify(output)
+            # Type ignore: jsonify returns object but runtime result is JSONSerializable
+            output = jsonify(output)  # type: ignore[assignment]
             if example.input:
-                span.set_attribute(INPUT_VALUE, example.input)  # type: ignore
+                # OpenTelemetry type hints are restrictive, but Arize's tracing layer
+                # accepts JSON-serializable structures which are auto-serialized
+                span.set_attribute(INPUT_VALUE, example.input)  # type: ignore[arg-type]
             else:
                 span.set_attribute(
                     INPUT_VALUE,
@@ -187,9 +202,9 @@ def run_experiment(
                 else datetime.now(tz=timezone.utc)
             ),
             dataset_example_id=example.id,
-            output=output,  # type:ignore
+            output=output,
             error=repr(error) if error else None,
-            trace_id=_str_trace_id(span.get_span_context().trace_id),  # type: ignore
+            trace_id=_str_trace_id(span.get_span_context().trace_id),
         )
 
     async def async_run_experiment(example: Example) -> ExperimentRun:
@@ -197,11 +212,12 @@ def run_experiment(
         error: BaseException | None = None
         status = Status(StatusCode.OK)
         with ExitStack() as stack:
-            span: Span = stack.enter_context(
+            # Type ignore: OpenTelemetry interface vs implementation type mismatch
+            span: Span = stack.enter_context(  # type: ignore[assignment]
                 cm=tracer.start_as_current_span(
                     name=root_span_name, context=Context()
                 )
-            )  # type:ignore
+            )
             stack.enter_context(capture_spans(resource))
             span.set_attribute(METADATA, json.dumps(md, ensure_ascii=False))
             try:
@@ -220,9 +236,12 @@ def run_experiment(
                 )
                 error = exc
                 _print_experiment_error(exc, example_id=example.id, kind="task")
-            output = jsonify(output)
+            # Type ignore: jsonify returns object but runtime result is JSONSerializable
+            output = jsonify(output)  # type: ignore[assignment]
             if example.input:
-                span.set_attribute(INPUT_VALUE, example.input)  # type: ignore
+                # OpenTelemetry type hints are restrictive, but Arize's tracing layer
+                # accepts JSON-serializable structures which are auto-serialized
+                span.set_attribute(INPUT_VALUE, example.input)  # type: ignore[arg-type]
             else:
                 span.set_attribute(
                     INPUT_VALUE,
@@ -261,14 +280,14 @@ def run_experiment(
                 else datetime.now(tz=timezone.utc)
             ),
             dataset_example_id=example.id,
-            output=output,  # type: ignore
+            output=output,
             error=repr(error) if error else None,
-            trace_id=_str_trace_id(span.get_span_context().trace_id),  # type: ignore
+            trace_id=_str_trace_id(span.get_span_context().trace_id),
         )
 
     _errors: tuple[type[BaseException], ...]
     if not isinstance(rate_limit_errors, Sequence):
-        _errors = (rate_limit_errors,)  # type: ignore
+        _errors = (rate_limit_errors,) if rate_limit_errors is not None else ()
     else:
         _errors = tuple(filter(None, rate_limit_errors))
     rate_limiters = [RateLimiter(rate_limit_error=rle) for rle in _errors]
@@ -284,8 +303,13 @@ def run_experiment(
     )
 
     executor = get_executor_on_sync_context(
-        sync_fn=rate_limited_sync_run_experiment,
-        async_fn=rate_limited_async_run_experiment,
+        sync_fn=cast(
+            "Callable[[object], Any]", rate_limited_sync_run_experiment
+        ),
+        async_fn=cast(
+            "Callable[[object], Coroutine[Any, Any, Any]]",
+            rate_limited_async_run_experiment,
+        ),
         max_retries=0,
         exit_on_error=exit_on_error,
         fallback_return_value=None,
@@ -295,20 +319,27 @@ def run_experiment(
     )
 
     runs, _ = executor.run(examples)
-    task_summary = _TaskSummary.from_task_runs(len(dataset), runs)
+    task_summary = _TaskSummary.from_task_runs(
+        len(dataset), cast("list[ExperimentRun | None]", runs)
+    )
 
     if exit_on_error and (None in runs):
         # When exit_on_error is True, the result of a failed task execution is None
         # If any task execution failed, raise an error to exit early
         raise RuntimeError("An error occurred during execution of tasks.")
 
+    # Filter out None values before accessing attributes
+    runs_filtered = [
+        r for r in cast("list[ExperimentRun | None]", runs) if r is not None
+    ]
+
     out_df = pd.DataFrame()
-    out_df["id"] = [run.id for run in runs]
-    out_df["example_id"] = [run.dataset_example_id for run in runs]
-    out_df["result"] = [run.output for run in runs]
-    out_df["result.trace.id"] = [run.trace_id for run in runs]
+    out_df["id"] = [run.id for run in runs_filtered]
+    out_df["example_id"] = [run.dataset_example_id for run in runs_filtered]
+    out_df["result"] = [run.output for run in runs_filtered]  # type: ignore[assignment]
+    out_df["result.trace.id"] = [run.trace_id for run in runs_filtered]
     out_df["result.trace.timestamp"] = [
-        int(run.start_time.timestamp() * 1e3) for run in runs
+        int(run.start_time.timestamp() * 1e3) for run in runs_filtered
     ]
     out_df.set_index("id", inplace=True, drop=False)
     logger.info(f"✅ Task runs completed.\n{task_summary}")
@@ -317,7 +348,7 @@ def run_experiment(
         eval_results = evaluate_experiment(
             experiment_name=experiment_name,
             examples=examples,
-            experiment_results=runs,
+            experiment_results=cast("Sequence[ExperimentRun]", runs),
             evaluators=evaluators,
             rate_limit_errors=rate_limit_errors,
             concurrency=concurrency,
@@ -333,7 +364,7 @@ def run_experiment(
             )
 
         # group evaluation results by name
-        eval_results_by_name = {}
+        eval_results_by_name: dict[str, list[ExperimentEvaluationRun]] = {}
         for r in eval_results:
             if r is None:
                 continue
@@ -355,7 +386,8 @@ def run_experiment(
             }
 
             for attr, getter in eval_data.items():
-                out_df[f"eval.{eval_name}.{attr}"] = out_df.index.map(
+                # Type ignore: pandas DataFrame column assignment type is overly restrictive
+                out_df[f"eval.{eval_name}.{attr}"] = out_df.index.map(  # type: ignore[assignment]
                     {r.experiment_run_id: getter(r) for r in eval_res}
                 )
             out_df = _add_metadata_to_output_df(out_df, eval_res, eval_name)
@@ -372,7 +404,7 @@ def evaluate_experiment(
     evaluators: Evaluators | None = None,
     rate_limit_errors: RateLimitErrors | None = None,
     concurrency: int = 3,
-    tracer: Tracer | None = None,
+    tracer: Tracer = _NOOP_TRACER,
     resource: Resource | None = None,
     exit_on_error: bool = False,
     timeout: int = 120,
@@ -386,7 +418,7 @@ def evaluate_experiment(
         evaluators (Evaluators): The evaluators to use for assessment.
         rate_limit_errors (RateLimitErrors | :obj:`None`): Optional rate limit errors.
         concurrency (int): The number of concurrent tasks to run. Default is 3.
-        tracer (Tracer | :obj:`None`): Optional tracer for tracing the evaluation.
+        tracer (Tracer): Tracer for tracing the evaluation. Defaults to NoOpTracer().
         resource (Resource | :obj:`None`): Optional resource for the evaluation.
         exit_on_error (bool): Whether to exit on error. Default is False.
         timeout (int): The timeout for each evaluation in seconds. Default is 120.
@@ -425,12 +457,16 @@ def evaluate_experiment(
         status = Status(StatusCode.OK)
         root_span_name = f"Evaluation: {evaluator.name}"
         with ExitStack() as stack:
-            span: Span = stack.enter_context(
-                tracer.start_as_current_span(  # type:ignore
-                    name=root_span_name, context=Context()
-                )
+            span: Span = cast(
+                "Span",
+                stack.enter_context(
+                    tracer.start_as_current_span(
+                        name=root_span_name, context=Context()
+                    )
+                ),
             )
-            stack.enter_context(capture_spans(resource))  # type:ignore
+            if resource is not None:
+                stack.enter_context(capture_spans(resource))
             span.set_attribute(METADATA, json.dumps(md, ensure_ascii=False))
             try:
                 result = evaluator.evaluate(
@@ -456,7 +492,15 @@ def evaluate_experiment(
                 )
             if result:
                 span.set_attributes(
-                    dict(flatten(jsonify(result), recurse_on_sequence=True))
+                    dict(
+                        flatten(
+                            cast(
+                                "Mapping[str, Any] | Iterable[Any]",
+                                jsonify(result),
+                            ),
+                            recurse_on_sequence=True,
+                        )
+                    )
                 )
             span.set_attribute(OPENINFERENCE_SPAN_KIND, root_span_kind)
             span.set_status(status)
@@ -473,7 +517,7 @@ def evaluate_experiment(
             annotator_kind=evaluator.kind,
             error=repr(error) if error else None,
             result=result,
-            trace_id=_str_trace_id(span.get_span_context().trace_id),  # type:ignore
+            trace_id=_str_trace_id(span.get_span_context().trace_id),
         )
 
     async def async_eval_run(
@@ -485,12 +529,16 @@ def evaluate_experiment(
         status = Status(StatusCode.OK)
         root_span_name = f"Evaluation: {evaluator.name}"
         with ExitStack() as stack:
-            span: Span = stack.enter_context(
-                tracer.start_as_current_span(  # type:ignore
-                    name=root_span_name, context=Context()
-                )
+            span: Span = cast(
+                "Span",
+                stack.enter_context(
+                    tracer.start_as_current_span(
+                        name=root_span_name, context=Context()
+                    )
+                ),
             )
-            stack.enter_context(capture_spans(resource))  # type:ignore
+            if resource is not None:
+                stack.enter_context(capture_spans(resource))
             span.set_attribute(METADATA, json.dumps(md, ensure_ascii=False))
             try:
                 result = await evaluator.async_evaluate(
@@ -516,7 +564,15 @@ def evaluate_experiment(
                 )
             if result:
                 span.set_attributes(
-                    dict(flatten(jsonify(result), recurse_on_sequence=True))
+                    dict(
+                        flatten(
+                            cast(
+                                "Mapping[str, Any] | Iterable[Any]",
+                                jsonify(result),
+                            ),
+                            recurse_on_sequence=True,
+                        )
+                    )
                 )
             span.set_attribute(OPENINFERENCE_SPAN_KIND, root_span_kind)
             span.set_status(status)
@@ -532,7 +588,7 @@ def evaluate_experiment(
             annotator_kind=evaluator.kind,
             error=repr(error) if error else None,
             result=result,
-            trace_id=_str_trace_id(span.get_span_context().trace_id),  # type:ignore
+            trace_id=_str_trace_id(span.get_span_context().trace_id),
         )
 
     _errors: tuple[type[BaseException], ...]
@@ -553,8 +609,11 @@ def evaluate_experiment(
     )
 
     executor = get_executor_on_sync_context(
-        rate_limited_sync_evaluate_run,
-        rate_limited_async_evaluate_run,
+        cast("Callable[[object], Any]", rate_limited_sync_evaluate_run),
+        cast(
+            "Callable[[object], Coroutine[Any, Any, Any]]",
+            rate_limited_async_evaluate_run,
+        ),
         max_retries=0,
         exit_on_error=exit_on_error,
         fallback_return_value=None,
@@ -565,14 +624,15 @@ def evaluate_experiment(
         timeout=timeout,
     )
     eval_runs, _ = executor.run(evaluation_input)
-    return eval_runs
+    # Cast: run returns list[Unset | object], but sync/async_eval_run guarantee ExperimentEvaluationRun
+    return cast("list[ExperimentEvaluationRun]", eval_runs)
 
 
 def _add_metadata_to_output_df(
     output_df: pd.DataFrame,
     eval_runs: list[ExperimentEvaluationRun],
     evaluator_name: str,
-) -> object:
+) -> pd.DataFrame:
     for eval_run in eval_runs:
         if eval_run.result is None:
             continue
@@ -603,7 +663,9 @@ def _dataframe_to_examples(dataset: pd.DataFrame) -> list[Example]:
     examples = []
 
     for _, row in dataset.iterrows():
-        example = Example(dataset_row=row.to_dict())
+        example = Example(
+            dataset_row=cast("Mapping[str, JSONSerializable]", row.to_dict())
+        )
         examples.append(example)
     return examples
 
@@ -770,7 +832,8 @@ def get_result_attr(r: object, attr: str, default: object = None) -> object:
     Returns:
         The attribute value if found, otherwise the default value.
     """
-    return getattr(r.result, attr, default) if r.result else default
+    # Type ignore: r typed as object but expected to have result attribute at runtime
+    return getattr(r.result, attr, default) if r.result else default  # type: ignore[attr-defined]
 
 
 def transform_to_experiment_format(

arize/experiments/tracing.py

@@ -57,7 +57,10 @@ _ACTIVE_MODIFIER: ContextVar[SpanModifier | None] = ContextVar(
 
 
 def override_span(
-    init: Callable[..., None], span: ReadableSpan, args: object, kwargs: object
+    init: Callable[..., None],
+    span: ReadableSpan,
+    args: tuple[object, ...],
+    kwargs: dict[str, object],
 ) -> None:
     """Override span initialization to apply active span modifiers.
 

arize/experiments/types.py

@@ -11,6 +11,7 @@ from datetime import datetime, timezone
 from importlib.metadata import version
 from random import getrandbits
 from typing import (
+    NoReturn,
     cast,
 )
 
@@ -90,11 +91,13 @@ class Example:
     def from_dict(cls, obj: Mapping[str, object]) -> Example:
         """Create an Example instance from a dictionary."""
         return cls(
-            id=obj["id"],
-            input=obj["input"],
-            output=obj["output"],
-            metadata=obj.get("metadata") or {},
-            updated_at=obj["updated_at"],
+            id=cast("str", obj["id"]),
+            input=cast("Mapping[str, JSONSerializable]", obj["input"]),
+            output=cast("Mapping[str, JSONSerializable]", obj["output"]),
+            metadata=cast(
+                "Mapping[str, JSONSerializable]", obj.get("metadata") or {}
+            ),
+            updated_at=cast("datetime", obj["updated_at"]),
         )
 
     def __repr__(self) -> str:
@@ -148,7 +151,7 @@ def _make_read_only(
     return obj
 
 
-class _ReadOnly(ObjectProxy):  # type: ignore[misc]
+class _ReadOnly(ObjectProxy):
     def __setitem__(self, *args: object, **kwargs: object) -> object:
         raise NotImplementedError
 
@@ -227,15 +230,15 @@ class ExperimentRun:
     def from_dict(cls, obj: Mapping[str, object]) -> ExperimentRun:
         """Create an ExperimentRun instance from a dictionary."""
         return cls(
-            start_time=obj["start_time"],
-            end_time=obj["end_time"],
-            experiment_id=obj["experiment_id"],
-            dataset_example_id=obj["dataset_example_id"],
-            repetition_number=obj.get("repetition_number") or 1,
-            output=_make_read_only(obj.get("output")),
-            error=obj.get("error"),
-            id=obj["id"],
-            trace_id=obj.get("trace_id"),
+            start_time=cast("datetime", obj["start_time"]),
+            end_time=cast("datetime", obj["end_time"]),
+            experiment_id=cast("str", obj["experiment_id"]),
+            dataset_example_id=cast("str", obj["dataset_example_id"]),
+            repetition_number=cast("int", obj.get("repetition_number") or 1),
+            output=cast("JSONSerializable", _make_read_only(obj.get("output"))),
+            error=cast("str | None", obj.get("error")),
+            id=cast("str", obj["id"]),
+            trace_id=cast("str | None", obj.get("trace_id")),
         )
 
     def __post_init__(self) -> None:
@@ -280,15 +283,17 @@ class ExperimentEvaluationRun:
     def from_dict(cls, obj: Mapping[str, object]) -> ExperimentEvaluationRun:
         """Create an ExperimentEvaluationRun instance from a dictionary."""
         return cls(
-            experiment_run_id=obj["experiment_run_id"],
-            start_time=obj["start_time"],
-            end_time=obj["end_time"],
-            name=obj["name"],
-            annotator_kind=obj["annotator_kind"],
-            error=obj.get("error"),
-            result=EvaluationResult.from_dict(obj.get("result")),
-            id=obj["id"],
-            trace_id=obj.get("trace_id"),
+            experiment_run_id=cast("str", obj["experiment_run_id"]),
+            start_time=cast("datetime", obj["start_time"]),
+            end_time=cast("datetime", obj["end_time"]),
+            name=cast("str", obj["name"]),
+            annotator_kind=cast("str", obj["annotator_kind"]),
+            error=cast("str | None", obj.get("error")),
+            result=EvaluationResult.from_dict(
+                cast("Mapping[str, object] | None", obj.get("result"))
+            ),
+            id=cast("str", obj["id"]),
+            trace_id=cast("str | None", obj.get("trace_id")),
         )
 
     def __post_init__(self) -> None:
@@ -334,7 +339,7 @@ class _HasStats:
             text = self.stats.__str__()
         else:
             text = self.stats.to_markdown(index=False)
-        return f"{self.title}\n{'-' * len(self.title)}\n" + text  # type: ignore
+        return f"{self.title}\n{'-' * len(self.title)}\n" + text
 
 
 @dataclass(frozen=True)
@@ -378,7 +383,7 @@ class _TaskSummary(_HasStats):
         return summary
 
     @classmethod
-    def __new__(cls, *args: object, **kwargs: object) -> object:
+    def __new__(cls, *args: object, **kwargs: object) -> NoReturn:
         # Direct instantiation by users is discouraged.
         raise NotImplementedError
 

arize/logging.py

@@ -6,8 +6,11 @@ import json
 import logging
 import os
 import sys
-from collections.abc import Iterable, Mapping
-from typing import Any, ClassVar
+from collections.abc import Iterable, Mapping, MutableMapping
+from typing import TYPE_CHECKING, Any, ClassVar
+
+if TYPE_CHECKING:
+    import requests
 
 from arize.config import _parse_bool
 from arize.constants.config import (
@@ -34,9 +37,18 @@ class CtxAdapter(logging.LoggerAdapter):
     """LoggerAdapter that merges bound context with per-call extras safely."""
 
     def process(
-        self, msg: object, kwargs: dict[str, object]
-    ) -> tuple[object, dict[str, object]]:
-        """Process the logging call by merging bound and call extras."""
+        self, msg: object, kwargs: MutableMapping[str, Any]
+    ) -> tuple[object, MutableMapping[str, Any]]:
+        """Process the logging call by merging bound and call extras.
+
+        Args:
+            msg: The log message to process.
+            kwargs: Keyword arguments from the logging call, may include 'extra' dict.
+
+        Returns:
+            tuple[object, dict[str, object]]: A tuple of (message, modified_kwargs) with
+                merged extra context.
+        """
         call_extra = _coerce_mapping(kwargs.pop("extra", None))
         bound_extra = _coerce_mapping(self.extra)
         merged = (
@@ -49,13 +61,24 @@ class CtxAdapter(logging.LoggerAdapter):
         return msg, kwargs
 
     def with_extra(self, **more: object) -> CtxAdapter:
-        """Return a copy of this adapter with additional bound extras."""
+        """Return a copy of this adapter with additional bound extras.
+
+        Args:
+            **more: Additional key-value pairs to merge into the bound extras.
+
+        Returns:
+            CtxAdapter: A new adapter instance with merged extra context.
+        """
         base = _coerce_mapping(self.extra)
         base.update(_coerce_mapping(more))
         return type(self)(self.logger, base)
 
     def without_extra(self) -> CtxAdapter:
-        """Return a copy of this adapter with *no* bound extras."""
+        """Return a copy of this adapter with *no* bound extras.
+
+        Returns:
+            CtxAdapter: A new adapter instance without any bound extra context.
+        """
         return type(self)(self.logger, None)
 
 
@@ -86,7 +109,14 @@ class CustomLogFormatter(logging.Formatter):
         super().__init__(fmt=fmt)
 
     def format(self, record: logging.LogRecord) -> str:
-        """Format the log record with color based on log level."""
+        """Format the log record with color based on log level.
+
+        Args:
+            record: The log record to format.
+
+        Returns:
+            str: Formatted and color-coded log message with any extra fields appended.
+        """
         # Build the base message without any color.
         base = super().format(record)
 
@@ -113,31 +143,22 @@ class JsonFormatter(logging.Formatter):
     """Minimal JSON formatter (one JSON object per line)."""
 
     # fields to skip copying from record.__dict__
-    _skip: ClassVar[set[str]] = {
-        # "name",
-        # "msg",
-        # "args",
-        # "levelname",
-        # "levelno",
-        # "pathname",
-        # "filename",
-        # "module",
-        # "exc_info",
-        # "exc_text",
-        # "stack_info",
-        # "lineno",
-        # "funcName",
-        # "created",
-        # "msecs",
-        # "relativeCreated",
-        # "thread",
-        # "threadName",
-        # "processName",
-        # "process",
-    }
+    _skip: ClassVar[set[str]] = set()
+    # Potential fields to skip:
+    # "name", "msg", "args", "levelname", "levelno", "pathname",
+    # "filename", "module", "exc_info", "exc_text", "stack_info",
+    # "lineno", "funcName", "created", "msecs", "relativeCreated",
+    # "thread", "threadName", "processName", "process"
 
     def format(self, record: logging.LogRecord) -> str:
-        """Format the log record as a JSON string."""
+        """Format the log record as a JSON string.
+
+        Args:
+            record: The log record to format.
+
+        Returns:
+            str: JSON-formatted log message as a single line with all fields and extras.
+        """
         payload: dict[str, object] = {
             # "time": self.formatTime(record, datefmt="%Y-%m-%dT%H:%M:%S%z"),
             # "logger": record.name,
@@ -277,7 +298,7 @@ def log_a_list(values: Iterable[Any] | None, join_word: str) -> str:
     )
 
 
-def get_arize_project_url(response: object) -> str:
+def get_arize_project_url(response: requests.Response) -> str:
     """Extract the Arize project URL from an API response.
 
     Args: