arize 8.0.0a2__py3-none-any.whl → 8.0.0a4__py3-none-any.whl

arize/_exporter/client.py

@@ -1,6 +1,4 @@
 # type: ignore[pb2]
-from __future__ import annotations
-
 import logging
 from dataclasses import dataclass
 from datetime import datetime
@@ -14,17 +12,14 @@ from google.protobuf.wrappers_pb2 import Int64Value
 from pyarrow import flight
 from tqdm import tqdm
 
-from arize._exporter.parsers.tracing_data_parser import (
-    OtelTracingDataTransformer,
-)
 from arize._exporter.validation import (
     validate_input_type,
-    validate_input_value,
     validate_start_end_time,
 )
 from arize._generated.protocol.flight import export_pb2
+from arize.logging import CtxAdapter
 from arize.types import Environments, SimilaritySearchParams
-from arize.utils.pandas import reset_dataframe_index
+from arize.utils.dataframe import reset_dataframe_index
 from arize.utils.proto import get_pb_flight_doput_request
 
 logger = logging.getLogger(__name__)
@@ -124,6 +119,10 @@ class ArizeExportClient:
         df.drop(null_columns, axis=1, inplace=True)
 
         if environment == Environments.TRACING:
+            from arize._exporter.parsers.tracing_data_parser import (
+                OtelTracingDataTransformer,
+            )
+
             # by default, transform the exported tracing data so that it's
             # easier to work with in Python
             df = OtelTracingDataTransformer().transform(df)
@@ -238,6 +237,27 @@ class ArizeExportClient:
         columns: List | None = None,
         stream_chunk_size: int | None = None,
     ) -> Tuple[flight.FlightStreamReader, int]:
+        # Bind common context for this operation
+        log = CtxAdapter(
+            logger,
+            {
+                "component": "exporter",
+                "operation": "export_to_df",
+                "space_id": space_id,
+                "model_id": model_id,
+                "environment": environment.name,
+                "model_version": model_version,
+                "batch_id": batch_id,
+                "include_actuals": include_actuals,
+                "where": where,
+                "columns": columns,
+                "similarity_search_params": similarity_search_params,
+                "stream_chunk_size": stream_chunk_size,
+                "start_time": start_time,
+                "end_time": end_time,
+            },
+        )
+        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)
@@ -281,7 +301,7 @@ class ArizeExportClient:
                     json_format.MessageToJson(query_descriptor)  # type: ignore
                 ),
             )
-            logger.debug("Fetching data...")
+            logger.info("Fetching data...")
 
             if flight_info.total_records == 0:
                 logger.warning("Query returns no data")

arize/_exporter/parsers/tracing_data_parser.py

@@ -1,9 +1,8 @@
+from __future__ import annotations
+
 import json
 import logging
-from typing import List
-
-import numpy as np
-import pandas as pd
+from typing import TYPE_CHECKING, List
 
 from arize.spans.columns import (
     SPAN_ATTRIBUTES_EMBEDDING_EMBEDDINGS_COL,
@@ -19,6 +18,10 @@ from arize.spans.columns import (
     SPAN_START_TIME_COL,
 )
 
+if TYPE_CHECKING:
+    import numpy as np
+    import pandas as pd
+
 logger = logging.getLogger(__name__)
 
 

arize/_exporter/validation.py

@@ -1,5 +1,9 @@
-from datetime import datetime
-from typing import Any
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from datetime import datetime
 
 
 def validate_input_type(
@@ -35,6 +39,6 @@ def validate_input_value(
     )
 
 
-def validate_start_end_time(start_time, end_time: datetime) -> None:
+def validate_start_end_time(start_time: datetime, end_time: datetime) -> None:
     if start_time >= end_time:
         raise ValueError("start_time must be before end_time")

arize/_flight/client.py

@@ -2,29 +2,28 @@
 from __future__ import annotations
 
 import base64
+import logging
 from dataclasses import dataclass, field
-from typing import Any, Dict, List, Tuple
+from typing import TYPE_CHECKING, Any, Dict, List, Tuple
 
-import pyarrow as pa
 from google.protobuf import json_format
 from pyarrow import flight
 
-from arize._flight.types import FlightRequestType
 from arize._generated.protocol.flight.ingest_pb2 import (
     WriteSpanAnnotationResponse,
     WriteSpanAttributesMetadataResponse,
     WriteSpanEvaluationResponse,
 )
 from arize.config import get_python_version
-from arize.constants.config import (
-    DEFAULT_FLIGHT_HOST,
-    DEFAULT_FLIGHT_PORT,
-    DEFAULT_FLIGHT_TRANSPORT_SCHEME,
-)
 from arize.logging import log_a_list
 from arize.utils.proto import get_pb_flight_doput_request, get_pb_schema_tracing
 from arize.version import __version__
 
+if TYPE_CHECKING:
+    import pyarrow as pa
+
+    from arize._flight.client import FlightRequestType
+
 BytesPair = Tuple[bytes, bytes]
 Headers = List[BytesPair]
 WriteSpanResponse = (
@@ -32,7 +31,6 @@ WriteSpanResponse = (
     | WriteSpanAnnotationResponse
     | WriteSpanAttributesMetadataResponse
 )
-import logging
 
 logger = logging.getLogger(__name__)
 
@@ -40,10 +38,10 @@ logger = logging.getLogger(__name__)
 @dataclass(frozen=True)
 class ArizeFlightClient:
     api_key: str = field(repr=False)
-    host: str = DEFAULT_FLIGHT_HOST
-    port: int = DEFAULT_FLIGHT_PORT
-    scheme: str = DEFAULT_FLIGHT_TRANSPORT_SCHEME
-    request_verify: bool = True
+    host: str
+    port: int
+    scheme: str
+    request_verify: bool
 
     # internal cache for the underlying FlightClient
     _client: flight.FlightClient | None = field(
@@ -54,7 +52,6 @@ class ArizeFlightClient:
 
     @property
     def headers(self) -> Headers:
-        # Keep the typing simple: (bytes, bytes)
         return [
             (b"origin", b"arize-logging-client"),
             (b"auth-token-bin", str(self.api_key).encode("utf-8")),

arize/_lazy.py

@@ -1,44 +1,13 @@
+# src/arize/_lazy.py
 from __future__ import annotations
 
 import sys
 import threading
 from importlib import import_module
-from typing import Any, Dict, Tuple
+from typing import TYPE_CHECKING, Any, Dict, Tuple
 
-from arize.config import SDKConfiguration
-
-
-class OptionalDependencyError(ImportError): ...
-
-
-def _ensure_imports(
-    extra_key: str | None,
-    required: Tuple[str, ...],
-    pkgname="arize",
-):
-    if not required:
-        return
-    missing = []
-    for p in required:
-        try:
-            import_module(p)
-        except Exception:
-            missing.append(p)
-    if missing:
-        raise OptionalDependencyError(
-            f"Missing optional dependencies: {', '.join(missing)}. "
-            f"Install via: pip install {pkgname}[{extra_key}]"
-        )
-
-
-def _dynamic_import(modname: str, retries: int = 2):
-    for i in range(retries):
-        try:
-            return import_module(modname)
-        except (ModuleNotFoundError, ImportError, KeyError):
-            sys.modules.pop(modname, None)
-            if i + 1 == retries:
-                raise
+if TYPE_CHECKING:
+    from arize.config import SDKConfiguration
 
 
 class LazySubclientsMixin:
@@ -63,7 +32,7 @@ class LazySubclientsMixin:
 
             module_path, class_name = subs[name]
             extra_key, required = self._EXTRAS.get(name, (None, ()))
-            _ensure_imports(extra_key, required, pkgname="arize")
+            require(extra_key, required)
 
             module = _dynamic_import(module_path)
             klass = getattr(module, class_name)
@@ -79,3 +48,36 @@ class LazySubclientsMixin:
 
     def __dir__(self):
         return sorted({*super().__dir__(), *self._SUBCLIENTS.keys()})
+
+
+class OptionalDependencyError(ImportError): ...
+
+
+def require(
+    extra_key: str | None,
+    required: Tuple[str, ...],
+    pkgname="arize",
+):
+    if not required:
+        return
+    missing = []
+    for p in required:
+        try:
+            import_module(p)
+        except Exception:
+            missing.append(p)
+    if missing:
+        raise OptionalDependencyError(
+            f"Missing optional dependencies: {', '.join(missing)}. "
+            f"Install via: pip install {pkgname}[{extra_key}]"
+        )
+
+
+def _dynamic_import(modname: str, retries: int = 2):
+    for i in range(retries):
+        try:
+            return import_module(modname)
+        except (ModuleNotFoundError, ImportError, KeyError):
+            sys.modules.pop(modname, None)
+            if i + 1 == retries:
+                raise

arize/client.py

@@ -1,4 +1,3 @@
-# src/arize/client.py
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
@@ -9,13 +8,16 @@ from arize.config import SDKConfiguration
 if TYPE_CHECKING:
     from arize.datasets.client import DatasetsClient
     from arize.experiments.client import ExperimentsClient
+    from arize.models.client import MLModelsClient
     from arize.spans.client import SpansClient
 
 
-# TODO(Kiko): Root client must have all options that the configuration has
-# TODO(Kiko): ArizeFlightClient options
+# TODO(Kiko): Go through main APIs and add CtxAdapter where missing
+# TODO(Kiko): Search and handle other TODOs
+# TODO(Kiko): Go over **every file** and do not import anything at runtime, use `if TYPE_CHECKING`
+# with `from __future__ import annotations` (must include for Python < 3.11)
+# TODO(Kiko): MIMIC Explainer not done
 # TODO(Kiko): Go over docstrings
-# TODO(Kiko): Missing a __repr__ method
 class ArizeClient(LazySubclientsMixin):
     """
     Root client for the Arize SDK. All parameters are optional. If not provided, they will be read
@@ -42,6 +44,7 @@ class ArizeClient(LazySubclientsMixin):
         "datasets": ("arize.datasets.client", "DatasetsClient"),
         "experiments": ("arize.experiments.client", "ExperimentsClient"),
         "spans": ("arize.spans.client", "SpansClient"),
+        "models": ("arize.models.client", "MLModelsClient"),
     }
     _EXTRAS = {
         # Gate only the generated-backed ones
@@ -65,6 +68,10 @@ class ArizeClient(LazySubclientsMixin):
                 "tqdm",
             ),
         ),
+        # Imports are gated in each method of the models client
+        # This is to allow for very lean package install if people only
+        # want to stream ML records
+        "models": (None, ()),
     }
 
     def __init__(
@@ -76,6 +83,8 @@ class ArizeClient(LazySubclientsMixin):
         flight_server_port: int | None = None,
         flight_scheme: str | None = None,
         request_verify: bool | None = None,
+        stream_max_workers: int | None = None,
+        stream_max_queue_bound: int | None = None,
     ):
         cfg_kwargs: dict = {}
         if api_key is not None:
@@ -92,6 +101,10 @@ class ArizeClient(LazySubclientsMixin):
             cfg_kwargs["flight_scheme"] = flight_scheme
         if request_verify is not None:
             cfg_kwargs["request_verify"] = request_verify
+        if stream_max_workers is not None:
+            cfg_kwargs["stream_max_workers"] = stream_max_workers
+        if stream_max_queue_bound is not None:
+            cfg_kwargs["stream_max_queue_bound"] = stream_max_queue_bound
 
         # Only the explicitly provided fields are passed; the rest use
         # SDKConfiguration’s default factories / defaults.
@@ -109,3 +122,22 @@ class ArizeClient(LazySubclientsMixin):
     @property
     def spans(self) -> SpansClient:
         return self.__getattr__("spans")
+
+    @property
+    def models(self) -> MLModelsClient:
+        return self.__getattr__("models")
+
+    def __repr__(self) -> str:
+        lines = [f"{self.__class__.__name__}("]
+        # Indent the SDKConfiguration repr
+        cfg_repr = repr(self.sdk_config).splitlines()
+        lines.append(f"  sdk_config={cfg_repr[0]}")
+        lines.extend("  " + line for line in cfg_repr[1:])
+        # Add subclient states
+        lines.append("  subclients={")
+        for name in self._SUBCLIENTS:
+            state = "loaded" if name in self._lazy_cache else "lazy"
+            lines.append(f"    {name!r}: {state},")
+        lines.append("  }")
+        lines.append(")")
+        return "\n".join(lines)

arize/config.py

@@ -11,6 +11,8 @@ from arize.constants.config import (
     DEFAULT_FLIGHT_PORT,
     DEFAULT_FLIGHT_TRANSPORT_SCHEME,
     DEFAULT_REQUEST_VERIFY,
+    DEFAULT_STREAM_MAX_QUEUE_BOUND,
+    DEFAULT_STREAM_MAX_WORKERS,
     ENV_API_HOST,
     ENV_API_INSECURE,
     ENV_API_KEY,
@@ -18,6 +20,8 @@ from arize.constants.config import (
     ENV_FLIGHT_PORT,
     ENV_FLIGHT_TRANSPORT_SCHEME,
     ENV_REQUEST_VERIFY,
+    ENV_STREAM_MAX_QUEUE_BOUND,
+    ENV_STREAM_MAX_WORKERS,
 )
 from arize.exceptions.auth import MissingAPIKeyError
 from arize.version import __version__
@@ -62,6 +66,16 @@ def _verify_factory() -> bool:
     return _parse_bool(os.getenv(ENV_REQUEST_VERIFY, DEFAULT_REQUEST_VERIFY))
 
 
+def _stream_max_workers_factory() -> int:
+    return int(os.getenv(ENV_STREAM_MAX_WORKERS, DEFAULT_STREAM_MAX_WORKERS))
+
+
+def _stream_max_queue_bound_factory() -> int:
+    return int(
+        os.getenv(ENV_STREAM_MAX_QUEUE_BOUND, DEFAULT_STREAM_MAX_QUEUE_BOUND)
+    )
+
+
 def _mask_secret(secret: str, N: int = 4) -> str:
     """Show first N chars then '***'; empty string if empty."""
     return f"{secret[:N]}***"
@@ -80,6 +94,10 @@ class SDKConfiguration:
     flight_server_port: int = field(default_factory=_flight_port_factory)
     flight_scheme: str = field(default_factory=_flight_scheme_factory)
     request_verify: bool = field(default_factory=_verify_factory)
+    stream_max_workers: int = field(default_factory=_stream_max_workers_factory)
+    stream_max_queue_bound: int = field(
+        default_factory=_stream_max_queue_bound_factory
+    )
 
     # Private, excluded from comparisons & repr
     _headers: Dict[str, str] = field(init=False, repr=False, compare=False)
@@ -97,6 +115,10 @@ class SDKConfiguration:
     def files_url(self) -> str:
         return _endpoint(self.api_scheme, self.api_host, "/v1/pandas_arrow")
 
+    @property
+    def records_url(self) -> str:
+        return _endpoint(self.api_scheme, self.api_host, "/v1/log")
+
     @property
     def headers(self) -> Dict[str, str]:
         # Create base headers
@@ -110,17 +132,29 @@ class SDKConfiguration:
             # "sync": "0",  # Defaults to async logging
         }
 
+    @property
+    def headers_grpc(self) -> Dict[str, str]:
+        return {
+            "authorization": self.api_key,
+            "Grpc-Metadata-sdk-language": "python",
+            "Grpc-Metadata-language-version": get_python_version(),
+            "Grpc-Metadata-sdk-version": __version__,
+            # "Grpc-Metadata-arize-space-id": space_id,
+            # "Grpc-Metadata-arize-interface": "stream",
+        }
+
     def __repr__(self) -> str:
         # Dynamically build repr for all fields
-        parts = []
+        lines = [f"{self.__class__.__name__}("]
         for f in fields(self):
             if not f.repr:
                 continue
             val = getattr(self, f.name)
             if f.name == "api_key":
                 val = _mask_secret(val, 6)
-            parts.append(f"{f.name}={val!r}")
-        return f"{self.__class__.__name__}({', '.join(parts)})"
+            lines.append(f"  {f.name}={val!r},")
+        lines.append(")")
+        return "\n".join(lines)
 
     # TODO(Kiko): This may not be well placed in this class
     def get_generated_client(self):

arize/constants/config.py

@@ -17,6 +17,12 @@ DEFAULT_FLIGHT_PORT = 443
 DEFAULT_FLIGHT_TRANSPORT_SCHEME = "grpc+tls"
 DEFAULT_REQUEST_VERIFY = True
 
+# ML Streaming configuration
+ENV_STREAM_MAX_WORKERS = "ARIZE_STREAM_MAX_WORKERS"
+ENV_STREAM_MAX_QUEUE_BOUND = "ARIZE_STREAM_MAX_QUEUE_BOUND"
+DEFAULT_STREAM_MAX_WORKERS = 8
+DEFAULT_STREAM_MAX_QUEUE_BOUND = 5000
+
 # Logging
 ENV_LOG_ENABLE = "ARIZE_LOG_ENABLE"
 ENV_LOG_LEVEL = "ARIZE_LOG_LEVEL"

arize/constants/ml.py

@@ -1,28 +1,31 @@
+import json
+from pathlib import Path
+
 # MAX_BYTES_PER_BULK_RECORD = 100000
 # MAX_DAYS_WITHIN_RANGE = 365
-# MIN_PREDICTION_ID_LEN = 1
-# MAX_PREDICTION_ID_LEN = 512
-# MIN_DOCUMENT_ID_LEN = 1
-# MAX_DOCUMENT_ID_LEN = 128
+MIN_PREDICTION_ID_LEN = 1
+MAX_PREDICTION_ID_LEN = 512
+MIN_DOCUMENT_ID_LEN = 1
+MAX_DOCUMENT_ID_LEN = 128
 # # The maximum number of character for tag values
 MAX_TAG_LENGTH = 20_000
-# MAX_TAG_LENGTH_TRUNCATION = 1_000
+MAX_TAG_LENGTH_TRUNCATION = 1_000
 # # The maximum number of character for embedding raw data
 MAX_RAW_DATA_CHARACTERS = 2_000_000
-# MAX_RAW_DATA_CHARACTERS_TRUNCATION = 5_000
-# # The maximum number of acceptable years in the past from current time for prediction_timestamps
+MAX_RAW_DATA_CHARACTERS_TRUNCATION = 5_000
+# The maximum number of acceptable years in the past from current time for prediction_timestamps
 MAX_PAST_YEARS_FROM_CURRENT_TIME = 5
-# # The maximum number of acceptable years in the future from current time for prediction_timestamps
+# The maximum number of acceptable years in the future from current time for prediction_timestamps
 MAX_FUTURE_YEARS_FROM_CURRENT_TIME = 1
 # # The maximum number of character for llm model name
-# MAX_LLM_MODEL_NAME_LENGTH = 20_000
-# MAX_LLM_MODEL_NAME_LENGTH_TRUNCATION = 50
+MAX_LLM_MODEL_NAME_LENGTH = 20_000
+MAX_LLM_MODEL_NAME_LENGTH_TRUNCATION = 50
 # # The maximum number of character for prompt template
-# MAX_PROMPT_TEMPLATE_LENGTH = 50_000
-# MAX_PROMPT_TEMPLATE_LENGTH_TRUNCATION = 5_000
+MAX_PROMPT_TEMPLATE_LENGTH = 50_000
+MAX_PROMPT_TEMPLATE_LENGTH_TRUNCATION = 5_000
 # # The maximum number of character for prompt template version
-# MAX_PROMPT_TEMPLATE_VERSION_LENGTH = 20_000
-# MAX_PROMPT_TEMPLATE_VERSION_LENGTH_TRUNCATION = 50
+MAX_PROMPT_TEMPLATE_VERSION_LENGTH = 20_000
+MAX_PROMPT_TEMPLATE_VERSION_LENGTH_TRUNCATION = 50
 # # The maximum number of embeddings
 MAX_NUMBER_OF_EMBEDDINGS = 30
 MAX_EMBEDDING_DIMENSIONALITY = 20_000
@@ -37,21 +40,20 @@ MAX_NUMBER_OF_SIMILARITY_REFERENCES = 10
 # GENERATED_LLM_PARAMS_JSON_COL = "arize_generated_llm_params_json"
 #
 # # reserved columns for LLM run metadata
-# LLM_RUN_METADATA_TOTAL_TOKEN_COUNT_TAG_NAME = "total_token_count"
-# LLM_RUN_METADATA_PROMPT_TOKEN_COUNT_TAG_NAME = "prompt_token_count"
-# LLM_RUN_METADATA_RESPONSE_TOKEN_COUNT_TAG_NAME = "response_token_count"
-# LLM_RUN_METADATA_RESPONSE_LATENCY_MS_TAG_NAME = "response_latency_ms"
-#
-# # all reserved tags
-# RESERVED_TAG_COLS = [
-#     LLM_RUN_METADATA_TOTAL_TOKEN_COUNT_TAG_NAME,
-#     LLM_RUN_METADATA_PROMPT_TOKEN_COUNT_TAG_NAME,
-#     LLM_RUN_METADATA_RESPONSE_TOKEN_COUNT_TAG_NAME,
-#     LLM_RUN_METADATA_RESPONSE_LATENCY_MS_TAG_NAME,
-# ]
-#
-#
-# path = Path(__file__).with_name("model_mapping.json")
-# with path.open("r") as f:
-#     MODEL_MAPPING_CONFIG = json.load(f)
+LLM_RUN_METADATA_TOTAL_TOKEN_COUNT_TAG_NAME = "total_token_count"
+LLM_RUN_METADATA_PROMPT_TOKEN_COUNT_TAG_NAME = "prompt_token_count"
+LLM_RUN_METADATA_RESPONSE_TOKEN_COUNT_TAG_NAME = "response_token_count"
+LLM_RUN_METADATA_RESPONSE_LATENCY_MS_TAG_NAME = "response_latency_ms"
 #
+# all reserved tags
+RESERVED_TAG_COLS = [
+    LLM_RUN_METADATA_TOTAL_TOKEN_COUNT_TAG_NAME,
+    LLM_RUN_METADATA_PROMPT_TOKEN_COUNT_TAG_NAME,
+    LLM_RUN_METADATA_RESPONSE_TOKEN_COUNT_TAG_NAME,
+    LLM_RUN_METADATA_RESPONSE_LATENCY_MS_TAG_NAME,
+]
+
+
+path = Path(__file__).with_name("model_mapping.json")
+with path.open("r") as f:
+    MODEL_MAPPING_CONFIG = json.load(f)