arize 8.0.0b2__tar.gz → 8.0.0b4__tar.gz

{arize-8.0.0b2 → arize-8.0.0b4}/.gitignore

@@ -41,3 +41,4 @@ docs/_build
 output/
 
 .workon
+.coverage.*

{arize-8.0.0b2 → arize-8.0.0b4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arize
-Version: 8.0.0b2
+Version: 8.0.0b4
 Summary: A helper library to interact with Arize AI APIs
 Project-URL: Homepage, https://arize.com
 Project-URL: Documentation, https://docs.arize.com/arize
@@ -44,8 +44,16 @@ Requires-Dist: typing-extensions<5,>=4.7.1
 Requires-Dist: urllib3<3,>=2.1.0
 Requires-Dist: wrapt<2.0.0,>=1.0.0
 Provides-Extra: dev
+Requires-Dist: mypy==1.19.1; extra == 'dev'
+Requires-Dist: pandas-stubs>=2.2.0; extra == 'dev'
+Requires-Dist: pytest-cov==6.0.0; extra == 'dev'
 Requires-Dist: pytest==8.4.2; extra == 'dev'
-Requires-Dist: ruff==0.13.2; extra == 'dev'
+Requires-Dist: ruff==0.14.9; extra == 'dev'
+Requires-Dist: taskipy<2,>=1.14.1; extra == 'dev'
+Requires-Dist: types-python-dateutil>=2.9.0; extra == 'dev'
+Requires-Dist: types-requests>=2.31.0; extra == 'dev'
+Requires-Dist: types-tabulate>=0.9.0; extra == 'dev'
+Requires-Dist: types-tqdm>=4.66.0; extra == 'dev'
 Provides-Extra: embeddings
 Requires-Dist: datasets!=2.14.*,<3,>=2.8; extra == 'embeddings'
 Requires-Dist: pillow<11,>=8.4.0; extra == 'embeddings'

arize-8.0.0b4/pyproject.toml

@@ -0,0 +1,290 @@
+[project]
+name = "arize"
+description = "A helper library to interact with Arize AI APIs"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { text = "Apache-2.0" }
+license-files = ["LICENSE", "NOTICE"]
+keywords = [
+    "Arize",
+    "Observability",
+    "Monitoring",
+    "Explainability",
+    "Tracing",
+    "LLM",
+    "Evaluations",
+]
+authors = [{ name = "Arize AI", email = "support@arize.com" }]
+maintainers = [{ name = "Arize AI", email = "support@arize.com" }]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Programming Language :: Python",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
+    "Intended Audience :: Developers",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+    "Topic :: System :: Logging",
+    "Topic :: System :: Monitoring",
+]
+dependencies = [
+    "numpy>=2.0.0",
+    "openinference-semantic-conventions>=0.1.25, <1",
+    "opentelemetry-exporter-otlp-proto-common>=1.38.0",
+    "opentelemetry-exporter-otlp-proto-grpc>=1.38.0",
+    "opentelemetry-sdk>=1.38.0",
+    "opentelemetry-semantic-conventions>=0.43b0, <1",
+    "pandas>=2.0.0, <3",
+    "protobuf>=4.21.0, <6",
+    "pyarrow>=0.15.0",
+    "pydantic>=2, <3",
+    "python-dateutil>=2.8.2, <3",
+    "requests>=2.0.0, <3",
+    "requests_futures>=1.0.0, <2",
+    "tqdm>4, <5",                                       # For progress bars
+    "typing-extensions>=4.7.1, <5",
+    "urllib3>=2.1.0, <3",
+    "wrapt>=1.0.0, <2.0.0",
+]
+dynamic = ["version"]
+
+[project.optional-dependencies]
+dev = [
+    "mypy==1.19.1",
+    "pandas-stubs>=2.2.0",
+    "pytest-cov==6.0.0",
+    "pytest==8.4.2",
+    "ruff==0.14.9",
+    "taskipy>=1.14.1, <2",
+    "types-python-dateutil>=2.9.0",
+    "types-requests>=2.31.0",
+    "types-tabulate>=0.9.0",
+    "types-tqdm>=4.66.0",
+]
+otel = [
+    "arize-otel>=0.11.0, <1",
+]
+mimic = [
+    "interpret-community[mimic]>=0.22.0,<1",
+]
+embeddings = [
+    "Pillow>=8.4.0, <11",
+    "datasets>=2.8, <3, !=2.14.*",
+    "tokenizers>=0.13, <1",
+    "torch>=1.13, <3",
+    "transformers>=4.25, <5",
+]
+
+[project.urls]
+Homepage = "https://arize.com"
+Documentation = "https://docs.arize.com/arize"
+Issues = "https://github.com/Arize-ai/client_python/issues"
+Source = "https://github.com/Arize-ai/client_python"
+Changelog = "https://github.com/Arize-ai/client_python/blob/main/CHANGELOG.md"
+
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.version]
+path = "src/arize/version.py"
+
+[tool.hatch.build]
+only-packages = true
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/arize"]
+
+[tool.hatch.build.targets.sdist]
+exclude = ["tests", "docs"]
+
+
+[tool.black]
+include = '\.pyi?$'
+exclude = '(_pb2\.py$|docs/source/.*\.py)'
+
+[tool.ruff]
+target-version = "py310"
+line-length = 80
+exclude = [
+    "dist/",
+    "testing_notebooks/",
+    "__pycache__",
+    "src/arize/_generated/**/*.py",
+    "*_pb2.py*",
+    "*_pb2_grpc.py*",
+    "*.pyi",
+    "docs/",
+]
+[tool.ruff.format]
+docstring-code-format = true
+line-ending = "native"
+
+[tool.ruff.lint]
+select = [
+    # Core correctness & style
+    # ------------------------
+    "E",   # pycodestyle errors (syntax, indentation, whitespace)
+    "W",   # pycodestyle warnings (less severe style issues)
+    "F",   # Pyflakes (unused imports, undefined names, etc.)
+    "UP",  # pyupgrade (modernize syntax for your Python version)
+    "B",   # flake8-bugbear (likely bugs and design problems)
+    "SIM", # flake8-simplify (simpler / clearer code patterns)
+    "I",   # isort (import ordering)
+
+    # Typing & documentation
+    # ------------------------
+    "ANN", # flake8-annotations (enforce type annotations)
+    "D",   # pydocstyle (docstring conventions for public APIs)
+    "TCH", # flake8-type-checking (TYPE_CHECKING import hygiene)
+
+    # Recommended for SDKs
+    # ------------------------
+    "RUF",  # Ruff-native rules (high-signal correctness & footguns)
+    "C4",   # flake8-comprehensions (cleaner comprehensions)
+    "PIE",  # flake8-pie (small correctness & readability improvements)
+    "PERF", # perflint (common performance pitfalls)
+    "DTZ",  # flake8-datetimez (timezone-aware datetime usage)
+    "TRY",  # tryceratops (better exception handling patterns)
+    "RET",  # flake8-return (return consistency)
+    "S",    # flake8-bandit (basic security issues; tune via ignores)
+]
+ignore = [
+    "TRY003", # Put long / detailed messages inside the exception class
+]
+
+[tool.ruff.lint.per-file-ignores]
+"tests/**/*.py" = [
+    "D104",   # Missing docstring in public package
+    "DTZ001", # Allow naive datetime usage in tests
+    "S101",   # Allow assert statements in tests
+    "S108",   # Allow insecure usage of temporary file or directory
+    "TRY301", # Allow generic exception catching in tests
+]
+
+[tool.ruff.lint.isort]
+force-wrap-aliases = true
+
+[tool.ruff.lint.pycodestyle]
+max-doc-length = 110
+max-line-length = 110
+
+[tool.ruff.lint.pydocstyle]
+convention = "google"
+
+[tool.ruff.lint.pyupgrade]
+keep-runtime-typing = true
+
+[tool.mypy]
+python_version = "3.10"
+warn_unused_configs = true
+exclude = ["testing_notebooks/", "src/arize/_generated/", "tests/"]
+# Incrementally enable strict type checking by uncommenting checks from bottom to top.
+# Fix foundational issues (imports, syntax) before type definitions, then signatures, then implementations.
+disable_error_code = []
+
+# Ignore missing imports for libraries without type stubs
+[[tool.mypy.overrides]]
+module = [
+    "pyarrow.*",
+    "google.protobuf.*",
+    "wrapt.*",
+    "datasets.*",
+    "requests_futures.*",
+    "interpret_community.*",
+    "sklearn.*",
+]
+ignore_missing_imports = true
+
+[tool.pytest.ini_options]
+# Directory where pytest will search for tests
+testpaths = ["tests"]
+# Patterns for discovering test files (files matching these patterns will be collected)
+python_files = ["test_*.py", "*_test.py"]
+# Patterns for discovering test classes (classes matching this pattern will be collected)
+python_classes = ["Test*"]
+# Patterns for discovering test functions (functions matching this pattern will be collected)
+python_functions = ["test_*"]
+# Directories and patterns to ignore during test collection
+norecursedirs = [
+    "testing_notebooks",
+    ".*",
+    "build",
+    "dist",
+    "*.egg"
+]
+# Additional command-line options to always apply when running pytest
+addopts = [
+    "-v",               # Verbose: show individual test names as they run
+    "--strict-markers", # Error on undefined markers (prevents typos like @pytest.mark.slo instead of @pytest.mark.slow)
+    "--tb=short",       # Shorter tracebacks: show only the failure point without full stack
+    "--durations=10",   # Show the 10 slowest tests at the end (helps identify performance bottlenecks)
+]
+# Custom markers that can be used to categorize tests (use with @pytest.mark.marker_name)
+markers = [
+    "slow: marks tests as slow (deselect with '-m \"not slow\"')",
+    "integration: marks tests as integration tests",
+    "unit: marks tests as unit tests",
+]
+
+[tool.coverage.run]
+source = ["src/arize"]
+omit = [
+    "*/tests/*",
+    "*/testing_notebooks/*",
+    "*/_generated/*",
+    "*/__pycache__/*",
+    "*/site-packages/*",
+    # Exclude other src/arize modules. TODO(Kiko): Remove them progressively.
+    "src/arize/constants/*",
+    "src/arize/datasets/*",
+    "src/arize/embeddings/*",
+    "src/arize/exceptions/*",
+    "src/arize/experiments/*",
+    "src/arize/ml/*",
+    "src/arize/projects/*",
+    "src/arize/spans/*",
+    "src/arize/utils/*",
+]
+# Enable branch coverage (measures whether all branches of if/else/try/except are tested)
+# Without this, only line coverage is measured
+branch = true
+
+[tool.coverage.report]
+# Number of decimal places to show in coverage percentages (e.g., 85.47% instead of 85%)
+precision = 2
+# Show line numbers for lines that are missing coverage in the report
+show_missing = true
+# If true, files with 100% coverage would be hidden from the report
+skip_covered = false
+# Minimum coverage percentage required (build fails if total coverage is below this)
+fail_under = 90
+# Lines matching these patterns are excluded from coverage measurement
+exclude_lines = [
+    "pragma: no cover",
+    "def __repr__",
+    "raise AssertionError",
+    "raise NotImplementedError",
+    "if __name__ == .__main__.:",
+    "if TYPE_CHECKING:",
+    "@abstractmethod",
+    "@abc.abstractmethod",
+]
+
+[tool.coverage.html]
+# Directory where HTML coverage reports will be generated (browsable in a web browser)
+directory = "htmlcov"
+
+[tool.taskipy.tasks]
+lint = { cmd = "ruff format . && ruff check --fix .", cwd = ".", help = "Run ruff to format and lint the package." }
+type-check = {cmd = "mypy --no-incremental --show-traceback --config-file pyproject.toml .", cwd = ".", help = "Run mypy for type checking."}
+test = {cmd = "pytest --cov .", cwd = ".", help = "Run the test suite with coverage."}
+# Exclusive for CI
+ci-format = { cmd = "ruff format --check .", cwd = ".", help = "Run ruff to check formatting of the package during CI." }
+ci-lint = { cmd = "ruff check .", cwd = ".", help = "Run ruff to check linting rules of the package during CI." }
+ci-type-check = {cmd = "mypy --no-incremental --show-traceback --config-file pyproject.toml .", cwd = ".", help = "Run mypy for type checking during CI."}
+ci-test = {cmd = "pytest --cov --cov-report=term-missing --cov-report=html .", cwd = ".", help = "Run the test suite with coverage during CI."}

{arize-8.0.0b2 → arize-8.0.0b4}/src/arize/__init__.py

@@ -2,6 +2,7 @@
 
 import logging
 from collections.abc import Mapping
+from typing import Literal, cast
 
 from arize._generated.api_client import models
 from arize.client import ArizeClient
@@ -82,7 +83,11 @@ def make_to_df(field_name: str) -> object:
 
         # Drop None/NaN columns if requested
         if exclude_none in ("any", "all", True):
-            drop_how = "all" if exclude_none is True else exclude_none
+            drop_how: Literal["any", "all"] = (
+                "all"
+                if exclude_none is True
+                else cast("Literal['any', 'all']", exclude_none)
+            )
             df.dropna(axis=1, how=drop_how, inplace=True)
 
         if convert_dtypes:
@@ -92,6 +97,8 @@ def make_to_df(field_name: str) -> object:
     return to_df
 
 
+# Monkey-patch convenience methods onto generated response models
+# Type ignore comments needed: mypy can't verify runtime attribute additions
 models.DatasetsList200Response.to_df = make_to_df("datasets")  # type: ignore[attr-defined]
 models.DatasetsExamplesList200Response.to_df = make_to_df("examples")  # type: ignore[attr-defined]
 models.ExperimentsList200Response.to_df = make_to_df("experiments")  # type: ignore[attr-defined]

{arize-8.0.0b2 → arize-8.0.0b4}/src/arize/_exporter/client.py

@@ -1,4 +1,3 @@
-# type: ignore[pb2]
 import logging
 from dataclasses import dataclass
 from datetime import datetime
@@ -41,7 +40,7 @@ class ArizeExportClient:
         batch_id: str = "",
         include_actuals: bool = False,
         stream_chunk_size: int | None = None,
-    ) -> object:
+    ) -> pd.DataFrame:
         """Exports data of a specific model in the Arize platform to a pandas dataframe.
 
         The export covers a defined time interval and model environment, and can
@@ -237,6 +236,22 @@ class ArizeExportClient:
         columns: list | None = None,
         stream_chunk_size: int | None = None,
     ) -> tuple[flight.FlightStreamReader | None, int]:
+        # Validate inputs first before creating logging context
+        validate_input_type(space_id, "space_id", str)
+        validate_input_type(model_id, "model_id", str)
+        validate_input_type(environment, "environment", Environments)
+        validate_input_type(include_actuals, "include_actuals", bool)
+        validate_input_type(start_time, "start_time", datetime)
+        validate_input_type(end_time, "end_time", datetime)
+        validate_input_type(model_version, "model_version", str)
+        validate_input_type(batch_id, "batch_id", str)
+        validate_input_type(where, "where", str)
+        validate_input_type(columns, "columns", list, allow_none=True)
+        validate_input_type(
+            stream_chunk_size, "stream_chunk_size", int, allow_none=True
+        )
+        validate_start_end_time(start_time, end_time)
+
         # Bind common context for this operation
         log = CtxAdapter(
             logger,
@@ -258,20 +273,6 @@ class ArizeExportClient:
             },
         )
         log.debug("Getting stream reader...")
-        validate_input_type(space_id, "space_id", str)
-        validate_input_type(model_id, "model_id", str)
-        validate_input_type(environment, "environment", Environments)
-        validate_input_type(include_actuals, "include_actuals", bool)
-        validate_input_type(start_time, "start_time", datetime)
-        validate_input_type(end_time, "end_time", datetime)
-        validate_input_type(model_version, "model_version", str)
-        validate_input_type(batch_id, "batch_id", str)
-        validate_input_type(where, "where", str)
-        validate_input_type(columns, "columns", list, allow_none=True)
-        validate_input_type(
-            stream_chunk_size, "stream_chunk_size", int, allow_none=True
-        )
-        validate_start_end_time(start_time, end_time)
 
         # Create query descriptor
         query_descriptor = flight_pb2.RecordQueryDescriptor(
@@ -300,7 +301,7 @@ class ArizeExportClient:
         try:
             flight_info = self.flight_client.get_flight_info(
                 flight.FlightDescriptor.for_command(
-                    json_format.MessageToJson(query_descriptor)  # type: ignore
+                    json_format.MessageToJson(query_descriptor)
                 ),
             )
             logger.info("Fetching data...")

{arize-8.0.0b2 → arize-8.0.0b4}/src/arize/_exporter/parsers/tracing_data_parser.py

@@ -1,5 +1,7 @@
 import json
 import logging
+from collections.abc import Callable
+from typing import Any, TypeGuard
 
 import numpy as np
 import pandas as pd
@@ -28,7 +30,10 @@ logger = logging.getLogger(__name__)
 # of the error is on the data; It should not prevent a user from continuing to use the data
 class OtelTracingDataTransformer:
     def _apply_column_transformation(
-        self, df: pd.DataFrame, col_name: str, transform_func: object
+        self,
+        df: pd.DataFrame,
+        col_name: str,
+        transform_func: Callable[[Any], Any],
     ) -> str | None:
         """Apply a transformation to a column and return error message if it fails."""
         try:
@@ -89,7 +94,7 @@ class OtelTracingDataTransformer:
             if col.name in df.columns
         ]
         for col_name in dirty_string_column_names:
-            df[col_name] = df[col_name].apply(self._clean_json_string)
+            df[col_name] = df[col_name].apply(self._clean_json_string)  # type: ignore[arg-type]
 
         # Convert timestamp columns to datetime objects
         timestamp_column_names: list[str] = [
@@ -102,7 +107,7 @@ class OtelTracingDataTransformer:
         ]
         for col_name in timestamp_column_names:
             df[col_name] = df[col_name].apply(
-                self._convert_timestamp_to_datetime
+                self._convert_timestamp_to_datetime  # type: ignore[arg-type]
             )
 
         for err in errors:
@@ -138,7 +143,7 @@ class OtelTracingDataTransformer:
             return None
         return None
 
-    def _is_non_empty_string(self, value: object) -> bool:
+    def _is_non_empty_string(self, value: object) -> TypeGuard[str]:
         return isinstance(value, str) and value != ""
 
     def _deserialize_json_string_to_dict(self, value: str) -> object:

{arize-8.0.0b2 → arize-8.0.0b4}/src/arize/_exporter/validation.py

@@ -35,7 +35,7 @@ def validate_input_value(
     if input in choices:
         return
     raise ValueError(
-        f"{input_name} is {input}, but must be one of {', '.join(choices)}"
+        f"{input_name} is {input}, but must be one of {', '.join(str(c) for c in choices)}"
     )
 
 

{arize-8.0.0b2 → arize-8.0.0b4}/src/arize/_flight/client.py

@@ -1,11 +1,10 @@
-# type: ignore[pb2]
 from __future__ import annotations
 
 import base64
 import logging
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, TypeAlias
 
 from google.protobuf import json_format
 from pyarrow import flight
@@ -20,15 +19,15 @@ from arize.version import __version__
 
 if TYPE_CHECKING:
     import types
-    from collections.abc import Iterable
+    from collections.abc import Iterator
 
     import pandas as pd
     import pyarrow as pa
 
 
-BytesPair = tuple[bytes, bytes]
-Headers = list[BytesPair]
-FlightPostArrowFileResponse = (
+BytesPair: TypeAlias = tuple[bytes, bytes]
+Headers: TypeAlias = list[BytesPair]
+FlightPostArrowFileResponse: TypeAlias = (
     flight_pb2.WriteSpanEvaluationResponse
     | flight_pb2.WriteSpanAnnotationResponse
     | flight_pb2.WriteSpanAttributesMetadataResponse
@@ -90,7 +89,11 @@ class ArizeFlightClient:
     # ---------- Connection management ----------
 
     def _ensure_client(self) -> flight.FlightClient:
-        """Lazily initialize and return the underlying Flight client connection."""
+        """Lazily initialize and return the underlying Flight client connection.
+
+        Returns:
+            flight.FlightClient: The initialized Apache Arrow Flight client.
+        """
         client = object.__getattribute__(self, "_client")
         if client is not None:
             return client
@@ -135,7 +138,11 @@ class ArizeFlightClient:
     # ---------- methods simple passthrough wrappers ----------
 
     def get_flight_info(self, *args: object, **kwargs: object) -> object:
-        """Get flight information. Passthrough to underlying Flight client with auth options."""
+        """Get flight information. Passthrough to underlying Flight client with auth options.
+
+        Returns:
+            object: FlightInfo object containing metadata about the requested data stream.
+        """
         client = self._ensure_client()
         kwargs.setdefault("options", self.call_options)
         return client.get_flight_info(*args, **kwargs)
@@ -146,6 +153,9 @@ class ArizeFlightClient:
         """Retrieve data stream via Flight DoGet.
 
         Passthrough to underlying Flight client with auth options.
+
+        Returns:
+            flight.FlightStreamReader: A stream reader for retrieving Arrow record batches.
         """
         client = self._ensure_client()
         kwargs.setdefault("options", self.call_options)
@@ -153,10 +163,15 @@ class ArizeFlightClient:
 
     def do_put(
         self, *args: object, **kwargs: object
-    ) -> [flight.FlightStreamWriter, flight.FlightMetadataReader]:
+    ) -> tuple[flight.FlightStreamWriter, flight.FlightMetadataReader]:
         """Upload data stream via Flight DoPut.
 
         Passthrough to underlying Flight client with auth options.
+
+        Returns:
+            tuple[flight.FlightStreamWriter, flight.FlightMetadataReader]: A tuple containing
+                a stream writer for uploading Arrow record batches and a metadata reader for
+                receiving server responses.
         """
         client = self._ensure_client()
         kwargs.setdefault("options", self.call_options)
@@ -164,10 +179,13 @@ class ArizeFlightClient:
 
     def do_action(
         self, *args: object, **kwargs: object
-    ) -> Iterable[flight.Result]:
+    ) -> Iterator[flight.Result]:
         """Execute an action via Flight DoAction.
 
         Passthrough to underlying Flight client with auth options.
+
+        Returns:
+            Iterable[flight.Result]: An iterable of Result objects from the action execution.
         """
         client = self._ensure_client()
         kwargs.setdefault("options", self.call_options)
@@ -210,6 +228,10 @@ class ArizeFlightClient:
             FlightRequestType.ANNOTATION,
             FlightRequestType.METADATA,
         ):
+            if project_name is None:
+                raise ValueError(
+                    f"project_name is required for {request_type.name} request type"
+                )
             proto_schema = get_pb_schema_tracing(project_name=project_name)
             base64_schema = base64.b64encode(proto_schema.SerializeToString())
             pa_schema = append_to_pyarrow_metadata(
@@ -260,8 +282,6 @@ class ArizeFlightClient:
             case FlightRequestType.LOG_EXPERIMENT_DATA:
                 res = flight_pb2.PostExperimentDataResponse()
                 res.ParseFromString(flight_response.to_pybytes())
-            case _:
-                raise ValueError(f"Unsupported request_type: {request_type}")
         return res
 
     # ---------- dataset methods ----------
@@ -271,7 +291,7 @@ class ArizeFlightClient:
         space_id: str,
         dataset_name: str,
         pa_table: pa.Table,
-    ) -> str:
+    ) -> str | None:
         """Create a new dataset via Flight DoPut.
 
         Args:

{arize-8.0.0b2 → arize-8.0.0b4}/src/arize/_lazy.py

@@ -52,7 +52,7 @@ class LazySubclientsMixin:
             # Determine which parameters this subclient needs
             # and build kwargs accordingly
             sig = inspect.signature(klass.__init__)
-            kwargs = {}
+            kwargs: dict[str, object] = {}
             if "sdk_config" in sig.parameters:
                 kwargs["sdk_config"] = self.sdk_config
             if "generated_client" in sig.parameters:
@@ -72,7 +72,14 @@ class OptionalDependencyError(ImportError): ...
 
 
 def _can_import(module_name: str) -> bool:
-    """Check if a module can be imported without raising an exception."""
+    """Check if a module can be imported without raising an exception.
+
+    Args:
+        module_name: The fully qualified module name to check (e.g., 'numpy', 'sklearn.preprocessing').
+
+    Returns:
+        bool: True if the module can be imported successfully, False otherwise.
+    """
     try:
         import_module(module_name)
     except Exception:
@@ -86,6 +93,18 @@ def require(
     required: tuple[str, ...],
     pkgname: str = "arize",
 ) -> None:
+    """Ensure required optional dependencies are installed, raising an error if missing.
+
+    Args:
+        extra_key: The extras group key for pip install (e.g., 'mimic', 'embeddings').
+            Used in the error message to guide users.
+        required: Tuple of required module names to check for availability.
+        pkgname: The package name for installation instructions. Defaults to 'arize'.
+
+    Raises:
+        OptionalDependencyError: If any of the required modules cannot be imported.
+            The error message includes pip install instructions with the extras group.
+    """
     if not required:
         return
     missing = [p for p in required if not _can_import(p)]
@@ -97,6 +116,22 @@ def require(
 
 
 def _dynamic_import(modname: str, retries: int = 2) -> types.ModuleType:
+    """Dynamically import a module with retry logic and sys.modules cleanup on failure.
+
+    Args:
+        modname: The fully qualified module name to import.
+        retries: Number of import attempts to make. Must be > 0. Defaults to 2.
+
+    Returns:
+        types.ModuleType: The successfully imported module.
+
+    Raises:
+        ValueError: If retries is <= 0.
+        ModuleNotFoundError: If the module cannot be found after all retry attempts.
+        ImportError: If the module import fails after all retry attempts.
+        KeyError: If a key error occurs during import after all retry attempts.
+    """
+
     def _attempt_import(remaining_attempts: int) -> types.ModuleType:
         try:
             return import_module(modname)