aeri-python 4.0.0__py3-none-any.whl

aeri/_client/constants.py

@@ -0,0 +1,65 @@
+"""Constants used by the Aeri OpenTelemetry integration.
+
+This module defines constants used throughout the Aeri OpenTelemetry integration.
+"""
+
+from typing import Any, List, Literal, Union, get_args
+
+from typing_extensions import TypeAlias
+
+AERI_TRACER_NAME = "aeri-sdk"
+
+AERI_SDK_EXPERIMENT_ENVIRONMENT = "sdk-experiment"
+
+"""Note: this type is used with .__args__ / get_args in some cases and therefore must remain flat"""
+ObservationTypeGenerationLike: TypeAlias = Literal[
+    "generation",
+    "embedding",
+]
+
+ObservationTypeSpanLike: TypeAlias = Literal[
+    "span",
+    "agent",
+    "tool",
+    "chain",
+    "retriever",
+    "evaluator",
+    "guardrail",
+]
+
+ObservationTypeLiteralNoEvent: TypeAlias = Union[
+    ObservationTypeGenerationLike,
+    ObservationTypeSpanLike,
+]
+
+"""Enumeration of valid observation types for Aeri tracing.
+
+This Literal defines all available observation types that can be used with the @observe
+decorator and other Aeri SDK methods.
+"""
+ObservationTypeLiteral: TypeAlias = Union[
+    ObservationTypeLiteralNoEvent, Literal["event"]
+]
+
+
+def get_observation_types_list(
+    literal_type: Any,
+) -> List[str]:
+    """Flattens the Literal type to provide a list of strings.
+
+    Args:
+        literal_type: A Literal type, TypeAlias, or union of Literals to flatten
+
+    Returns:
+        Flat list of all string values contained in the Literal type
+    """
+    result = []
+    args = get_args(literal_type)
+
+    for arg in args:
+        if hasattr(arg, "__args__"):
+            result.extend(get_observation_types_list(arg))
+        else:
+            result.append(arg)
+
+    return result

aeri/_client/datasets.py

@@ -0,0 +1,302 @@
+import datetime as dt
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+from aeri.api import (
+    Dataset,
+    DatasetItem,
+)
+from aeri.batch_evaluation import CompositeEvaluatorFunction
+from aeri.experiment import (
+    EvaluatorFunction,
+    ExperimentResult,
+    RunEvaluatorFunction,
+    TaskFunction,
+)
+
+if TYPE_CHECKING:
+    from aeri._client.client import Aeri
+
+
+class DatasetClient:
+    """Class for managing datasets in Aeri.
+
+    Attributes:
+        id (str): Unique identifier of the dataset.
+        name (str): Name of the dataset.
+        description (Optional[str]): Description of the dataset.
+        metadata (Optional[typing.Any]): Additional metadata of the dataset.
+        project_id (str): Identifier of the project to which the dataset belongs.
+        created_at (datetime): Timestamp of dataset creation.
+        updated_at (datetime): Timestamp of the last update to the dataset.
+        items (List[DatasetItem]): List of dataset items associated with the dataset.
+        version (Optional[datetime]): Timestamp of the dataset version.
+    Example:
+        Print the input of each dataset item in a dataset.
+        ```python
+        from aeri import Aeri
+
+        aeri = Aeri()
+
+        dataset = aeri.get_dataset("<dataset_name>")
+
+        for item in dataset.items:
+            print(item.input)
+        ```
+    """
+
+    id: str
+    name: str
+    description: Optional[str]
+    project_id: str
+    metadata: Optional[Any]
+    created_at: dt.datetime
+    updated_at: dt.datetime
+    input_schema: Optional[Any]
+    expected_output_schema: Optional[Any]
+    items: List[DatasetItem]
+    version: Optional[dt.datetime]
+
+    def __init__(
+        self,
+        *,
+        dataset: Dataset,
+        items: List[DatasetItem],
+        aeri_client: "Aeri",
+        version: Optional[dt.datetime] = None,
+    ):
+        """Initialize the DatasetClient."""
+        self.id = dataset.id
+        self.name = dataset.name
+        self.description = dataset.description
+        self.project_id = dataset.project_id
+        self.metadata = dataset.metadata
+        self.created_at = dataset.created_at
+        self.updated_at = dataset.updated_at
+        self.input_schema = dataset.input_schema
+        self.expected_output_schema = dataset.expected_output_schema
+        self.items = items
+        self.version = version
+        self._aeri_client: "Aeri" = aeri_client
+
+    def run_experiment(
+        self,
+        *,
+        name: str,
+        run_name: Optional[str] = None,
+        description: Optional[str] = None,
+        task: TaskFunction,
+        evaluators: List[EvaluatorFunction] = [],
+        composite_evaluator: Optional[CompositeEvaluatorFunction] = None,
+        run_evaluators: List[RunEvaluatorFunction] = [],
+        max_concurrency: int = 50,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> ExperimentResult:
+        """Run an experiment on this Aeri dataset with automatic tracking.
+
+        This is a convenience method that runs an experiment using all items in this
+        dataset. It automatically creates a dataset run in Aeri for tracking and
+        comparison purposes, linking all experiment results to the dataset.
+
+        Key benefits of using dataset.run_experiment():
+        - Automatic dataset run creation and linking in Aeri UI
+        - Built-in experiment tracking and versioning
+        - Easy comparison between different experiment runs
+        - Direct access to dataset items with their metadata and expected outputs
+        - Automatic URL generation for viewing results in Aeri dashboard
+
+        Args:
+            name: Human-readable name for the experiment run. This will be used as
+                the dataset run name in Aeri for tracking and identification.
+            run_name: Optional exact name for the dataset run. If provided, this will be
+                used as the exact dataset run name in Aeri. If not provided, this will
+                default to the experiment name appended with an ISO timestamp.
+            description: Optional description of the experiment's purpose, methodology,
+                or what you're testing. Appears in the Aeri UI for context.
+            task: Function that processes each dataset item and returns output.
+                The function will receive DatasetItem objects with .input, .expected_output,
+                .metadata attributes. Signature should be: task(*, item, **kwargs) -> Any
+            evaluators: List of functions to evaluate each item's output individually.
+                These will have access to the item's expected_output for comparison.
+            composite_evaluator: Optional function that creates composite scores from item-level evaluations.
+                Receives the same inputs as item-level evaluators (input, output, expected_output, metadata)
+                plus the list of evaluations from item-level evaluators. Useful for weighted averages,
+                pass/fail decisions based on multiple criteria, or custom scoring logic combining multiple metrics.
+            run_evaluators: List of functions to evaluate the entire experiment run.
+                Useful for computing aggregate statistics across all dataset items.
+            max_concurrency: Maximum number of concurrent task executions (default: 50).
+                Adjust based on API rate limits and system resources.
+            metadata: Optional metadata to attach to the experiment run and all traces.
+                Will be combined with individual item metadata.
+
+        Returns:
+            ExperimentResult object containing:
+            - name: The experiment name.
+            - run_name: The experiment run name (equivalent to the dataset run name).
+            - description: Optional experiment description.
+            - item_results: Results for each dataset item with outputs and evaluations.
+            - run_evaluations: Aggregate evaluation results for the entire run.
+            - dataset_run_id: ID of the created dataset run in Aeri.
+            - dataset_run_url: Direct URL to view the experiment results in Aeri UI.
+
+            The result object provides a format() method for human-readable output:
+            ```python
+            result = dataset.run_experiment(...)
+            print(result.format())  # Summary view
+            print(result.format(include_item_results=True))  # Detailed view
+            ```
+
+        Raises:
+            ValueError: If the dataset has no items or no Aeri client is available.
+
+        Examples:
+            Basic dataset experiment:
+            ```python
+            dataset = aeri.get_dataset("qa-evaluation-set")
+
+            def answer_questions(*, item, **kwargs):
+                # item is a DatasetItem with .input, .expected_output, .metadata
+                question = item.input
+                return my_qa_system.answer(question)
+
+            def accuracy_evaluator(*, input, output, expected_output=None, **kwargs):
+                if not expected_output:
+                    return {"name": "accuracy", "value": 0, "comment": "No expected output"}
+
+                is_correct = output.strip().lower() == expected_output.strip().lower()
+                return {
+                    "name": "accuracy",
+                    "value": 1.0 if is_correct else 0.0,
+                    "comment": "Correct" if is_correct else "Incorrect"
+                }
+
+            result = dataset.run_experiment(
+                name="QA System v2.0 Evaluation",
+                description="Testing improved QA system on curated question set",
+                task=answer_questions,
+                evaluators=[accuracy_evaluator]
+            )
+
+            print(f"Evaluated {len(result['item_results'])} questions")
+            print(f"View detailed results: {result['dataset_run_url']}")
+            ```
+
+            Advanced experiment with multiple evaluators and run-level analysis:
+            ```python
+            dataset = aeri.get_dataset("content-generation-benchmark")
+
+            async def generate_content(*, item, **kwargs):
+                prompt = item.input
+                response = await openai_client.chat.completions.create(
+                    model="gpt-4",
+                    messages=[{"role": "user", "content": prompt}],
+                    temperature=0.7
+                )
+                return response.choices[0].message.content
+
+            def quality_evaluator(*, input, output, expected_output=None, metadata=None, **kwargs):
+                # Use metadata for context-aware evaluation
+                content_type = metadata.get("type", "general") if metadata else "general"
+
+                # Basic quality checks
+                word_count = len(output.split())
+                min_words = {"blog": 300, "tweet": 10, "summary": 100}.get(content_type, 50)
+
+                return [
+                    {
+                        "name": "word_count",
+                        "value": word_count,
+                        "comment": f"Generated {word_count} words"
+                    },
+                    {
+                        "name": "meets_length_requirement",
+                        "value": word_count >= min_words,
+                        "comment": f"{'Meets' if word_count >= min_words else 'Below'} minimum {min_words} words for {content_type}"
+                    }
+                ]
+
+            def content_diversity(*, item_results, **kwargs):
+                # Analyze diversity across all generated content
+                all_outputs = [result["output"] for result in item_results]
+                unique_words = set()
+                total_words = 0
+
+                for output in all_outputs:
+                    words = output.lower().split()
+                    unique_words.update(words)
+                    total_words += len(words)
+
+                diversity_ratio = len(unique_words) / total_words if total_words > 0 else 0
+
+                return {
+                    "name": "vocabulary_diversity",
+                    "value": diversity_ratio,
+                    "comment": f"Used {len(unique_words)} unique words out of {total_words} total ({diversity_ratio:.2%} diversity)"
+                }
+
+            result = dataset.run_experiment(
+                name="Content Generation Diversity Test",
+                description="Evaluating content quality and vocabulary diversity across different content types",
+                task=generate_content,
+                evaluators=[quality_evaluator],
+                run_evaluators=[content_diversity],
+                max_concurrency=3,  # Limit API calls
+                metadata={"model": "gpt-4", "temperature": 0.7}
+            )
+
+            # Results are automatically linked to dataset in Aeri
+            print(f"Experiment completed! View in Aeri: {result['dataset_run_url']}")
+
+            # Access individual results
+            for i, item_result in enumerate(result["item_results"]):
+                print(f"Item {i+1}: {item_result['evaluations']}")
+            ```
+
+            Comparing different model versions:
+            ```python
+            # Run multiple experiments on the same dataset for comparison
+            dataset = aeri.get_dataset("model-benchmark")
+
+            # Experiment 1: GPT-4
+            result_gpt4 = dataset.run_experiment(
+                name="GPT-4 Baseline",
+                description="Baseline performance with GPT-4",
+                task=lambda *, item, **kwargs: gpt4_model.generate(item.input),
+                evaluators=[accuracy_evaluator, fluency_evaluator]
+            )
+
+            # Experiment 2: Custom model
+            result_custom = dataset.run_experiment(
+                name="Custom Model v1.2",
+                description="Testing our fine-tuned model",
+                task=lambda *, item, **kwargs: custom_model.generate(item.input),
+                evaluators=[accuracy_evaluator, fluency_evaluator]
+            )
+
+            # Both experiments are now visible in Aeri for easy comparison
+            print("Compare results in Aeri:")
+            print(f"GPT-4: {result_gpt4.dataset_run_url}")
+            print(f"Custom: {result_custom.dataset_run_url}")
+            ```
+
+        Note:
+            - All experiment results are automatically tracked in Aeri as dataset runs
+            - Dataset items provide .input, .expected_output, and .metadata attributes
+            - Results can be easily compared across different experiment runs in the UI
+            - The dataset_run_url provides direct access to detailed results and analysis
+            - Failed items are handled gracefully and logged without stopping the experiment
+            - This method works in both sync and async contexts (Jupyter notebooks, web apps, etc.)
+            - Async execution is handled automatically with smart event loop detection
+        """
+        return self._aeri_client.run_experiment(
+            name=name,
+            run_name=run_name,
+            description=description,
+            data=self.items,
+            task=task,
+            evaluators=evaluators,
+            composite_evaluator=composite_evaluator,
+            run_evaluators=run_evaluators,
+            max_concurrency=max_concurrency,
+            metadata=metadata,
+            _dataset_version=self.version,
+        )

aeri/_client/environment_variables.py

@@ -0,0 +1,158 @@
+"""Environment variable definitions for Aeri OpenTelemetry integration.
+
+This module defines environment variables used to configure the Aeri OpenTelemetry integration.
+Each environment variable includes documentation on its purpose, expected values, and defaults.
+"""
+
+AERI_TRACING_ENVIRONMENT = "AERI_TRACING_ENVIRONMENT"
+"""
+.. envvar:: AERI_TRACING_ENVIRONMENT
+
+The tracing environment. Can be any lowercase alphanumeric string with hyphens and underscores that does not start with 'aeri'.
+
+**Default value:** ``"default"``
+"""
+
+AERI_RELEASE = "AERI_RELEASE"
+"""
+.. envvar:: AERI_RELEASE
+
+Release number/hash of the application to provide analytics grouped by release.
+"""
+
+
+AERI_PUBLIC_KEY = "AERI_PUBLIC_KEY"
+"""
+.. envvar:: AERI_PUBLIC_KEY
+
+Public API key of Aeri project
+"""
+
+AERI_SECRET_KEY = "AERI_SECRET_KEY"
+"""
+.. envvar:: AERI_SECRET_KEY
+
+Secret API key of Aeri project
+"""
+
+AERI_BASE_URL = "AERI_BASE_URL"
+"""
+.. envvar:: AERI_BASE_URL
+
+Base URL of Aeri API. Can be set via `AERI_BASE_URL` environment variable.
+
+**Default value:** ``"https://cloud.aeri.com"``
+"""
+
+AERI_HOST = "AERI_HOST"
+"""
+.. envvar:: AERI_HOST
+
+Deprecated. Use AERI_BASE_URL instead. Host of Aeri API. Can be set via `AERI_HOST` environment variable.
+
+**Default value:** ``"https://cloud.aeri.com"``
+"""
+
+AERI_OTEL_TRACES_EXPORT_PATH = "AERI_OTEL_TRACES_EXPORT_PATH"
+"""
+.. envvar:: AERI_OTEL_TRACES_EXPORT_PATH
+
+URL path on the configured host to export traces to.
+
+**Default value:** ``/api/public/otel/v1/traces``
+"""
+
+AERI_DEBUG = "AERI_DEBUG"
+"""
+.. envvar:: AERI_DEBUG
+
+Enables debug mode for more verbose logging.
+
+**Default value:** ``"False"``
+"""
+
+AERI_TRACING_ENABLED = "AERI_TRACING_ENABLED"
+"""
+.. envvar:: AERI_TRACING_ENABLED
+
+Enables or disables the Aeri client. If disabled, all observability calls to the backend will be no-ops. Default is True. Set to `False` to disable tracing.
+
+**Default value:** ``"True"``
+"""
+
+AERI_MEDIA_UPLOAD_THREAD_COUNT = "AERI_MEDIA_UPLOAD_THREAD_COUNT"
+"""
+.. envvar:: AERI_MEDIA_UPLOAD_THREAD_COUNT 
+
+Number of background threads to handle media uploads from trace ingestion.
+
+**Default value:** ``1``
+"""
+
+AERI_FLUSH_AT = "AERI_FLUSH_AT"
+"""
+.. envvar:: AERI_FLUSH_AT
+
+Max batch size until a new ingestion batch is sent to the API.
+**Default value:** same as OTEL ``OTEL_BSP_MAX_EXPORT_BATCH_SIZE``
+"""
+
+AERI_FLUSH_INTERVAL = "AERI_FLUSH_INTERVAL"
+"""
+.. envvar:: AERI_FLUSH_INTERVAL
+
+Max delay in seconds until a new ingestion batch is sent to the API.
+**Default value:** same as OTEL ``OTEL_BSP_SCHEDULE_DELAY``
+"""
+
+AERI_SAMPLE_RATE = "AERI_SAMPLE_RATE"
+"""
+.. envvar: AERI_SAMPLE_RATE
+
+Float between 0 and 1 indicating the sample rate of traces to bet sent to Aeri servers.
+
+**Default value**: ``1.0``
+
+"""
+AERI_OBSERVE_DECORATOR_IO_CAPTURE_ENABLED = (
+    "AERI_OBSERVE_DECORATOR_IO_CAPTURE_ENABLED"
+)
+"""
+.. envvar: AERI_OBSERVE_DECORATOR_IO_CAPTURE_ENABLED
+
+Default capture of function args, kwargs and return value when using the @observe decorator.
+
+Having default IO capture enabled for observe decorated function may have a performance impact on your application
+if large or deeply nested objects are attempted to be serialized. Set this value to `False` and use manual
+input/output setting on your observation to avoid this.
+
+**Default value**: ``True``
+"""
+
+AERI_MEDIA_UPLOAD_ENABLED = "AERI_MEDIA_UPLOAD_ENABLED"
+"""
+.. envvar: AERI_MEDIA_UPLOAD_ENABLED
+
+Controls whether media detection and upload is attempted by the SDK.
+
+**Default value**: ``True``
+"""
+
+AERI_TIMEOUT = "AERI_TIMEOUT"
+"""
+.. envvar: AERI_TIMEOUT
+
+Controls the timeout for all API requests in seconds
+
+**Default value**: ``5``
+"""
+
+AERI_PROMPT_CACHE_DEFAULT_TTL_SECONDS = "AERI_PROMPT_CACHE_DEFAULT_TTL_SECONDS"
+"""
+.. envvar: AERI_PROMPT_CACHE_DEFAULT_TTL_SECONDS
+
+Controls the default time-to-live (TTL) in seconds for cached prompts.
+This setting determines how long prompt responses are cached before they expire.
+
+**Default value**: ``60``
+"""

aeri/_client/get_client.py

@@ -0,0 +1,149 @@
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import Iterator, Optional
+
+from aeri._client.client import Aeri
+from aeri._client.resource_manager import AeriResourceManager
+from aeri.logger import aeri_logger
+
+# Context variable to track the current aeri_public_key in execution context
+_current_public_key: ContextVar[Optional[str]] = ContextVar(
+    "aeri_public_key", default=None
+)
+
+
+@contextmanager
+def _set_current_public_key(public_key: Optional[str]) -> Iterator[None]:
+    """Context manager to set and restore the current public key in execution context.
+
+    Args:
+        public_key: The public key to set in context. If None, context is not modified.
+
+    Yields:
+        None
+    """
+    if public_key is None:
+        yield  # Don't modify context if no key provided
+        return
+
+    token = _current_public_key.set(public_key)
+    try:
+        yield
+    finally:
+        _current_public_key.reset(token)
+
+
+def _create_client_from_instance(
+    instance: "AeriResourceManager", public_key: Optional[str] = None
+) -> Aeri:
+    """Create a Aeri client from a resource manager instance with all settings preserved."""
+    return Aeri(
+        public_key=public_key or instance.public_key,
+        secret_key=instance.secret_key,
+        base_url=instance.base_url,
+        tracing_enabled=instance.tracing_enabled,
+        environment=instance.environment,
+        timeout=instance.timeout,
+        flush_at=instance.flush_at,
+        flush_interval=instance.flush_interval,
+        release=instance.release,
+        media_upload_thread_count=instance.media_upload_thread_count,
+        sample_rate=instance.sample_rate,
+        mask=instance.mask,
+        blocked_instrumentation_scopes=instance.blocked_instrumentation_scopes,
+        should_export_span=instance.should_export_span,
+        additional_headers=instance.additional_headers,
+        tracer_provider=instance.tracer_provider,
+        httpx_client=instance.httpx_client,
+    )
+
+
+def get_client(*, public_key: Optional[str] = None) -> Aeri:
+    """Get or create a Aeri client instance.
+
+    Returns an existing Aeri client or creates a new one if none exists. In multi-project setups,
+    providing a public_key is required. Multi-project support is experimental - see Aeri docs.
+
+    Behavior:
+    - Single project: Returns existing client or creates new one
+    - Multi-project: Requires public_key to return specific client
+    - No public_key in multi-project: Returns disabled client to prevent data leakage
+
+    The function uses a singleton pattern per public_key to conserve resources and maintain state.
+
+    Args:
+        public_key (Optional[str]): Project identifier
+            - With key: Returns client for that project
+            - Without key: Returns single client or disabled client if multiple exist
+
+    Returns:
+        Aeri: Client instance in one of three states:
+            1. Client for specified public_key
+            2. Default client for single-project setup
+            3. Disabled client when multiple projects exist without key
+
+    Security:
+        Disables tracing when multiple projects exist without explicit key to prevent
+        cross-project data leakage. Multi-project setups are experimental.
+
+    Example:
+        ```python
+        # Single project
+        client = get_client()  # Default client
+
+        # In multi-project usage:
+        client_a = get_client(public_key="project_a_key")  # Returns project A's client
+        client_b = get_client(public_key="project_b_key")  # Returns project B's client
+
+        # Without specific key in multi-project setup:
+        client = get_client()  # Returns disabled client for safety
+        ```
+    """
+    with AeriResourceManager._lock:
+        active_instances = AeriResourceManager._instances
+
+        # If no explicit public_key provided, check execution context
+        if not public_key:
+            public_key = _current_public_key.get(None)
+
+        if not public_key:
+            if len(active_instances) == 0:
+                # No clients initialized yet, create default instance
+                return Aeri()
+
+            if len(active_instances) == 1:
+                # Only one client exists, safe to use without specifying key
+                instance = list(active_instances.values())[0]
+
+                # Initialize with the credentials bound to the instance
+                # This is important if the original instance was instantiated
+                # via constructor arguments
+                return _create_client_from_instance(instance)
+
+            else:
+                # Multiple clients exist but no key specified - disable tracing
+                # to prevent cross-project data leakage
+                aeri_logger.warning(
+                    "No 'aeri_public_key' passed to decorated function, but multiple aeri clients are instantiated in current process. Skipping tracing for this function to avoid cross-project leakage."
+                )
+                return Aeri(
+                    tracing_enabled=False, public_key="fake", secret_key="fake"
+                )
+
+        else:
+            # Specific key provided, look up existing instance
+            target_instance: Optional[AeriResourceManager] = active_instances.get(
+                public_key, None
+            )
+
+            if target_instance is None:
+                # No instance found with this key - client not initialized properly
+                aeri_logger.warning(
+                    f"No Aeri client with public key {public_key} has been initialized. Skipping tracing for decorated function."
+                )
+                return Aeri(
+                    tracing_enabled=False, public_key="fake", secret_key="fake"
+                )
+
+            # target_instance is guaranteed to be not None at this point
+            return _create_client_from_instance(target_instance, public_key)