"""ABV OpenTelemetry integration module.

This module implements ABV's core observability functionality on top of the OpenTelemetry (OTel) standard.
"""

import logging
import os
import warnings
import re
import urllib.parse
from datetime import datetime
from hashlib import sha256
from time import time_ns
from typing import (
    Any,
    Dict,
    List,
    Literal,
    Optional,
    Union,
    Type,
    cast,
    overload,
)

import backoff
import httpx
from opentelemetry import trace
from opentelemetry import trace as otel_trace_api
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.id_generator import RandomIdGenerator
from opentelemetry.util._decorator import (
    _AgnosticContextManager,
    _agnosticcontextmanager,
)
from packaging.version import Version

from abvdev._client.attributes import ABVOtelSpanAttributes
from abvdev._client.datasets import DatasetClient, DatasetItemClient
from abvdev._client.environment_variables import (
    ABV_DEBUG,
    ABV_HOST,
    ABV_SAMPLE_RATE,
    ABV_TIMEOUT,
    ABV_TRACING_ENABLED,
    ABV_TRACING_ENVIRONMENT,
)
from abvdev._client.constants import (
    ObservationTypeLiteral,
    ObservationTypeLiteralNoEvent,
    ObservationTypeGenerationLike,
    ObservationTypeSpanLike,
    get_observation_types_list,
)
from abvdev._client.resource_manager import ABVResourceManager
from abvdev._client.span import (
    ABVEvent,
    ABVGeneration,
    ABVSpan,
    ABVAgent,
    ABVTool,
    ABVChain,
    ABVRetriever,
    ABVEvaluator,
    ABVEmbedding,
    ABVGuardrail,
)
from abvdev._utils import _get_timestamp
from abvdev._utils.parse_error import handle_fern_exception
from abvdev._utils.prompt_cache import PromptCache
from abvdev.api.resources.commons.errors.error import Error
from abvdev.api.resources.ingestion.types.score_body import ScoreBody
from abvdev.api.resources.prompts.types import (
    CreatePromptRequest_Chat,
    CreatePromptRequest_Text,
    Prompt_Chat,
    Prompt_Text,
)
from abvdev.logger import abv_logger
from abvdev.media import ABVMedia
from abvdev.model import (
    ChatMessageDict,
    ChatMessageWithPlaceholdersDict,
    ChatPromptClient,
    CreateDatasetItemRequest,
    CreateDatasetRequest,
    Dataset,
    DatasetItem,
    DatasetStatus,
    MapValue,
    PromptClient,
    TextPromptClient,
)
from abvdev.types import MaskFunction, ScoreDataType, SpanLevel, TraceContext


class ABV:
    """Main client for ABV tracing and platform features.

    This class provides an interface for creating and managing traces, spans,
    and generations in ABV as well as interacting with the ABV API.

    The client features a thread-safe singleton pattern for each unique public API key,
    ensuring consistent trace context propagation across your application. It implements
    efficient batching of spans with configurable flush settings and includes background
    thread management for media uploads and score ingestion.

    Configuration is flexible through either direct parameters or environment variables,
    with graceful fallbacks and runtime configuration updates.

    Attributes:
        api: Synchronous API client for ABV backend communication
        async_api: Asynchronous API client for ABV backend communication
        abv_tracer: Internal ABVTracer instance managing OpenTelemetry components

    Parameters:
        api_key (Optional[str]): Your ABV API key. Can also be set via ABV_API_KEY environment variable.
        region (Optional[Literal["us", "eu"]]): Data region for ABV API and Gateway. Determines which regional endpoints to use. Defaults to "us".
            - "us": Routes to app.abv.dev and gateway.abv.dev
            - "eu": Routes to eu.app.abv.dev and eu.gateway.abv.dev
        host (Optional[str]): The ABV API host URL. If specified, this overrides the region-based URL. Can also be set via ABV_HOST environment variable. Defaults to "https://app.abv.dev" (US) or "https://eu.app.abv.dev" (EU).
        gateway_base_url (Optional[str]): Base URL for the ABV AI Gateway. If specified, this overrides the region-based gateway URL. Defaults to "https://gateway.abv.dev" (US) or "https://eu.gateway.abv.dev" (EU).
        timeout (Optional[int]): Timeout in seconds for API requests. Defaults to 5 seconds.
        httpx_client (Optional[httpx.Client]): Custom httpx client for making non-tracing HTTP requests. If not provided, a default client will be created.
        debug (bool): Enable debug logging. Defaults to False. Can also be set via ABV_DEBUG environment variable.
        tracing_enabled (Optional[bool]): Enable or disable tracing. Defaults to True. Can also be set via ABV_TRACING_ENABLED environment variable.
        flush_at (Optional[int]): Number of spans to batch before sending to the API. Defaults to 512. Can also be set via ABV_FLUSH_AT environment variable.
        flush_interval (Optional[float]): Time in seconds between batch flushes. Defaults to 5 seconds. Can also be set via ABV_FLUSH_INTERVAL environment variable.
        environment (Optional[str]): Environment name for tracing. Default is 'default'. Can also be set via ABV_TRACING_ENVIRONMENT environment variable. Can be any lowercase alphanumeric string with hyphens and underscores that does not start with 'abv'.
        release (Optional[str]): Release version/hash of your application. Used for grouping analytics by release.
        media_upload_thread_count (Optional[int]): Number of background threads for handling media uploads. Defaults to 1. Can also be set via ABV_MEDIA_UPLOAD_THREAD_COUNT environment variable.
        sample_rate (Optional[float]): Sampling rate for traces (0.0 to 1.0). Defaults to 1.0 (100% of traces are sampled). Can also be set via ABV_SAMPLE_RATE environment variable.
        mask (Optional[MaskFunction]): Function to mask sensitive data in traces before sending to the API.
        blocked_instrumentation_scopes (Optional[List[str]]): List of instrumentation scope names to block from being exported to ABV. Spans from these scopes will be filtered out before being sent to the API. Useful for filtering out spans from specific libraries or frameworks. For exported spans, you can see the instrumentation scope name in the span metadata in ABV (`metadata.scope.name`)
        additional_headers (Optional[Dict[str, str]]): Additional headers to include in all API requests and OTLPSpanExporter requests. These headers will be merged with default headers. Note: If httpx_client is provided, additional_headers must be set directly on your custom httpx_client as well.
        tracer_provider(Optional[TracerProvider]): OpenTelemetry TracerProvider to use for ABV. This can be useful to set to have disconnected tracing between ABV and other OpenTelemetry-span emitting libraries. Note: To track active spans, the context is still shared between TracerProviders. This may lead to broken trace trees.

    Example:
        ```python
        from abvdev.otel import ABV

        # Initialize the client (reads from env vars if not provided)
        abv = ABV(
            api_key="your-api-key",
            host="https://app.abv.dev",  # Optional, default shown
        )

        # Create a trace span
        with abv.start_as_current_span(name="process-query") as span:
            # Your application code here

            # Create a nested generation span for an LLM call
            with span.start_as_current_generation(
                name="generate-response",
                model="gpt-4",
                input={"query": "Tell me about AI"},
                model_parameters={"temperature": 0.7, "max_tokens": 500}
            ) as generation:
                # Generate response here
                response = "AI is a field of computer science..."

                generation.update(
                    output=response,
                    usage_details={"prompt_tokens": 10, "completion_tokens": 50},
                    cost_details={"total_cost": 0.0023}
                )

                # Score the generation (supports NUMERIC, BOOLEAN, CATEGORICAL)
                generation.score(name="relevance", value=0.95, data_type="NUMERIC")
        ```
    """

    _resources: Optional[ABVResourceManager] = None
    _mask: Optional[MaskFunction] = None
    _otel_tracer: otel_trace_api.Tracer

    def __init__(
        self,
        *,
        api_key: Optional[str] = None,
        region: Optional[Literal["us", "eu"]] = None,
        host: Optional[str] = None,
        gateway_base_url: Optional[str] = None,
        timeout: Optional[int] = None,
        httpx_client: Optional[httpx.Client] = None,
        debug: bool = False,
        tracing_enabled: Optional[bool] = True,
        flush_at: Optional[int] = None,
        flush_interval: Optional[float] = None,
        environment: Optional[str] = None,
        release: Optional[str] = None,
        media_upload_thread_count: Optional[int] = None,
        sample_rate: Optional[float] = None,
        mask: Optional[MaskFunction] = None,
        blocked_instrumentation_scopes: Optional[List[str]] = None,
        additional_headers: Optional[Dict[str, str]] = None,
        tracer_provider: Optional[TracerProvider] = None,
    ):
        # Determine host based on region (unless explicitly provided)
        region = region or "us"
        region_host = "https://eu.app.abv.dev" if region == "eu" else "https://app.abv.dev"
        self._host = host or os.environ.get(ABV_HOST) or region_host

        self._environment = environment or cast(
            str, os.environ.get(ABV_TRACING_ENVIRONMENT)
        )
        self._project_id: Optional[str] = None
        self._region = region
        self._gateway_base_url = gateway_base_url
        sample_rate = sample_rate or float(os.environ.get(ABV_SAMPLE_RATE, 1.0))
        if not 0.0 <= sample_rate <= 1.0:
            raise ValueError(
                f"Sample rate must be between 0.0 and 1.0, got {sample_rate}"
            )

        timeout = timeout or int(os.environ.get(ABV_TIMEOUT, 5))

        self._tracing_enabled = (
            tracing_enabled
            and os.environ.get(ABV_TRACING_ENABLED, "true").lower() != "false"
        )
        if not self._tracing_enabled:
            abv_logger.info(
                "Configuration: ABV tracing is explicitly disabled. No data will be sent to the ABV API."
            )

        debug = (
            debug if debug else (os.getenv(ABV_DEBUG, "false").lower() == "true")
        )
        if debug:
            logging.basicConfig(
                format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
            )
            abv_logger.setLevel(logging.DEBUG)

        api_key = api_key or os.environ.get("ABV_API_KEY")
        if api_key is None:
            abv_logger.warning(
                "Authentication error: ABV client initialized without api_key. Client will be disabled. "
                "Provide an api_key parameter or set ABV_API_KEY environment variable. "
            )
            self._otel_tracer = otel_trace_api.NoOpTracer()
            return

        if os.environ.get("OTEL_SDK_DISABLED", "false").lower() == "true":
            abv_logger.warning(
                "OTEL_SDK_DISABLED is set. ABV tracing will be disabled and no traces will appear in the UI."
            )

        # Initialize api and tracer if requirements are met
        self._resources = ABVResourceManager(
            api_key=api_key,
            host=self._host,
            timeout=timeout,
            environment=environment,
            release=release,
            flush_at=flush_at,
            flush_interval=flush_interval,
            httpx_client=httpx_client,
            media_upload_thread_count=media_upload_thread_count,
            sample_rate=sample_rate,
            mask=mask,
            tracing_enabled=self._tracing_enabled,
            blocked_instrumentation_scopes=blocked_instrumentation_scopes,
            additional_headers=additional_headers,
            tracer_provider=tracer_provider,
        )
        self._mask = self._resources.mask

        self._otel_tracer = (
            self._resources.tracer
            if self._tracing_enabled and self._resources.tracer is not None
            else otel_trace_api.NoOpTracer()
        )
        self.api = self._resources.api
        self.async_api = self._resources.async_api

        # Always initialize gateway if we have an API key
        self.gateway = None
        if api_key:
            from abvdev._client.gateway import GatewayClient

            self.gateway = GatewayClient(
                api_key=api_key,
                region=self._region,
                gateway_base_url=self._gateway_base_url,
                abv_client=self,
            )

            abv_logger.debug(
                f"Initialized ABV Gateway with region={self._region}, gateway_base_url={self._gateway_base_url}"
            )

    def start_span(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
    ) -> ABVSpan:
        """Create a new span for tracing a unit of work.

        This method creates a new span but does not set it as the current span in the
        context. To create and use a span within a context, use start_as_current_span().

        The created span will be the child of the current span in the context.

        Args:
            trace_context: Optional context for connecting to an existing trace
            name: Name of the span (e.g., function or operation name)
            input: Input data for the operation (can be any JSON-serializable object)
            output: Output data from the operation (can be any JSON-serializable object)
            metadata: Additional metadata to associate with the span
            version: Version identifier for the code or component
            level: Importance level of the span (info, warning, error)
            status_message: Optional status message for the span

        Returns:
            A ABVSpan object that must be ended with .end() when the operation completes

        Example:
            ```python
            span = abv.start_span(name="process-data")
            try:
                # Do work
                span.update(output="result")
            finally:
                span.end()
            ```
        """
        return self.start_observation(
            trace_context=trace_context,
            name=name,
            as_type="span",
            input=input,
            output=output,
            metadata=metadata,
            version=version,
            level=level,
            status_message=status_message,
        )

    def start_as_current_span(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
        end_on_exit: Optional[bool] = None,
    ) -> _AgnosticContextManager[ABVSpan]:
        """Create a new span and set it as the current span in a context manager.

        This method creates a new span and sets it as the current span within a context
        manager. Use this method with a 'with' statement to automatically handle span
        lifecycle within a code block.

        The created span will be the child of the current span in the context.

        Args:
            trace_context: Optional context for connecting to an existing trace
            name: Name of the span (e.g., function or operation name)
            input: Input data for the operation (can be any JSON-serializable object)
            output: Output data from the operation (can be any JSON-serializable object)
            metadata: Additional metadata to associate with the span
            version: Version identifier for the code or component
            level: Importance level of the span (info, warning, error)
            status_message: Optional status message for the span
            end_on_exit (default: True): Whether to end the span automatically when leaving the context manager. If False, the span must be manually ended to avoid memory leaks.

        Returns:
            A context manager that yields a ABVSpan

        Example:
            ```python
            with abv.start_as_current_span(name="process-query") as span:
                # Do work
                result = process_data()
                span.update(output=result)

                # Create a child span automatically
                with span.start_as_current_span(name="sub-operation") as child_span:
                    # Do sub-operation work
                    child_span.update(output="sub-result")
            ```
        """
        return self.start_as_current_observation(
            trace_context=trace_context,
            name=name,
            as_type="span",
            input=input,
            output=output,
            metadata=metadata,
            version=version,
            level=level,
            status_message=status_message,
            end_on_exit=end_on_exit,
        )

    @overload
    def start_observation(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        as_type: Literal["generation"],
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
        completion_start_time: Optional[datetime] = None,
        model: Optional[str] = None,
        model_parameters: Optional[Dict[str, MapValue]] = None,
        usage_details: Optional[Dict[str, int]] = None,
        cost_details: Optional[Dict[str, float]] = None,
        prompt: Optional[PromptClient] = None,
    ) -> ABVGeneration: ...

    @overload
    def start_observation(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        as_type: Literal["span"] = "span",
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
    ) -> ABVSpan: ...

    @overload
    def start_observation(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        as_type: Literal["agent"],
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
    ) -> ABVAgent: ...

    @overload
    def start_observation(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        as_type: Literal["tool"],
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
    ) -> ABVTool: ...

    @overload
    def start_observation(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        as_type: Literal["chain"],
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
    ) -> ABVChain: ...

    @overload
    def start_observation(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        as_type: Literal["retriever"],
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
    ) -> ABVRetriever: ...

    @overload
    def start_observation(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        as_type: Literal["evaluator"],
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
    ) -> ABVEvaluator: ...

    @overload
    def start_observation(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        as_type: Literal["embedding"],
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
        completion_start_time: Optional[datetime] = None,
        model: Optional[str] = None,
        model_parameters: Optional[Dict[str, MapValue]] = None,
        usage_details: Optional[Dict[str, int]] = None,
        cost_details: Optional[Dict[str, float]] = None,
        prompt: Optional[PromptClient] = None,
    ) -> ABVEmbedding: ...

    @overload
    def start_observation(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        as_type: Literal["guardrail"],
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
    ) -> ABVGuardrail: ...

    def start_observation(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        as_type: ObservationTypeLiteralNoEvent = "span",
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
        completion_start_time: Optional[datetime] = None,
        model: Optional[str] = None,
        model_parameters: Optional[Dict[str, MapValue]] = None,
        usage_details: Optional[Dict[str, int]] = None,
        cost_details: Optional[Dict[str, float]] = None,
        prompt: Optional[PromptClient] = None,
    ) -> Union[
        ABVSpan,
        ABVGeneration,
        ABVAgent,
        ABVTool,
        ABVChain,
        ABVRetriever,
        ABVEvaluator,
        ABVEmbedding,
        ABVGuardrail,
    ]:
        """Create a new observation of the specified type.

        This method creates a new observation but does not set it as the current span in the
        context. To create and use an observation within a context, use start_as_current_observation().

        Args:
            trace_context: Optional context for connecting to an existing trace
            name: Name of the observation
            as_type: Type of observation to create (defaults to "span")
            input: Input data for the operation
            output: Output data from the operation
            metadata: Additional metadata to associate with the observation
            version: Version identifier for the code or component
            level: Importance level of the observation
            status_message: Optional status message for the observation
            completion_start_time: When the model started generating (for generation types)
            model: Name/identifier of the AI model used (for generation types)
            model_parameters: Parameters used for the model (for generation types)
            usage_details: Token usage information (for generation types)
            cost_details: Cost information (for generation types)
            prompt: Associated prompt template (for generation types)

        Returns:
            An observation object of the appropriate type that must be ended with .end()
        """
        if trace_context:
            trace_id = trace_context.get("trace_id", None)
            parent_span_id = trace_context.get("parent_span_id", None)

            if trace_id:
                remote_parent_span = self._create_remote_parent_span(
                    trace_id=trace_id, parent_span_id=parent_span_id
                )

                with otel_trace_api.use_span(
                    cast(otel_trace_api.Span, remote_parent_span)
                ):
                    otel_span = self._otel_tracer.start_span(name=name)
                    otel_span.set_attribute(ABVOtelSpanAttributes.AS_ROOT, True)

                    return self._create_observation_from_otel_span(
                        otel_span=otel_span,
                        as_type=as_type,
                        input=input,
                        output=output,
                        metadata=metadata,
                        version=version,
                        level=level,
                        status_message=status_message,
                        completion_start_time=completion_start_time,
                        model=model,
                        model_parameters=model_parameters,
                        usage_details=usage_details,
                        cost_details=cost_details,
                        prompt=prompt,
                    )

        otel_span = self._otel_tracer.start_span(name=name)

        return self._create_observation_from_otel_span(
            otel_span=otel_span,
            as_type=as_type,
            input=input,
            output=output,
            metadata=metadata,
            version=version,
            level=level,
            status_message=status_message,
            completion_start_time=completion_start_time,
            model=model,
            model_parameters=model_parameters,
            usage_details=usage_details,
            cost_details=cost_details,
            prompt=prompt,
        )

    def _create_observation_from_otel_span(
        self,
        *,
        otel_span: otel_trace_api.Span,
        as_type: ObservationTypeLiteralNoEvent,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
        completion_start_time: Optional[datetime] = None,
        model: Optional[str] = None,
        model_parameters: Optional[Dict[str, MapValue]] = None,
        usage_details: Optional[Dict[str, int]] = None,
        cost_details: Optional[Dict[str, float]] = None,
        prompt: Optional[PromptClient] = None,
    ) -> Union[
        ABVSpan,
        ABVGeneration,
        ABVAgent,
        ABVTool,
        ABVChain,
        ABVRetriever,
        ABVEvaluator,
        ABVEmbedding,
        ABVGuardrail,
    ]:
        """Create the appropriate observation type from an OTEL span."""
        if as_type in get_observation_types_list(ObservationTypeGenerationLike):
            observation_class = self._get_span_class(as_type)
            # Type ignore to prevent overloads of internal _get_span_class function,
            # issue is that ABVEvent could be returned and that classes have diff. args
            return observation_class(  # type: ignore[return-value,call-arg]
                otel_span=otel_span,
                abv_client=self,
                environment=self._environment,
                input=input,
                output=output,
                metadata=metadata,
                version=version,
                level=level,
                status_message=status_message,
                completion_start_time=completion_start_time,
                model=model,
                model_parameters=model_parameters,
                usage_details=usage_details,
                cost_details=cost_details,
                prompt=prompt,
            )
        else:
            # For other types (e.g. span, guardrail), create appropriate class without generation properties
            observation_class = self._get_span_class(as_type)
            # Type ignore to prevent overloads of internal _get_span_class function,
            # issue is that ABVEvent could be returned and that classes have diff. args
            return observation_class(  # type: ignore[return-value,call-arg]
                otel_span=otel_span,
                abv_client=self,
                environment=self._environment,
                input=input,
                output=output,
                metadata=metadata,
                version=version,
                level=level,
                status_message=status_message,
            )
            # span._observation_type = as_type
            # span._otel_span.set_attribute("abv.observation.type", as_type)
            # return span

    def start_generation(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
        completion_start_time: Optional[datetime] = None,
        model: Optional[str] = None,
        model_parameters: Optional[Dict[str, MapValue]] = None,
        usage_details: Optional[Dict[str, int]] = None,
        cost_details: Optional[Dict[str, float]] = None,
        prompt: Optional[PromptClient] = None,
    ) -> ABVGeneration:
        """[DEPRECATED] Create a new generation span for model generations.

        DEPRECATED: This method is deprecated and will be removed in a future version.
        Use start_observation(as_type='generation') instead.

        This method creates a specialized span for tracking model generations.
        It includes additional fields specific to model generations such as model name,
        token usage, and cost details.

        The created generation span will be the child of the current span in the context.

        Args:
            trace_context: Optional context for connecting to an existing trace
            name: Name of the generation operation
            input: Input data for the model (e.g., prompts)
            output: Output from the model (e.g., completions)
            metadata: Additional metadata to associate with the generation
            version: Version identifier for the model or component
            level: Importance level of the generation (info, warning, error)
            status_message: Optional status message for the generation
            completion_start_time: When the model started generating the response
            model: Name/identifier of the AI model used (e.g., "gpt-4")
            model_parameters: Parameters used for the model (e.g., temperature, max_tokens)
            usage_details: Token usage information (e.g., prompt_tokens, completion_tokens)
            cost_details: Cost information for the model call
            prompt: Associated prompt template from ABV prompt management

        Returns:
            A ABVGeneration object that must be ended with .end() when complete

        Example:
            ```python
            generation = abv.start_generation(
                name="answer-generation",
                model="gpt-4",
                input={"prompt": "Explain quantum computing"},
                model_parameters={"temperature": 0.7}
            )
            try:
                # Call model API
                response = llm.generate(...)

                generation.update(
                    output=response.text,
                    usage_details={
                        "prompt_tokens": response.usage.prompt_tokens,
                        "completion_tokens": response.usage.completion_tokens
                    }
                )
            finally:
                generation.end()
            ```
        """
        warnings.warn(
            "start_generation is deprecated and will be removed in a future version. "
            "Use start_observation(as_type='generation') instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        return self.start_observation(
            trace_context=trace_context,
            name=name,
            as_type="generation",
            input=input,
            output=output,
            metadata=metadata,
            version=version,
            level=level,
            status_message=status_message,
            completion_start_time=completion_start_time,
            model=model,
            model_parameters=model_parameters,
            usage_details=usage_details,
            cost_details=cost_details,
            prompt=prompt,
        )

    def start_as_current_generation(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
        completion_start_time: Optional[datetime] = None,
        model: Optional[str] = None,
        model_parameters: Optional[Dict[str, MapValue]] = None,
        usage_details: Optional[Dict[str, int]] = None,
        cost_details: Optional[Dict[str, float]] = None,
        prompt: Optional[PromptClient] = None,
        end_on_exit: Optional[bool] = None,
    ) -> _AgnosticContextManager[ABVGeneration]:
        """[DEPRECATED] Create a new generation span and set it as the current span in a context manager.

        DEPRECATED: This method is deprecated and will be removed in a future version.
        Use start_as_current_observation(as_type='generation') instead.

        This method creates a specialized span for model generations and sets it as the
        current span within a context manager. Use this method with a 'with' statement to
        automatically handle the generation span lifecycle within a code block.

        The created generation span will be the child of the current span in the context.

        Args:
            trace_context: Optional context for connecting to an existing trace
            name: Name of the generation operation
            input: Input data for the model (e.g., prompts)
            output: Output from the model (e.g., completions)
            metadata: Additional metadata to associate with the generation
            version: Version identifier for the model or component
            level: Importance level of the generation (info, warning, error)
            status_message: Optional status message for the generation
            completion_start_time: When the model started generating the response
            model: Name/identifier of the AI model used (e.g., "gpt-4")
            model_parameters: Parameters used for the model (e.g., temperature, max_tokens)
            usage_details: Token usage information (e.g., prompt_tokens, completion_tokens)
            cost_details: Cost information for the model call
            prompt: Associated prompt template from ABV prompt management
            end_on_exit (default: True): Whether to end the span automatically when leaving the context manager. If False, the span must be manually ended to avoid memory leaks.

        Returns:
            A context manager that yields a ABVGeneration

        Example:
            ```python
            with abv.start_as_current_generation(
                name="answer-generation",
                model="gpt-4",
                input={"prompt": "Explain quantum computing"}
            ) as generation:
                # Call model API
                response = llm.generate(...)

                # Update with results
                generation.update(
                    output=response.text,
                    usage_details={
                        "prompt_tokens": response.usage.prompt_tokens,
                        "completion_tokens": response.usage.completion_tokens
                    }
                )
            ```
        """
        warnings.warn(
            "start_as_current_generation is deprecated and will be removed in a future version. "
            "Use start_as_current_observation(as_type='generation') instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        return self.start_as_current_observation(
            trace_context=trace_context,
            name=name,
            as_type="generation",
            input=input,
            output=output,
            metadata=metadata,
            version=version,
            level=level,
            status_message=status_message,
            completion_start_time=completion_start_time,
            model=model,
            model_parameters=model_parameters,
            usage_details=usage_details,
            cost_details=cost_details,
            prompt=prompt,
            end_on_exit=end_on_exit,
        )

    @overload
    def start_as_current_observation(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        as_type: Literal["generation"],
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
        completion_start_time: Optional[datetime] = None,
        model: Optional[str] = None,
        model_parameters: Optional[Dict[str, MapValue]] = None,
        usage_details: Optional[Dict[str, int]] = None,
        cost_details: Optional[Dict[str, float]] = None,
        prompt: Optional[PromptClient] = None,
        end_on_exit: Optional[bool] = None,
    ) -> _AgnosticContextManager[ABVGeneration]: ...

    @overload
    def start_as_current_observation(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        as_type: Literal["span"] = "span",
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
        end_on_exit: Optional[bool] = None,
    ) -> _AgnosticContextManager[ABVSpan]: ...

    @overload
    def start_as_current_observation(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        as_type: Literal["agent"],
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
        end_on_exit: Optional[bool] = None,
    ) -> _AgnosticContextManager[ABVAgent]: ...

    @overload
    def start_as_current_observation(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        as_type: Literal["tool"],
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
        end_on_exit: Optional[bool] = None,
    ) -> _AgnosticContextManager[ABVTool]: ...

    @overload
    def start_as_current_observation(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        as_type: Literal["chain"],
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
        end_on_exit: Optional[bool] = None,
    ) -> _AgnosticContextManager[ABVChain]: ...

    @overload
    def start_as_current_observation(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        as_type: Literal["retriever"],
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
        end_on_exit: Optional[bool] = None,
    ) -> _AgnosticContextManager[ABVRetriever]: ...

    @overload
    def start_as_current_observation(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        as_type: Literal["evaluator"],
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
        end_on_exit: Optional[bool] = None,
    ) -> _AgnosticContextManager[ABVEvaluator]: ...

    @overload
    def start_as_current_observation(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        as_type: Literal["embedding"],
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
        completion_start_time: Optional[datetime] = None,
        model: Optional[str] = None,
        model_parameters: Optional[Dict[str, MapValue]] = None,
        usage_details: Optional[Dict[str, int]] = None,
        cost_details: Optional[Dict[str, float]] = None,
        prompt: Optional[PromptClient] = None,
        end_on_exit: Optional[bool] = None,
    ) -> _AgnosticContextManager[ABVEmbedding]: ...

    @overload
    def start_as_current_observation(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        as_type: Literal["guardrail"],
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
        end_on_exit: Optional[bool] = None,
    ) -> _AgnosticContextManager[ABVGuardrail]: ...

    def start_as_current_observation(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        as_type: ObservationTypeLiteralNoEvent = "span",
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
        completion_start_time: Optional[datetime] = None,
        model: Optional[str] = None,
        model_parameters: Optional[Dict[str, MapValue]] = None,
        usage_details: Optional[Dict[str, int]] = None,
        cost_details: Optional[Dict[str, float]] = None,
        prompt: Optional[PromptClient] = None,
        end_on_exit: Optional[bool] = None,
    ) -> Union[
        _AgnosticContextManager[ABVGeneration],
        _AgnosticContextManager[ABVSpan],
        _AgnosticContextManager[ABVAgent],
        _AgnosticContextManager[ABVTool],
        _AgnosticContextManager[ABVChain],
        _AgnosticContextManager[ABVRetriever],
        _AgnosticContextManager[ABVEvaluator],
        _AgnosticContextManager[ABVEmbedding],
        _AgnosticContextManager[ABVGuardrail],
    ]:
        """Create a new observation and set it as the current span in a context manager.

        This method creates a new observation of the specified type and sets it as the
        current span within a context manager. Use this method with a 'with' statement to
        automatically handle the observation lifecycle within a code block.

        The created observation will be the child of the current span in the context.

        Args:
            trace_context: Optional context for connecting to an existing trace
            name: Name of the observation (e.g., function or operation name)
            as_type: Type of observation to create (defaults to "span")
            input: Input data for the operation (can be any JSON-serializable object)
            output: Output data from the operation (can be any JSON-serializable object)
            metadata: Additional metadata to associate with the observation
            version: Version identifier for the code or component
            level: Importance level of the observation (info, warning, error)
            status_message: Optional status message for the observation
            end_on_exit (default: True): Whether to end the span automatically when leaving the context manager. If False, the span must be manually ended to avoid memory leaks.

            The following parameters are available when as_type is: "generation" or "embedding".
            completion_start_time: When the model started generating the response
            model: Name/identifier of the AI model used (e.g., "gpt-4")
            model_parameters: Parameters used for the model (e.g., temperature, max_tokens)
            usage_details: Token usage information (e.g., prompt_tokens, completion_tokens)
            cost_details: Cost information for the model call
            prompt: Associated prompt template from ABV prompt management

        Returns:
            A context manager that yields the appropriate observation type based on as_type

        Example:
            ```python
            # Create a span
            with abv.start_as_current_observation(name="process-query", as_type="span") as span:
                # Do work
                result = process_data()
                span.update(output=result)

                # Create a child span automatically
                with span.start_as_current_span(name="sub-operation") as child_span:
                    # Do sub-operation work
                    child_span.update(output="sub-result")

            # Create a tool observation
            with abv.start_as_current_observation(name="web-search", as_type="tool") as tool:
                # Do tool work
                results = search_web(query)
                tool.update(output=results)

            # Create a generation observation
            with abv.start_as_current_observation(
                name="answer-generation",
                as_type="generation",
                model="gpt-4"
            ) as generation:
                # Generate answer
                response = llm.generate(...)
                generation.update(output=response)
            ```
        """
        if as_type in get_observation_types_list(ObservationTypeGenerationLike):
            if trace_context:
                trace_id = trace_context.get("trace_id", None)
                parent_span_id = trace_context.get("parent_span_id", None)

                if trace_id:
                    remote_parent_span = self._create_remote_parent_span(
                        trace_id=trace_id, parent_span_id=parent_span_id
                    )

                    return cast(
                        Union[
                            _AgnosticContextManager[ABVGeneration],
                            _AgnosticContextManager[ABVEmbedding],
                        ],
                        self._create_span_with_parent_context(
                            as_type=as_type,
                            name=name,
                            remote_parent_span=remote_parent_span,
                            parent=None,
                            end_on_exit=end_on_exit,
                            input=input,
                            output=output,
                            metadata=metadata,
                            version=version,
                            level=level,
                            status_message=status_message,
                            completion_start_time=completion_start_time,
                            model=model,
                            model_parameters=model_parameters,
                            usage_details=usage_details,
                            cost_details=cost_details,
                            prompt=prompt,
                        ),
                    )

            return cast(
                Union[
                    _AgnosticContextManager[ABVGeneration],
                    _AgnosticContextManager[ABVEmbedding],
                ],
                self._start_as_current_otel_span_with_processed_media(
                    as_type=as_type,
                    name=name,
                    end_on_exit=end_on_exit,
                    input=input,
                    output=output,
                    metadata=metadata,
                    version=version,
                    level=level,
                    status_message=status_message,
                    completion_start_time=completion_start_time,
                    model=model,
                    model_parameters=model_parameters,
                    usage_details=usage_details,
                    cost_details=cost_details,
                    prompt=prompt,
                ),
            )

        if as_type in get_observation_types_list(ObservationTypeSpanLike):
            if trace_context:
                trace_id = trace_context.get("trace_id", None)
                parent_span_id = trace_context.get("parent_span_id", None)

                if trace_id:
                    remote_parent_span = self._create_remote_parent_span(
                        trace_id=trace_id, parent_span_id=parent_span_id
                    )

                    return cast(
                        Union[
                            _AgnosticContextManager[ABVSpan],
                            _AgnosticContextManager[ABVAgent],
                            _AgnosticContextManager[ABVTool],
                            _AgnosticContextManager[ABVChain],
                            _AgnosticContextManager[ABVRetriever],
                            _AgnosticContextManager[ABVEvaluator],
                            _AgnosticContextManager[ABVGuardrail],
                        ],
                        self._create_span_with_parent_context(
                            as_type=as_type,
                            name=name,
                            remote_parent_span=remote_parent_span,
                            parent=None,
                            end_on_exit=end_on_exit,
                            input=input,
                            output=output,
                            metadata=metadata,
                            version=version,
                            level=level,
                            status_message=status_message,
                        ),
                    )

            return cast(
                Union[
                    _AgnosticContextManager[ABVSpan],
                    _AgnosticContextManager[ABVAgent],
                    _AgnosticContextManager[ABVTool],
                    _AgnosticContextManager[ABVChain],
                    _AgnosticContextManager[ABVRetriever],
                    _AgnosticContextManager[ABVEvaluator],
                    _AgnosticContextManager[ABVGuardrail],
                ],
                self._start_as_current_otel_span_with_processed_media(
                    as_type=as_type,
                    name=name,
                    end_on_exit=end_on_exit,
                    input=input,
                    output=output,
                    metadata=metadata,
                    version=version,
                    level=level,
                    status_message=status_message,
                ),
            )

        # This should never be reached since all valid types are handled above
        abv_logger.warning(
            f"Unknown observation type: {as_type}, falling back to span"
        )
        return self._start_as_current_otel_span_with_processed_media(
            as_type="span",
            name=name,
            end_on_exit=end_on_exit,
            input=input,
            output=output,
            metadata=metadata,
            version=version,
            level=level,
            status_message=status_message,
        )

    def _get_span_class(
        self,
        as_type: ObservationTypeLiteral,
    ) -> Union[
        Type[ABVAgent],
        Type[ABVTool],
        Type[ABVChain],
        Type[ABVRetriever],
        Type[ABVEvaluator],
        Type[ABVEmbedding],
        Type[ABVGuardrail],
        Type[ABVGeneration],
        Type[ABVEvent],
        Type[ABVSpan],
    ]:
        """Get the appropriate span class based on as_type."""
        normalized_type = as_type.lower()

        if normalized_type == "agent":
            return ABVAgent
        elif normalized_type == "tool":
            return ABVTool
        elif normalized_type == "chain":
            return ABVChain
        elif normalized_type == "retriever":
            return ABVRetriever
        elif normalized_type == "evaluator":
            return ABVEvaluator
        elif normalized_type == "embedding":
            return ABVEmbedding
        elif normalized_type == "guardrail":
            return ABVGuardrail
        elif normalized_type == "generation":
            return ABVGeneration
        elif normalized_type == "event":
            return ABVEvent
        elif normalized_type == "span":
            return ABVSpan
        else:
            return ABVSpan

    @_agnosticcontextmanager
    def _create_span_with_parent_context(
        self,
        *,
        name: str,
        parent: Optional[otel_trace_api.Span] = None,
        remote_parent_span: Optional[otel_trace_api.Span] = None,
        as_type: ObservationTypeLiteralNoEvent,
        end_on_exit: Optional[bool] = None,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
        completion_start_time: Optional[datetime] = None,
        model: Optional[str] = None,
        model_parameters: Optional[Dict[str, MapValue]] = None,
        usage_details: Optional[Dict[str, int]] = None,
        cost_details: Optional[Dict[str, float]] = None,
        prompt: Optional[PromptClient] = None,
    ) -> Any:
        parent_span = parent or cast(otel_trace_api.Span, remote_parent_span)

        with otel_trace_api.use_span(parent_span):
            with self._start_as_current_otel_span_with_processed_media(
                name=name,
                as_type=as_type,
                end_on_exit=end_on_exit,
                input=input,
                output=output,
                metadata=metadata,
                version=version,
                level=level,
                status_message=status_message,
                completion_start_time=completion_start_time,
                model=model,
                model_parameters=model_parameters,
                usage_details=usage_details,
                cost_details=cost_details,
                prompt=prompt,
            ) as abv_span:
                if remote_parent_span is not None:
                    abv_span._otel_span.set_attribute(
                        ABVOtelSpanAttributes.AS_ROOT, True
                    )

                yield abv_span

    @_agnosticcontextmanager
    def _start_as_current_otel_span_with_processed_media(
        self,
        *,
        name: str,
        as_type: Optional[ObservationTypeLiteralNoEvent] = None,
        end_on_exit: Optional[bool] = None,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
        completion_start_time: Optional[datetime] = None,
        model: Optional[str] = None,
        model_parameters: Optional[Dict[str, MapValue]] = None,
        usage_details: Optional[Dict[str, int]] = None,
        cost_details: Optional[Dict[str, float]] = None,
        prompt: Optional[PromptClient] = None,
    ) -> Any:
        with self._otel_tracer.start_as_current_span(
            name=name,
            end_on_exit=end_on_exit if end_on_exit is not None else True,
        ) as otel_span:
            span_class = self._get_span_class(
                as_type or "generation"
            )  # default was "generation"
            common_args = {
                "otel_span": otel_span,
                "abv_client": self,
                "environment": self._environment,
                "input": input,
                "output": output,
                "metadata": metadata,
                "version": version,
                "level": level,
                "status_message": status_message,
            }

            if span_class in [
                ABVGeneration,
                ABVEmbedding,
            ]:
                common_args.update(
                    {
                        "completion_start_time": completion_start_time,
                        "model": model,
                        "model_parameters": model_parameters,
                        "usage_details": usage_details,
                        "cost_details": cost_details,
                        "prompt": prompt,
                    }
                )
            # For span-like types (span, agent, tool, chain, retriever, evaluator, guardrail), no generation properties needed

            yield span_class(**common_args)  # type: ignore[arg-type]

    def _get_current_otel_span(self) -> Optional[otel_trace_api.Span]:
        current_span = otel_trace_api.get_current_span()

        if current_span is otel_trace_api.INVALID_SPAN:
            abv_logger.warning(
                "Context error: No active span in current context. Operations that depend on an active span will be skipped. "
                "Ensure spans are created with start_as_current_span() or that you're operating within an active span context."
            )
            return None

        return current_span

    def update_current_generation(
        self,
        *,
        name: Optional[str] = None,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
        completion_start_time: Optional[datetime] = None,
        model: Optional[str] = None,
        model_parameters: Optional[Dict[str, MapValue]] = None,
        usage_details: Optional[Dict[str, int]] = None,
        cost_details: Optional[Dict[str, float]] = None,
        prompt: Optional[PromptClient] = None,
    ) -> None:
        """Update the current active generation span with new information.

        This method updates the current generation span in the active context with
        additional information. It's useful for adding output, usage stats, or other
        details that become available during or after model generation.

        Args:
            name: The generation name
            input: Updated input data for the model
            output: Output from the model (e.g., completions)
            metadata: Additional metadata to associate with the generation
            version: Version identifier for the model or component
            level: Importance level of the generation (info, warning, error)
            status_message: Optional status message for the generation
            completion_start_time: When the model started generating the response
            model: Name/identifier of the AI model used (e.g., "gpt-4")
            model_parameters: Parameters used for the model (e.g., temperature, max_tokens)
            usage_details: Token usage information (e.g., prompt_tokens, completion_tokens)
            cost_details: Cost information for the model call
            prompt: Associated prompt template from ABV prompt management

        Example:
            ```python
            with abv.start_as_current_generation(name="answer-query") as generation:
                # Initial setup and API call
                response = llm.generate(...)

                # Update with results that weren't available at creation time
                abv.update_current_generation(
                    output=response.text,
                    usage_details={
                        "prompt_tokens": response.usage.prompt_tokens,
                        "completion_tokens": response.usage.completion_tokens
                    }
                )
            ```
        """
        if not self._tracing_enabled:
            abv_logger.debug(
                "Operation skipped: update_current_generation - Tracing is disabled or client is in no-op mode."
            )
            return

        current_otel_span = self._get_current_otel_span()

        if current_otel_span is not None:
            generation = ABVGeneration(
                otel_span=current_otel_span, abv_client=self
            )

            if name:
                current_otel_span.update_name(name)

            generation.update(
                input=input,
                output=output,
                metadata=metadata,
                version=version,
                level=level,
                status_message=status_message,
                completion_start_time=completion_start_time,
                model=model,
                model_parameters=model_parameters,
                usage_details=usage_details,
                cost_details=cost_details,
                prompt=prompt,
            )

    def update_current_span(
        self,
        *,
        name: Optional[str] = None,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
    ) -> None:
        """Update the current active span with new information.

        This method updates the current span in the active context with
        additional information. It's useful for adding outputs or metadata
        that become available during execution.

        Args:
            name: The span name
            input: Updated input data for the operation
            output: Output data from the operation
            metadata: Additional metadata to associate with the span
            version: Version identifier for the code or component
            level: Importance level of the span (info, warning, error)
            status_message: Optional status message for the span

        Example:
            ```python
            with abv.start_as_current_span(name="process-data") as span:
                # Initial processing
                result = process_first_part()

                # Update with intermediate results
                abv.update_current_span(metadata={"intermediate_result": result})

                # Continue processing
                final_result = process_second_part(result)

                # Final update
                abv.update_current_span(output=final_result)
            ```
        """
        if not self._tracing_enabled:
            abv_logger.debug(
                "Operation skipped: update_current_span - Tracing is disabled or client is in no-op mode."
            )
            return

        current_otel_span = self._get_current_otel_span()

        if current_otel_span is not None:
            span = ABVSpan(
                otel_span=current_otel_span,
                abv_client=self,
                environment=self._environment,
            )

            if name:
                current_otel_span.update_name(name)

            span.update(
                input=input,
                output=output,
                metadata=metadata,
                version=version,
                level=level,
                status_message=status_message,
            )

    def update_current_trace(
        self,
        *,
        name: Optional[str] = None,
        user_id: Optional[str] = None,
        session_id: Optional[str] = None,
        version: Optional[str] = None,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        tags: Optional[List[str]] = None,
        public: Optional[bool] = None,
    ) -> None:
        """Update the current trace with additional information.

        This method updates the ABV trace that the current span belongs to. It's useful for
        adding trace-level metadata like user ID, session ID, or tags that apply to
        the entire ABV trace rather than just a single observation.

        Args:
            name: Updated name for the ABV trace
            user_id: ID of the user who initiated the ABV trace
            session_id: Session identifier for grouping related ABV traces
            version: Version identifier for the application or service
            input: Input data for the overall ABV trace
            output: Output data from the overall ABV trace
            metadata: Additional metadata to associate with the ABV trace
            tags: List of tags to categorize the ABV trace
            public: Whether the ABV trace should be publicly accessible

        Example:
            ```python
            with abv.start_as_current_span(name="handle-request") as span:
                # Get user information
                user = authenticate_user(request)

                # Update trace with user context
                abv.update_current_trace(
                    user_id=user.id,
                    session_id=request.session_id,
                    tags=["production", "web-app"]
                )

                # Continue processing
                response = process_request(request)

                # Update span with results
                span.update(output=response)
            ```
        """
        if not self._tracing_enabled:
            abv_logger.debug(
                "Operation skipped: update_current_trace - Tracing is disabled or client is in no-op mode."
            )
            return

        current_otel_span = self._get_current_otel_span()

        if current_otel_span is not None:
            existing_observation_type = current_otel_span.attributes.get(  # type: ignore[attr-defined]
                ABVOtelSpanAttributes.OBSERVATION_TYPE, "span"
            )
            # We need to preserve the class to keep the corret observation type
            span_class = self._get_span_class(existing_observation_type)
            span = span_class(
                otel_span=current_otel_span,
                abv_client=self,
                environment=self._environment,
            )

            span.update_trace(
                name=name,
                user_id=user_id,
                session_id=session_id,
                version=version,
                input=input,
                output=output,
                metadata=metadata,
                tags=tags,
                public=public,
            )

    def create_event(
        self,
        *,
        trace_context: Optional[TraceContext] = None,
        name: str,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[SpanLevel] = None,
        status_message: Optional[str] = None,
    ) -> ABVEvent:
        """Create a new ABV observation of type 'EVENT'.

        The created ABV Event observation will be the child of the current span in the context.

        Args:
            trace_context: Optional context for connecting to an existing trace
            name: Name of the span (e.g., function or operation name)
            input: Input data for the operation (can be any JSON-serializable object)
            output: Output data from the operation (can be any JSON-serializable object)
            metadata: Additional metadata to associate with the span
            version: Version identifier for the code or component
            level: Importance level of the span (info, warning, error)
            status_message: Optional status message for the span

        Returns:
            The ABV Event object

        Example:
            ```python
            event = abv.create_event(name="process-event")
            ```
        """
        timestamp = time_ns()

        if trace_context:
            trace_id = trace_context.get("trace_id", None)
            parent_span_id = trace_context.get("parent_span_id", None)

            if trace_id:
                remote_parent_span = self._create_remote_parent_span(
                    trace_id=trace_id, parent_span_id=parent_span_id
                )

                with otel_trace_api.use_span(
                    cast(otel_trace_api.Span, remote_parent_span)
                ):
                    otel_span = self._otel_tracer.start_span(
                        name=name, start_time=timestamp
                    )
                    otel_span.set_attribute(ABVOtelSpanAttributes.AS_ROOT, True)

                    return cast(
                        ABVEvent,
                        ABVEvent(
                            otel_span=otel_span,
                            abv_client=self,
                            environment=self._environment,
                            input=input,
                            output=output,
                            metadata=metadata,
                            version=version,
                            level=level,
                            status_message=status_message,
                        ).end(end_time=timestamp),
                    )

        otel_span = self._otel_tracer.start_span(name=name, start_time=timestamp)

        return cast(
            ABVEvent,
            ABVEvent(
                otel_span=otel_span,
                abv_client=self,
                environment=self._environment,
                input=input,
                output=output,
                metadata=metadata,
                version=version,
                level=level,
                status_message=status_message,
            ).end(end_time=timestamp),
        )

    def _create_remote_parent_span(
        self, *, trace_id: str, parent_span_id: Optional[str]
    ) -> Any:
        if not self._is_valid_trace_id(trace_id):
            abv_logger.warning(
                f"Passed trace ID '{trace_id}' is not a valid 32 lowercase hex char ABV trace id. Ignoring trace ID."
            )

        if parent_span_id and not self._is_valid_span_id(parent_span_id):
            abv_logger.warning(
                f"Passed span ID '{parent_span_id}' is not a valid 16 lowercase hex char ABV span id. Ignoring parent span ID."
            )

        int_trace_id = int(trace_id, 16)
        int_parent_span_id = (
            int(parent_span_id, 16)
            if parent_span_id
            else RandomIdGenerator().generate_span_id()
        )

        span_context = otel_trace_api.SpanContext(
            trace_id=int_trace_id,
            span_id=int_parent_span_id,
            trace_flags=otel_trace_api.TraceFlags(0x01),  # mark span as sampled
            is_remote=False,
        )

        return trace.NonRecordingSpan(span_context)

    def _is_valid_trace_id(self, trace_id: str) -> bool:
        pattern = r"^[0-9a-f]{32}$"

        return bool(re.match(pattern, trace_id))

    def _is_valid_span_id(self, span_id: str) -> bool:
        pattern = r"^[0-9a-f]{16}$"

        return bool(re.match(pattern, span_id))

    def _create_observation_id(self, *, seed: Optional[str] = None) -> str:
        """Create a unique observation ID for use with ABV.

        This method generates a unique observation ID (span ID in OpenTelemetry terms)
        for use with various ABV APIs. It can either generate a random ID or
        create a deterministic ID based on a seed string.

        Observation IDs must be 16 lowercase hexadecimal characters, representing 8 bytes.
        This method ensures the generated ID meets this requirement. If you need to
        correlate an external ID with a ABV observation ID, use the external ID as
        the seed to get a valid, deterministic observation ID.

        Args:
            seed: Optional string to use as a seed for deterministic ID generation.
                 If provided, the same seed will always produce the same ID.
                 If not provided, a random ID will be generated.

        Returns:
            A 16-character lowercase hexadecimal string representing the observation ID.

        Example:
            ```python
            # Generate a random observation ID
            obs_id = abv.create_observation_id()

            # Generate a deterministic ID based on a seed
            user_obs_id = abv.create_observation_id(seed="user-123-feedback")

            # Correlate an external item ID with a ABV observation ID
            item_id = "item-789012"
            correlated_obs_id = abv.create_observation_id(seed=item_id)

            # Use the ID with ABV APIs
            abv.create_score(
                name="relevance",
                value=0.95,
                trace_id=trace_id,
                observation_id=obs_id
            )
            ```
        """
        if not seed:
            span_id_int = RandomIdGenerator().generate_span_id()

            return self._format_otel_span_id(span_id_int)

        return sha256(seed.encode("utf-8")).digest()[:8].hex()

    @staticmethod
    def create_trace_id(*, seed: Optional[str] = None) -> str:
        """Create a unique trace ID for use with ABV.

        This method generates a unique trace ID for use with various ABV APIs.
        It can either generate a random ID or create a deterministic ID based on
        a seed string.

        Trace IDs must be 32 lowercase hexadecimal characters, representing 16 bytes.
        This method ensures the generated ID meets this requirement. If you need to
        correlate an external ID with a ABV trace ID, use the external ID as the
        seed to get a valid, deterministic ABV trace ID.

        Args:
            seed: Optional string to use as a seed for deterministic ID generation.
                 If provided, the same seed will always produce the same ID.
                 If not provided, a random ID will be generated.

        Returns:
            A 32-character lowercase hexadecimal string representing the ABV trace ID.

        Example:
            ```python
            # Generate a random trace ID
            trace_id = abv.create_trace_id()

            # Generate a deterministic ID based on a seed
            session_trace_id = abv.create_trace_id(seed="session-456")

            # Correlate an external ID with a ABV trace ID
            external_id = "external-system-123456"
            correlated_trace_id = abv.create_trace_id(seed=external_id)

            # Use the ID with trace context
            with abv.start_as_current_span(
                name="process-request",
                trace_context={"trace_id": trace_id}
            ) as span:
                # Operation will be part of the specific trace
                pass
            ```
        """
        if not seed:
            trace_id_int = RandomIdGenerator().generate_trace_id()

            return ABV._format_otel_trace_id(trace_id_int)

        return sha256(seed.encode("utf-8")).digest()[:16].hex()

    def _get_otel_trace_id(self, otel_span: otel_trace_api.Span) -> str:
        span_context = otel_span.get_span_context()

        return self._format_otel_trace_id(span_context.trace_id)

    def _get_otel_span_id(self, otel_span: otel_trace_api.Span) -> str:
        span_context = otel_span.get_span_context()

        return self._format_otel_span_id(span_context.span_id)

    @staticmethod
    def _format_otel_span_id(span_id_int: int) -> str:
        """Format an integer span ID to a 16-character lowercase hex string.

        Internal method to convert an OpenTelemetry integer span ID to the standard
        W3C Trace Context format (16-character lowercase hex string).

        Args:
            span_id_int: 64-bit integer representing a span ID

        Returns:
            A 16-character lowercase hexadecimal string
        """
        return format(span_id_int, "016x")

    @staticmethod
    def _format_otel_trace_id(trace_id_int: int) -> str:
        """Format an integer trace ID to a 32-character lowercase hex string.

        Internal method to convert an OpenTelemetry integer trace ID to the standard
        W3C Trace Context format (32-character lowercase hex string).

        Args:
            trace_id_int: 128-bit integer representing a trace ID

        Returns:
            A 32-character lowercase hexadecimal string
        """
        return format(trace_id_int, "032x")

    @overload
    def create_score(
        self,
        *,
        name: str,
        value: float,
        session_id: Optional[str] = None,
        dataset_run_id: Optional[str] = None,
        trace_id: Optional[str] = None,
        observation_id: Optional[str] = None,
        score_id: Optional[str] = None,
        data_type: Optional[Literal["NUMERIC", "BOOLEAN"]] = None,
        comment: Optional[str] = None,
        config_id: Optional[str] = None,
        metadata: Optional[Any] = None,
    ) -> None: ...

    @overload
    def create_score(
        self,
        *,
        name: str,
        value: str,
        session_id: Optional[str] = None,
        dataset_run_id: Optional[str] = None,
        trace_id: Optional[str] = None,
        score_id: Optional[str] = None,
        observation_id: Optional[str] = None,
        data_type: Optional[Literal["CATEGORICAL"]] = "CATEGORICAL",
        comment: Optional[str] = None,
        config_id: Optional[str] = None,
        metadata: Optional[Any] = None,
    ) -> None: ...

    def create_score(
        self,
        *,
        name: str,
        value: Union[float, str],
        session_id: Optional[str] = None,
        dataset_run_id: Optional[str] = None,
        trace_id: Optional[str] = None,
        observation_id: Optional[str] = None,
        score_id: Optional[str] = None,
        data_type: Optional[ScoreDataType] = None,
        comment: Optional[str] = None,
        config_id: Optional[str] = None,
        metadata: Optional[Any] = None,
    ) -> None:
        """Create a score for a specific trace or observation.

        This method creates a score for evaluating a ABV trace or observation. Scores can be
        used to track quality metrics, user feedback, or automated evaluations.

        Args:
            name: Name of the score (e.g., "relevance", "accuracy")
            value: Score value (can be numeric for NUMERIC/BOOLEAN types or string for CATEGORICAL)
            session_id: ID of the ABV session to associate the score with
            dataset_run_id: ID of the ABV dataset run to associate the score with
            trace_id: ID of the ABV trace to associate the score with
            observation_id: Optional ID of the specific observation to score. Trace ID must be provided too.
            score_id: Optional custom ID for the score (auto-generated if not provided)
            data_type: Type of score (NUMERIC, BOOLEAN, or CATEGORICAL)
            comment: Optional comment or explanation for the score
            config_id: Optional ID of a score config defined in ABV
            metadata: Optional metadata to be attached to the score

        Example:
            ```python
            # Create a numeric score for accuracy
            abv.create_score(
                name="accuracy",
                value=0.92,
                trace_id="abcdef1234567890abcdef1234567890",
                data_type="NUMERIC",
                comment="High accuracy with minor irrelevant details"
            )

            # Create a categorical score for sentiment
            abv.create_score(
                name="sentiment",
                value="positive",
                trace_id="abcdef1234567890abcdef1234567890",
                observation_id="abcdef1234567890",
                data_type="CATEGORICAL"
            )
            ```
        """
        if not self._tracing_enabled:
            return

        score_id = score_id or self._create_observation_id()

        try:
            new_body = ScoreBody(
                id=score_id,
                sessionId=session_id,
                datasetRunId=dataset_run_id,
                traceId=trace_id,
                observationId=observation_id,
                name=name,
                value=value,
                dataType=data_type,  # type: ignore
                comment=comment,
                configId=config_id,
                environment=self._environment,
                metadata=metadata,
            )

            event = {
                "id": self.create_trace_id(),
                "type": "score-create",
                "timestamp": _get_timestamp(),
                "body": new_body,
            }

            if self._resources is not None:
                # Force the score to be in sample if it was for a legacy trace ID, i.e. non-32 hexchar
                force_sample = (
                    not self._is_valid_trace_id(trace_id) if trace_id else True
                )

                self._resources.add_score_task(
                    event,
                    force_sample=force_sample,
                )

        except Exception as e:
            abv_logger.exception(
                f"Error creating score: Failed to process score event for trace_id={trace_id}, name={name}. Error: {e}"
            )

    @overload
    def score_current_span(
        self,
        *,
        name: str,
        value: float,
        score_id: Optional[str] = None,
        data_type: Optional[Literal["NUMERIC", "BOOLEAN"]] = None,
        comment: Optional[str] = None,
        config_id: Optional[str] = None,
    ) -> None: ...

    @overload
    def score_current_span(
        self,
        *,
        name: str,
        value: str,
        score_id: Optional[str] = None,
        data_type: Optional[Literal["CATEGORICAL"]] = "CATEGORICAL",
        comment: Optional[str] = None,
        config_id: Optional[str] = None,
    ) -> None: ...

    def score_current_span(
        self,
        *,
        name: str,
        value: Union[float, str],
        score_id: Optional[str] = None,
        data_type: Optional[ScoreDataType] = None,
        comment: Optional[str] = None,
        config_id: Optional[str] = None,
    ) -> None:
        """Create a score for the current active span.

        This method scores the currently active span in the context. It's a convenient
        way to score the current operation without needing to know its trace and span IDs.

        Args:
            name: Name of the score (e.g., "relevance", "accuracy")
            value: Score value (can be numeric for NUMERIC/BOOLEAN types or string for CATEGORICAL)
            score_id: Optional custom ID for the score (auto-generated if not provided)
            data_type: Type of score (NUMERIC, BOOLEAN, or CATEGORICAL)
            comment: Optional comment or explanation for the score
            config_id: Optional ID of a score config defined in ABV

        Example:
            ```python
            with abv.start_as_current_generation(name="answer-query") as generation:
                # Generate answer
                response = generate_answer(...)
                generation.update(output=response)

                # Score the generation
                abv.score_current_span(
                    name="relevance",
                    value=0.85,
                    data_type="NUMERIC",
                    comment="Mostly relevant but contains some tangential information"
                )
            ```
        """
        current_span = self._get_current_otel_span()

        if current_span is not None:
            trace_id = self._get_otel_trace_id(current_span)
            observation_id = self._get_otel_span_id(current_span)

            abv_logger.info(
                f"Score: Creating score name='{name}' value={value} for current span ({observation_id}) in trace {trace_id}"
            )

            self.create_score(
                trace_id=trace_id,
                observation_id=observation_id,
                name=name,
                value=cast(str, value),
                score_id=score_id,
                data_type=cast(Literal["CATEGORICAL"], data_type),
                comment=comment,
                config_id=config_id,
            )

    @overload
    def score_current_trace(
        self,
        *,
        name: str,
        value: float,
        score_id: Optional[str] = None,
        data_type: Optional[Literal["NUMERIC", "BOOLEAN"]] = None,
        comment: Optional[str] = None,
        config_id: Optional[str] = None,
    ) -> None: ...

    @overload
    def score_current_trace(
        self,
        *,
        name: str,
        value: str,
        score_id: Optional[str] = None,
        data_type: Optional[Literal["CATEGORICAL"]] = "CATEGORICAL",
        comment: Optional[str] = None,
        config_id: Optional[str] = None,
    ) -> None: ...

    def score_current_trace(
        self,
        *,
        name: str,
        value: Union[float, str],
        score_id: Optional[str] = None,
        data_type: Optional[ScoreDataType] = None,
        comment: Optional[str] = None,
        config_id: Optional[str] = None,
    ) -> None:
        """Create a score for the current trace.

        This method scores the trace of the currently active span. Unlike score_current_span,
        this method associates the score with the entire trace rather than a specific span.
        It's useful for scoring overall performance or quality of the entire operation.

        Args:
            name: Name of the score (e.g., "user_satisfaction", "overall_quality")
            value: Score value (can be numeric for NUMERIC/BOOLEAN types or string for CATEGORICAL)
            score_id: Optional custom ID for the score (auto-generated if not provided)
            data_type: Type of score (NUMERIC, BOOLEAN, or CATEGORICAL)
            comment: Optional comment or explanation for the score
            config_id: Optional ID of a score config defined in ABV

        Example:
            ```python
            with abv.start_as_current_span(name="process-user-request") as span:
                # Process request
                result = process_complete_request()
                span.update(output=result)

                # Score the overall trace
                abv.score_current_trace(
                    name="overall_quality",
                    value=0.95,
                    data_type="NUMERIC",
                    comment="High quality end-to-end response"
                )
            ```
        """
        current_span = self._get_current_otel_span()

        if current_span is not None:
            trace_id = self._get_otel_trace_id(current_span)

            abv_logger.info(
                f"Score: Creating score name='{name}' value={value} for entire trace {trace_id}"
            )

            self.create_score(
                trace_id=trace_id,
                name=name,
                value=cast(str, value),
                score_id=score_id,
                data_type=cast(Literal["CATEGORICAL"], data_type),
                comment=comment,
                config_id=config_id,
            )

    def flush(self) -> None:
        """Force flush all pending spans and events to the ABV API.

        This method manually flushes any pending spans, scores, and other events to the
        ABV API. It's useful in scenarios where you want to ensure all data is sent
        before proceeding, without waiting for the automatic flush interval.

        Example:
            ```python
            # Record some spans and scores
            with abv.start_as_current_span(name="operation") as span:
                # Do work...
                pass

            # Ensure all data is sent to ABV before proceeding
            abv.flush()

            # Continue with other work
            ```
        """
        if self._resources is not None:
            self._resources.flush()

    def shutdown(self) -> None:
        """Shut down the ABV client and flush all pending data.

        This method cleanly shuts down the ABV client, ensuring all pending data
        is flushed to the API and all background threads are properly terminated.

        It's important to call this method when your application is shutting down to
        prevent data loss and resource leaks. For most applications, using the client
        as a context manager or relying on the automatic shutdown via atexit is sufficient.

        Example:
            ```python
            # Initialize ABV
            abv = ABV(api_key="...")

            # Use ABV throughout your application
            # ...

            # When application is shutting down
            abv.shutdown()
            ```
        """
        if self._resources is not None:
            self._resources.shutdown()

    def get_current_trace_id(self) -> Optional[str]:
        """Get the trace ID of the current active span.

        This method retrieves the trace ID from the currently active span in the context.
        It can be used to get the trace ID for referencing in logs, external systems,
        or for creating related operations.

        Returns:
            The current trace ID as a 32-character lowercase hexadecimal string,
            or None if there is no active span.

        Example:
            ```python
            with abv.start_as_current_span(name="process-request") as span:
                # Get the current trace ID for reference
                trace_id = abv.get_current_trace_id()

                # Use it for external correlation
                log.info(f"Processing request with trace_id: {trace_id}")

                # Or pass to another system
                external_system.process(data, trace_id=trace_id)
            ```
        """
        if not self._tracing_enabled:
            abv_logger.debug(
                "Operation skipped: get_current_trace_id - Tracing is disabled or client is in no-op mode."
            )
            return None

        current_otel_span = self._get_current_otel_span()

        return self._get_otel_trace_id(current_otel_span) if current_otel_span else None

    def get_current_observation_id(self) -> Optional[str]:
        """Get the observation ID (span ID) of the current active span.

        This method retrieves the observation ID from the currently active span in the context.
        It can be used to get the observation ID for referencing in logs, external systems,
        or for creating scores or other related operations.

        Returns:
            The current observation ID as a 16-character lowercase hexadecimal string,
            or None if there is no active span.

        Example:
            ```python
            with abv.start_as_current_span(name="process-user-query") as span:
                # Get the current observation ID
                observation_id = abv.get_current_observation_id()

                # Store it for later reference
                cache.set(f"query_{query_id}_observation", observation_id)

                # Process the query...
            ```
        """
        if not self._tracing_enabled:
            abv_logger.debug(
                "Operation skipped: get_current_observation_id - Tracing is disabled or client is in no-op mode."
            )
            return None

        current_otel_span = self._get_current_otel_span()

        return self._get_otel_span_id(current_otel_span) if current_otel_span else None

    def _get_project_id(self) -> Optional[str]:
        """Fetch and return the current project id. Persisted across requests. Returns None if no project id is found for api keys."""
        if not self._project_id:
            proj = self.api.projects.get()
            if not proj.data or not proj.data[0].id:
                return None

            self._project_id = proj.data[0].id

        return self._project_id

    def get_trace_url(self, *, trace_id: Optional[str] = None) -> Optional[str]:
        """Get the URL to view a trace in the ABV UI.

        This method generates a URL that links directly to a trace in the ABV UI.
        It's useful for providing links in logs, notifications, or debugging tools.

        Args:
            trace_id: Optional trace ID to generate a URL for. If not provided,
                     the trace ID of the current active span will be used.

        Returns:
            A URL string pointing to the trace in the ABV UI,
            or None if the project ID couldn't be retrieved or no trace ID is available.

        Example:
            ```python
            # Get URL for the current trace
            with abv.start_as_current_span(name="process-request") as span:
                trace_url = abv.get_trace_url()
                log.info(f"Processing trace: {trace_url}")

            # Get URL for a specific trace
            specific_trace_url = abv.get_trace_url(trace_id="1234567890abcdef1234567890abcdef")
            send_notification(f"Review needed for trace: {specific_trace_url}")
            ```
        """
        project_id = self._get_project_id()
        final_trace_id = trace_id or self.get_current_trace_id()

        return (
            f"{self._host}/project/{project_id}/traces/{final_trace_id}"
            if project_id and final_trace_id
            else None
        )

    def get_dataset(
        self, name: str, *, fetch_items_page_size: Optional[int] = 50
    ) -> "DatasetClient":
        """Fetch a dataset by its name.

        Args:
            name (str): The name of the dataset to fetch.
            fetch_items_page_size (Optional[int]): All items of the dataset will be fetched in chunks of this size. Defaults to 50.

        Returns:
            DatasetClient: The dataset with the given name.
        """
        try:
            abv_logger.debug(f"Getting datasets {name}")
            dataset = self.api.datasets.get(dataset_name=name)

            dataset_items = []
            page = 1

            while True:
                new_items = self.api.dataset_items.list(
                    dataset_name=self._url_encode(name, is_url_param=True),
                    page=page,
                    limit=fetch_items_page_size,
                )
                dataset_items.extend(new_items.data)

                if new_items.meta.total_pages <= page:
                    break

                page += 1

            items = [DatasetItemClient(i, abv=self) for i in dataset_items]

            return DatasetClient(dataset, items=items)

        except Error as e:
            handle_fern_exception(e)
            raise e

    def auth_check(self) -> bool:
        """Check if the provided credentials (public and secret key) are valid.

        Raises:
            Exception: If no projects were found for the provided credentials.

        Note:
            This method is blocking. It is discouraged to use it in production code.
        """
        try:
            projects = self.api.projects.get()
            abv_logger.debug(
                f"Auth check successful, found {len(projects.data)} projects"
            )
            if len(projects.data) == 0:
                raise Exception(
                    "Auth check failed, no project found for the keys provided."
                )
            return True

        except AttributeError as e:
            abv_logger.warning(
                f"Auth check failed: Client not properly initialized. Error: {e}"
            )
            return False

        except Error as e:
            handle_fern_exception(e)
            raise e

    def create_dataset(
        self,
        *,
        name: str,
        description: Optional[str] = None,
        metadata: Optional[Any] = None,
    ) -> Dataset:
        """Create a dataset with the given name on ABV.

        Args:
            name: Name of the dataset to create.
            description: Description of the dataset. Defaults to None.
            metadata: Additional metadata. Defaults to None.

        Returns:
            Dataset: The created dataset as returned by the ABV API.
        """
        try:
            body = CreateDatasetRequest(
                name=name, description=description, metadata=metadata
            )
            abv_logger.debug(f"Creating datasets {body}")

            return self.api.datasets.create(request=body)

        except Error as e:
            handle_fern_exception(e)
            raise e

    def create_dataset_item(
        self,
        *,
        dataset_name: str,
        input: Optional[Any] = None,
        expected_output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        source_trace_id: Optional[str] = None,
        source_observation_id: Optional[str] = None,
        status: Optional[DatasetStatus] = None,
        id: Optional[str] = None,
    ) -> DatasetItem:
        """Create a dataset item.

        Upserts if an item with id already exists.

        Args:
            dataset_name: Name of the dataset in which the dataset item should be created.
            input: Input data. Defaults to None. Can contain any dict, list or scalar.
            expected_output: Expected output data. Defaults to None. Can contain any dict, list or scalar.
            metadata: Additional metadata. Defaults to None. Can contain any dict, list or scalar.
            source_trace_id: Id of the source trace. Defaults to None.
            source_observation_id: Id of the source observation. Defaults to None.
            status: Status of the dataset item. Defaults to ACTIVE for newly created items.
            id: Id of the dataset item. Defaults to None. Provide your own id if you want to dedupe dataset items. Id needs to be globally unique and cannot be reused across datasets.

        Returns:
            DatasetItem: The created dataset item as returned by the ABV API.

        Example:
            ```python
            from abvdev import ABV

            abv = ABV()

            # Uploading items to the ABV dataset named "capital_cities"
            abv.create_dataset_item(
                dataset_name="capital_cities",
                input={"input": {"country": "Italy"}},
                expected_output={"expected_output": "Rome"},
                metadata={"foo": "bar"}
            )
            ```
        """
        try:
            body = CreateDatasetItemRequest(
                datasetName=dataset_name,
                input=input,
                expectedOutput=expected_output,
                metadata=metadata,
                sourceTraceId=source_trace_id,
                sourceObservationId=source_observation_id,
                status=status,
                id=id,
            )
            abv_logger.debug(f"Creating dataset item {body}")
            return self.api.dataset_items.create(request=body)
        except Error as e:
            handle_fern_exception(e)
            raise e

    def resolve_media_references(
        self,
        *,
        obj: Any,
        resolve_with: Literal["base64_data_uri"],
        max_depth: int = 10,
        content_fetch_timeout_seconds: int = 5,
    ) -> Any:
        """Replace media reference strings in an object with base64 data URIs.

        This method recursively traverses an object (up to max_depth) looking for media reference strings
        in the format "@@@abvMedia:...@@@". When found, it (synchronously) fetches the actual media content using
        the provided ABV client and replaces the reference string with a base64 data URI.

        If fetching media content fails for a reference string, a warning is logged and the reference
        string is left unchanged.

        Args:
            obj: The object to process. Can be a primitive value, array, or nested object.
                If the object has a __dict__ attribute, a dict will be returned instead of the original object type.
            resolve_with: The representation of the media content to replace the media reference string with.
                Currently only "base64_data_uri" is supported.
            max_depth: int: The maximum depth to traverse the object. Default is 10.
            content_fetch_timeout_seconds: int: The timeout in seconds for fetching media content. Default is 5.

        Returns:
            A deep copy of the input object with all media references replaced with base64 data URIs where possible.
            If the input object has a __dict__ attribute, a dict will be returned instead of the original object type.

        Example:
            obj = {
                "image": "@@@abvMedia:type=image/jpeg|id=123|source=bytes@@@",
                "nested": {
                    "pdf": "@@@abvMedia:type=application/pdf|id=456|source=bytes@@@"
                }
            }

            result = await ABVMedia.resolve_media_references(obj, abv_client)

            # Result:
            # {
            #     "image": "data:image/jpeg;base64,/9j/4AAQSkZJRg...",
            #     "nested": {
            #         "pdf": "data:application/pdf;base64,JVBERi0xLjcK..."
            #     }
            # }
        """
        return ABVMedia.resolve_media_references(
            abv_client=self,
            obj=obj,
            resolve_with=resolve_with,
            max_depth=max_depth,
            content_fetch_timeout_seconds=content_fetch_timeout_seconds,
        )

    @overload
    def get_prompt(
        self,
        name: str,
        *,
        version: Optional[int] = None,
        label: Optional[str] = None,
        type: Literal["chat"],
        cache_ttl_seconds: Optional[int] = None,
        fallback: Optional[List[ChatMessageDict]] = None,
        max_retries: Optional[int] = None,
        fetch_timeout_seconds: Optional[int] = None,
    ) -> ChatPromptClient: ...

    @overload
    def get_prompt(
        self,
        name: str,
        *,
        version: Optional[int] = None,
        label: Optional[str] = None,
        type: Literal["text"] = "text",
        cache_ttl_seconds: Optional[int] = None,
        fallback: Optional[str] = None,
        max_retries: Optional[int] = None,
        fetch_timeout_seconds: Optional[int] = None,
    ) -> TextPromptClient: ...

    def get_prompt(
        self,
        name: str,
        *,
        version: Optional[int] = None,
        label: Optional[str] = None,
        type: Literal["chat", "text"] = "text",
        cache_ttl_seconds: Optional[int] = None,
        fallback: Union[Optional[List[ChatMessageDict]], Optional[str]] = None,
        max_retries: Optional[int] = None,
        fetch_timeout_seconds: Optional[int] = None,
    ) -> PromptClient:
        """Get a prompt.

        This method attempts to fetch the requested prompt from the local cache. If the prompt is not found
        in the cache or if the cached prompt has expired, it will try to fetch the prompt from the server again
        and update the cache. If fetching the new prompt fails, and there is an expired prompt in the cache, it will
        return the expired prompt as a fallback.

        Args:
            name (str): The name of the prompt to retrieve.

        Keyword Args:
            version (Optional[int]): The version of the prompt to retrieve. If no label and version is specified, the `production` label is returned. Specify either version or label, not both.
            label: Optional[str]: The label of the prompt to retrieve. If no label and version is specified, the `production` label is returned. Specify either version or label, not both.
            cache_ttl_seconds: Optional[int]: Time-to-live in seconds for caching the prompt. Must be specified as a
            keyword argument. If not set, defaults to 60 seconds. Disables caching if set to 0.
            type: Literal["chat", "text"]: The type of the prompt to retrieve. Defaults to "text".
            fallback: Union[Optional[List[ChatMessageDict]], Optional[str]]: The prompt string to return if fetching the prompt fails. Important on the first call where no cached prompt is available. Follows ABV prompt formatting with double curly braces for variables. Defaults to None.
            max_retries: Optional[int]: The maximum number of retries in case of API/network errors. Defaults to 2. The maximum value is 4. Retries have an exponential backoff with a maximum delay of 10 seconds.
            fetch_timeout_seconds: Optional[int]: The timeout in milliseconds for fetching the prompt. Defaults to the default timeout set on the SDK, which is 5 seconds per default.

        Returns:
            The prompt object retrieved from the cache or directly fetched if not cached or expired of type
            - TextPromptClient, if type argument is 'text'.
            - ChatPromptClient, if type argument is 'chat'.

        Raises:
            Exception: Propagates any exceptions raised during the fetching of a new prompt, unless there is an
            expired prompt in the cache, in which case it logs a warning and returns the expired prompt.
        """
        if self._resources is None:
            raise Error(
                "SDK is not correctly initalized. Check the init logs for more details."
            )
        if version is not None and label is not None:
            raise ValueError("Cannot specify both version and label at the same time.")

        if not name:
            raise ValueError("Prompt name cannot be empty.")

        cache_key = PromptCache.generate_cache_key(name, version=version, label=label)
        bounded_max_retries = self._get_bounded_max_retries(
            max_retries, default_max_retries=2, max_retries_upper_bound=4
        )

        abv_logger.debug(f"Getting prompt '{cache_key}'")
        cached_prompt = self._resources.prompt_cache.get(cache_key)

        if cached_prompt is None or cache_ttl_seconds == 0:
            abv_logger.debug(
                f"Prompt '{cache_key}' not found in cache or caching disabled."
            )
            try:
                return self._fetch_prompt_and_update_cache(
                    name,
                    version=version,
                    label=label,
                    ttl_seconds=cache_ttl_seconds,
                    max_retries=bounded_max_retries,
                    fetch_timeout_seconds=fetch_timeout_seconds,
                )
            except Exception as e:
                if fallback:
                    abv_logger.warning(
                        f"Returning fallback prompt for '{cache_key}' due to fetch error: {e}"
                    )

                    fallback_client_args: Dict[str, Any] = {
                        "name": name,
                        "prompt": fallback,
                        "type": type,
                        "version": version or 0,
                        "config": {},
                        "labels": [label] if label else [],
                        "tags": [],
                    }

                    if type == "text":
                        return TextPromptClient(
                            prompt=Prompt_Text(**fallback_client_args),
                            is_fallback=True,
                        )

                    if type == "chat":
                        return ChatPromptClient(
                            prompt=Prompt_Chat(**fallback_client_args),
                            is_fallback=True,
                        )

                raise e

        if cached_prompt.is_expired():
            abv_logger.debug(f"Stale prompt '{cache_key}' found in cache.")
            try:
                # refresh prompt in background thread, refresh_prompt deduplicates tasks
                abv_logger.debug(f"Refreshing prompt '{cache_key}' in background.")

                def refresh_task() -> None:
                    self._fetch_prompt_and_update_cache(
                        name,
                        version=version,
                        label=label,
                        ttl_seconds=cache_ttl_seconds,
                        max_retries=bounded_max_retries,
                        fetch_timeout_seconds=fetch_timeout_seconds,
                    )

                self._resources.prompt_cache.add_refresh_prompt_task(
                    cache_key,
                    refresh_task,
                )
                abv_logger.debug(
                    f"Returning stale prompt '{cache_key}' from cache."
                )
                # return stale prompt
                return cached_prompt.value

            except Exception as e:
                abv_logger.warning(
                    f"Error when refreshing cached prompt '{cache_key}', returning cached version. Error: {e}"
                )
                # creation of refresh prompt task failed, return stale prompt
                return cached_prompt.value

        return cached_prompt.value

    def _fetch_prompt_and_update_cache(
        self,
        name: str,
        *,
        version: Optional[int] = None,
        label: Optional[str] = None,
        ttl_seconds: Optional[int] = None,
        max_retries: int,
        fetch_timeout_seconds: Optional[int],
    ) -> PromptClient:
        cache_key = PromptCache.generate_cache_key(name, version=version, label=label)
        abv_logger.debug(f"Fetching prompt '{cache_key}' from server...")

        try:

            @backoff.on_exception(
                backoff.constant, Exception, max_tries=max_retries + 1, logger=None
            )
            def fetch_prompts() -> Any:
                return self.api.prompts.get(
                    self._url_encode(name),
                    version=version,
                    label=label,
                    request_options={
                        "timeout_in_seconds": fetch_timeout_seconds,
                    }
                    if fetch_timeout_seconds is not None
                    else None,
                )

            prompt_response = fetch_prompts()

            prompt: PromptClient
            if prompt_response.type == "chat":
                prompt = ChatPromptClient(prompt_response)
            else:
                prompt = TextPromptClient(prompt_response)

            if self._resources is not None:
                self._resources.prompt_cache.set(cache_key, prompt, ttl_seconds)

            return prompt

        except Exception as e:
            abv_logger.error(
                f"Error while fetching prompt '{cache_key}': {str(e)}"
            )
            raise e

    def _get_bounded_max_retries(
        self,
        max_retries: Optional[int],
        *,
        default_max_retries: int = 2,
        max_retries_upper_bound: int = 4,
    ) -> int:
        if max_retries is None:
            return default_max_retries

        bounded_max_retries = min(
            max(max_retries, 0),
            max_retries_upper_bound,
        )

        return bounded_max_retries

    @overload
    def create_prompt(
        self,
        *,
        name: str,
        prompt: List[Union[ChatMessageDict, ChatMessageWithPlaceholdersDict]],
        labels: List[str] = [],
        tags: Optional[List[str]] = None,
        type: Optional[Literal["chat"]],
        config: Optional[Any] = None,
        commit_message: Optional[str] = None,
    ) -> ChatPromptClient: ...

    @overload
    def create_prompt(
        self,
        *,
        name: str,
        prompt: str,
        labels: List[str] = [],
        tags: Optional[List[str]] = None,
        type: Optional[Literal["text"]] = "text",
        config: Optional[Any] = None,
        commit_message: Optional[str] = None,
    ) -> TextPromptClient: ...

    def create_prompt(
        self,
        *,
        name: str,
        prompt: Union[
            str, List[Union[ChatMessageDict, ChatMessageWithPlaceholdersDict]]
        ],
        labels: List[str] = [],
        tags: Optional[List[str]] = None,
        type: Optional[Literal["chat", "text"]] = "text",
        config: Optional[Any] = None,
        commit_message: Optional[str] = None,
    ) -> PromptClient:
        """Create a new prompt in ABV.

        Keyword Args:
            name : The name of the prompt to be created.
            prompt : The content of the prompt to be created.
            is_active [DEPRECATED] : A flag indicating whether the prompt is active or not. This is deprecated and will be removed in a future release. Please use the 'production' label instead.
            labels: The labels of the prompt. Defaults to None. To create a default-served prompt, add the 'production' label.
            tags: The tags of the prompt. Defaults to None. Will be applied to all versions of the prompt.
            config: Additional structured data to be saved with the prompt. Defaults to None.
            type: The type of the prompt to be created. "chat" vs. "text". Defaults to "text".
            commit_message: Optional string describing the change.

        Returns:
            TextPromptClient: The prompt if type argument is 'text'.
            ChatPromptClient: The prompt if type argument is 'chat'.
        """
        try:
            abv_logger.debug(f"Creating prompt {name=}, {labels=}")

            if type == "chat":
                if not isinstance(prompt, list):
                    raise ValueError(
                        "For 'chat' type, 'prompt' must be a list of chat messages with role and content attributes."
                    )
                request: Union[CreatePromptRequest_Chat, CreatePromptRequest_Text] = (
                    CreatePromptRequest_Chat(
                        name=name,
                        prompt=cast(Any, prompt),
                        labels=labels,
                        tags=tags,
                        config=config or {},
                        commitMessage=commit_message,
                        type="chat",
                    )
                )
                server_prompt = self.api.prompts.create(request=request)

                if self._resources is not None:
                    self._resources.prompt_cache.invalidate(name)

                return ChatPromptClient(prompt=cast(Prompt_Chat, server_prompt))

            if not isinstance(prompt, str):
                raise ValueError("For 'text' type, 'prompt' must be a string.")

            request = CreatePromptRequest_Text(
                name=name,
                prompt=prompt,
                labels=labels,
                tags=tags,
                config=config or {},
                commitMessage=commit_message,
                type="text",
            )

            server_prompt = self.api.prompts.create(request=request)

            if self._resources is not None:
                self._resources.prompt_cache.invalidate(name)

            return TextPromptClient(prompt=cast(Prompt_Text, server_prompt))

        except Error as e:
            handle_fern_exception(e)
            raise e

    def update_prompt(
        self,
        *,
        name: str,
        version: int,
        new_labels: List[str] = [],
    ) -> Any:
        """Update an existing prompt version in ABV. The ABV SDK prompt cache is invalidated for all prompts witht he specified name.

        Args:
            name (str): The name of the prompt to update.
            version (int): The version number of the prompt to update.
            new_labels (List[str], optional): New labels to assign to the prompt version. Labels are unique across versions. The "latest" label is reserved and managed by ABV. Defaults to [].

        Returns:
            Prompt: The updated prompt from the ABV API.

        """
        updated_prompt = self.api.prompt_version.update(
            name=name,
            version=version,
            new_labels=new_labels,
        )

        if self._resources is not None:
            self._resources.prompt_cache.invalidate(name)

        return updated_prompt

    def _url_encode(self, url: str, *, is_url_param: Optional[bool] = False) -> str:
        # httpx ≥ 0.28 does its own WHATWG-compliant quoting (eg. encodes bare
        # “%”, “?”, “#”, “|”, … in query/path parts).  Re-quoting here would
        # double-encode, so we skip when the value is about to be sent straight
        # to httpx (`is_url_param=True`) and the installed version is ≥ 0.28.
        if is_url_param and Version(httpx.__version__) >= Version("0.28.0"):
            return url

        # urllib.parse.quote does not escape slashes "/" by default; we need to add safe="" to force escaping
        # we need add safe="" to force escaping of slashes
        # This is necessary for prompts in prompt folders
        return urllib.parse.quote(url, safe="")
