import csv
from dataclasses import dataclass
from dataclasses import field
import inspect
import json
import os
import time
from typing import Any
from typing import Callable
from typing import Dict
from typing import List
from typing import Literal
from typing import Optional
from typing import Set
from typing import Tuple
from typing import Union
from typing import cast

import ddtrace
from ddtrace import config
from ddtrace import patch
from ddtrace._trace.apm_filter import APMTracingEnabledFilter
from ddtrace._trace.context import Context
from ddtrace._trace.span import Span
from ddtrace._trace.tracer import Tracer
from ddtrace.constants import ERROR_MSG
from ddtrace.constants import ERROR_STACK
from ddtrace.constants import ERROR_TYPE
from ddtrace.ext import SpanTypes
from ddtrace.internal import atexit
from ddtrace.internal import core
from ddtrace.internal import forksafe
from ddtrace.internal._rand import rand64bits
from ddtrace.internal._rand import rand128bits
from ddtrace.internal.compat import ensure_text
from ddtrace.internal.logger import get_logger
from ddtrace.internal.remoteconfig.worker import remoteconfig_poller
from ddtrace.internal.service import Service
from ddtrace.internal.service import ServiceStatusError
from ddtrace.internal.telemetry import telemetry_writer
from ddtrace.internal.telemetry.constants import TELEMETRY_APM_PRODUCT
from ddtrace.internal.utils.formats import asbool
from ddtrace.internal.utils.formats import format_trace_id
from ddtrace.internal.utils.formats import parse_tags_str
from ddtrace.llmobs import _constants as constants
from ddtrace.llmobs import _telemetry as telemetry
from ddtrace.llmobs._constants import AGENT_MANIFEST
from ddtrace.llmobs._constants import ANNOTATIONS_CONTEXT_ID
from ddtrace.llmobs._constants import DECORATOR
from ddtrace.llmobs._constants import DEFAULT_PROJECT_NAME
from ddtrace.llmobs._constants import DISPATCH_ON_GUARDRAIL_SPAN_START
from ddtrace.llmobs._constants import DISPATCH_ON_LLM_SPAN_FINISH
from ddtrace.llmobs._constants import DISPATCH_ON_LLM_TOOL_CHOICE
from ddtrace.llmobs._constants import DISPATCH_ON_OPENAI_AGENT_SPAN_FINISH
from ddtrace.llmobs._constants import DISPATCH_ON_TOOL_CALL
from ddtrace.llmobs._constants import DISPATCH_ON_TOOL_CALL_OUTPUT_USED
from ddtrace.llmobs._constants import EXPERIMENT_CSV_FIELD_MAX_SIZE
from ddtrace.llmobs._constants import EXPERIMENT_EXPECTED_OUTPUT
from ddtrace.llmobs._constants import EXPERIMENT_ID_KEY
from ddtrace.llmobs._constants import EXPERIMENT_RUN_ID_KEY
from ddtrace.llmobs._constants import EXPERIMENT_RUN_ITERATION_KEY
from ddtrace.llmobs._constants import EXPERIMENTS_INPUT
from ddtrace.llmobs._constants import EXPERIMENTS_OUTPUT
from ddtrace.llmobs._constants import INPUT_DOCUMENTS
from ddtrace.llmobs._constants import INPUT_MESSAGES
from ddtrace.llmobs._constants import INPUT_PROMPT
from ddtrace.llmobs._constants import INPUT_VALUE
from ddtrace.llmobs._constants import INTEGRATION
from ddtrace.llmobs._constants import LLMOBS_TRACE_ID
from ddtrace.llmobs._constants import METADATA
from ddtrace.llmobs._constants import METRICS
from ddtrace.llmobs._constants import ML_APP
from ddtrace.llmobs._constants import MODEL_NAME
from ddtrace.llmobs._constants import MODEL_PROVIDER
from ddtrace.llmobs._constants import OUTPUT_DOCUMENTS
from ddtrace.llmobs._constants import OUTPUT_MESSAGES
from ddtrace.llmobs._constants import OUTPUT_VALUE
from ddtrace.llmobs._constants import PARENT_ID_KEY
from ddtrace.llmobs._constants import PROPAGATED_LLMOBS_TRACE_ID_KEY
from ddtrace.llmobs._constants import PROPAGATED_ML_APP_KEY
from ddtrace.llmobs._constants import PROPAGATED_PARENT_ID_KEY
from ddtrace.llmobs._constants import ROOT_PARENT_ID
from ddtrace.llmobs._constants import SESSION_ID
from ddtrace.llmobs._constants import SPAN_KIND
from ddtrace.llmobs._constants import SPAN_LINKS
from ddtrace.llmobs._constants import SPAN_START_WHILE_DISABLED_WARNING
from ddtrace.llmobs._constants import TAGS
from ddtrace.llmobs._constants import TOOL_DEFINITIONS
from ddtrace.llmobs._context import LLMObsContextProvider
from ddtrace.llmobs._evaluators.runner import EvaluatorRunner
from ddtrace.llmobs._experiment import Dataset
from ddtrace.llmobs._experiment import DatasetRecord
from ddtrace.llmobs._experiment import DatasetRecordInputType
from ddtrace.llmobs._experiment import Experiment
from ddtrace.llmobs._experiment import ExperimentConfigType
from ddtrace.llmobs._experiment import JSONType
from ddtrace.llmobs._experiment import Project
from ddtrace.llmobs._utils import AnnotationContext
from ddtrace.llmobs._utils import LinkTracker
from ddtrace.llmobs._utils import _get_ml_app
from ddtrace.llmobs._utils import _get_nearest_llmobs_ancestor
from ddtrace.llmobs._utils import _get_session_id
from ddtrace.llmobs._utils import _get_span_name
from ddtrace.llmobs._utils import _is_evaluation_span
from ddtrace.llmobs._utils import _validate_prompt
from ddtrace.llmobs._utils import enforce_message_role
from ddtrace.llmobs._utils import safe_json
from ddtrace.llmobs._writer import LLMObsEvalMetricWriter
from ddtrace.llmobs._writer import LLMObsEvaluationMetricEvent
from ddtrace.llmobs._writer import LLMObsExperimentsClient
from ddtrace.llmobs._writer import LLMObsSpanEvent
from ddtrace.llmobs._writer import LLMObsSpanWriter
from ddtrace.llmobs._writer import should_use_agentless
from ddtrace.llmobs.types import ExportedLLMObsSpan
from ddtrace.llmobs.types import Message
from ddtrace.llmobs.types import Prompt
from ddtrace.llmobs.types import _ErrorField
from ddtrace.llmobs.types import _Meta
from ddtrace.llmobs.types import _MetaIO
from ddtrace.llmobs.types import _SpanField
from ddtrace.llmobs.utils import Documents
from ddtrace.llmobs.utils import Messages
from ddtrace.llmobs.utils import extract_tool_definitions
from ddtrace.propagation.http import HTTPPropagator


log = get_logger(__name__)


SUPPORTED_LLMOBS_INTEGRATIONS = {
    "anthropic": "anthropic",
    "bedrock": "botocore",
    "openai": "openai",
    "langchain": "langchain",
    "google_adk": "google_adk",
    "google_generativeai": "google_generativeai",
    "google_genai": "google_genai",
    "vertexai": "vertexai",
    "langgraph": "langgraph",
    "litellm": "litellm",
    "crewai": "crewai",
    "openai_agents": "openai_agents",
    "mcp": "mcp",
    "pydantic_ai": "pydantic_ai",
    # requests/concurrent frameworks for distributed injection/extraction
    "requests": "requests",
    "httpx": "httpx",
    "urllib3": "urllib3",
    "grpc": "grpc",
    "flask": "flask",
    "starlette": "starlette",
    "fastapi": "fastapi",
    "aiohttp": "aiohttp",
    "asyncio": "asyncio",
    "futures": "futures",
}


@dataclass
class LLMObsSpan:
    """LLMObs span object.

    Passed to the `span_processor` function in the `enable` or `register_processor` methods.

    Example::
        def span_processor(span: LLMObsSpan) -> Optional[LLMObsSpan]:
            # Modify input/output
            if span.get_tag("omit_span") == "1":
                return None
            if span.get_tag("no_input") == "1":
                span.input = []
            return span
    """

    input: List[Message] = field(default_factory=list)
    output: List[Message] = field(default_factory=list)
    _tags: Dict[str, str] = field(default_factory=dict)

    def get_tag(self, key: str) -> Optional[str]:
        """Get a tag from the span.

        :param str key: The key of the tag to get.
        :return: The value of the tag or None if the tag does not exist.
        :rtype: Optional[str]
        """
        return self._tags.get(key)


class LLMObs(Service):
    _instance = None  # type: LLMObs
    enabled = False
    _app_key: str = os.getenv("DD_APP_KEY", "")
    _project_name: str = os.getenv("DD_LLMOBS_PROJECT_NAME", DEFAULT_PROJECT_NAME)

    def __init__(
        self,
        tracer: Optional[Tracer] = None,
        span_processor: Optional[Callable[[LLMObsSpan], Optional[LLMObsSpan]]] = None,
    ) -> None:
        super(LLMObs, self).__init__()
        self.tracer = tracer or ddtrace.tracer
        self._llmobs_context_provider = LLMObsContextProvider()
        self._user_span_processor = span_processor
        agentless_enabled = config._llmobs_agentless_enabled if config._llmobs_agentless_enabled is not None else True
        self._llmobs_span_writer = LLMObsSpanWriter(
            interval=float(os.getenv("_DD_LLMOBS_WRITER_INTERVAL", 1.0)),
            timeout=float(os.getenv("_DD_LLMOBS_WRITER_TIMEOUT", 5.0)),
            is_agentless=agentless_enabled,
        )
        self._llmobs_eval_metric_writer = LLMObsEvalMetricWriter(
            interval=float(os.getenv("_DD_LLMOBS_WRITER_INTERVAL", 1.0)),
            timeout=float(os.getenv("_DD_LLMOBS_WRITER_TIMEOUT", 5.0)),
            is_agentless=agentless_enabled,
        )
        self._evaluator_runner = EvaluatorRunner(
            interval=float(os.getenv("_DD_LLMOBS_EVALUATOR_INTERVAL", 1.0)),
            llmobs_service=self,
        )
        self._dne_client = LLMObsExperimentsClient(
            interval=float(os.getenv("_DD_LLMOBS_WRITER_INTERVAL", 1.0)),
            timeout=float(os.getenv("_DD_LLMOBS_WRITER_TIMEOUT", 5.0)),
            _app_key=self._app_key,
            _default_project=Project(name=self._project_name, _id=""),
            is_agentless=True,  # agent proxy doesn't seem to work for experiments
        )

        forksafe.register(self._child_after_fork)

        self._link_tracker = LinkTracker()
        self._annotations: List[Tuple[str, str, Dict[str, Any]]] = []
        self._annotation_context_lock = forksafe.RLock()

    def _on_span_start(self, span: Span) -> None:
        if self.enabled and span.span_type == SpanTypes.LLM:
            self._activate_llmobs_span(span)
            telemetry.record_span_started()
            self._do_annotations(span)

    def _on_span_finish(self, span: Span) -> None:
        if self.enabled and span.span_type == SpanTypes.LLM:
            self._submit_llmobs_span(span)
            telemetry.record_span_created(span)

    def _submit_llmobs_span(self, span: Span) -> None:
        """Generate and submit an LLMObs span event to be sent to LLMObs."""
        span_event = None
        try:
            span_event = self._llmobs_span_event(span)
            if span_event is None:
                return
            self._llmobs_span_writer.enqueue(span_event)
        except (KeyError, TypeError, ValueError):
            log.error(
                "Error generating LLMObs span event for span %s, likely due to malformed span", span, exc_info=True
            )
        finally:
            if span_event and span._get_ctx_item(SPAN_KIND) == "llm" and not _is_evaluation_span(span):
                if self._evaluator_runner:
                    self._evaluator_runner.enqueue(span_event, span)

    def _llmobs_span_event(self, span: Span) -> Optional[LLMObsSpanEvent]:
        """Span event object structure."""
        span_kind = span._get_ctx_item(SPAN_KIND)
        if not span_kind:
            raise KeyError("Span kind not found in span context")

        if span_kind == "llm":
            core.dispatch(DISPATCH_ON_LLM_SPAN_FINISH, (span,))

        llmobs_span = LLMObsSpan()
        _dd_attrs = {
            "span_id": str(span.span_id),
            "trace_id": format_trace_id(span.trace_id),
            "apm_trace_id": format_trace_id(span.trace_id),
        }

        meta = _Meta(span=_SpanField(kind=span_kind), input=_MetaIO(), output=_MetaIO())
        if span_kind in ("llm", "embedding") and span._get_ctx_item(MODEL_NAME) is not None:
            meta["model_name"] = span._get_ctx_item(MODEL_NAME) or ""
            meta["model_provider"] = (span._get_ctx_item(MODEL_PROVIDER) or "custom").lower()
        metadata = span._get_ctx_item(METADATA) or {}
        if span_kind == "agent" and span._get_ctx_item(AGENT_MANIFEST) is not None:
            metadata["agent_manifest"] = span._get_ctx_item(AGENT_MANIFEST)
        meta["metadata"] = metadata

        input_type: Literal["value", "messages", ""] = ""
        output_type: Literal["value", "messages", ""] = ""
        if span._get_ctx_item(INPUT_VALUE) is not None:
            input_type = "value"
            llmobs_span.input = [
                Message(content=safe_json(span._get_ctx_item(INPUT_VALUE), ensure_ascii=False) or "", role="")
            ]

        if span.context.get_baggage_item(EXPERIMENT_ID_KEY):
            _dd_attrs["scope"] = "experiments"
            if span_kind == "experiment":
                expected_output = span._get_ctx_item(EXPERIMENT_EXPECTED_OUTPUT)
                if expected_output:
                    meta["expected_output"] = expected_output

                input_data = span._get_ctx_item(EXPERIMENTS_INPUT)
                if input_data:
                    meta["input"] = input_data

                output_data = span._get_ctx_item(EXPERIMENTS_OUTPUT)
                if output_data:
                    meta["output"] = output_data

        input_messages = span._get_ctx_item(INPUT_MESSAGES)
        if span_kind == "llm" and input_messages is not None:
            input_type = "messages"
            llmobs_span.input = cast(List[Message], enforce_message_role(input_messages))

        if span._get_ctx_item(OUTPUT_VALUE) is not None:
            output_type = "value"
            llmobs_span.output = [
                Message(content=safe_json(span._get_ctx_item(OUTPUT_VALUE), ensure_ascii=False) or "", role="")
            ]

        output_messages = span._get_ctx_item(OUTPUT_MESSAGES)
        if span_kind == "llm" and output_messages is not None:
            output_type = "messages"
            llmobs_span.output = cast(List[Message], enforce_message_role(output_messages))

        if span_kind == "embedding" and span._get_ctx_item(INPUT_DOCUMENTS) is not None:
            meta["input"]["documents"] = span._get_ctx_item(INPUT_DOCUMENTS) or []
        if span_kind == "retrieval" and span._get_ctx_item(OUTPUT_DOCUMENTS) is not None:
            meta["output"]["documents"] = span._get_ctx_item(OUTPUT_DOCUMENTS) or []

        if span._get_ctx_item(INPUT_PROMPT) is not None:
            prompt_json_str = span._get_ctx_item(INPUT_PROMPT)
            if span_kind != "llm":
                log.warning(
                    "Dropping prompt on non-LLM span kind, annotating prompts is only supported for LLM span kinds."
                )
            else:
                prompt_dict = cast(Prompt, prompt_json_str)
                meta["input"]["prompt"] = prompt_dict
        elif span_kind == "llm":
            parent_span = _get_nearest_llmobs_ancestor(span)
            if parent_span is not None:
                parent_prompt = parent_span._get_ctx_item(INPUT_PROMPT)
                if parent_prompt is not None:
                    meta["input"]["prompt"] = parent_prompt

        if span._get_ctx_item(TOOL_DEFINITIONS) is not None:
            meta["tool_definitions"] = span._get_ctx_item(TOOL_DEFINITIONS) or []
        if span.error:
            meta["error"] = _ErrorField(
                message=span.get_tag(ERROR_MSG) or "",
                stack=span.get_tag(ERROR_STACK) or "",
                type=span.get_tag(ERROR_TYPE) or "",
            )

        if self._user_span_processor:
            error = False
            try:
                llmobs_span._tags = cast(Dict[str, str], span._get_ctx_item(TAGS))
                user_llmobs_span = self._user_span_processor(llmobs_span)
                if user_llmobs_span is None:
                    return None
                if not isinstance(user_llmobs_span, LLMObsSpan):
                    raise TypeError(
                        "User span processor must return an LLMObsSpan or None, got %r" % type(user_llmobs_span)
                    )
                llmobs_span = user_llmobs_span
            except Exception as e:
                log.error("Error in LLMObs span processor (%r): %r", self._user_span_processor, e)
                error = True
            finally:
                telemetry.record_llmobs_user_processor_called(error)

        if llmobs_span.input is not None:
            if input_type == "messages":
                meta["input"]["messages"] = llmobs_span.input
            elif input_type == "value":
                meta["input"]["value"] = llmobs_span.input[0].get("content", "")
        if llmobs_span.output is not None:
            if output_type == "messages":
                meta["output"]["messages"] = llmobs_span.output
            elif output_type == "value":
                meta["output"]["value"] = llmobs_span.output[0].get("content", "")

        if not meta["input"]:
            meta.pop("input")
        if not meta["output"]:
            meta.pop("output")
        metrics = span._get_ctx_item(METRICS) or {}
        ml_app = _get_ml_app(span)

        if ml_app is None:
            raise ValueError(
                "ML app is required for sending LLM Observability data. "
                "Ensure this configuration is set before running your application."
            )

        span._set_ctx_item(ML_APP, ml_app)
        parent_id = span._get_ctx_item(PARENT_ID_KEY) or ROOT_PARENT_ID

        llmobs_trace_id = span._get_ctx_item(LLMOBS_TRACE_ID)
        if llmobs_trace_id is None:
            raise ValueError("Failed to extract LLMObs trace ID from span context.")

        llmobs_span_event: LLMObsSpanEvent = {
            "trace_id": format_trace_id(llmobs_trace_id),
            "span_id": str(span.span_id),
            "parent_id": parent_id,
            "name": _get_span_name(span),
            "start_ns": span.start_ns,
            "duration": cast(int, span.duration_ns),
            "status": "error" if span.error else "ok",
            "meta": meta,
            "metrics": metrics,
            "tags": [],
            "_dd": _dd_attrs,
        }
        session_id = _get_session_id(span)
        if session_id is not None:
            span._set_ctx_item(SESSION_ID, session_id)
            llmobs_span_event["session_id"] = session_id

        llmobs_span_event["tags"] = self._llmobs_tags(span, ml_app, session_id)

        span_links = span._get_ctx_item(SPAN_LINKS)
        if isinstance(span_links, list) and span_links:
            llmobs_span_event["span_links"] = span_links

        return llmobs_span_event

    @staticmethod
    def _llmobs_tags(span: Span, ml_app: str, session_id: Optional[str] = None) -> List[str]:
        dd_tags = config.tags
        tags = {
            **dd_tags,
            "version": config.version or "",
            "env": config.env or "",
            "service": span.service or "",
            "source": "integration",
            "ml_app": ml_app,
            "ddtrace.version": ddtrace.__version__,
            "language": "python",
            "error": span.error,
        }
        err_type = span.get_tag(ERROR_TYPE)
        if err_type:
            tags["error_type"] = err_type
        if session_id:
            tags["session_id"] = session_id
        if span._get_ctx_item(INTEGRATION):
            tags["integration"] = span._get_ctx_item(INTEGRATION)
        if _is_evaluation_span(span):
            tags[constants.RUNNER_IS_INTEGRATION_SPAN_TAG] = "ragas"
        existing_tags = span._get_ctx_item(TAGS)
        if existing_tags is not None:
            tags.update(existing_tags)

        # set experiment tags on children spans if the tags do not already exist
        experiment_id = span.context.get_baggage_item(EXPERIMENT_ID_KEY)
        if experiment_id and "experiment_id" not in tags:
            tags["experiment_id"] = experiment_id

        run_id = span.context.get_baggage_item(EXPERIMENT_RUN_ID_KEY)
        if run_id and "run_id" not in tags:
            tags["run_id"] = run_id

        run_iteration = span.context.get_baggage_item(EXPERIMENT_RUN_ITERATION_KEY)
        if run_iteration and "run_iteration" not in tags:
            tags["run_iteration"] = run_iteration

        return ["{}:{}".format(k, v) for k, v in tags.items()]

    def _do_annotations(self, span: Span) -> None:
        # get the current span context
        # only do the annotations if it matches the context
        if span.span_type != SpanTypes.LLM:  # do this check to avoid the warning log in `annotate`
            return
        current_context = self._instance.tracer.current_trace_context()
        if current_context is None:
            return
        current_context_id = current_context.get_baggage_item(ANNOTATIONS_CONTEXT_ID)
        with self._annotation_context_lock:
            for _, context_id, annotation_kwargs in self._instance._annotations:
                if current_context_id == context_id:
                    self.annotate(span, **annotation_kwargs)

    def _child_after_fork(self) -> None:
        self._llmobs_span_writer = self._llmobs_span_writer.recreate()
        self._llmobs_eval_metric_writer = self._llmobs_eval_metric_writer.recreate()
        self._evaluator_runner = self._evaluator_runner.recreate()
        if self.enabled:
            self._start_service()

    def _start_service(self) -> None:
        try:
            self._llmobs_span_writer.start()
            self._llmobs_eval_metric_writer.start()
        except ServiceStatusError:
            log.debug("Error starting LLMObs writers")

        try:
            self._evaluator_runner.start()
        except ServiceStatusError:
            log.debug("Error starting evaluator runner")

    def _stop_service(self) -> None:
        try:
            self._evaluator_runner.stop()
            # flush remaining evaluation spans & evaluations
            self._instance._llmobs_span_writer.periodic()
            self._instance._llmobs_eval_metric_writer.periodic()
        except ServiceStatusError:
            log.debug("Error stopping evaluator runner")

        try:
            self._llmobs_span_writer.stop()
            self._llmobs_eval_metric_writer.stop()
        except ServiceStatusError:
            log.debug("Error stopping LLMObs writers")

        # Remove listener hooks for span events
        core.reset_listeners("trace.span_start", self._on_span_start)
        core.reset_listeners("trace.span_finish", self._on_span_finish)
        core.reset_listeners("http.span_inject", self._inject_llmobs_context)
        core.reset_listeners("http.activate_distributed_headers", self._activate_llmobs_distributed_context)
        core.reset_listeners("threading.submit", self._current_trace_context)
        core.reset_listeners("threading.execution", self._llmobs_context_provider.activate)
        core.reset_listeners("asyncio.create_task", self._on_asyncio_create_task)
        core.reset_listeners("asyncio.execute_task", self._on_asyncio_execute_task)

        core.reset_listeners(DISPATCH_ON_LLM_TOOL_CHOICE, self._link_tracker.on_llm_tool_choice)
        core.reset_listeners(DISPATCH_ON_TOOL_CALL, self._link_tracker.on_tool_call)
        core.reset_listeners(DISPATCH_ON_TOOL_CALL_OUTPUT_USED, self._link_tracker.on_tool_call_output_used)

        core.reset_listeners(DISPATCH_ON_GUARDRAIL_SPAN_START, self._link_tracker.on_guardrail_span_start)
        core.reset_listeners(DISPATCH_ON_LLM_SPAN_FINISH, self._link_tracker.on_llm_span_finish)
        core.reset_listeners(DISPATCH_ON_OPENAI_AGENT_SPAN_FINISH, self._link_tracker.on_openai_agent_span_finish)

        forksafe.unregister(self._child_after_fork)

    @classmethod
    def enable(
        cls,
        ml_app: Optional[str] = None,
        integrations_enabled: bool = True,
        agentless_enabled: Optional[bool] = None,
        instrumented_proxy_urls: Optional[Set[str]] = None,
        site: Optional[str] = None,
        api_key: Optional[str] = None,
        app_key: Optional[str] = None,
        project_name: Optional[str] = None,
        env: Optional[str] = None,
        service: Optional[str] = None,
        span_processor: Optional[Callable[[LLMObsSpan], Optional[LLMObsSpan]]] = None,
        _tracer: Optional[Tracer] = None,
        _auto: bool = False,
    ) -> None:
        """
        Enable LLM Observability tracing.

        :param str ml_app: The name of your ml application.
        :param bool integrations_enabled: Set to `true` to enable LLM integrations.
        :param bool agentless_enabled: Set to `true` to disable sending data that requires a Datadog Agent.
        :param set[str] instrumented_proxy_urls: A set of instrumented proxy URLs to help detect when to emit LLM spans.
        :param str site: Your datadog site.
        :param str api_key: Your datadog api key.
        :param str app_key: Your datadog application key.
        :param str project_name: Your project name used for experiments.
        :param str env: Your environment name.
        :param str service: Your service name.
        :param Callable[[LLMObsSpan], Optional[LLMObsSpan]] span_processor: A function that takes an LLMObsSpan and
            returns an LLMObsSpan or None. If None is returned, the span will be omitted and not sent to LLMObs.
        """
        if cls.enabled:
            log.debug("%s already enabled", cls.__name__)
            return

        if os.getenv("DD_LLMOBS_ENABLED") and not asbool(os.getenv("DD_LLMOBS_ENABLED")):
            log.debug("LLMObs.enable() called when DD_LLMOBS_ENABLED is set to false or 0, not starting LLMObs service")
            return
        # grab required values for LLMObs
        config._dd_site = site or config._dd_site
        config._dd_api_key = api_key or config._dd_api_key
        cls._app_key = app_key or cls._app_key
        cls._project_name = project_name or cls._project_name or DEFAULT_PROJECT_NAME
        config.env = env or config.env
        config.service = service or config.service
        config._llmobs_ml_app = ml_app or config._llmobs_ml_app
        config._llmobs_instrumented_proxy_urls = instrumented_proxy_urls or config._llmobs_instrumented_proxy_urls

        error = None
        start_ns = time.time_ns()
        try:
            config._llmobs_agentless_enabled = should_use_agentless(
                user_defined_agentless_enabled=(
                    agentless_enabled if agentless_enabled is not None else config._llmobs_agentless_enabled
                )
            )

            if config._llmobs_agentless_enabled:
                # validate required values for agentless LLMObs
                if not config._dd_api_key:
                    error = "missing_api_key"
                    raise ValueError(
                        "DD_API_KEY is required for sending LLMObs data when agentless mode is enabled. "
                        "Ensure this configuration is set before running your application."
                    )
                if not config._dd_site:
                    error = "missing_site"
                    raise ValueError(
                        "DD_SITE is required for sending LLMObs data when agentless mode is enabled. "
                        "Ensure this configuration is set before running your application."
                    )
                if not os.getenv("DD_REMOTE_CONFIG_ENABLED"):
                    config._remote_config_enabled = False
                    log.debug("Remote configuration disabled because DD_LLMOBS_AGENTLESS_ENABLED is set to true.")
                    remoteconfig_poller.disable()

                # Since the API key can be set programmatically and TelemetryWriter is already initialized by now,
                # we need to force telemetry to use agentless configuration
                telemetry_writer.enable_agentless_client(True)

            if integrations_enabled:
                cls._patch_integrations()

            # override the default _instance with a new tracer
            cls._instance = cls(tracer=_tracer, span_processor=span_processor)

            # Add APM trace filter to drop all APM traces when DD_APM_TRACING_ENABLED is falsy
            apm_filter = APMTracingEnabledFilter()
            cls._instance.tracer._span_aggregator.dd_processors.append(apm_filter)

            cls.enabled = True
            cls._instance.start()

            # Register hooks for span events
            core.on("trace.span_start", cls._instance._on_span_start)
            core.on("trace.span_finish", cls._instance._on_span_finish)
            core.on("http.span_inject", cls._inject_llmobs_context)
            core.on("http.activate_distributed_headers", cls._activate_llmobs_distributed_context)
            core.on("threading.submit", cls._instance._current_trace_context, "llmobs_ctx")
            core.on("threading.execution", cls._instance._llmobs_context_provider.activate)
            core.on("asyncio.create_task", cls._instance._on_asyncio_create_task)
            core.on("asyncio.execute_task", cls._instance._on_asyncio_execute_task)

            core.on(DISPATCH_ON_LLM_TOOL_CHOICE, cls._instance._link_tracker.on_llm_tool_choice)
            core.on(DISPATCH_ON_TOOL_CALL, cls._instance._link_tracker.on_tool_call)
            core.on(DISPATCH_ON_TOOL_CALL_OUTPUT_USED, cls._instance._link_tracker.on_tool_call_output_used)

            core.on(DISPATCH_ON_GUARDRAIL_SPAN_START, cls._instance._link_tracker.on_guardrail_span_start)
            core.on(DISPATCH_ON_LLM_SPAN_FINISH, cls._instance._link_tracker.on_llm_span_finish)
            core.on(DISPATCH_ON_OPENAI_AGENT_SPAN_FINISH, cls._instance._link_tracker.on_openai_agent_span_finish)

            atexit.register(cls.disable)
            telemetry_writer.product_activated(TELEMETRY_APM_PRODUCT.LLMOBS, True)

            log.debug("%s enabled; instrumented_proxy_urls: %s", cls.__name__, config._llmobs_instrumented_proxy_urls)

            llmobs_info = {
                "llmobs_enabled": cls.enabled,
                "llmobs_ml_app": config._llmobs_ml_app,
                "integrations_enabled": integrations_enabled,
                "llmobs_agentless_enabled": config._llmobs_agentless_enabled,
            }
            log.debug("LLMObs configurations: %s", llmobs_info)

        finally:
            telemetry.record_llmobs_enabled(
                error,
                config._llmobs_agentless_enabled,
                config._dd_site,
                start_ns,
                _auto,
                config._llmobs_instrumented_proxy_urls,
                config._llmobs_ml_app,
            )

    def _on_asyncio_create_task(self, task_data: Dict[str, Any]) -> None:
        """Propagates llmobs active trace context across asyncio tasks."""
        task_data["llmobs_ctx"] = self._current_trace_context()

    def _on_asyncio_execute_task(self, task_data: Dict[str, Any]) -> None:
        """Activates llmobs active trace context across asyncio task execution."""
        llmobs_ctx = task_data.get("llmobs_ctx")
        if llmobs_ctx is not None:
            self._llmobs_context_provider.activate(llmobs_ctx)

    @classmethod
    def pull_dataset(
        cls, dataset_name: str, project_name: Optional[str] = None, version: Optional[int] = None
    ) -> Dataset:
        ds = cls._instance._dne_client.dataset_get_with_records(
            dataset_name, (project_name or cls._project_name), version
        )
        return ds

    @classmethod
    def create_dataset(
        cls,
        dataset_name: str,
        project_name: Optional[str] = None,
        description: str = "",
        records: Optional[List[DatasetRecord]] = None,
    ) -> Dataset:
        if records is None:
            records = []
        ds = cls._instance._dne_client.dataset_create(dataset_name, project_name, description)
        for r in records:
            ds.append(r)
        if len(records) > 0:
            ds.push()
        return ds

    @classmethod
    def create_dataset_from_csv(
        cls,
        csv_path: str,
        dataset_name: str,
        input_data_columns: List[str],
        expected_output_columns: Optional[List[str]] = None,
        metadata_columns: Optional[List[str]] = None,
        csv_delimiter: str = ",",
        description: str = "",
        project_name: Optional[str] = None,
    ) -> Dataset:
        if expected_output_columns is None:
            expected_output_columns = []
        if metadata_columns is None:
            metadata_columns = []

        # Store the original field size limit to restore it later
        original_field_size_limit = csv.field_size_limit()

        csv.field_size_limit(EXPERIMENT_CSV_FIELD_MAX_SIZE)  # 10mb

        records = []
        try:
            with open(csv_path, mode="r") as csvfile:
                content = csvfile.readline().strip()
                if not content:
                    raise ValueError("CSV file appears to be empty or header is missing.")

                csvfile.seek(0)

                rows = csv.DictReader(csvfile, delimiter=csv_delimiter)

                if rows.fieldnames is None:
                    raise ValueError("CSV file appears to be empty or header is missing.")

                header_columns = rows.fieldnames
                missing_input_columns = [col for col in input_data_columns if col not in header_columns]
                missing_output_columns = [col for col in expected_output_columns if col not in header_columns]
                missing_metadata_columns = [col for col in metadata_columns if col not in metadata_columns]

                if any(col not in header_columns for col in input_data_columns):
                    raise ValueError(f"Input columns not found in CSV header: {missing_input_columns}")
                if any(col not in header_columns for col in expected_output_columns):
                    raise ValueError(f"Expected output columns not found in CSV header: {missing_output_columns}")
                if any(col not in header_columns for col in metadata_columns):
                    raise ValueError(f"Metadata columns not found in CSV header: {missing_metadata_columns}")

                for row in rows:
                    records.append(
                        DatasetRecord(
                            input_data={col: row[col] for col in input_data_columns},
                            expected_output={col: row[col] for col in expected_output_columns},
                            metadata={col: row[col] for col in metadata_columns},
                            record_id="",
                        )
                    )

        finally:
            # Always restore the original field size limit
            csv.field_size_limit(original_field_size_limit)

        ds = cls._instance._dne_client.dataset_create(dataset_name, project_name, description)
        for r in records:
            ds.append(r)
        if len(ds) > 0:
            cls._instance._dne_client.dataset_bulk_upload(ds._id, ds._records)
        return ds

    @classmethod
    def _delete_dataset(cls, dataset_id: str) -> None:
        return cls._instance._dne_client.dataset_delete(dataset_id)

    @classmethod
    def experiment(
        cls,
        name: str,
        task: Callable[[DatasetRecordInputType, Optional[ExperimentConfigType]], JSONType],
        dataset: Dataset,
        evaluators: List[Callable[[DatasetRecordInputType, JSONType, JSONType], JSONType]],
        description: str = "",
        project_name: Optional[str] = None,
        tags: Optional[Dict[str, str]] = None,
        config: Optional[ExperimentConfigType] = None,
        summary_evaluators: Optional[
            List[
                Callable[
                    [List[DatasetRecordInputType], List[JSONType], List[JSONType], Dict[str, List[JSONType]]], JSONType
                ]
            ]
        ] = None,
        runs: Optional[int] = 1,
    ) -> Experiment:
        """Initializes an Experiment to run a task on a Dataset and evaluators.

        :param name: The name of the experiment.
        :param task: The task function to run. Must accept parameters ``input_data`` and ``config``.
        :param dataset: The dataset to run the experiment on, created with LLMObs.pull/create_dataset().
        :param evaluators: A list of evaluator functions to evaluate the task output.
                           Must accept parameters ``input_data``, ``output_data``, and ``expected_output``.
        :param project_name: The name of the project to save the experiment to.
        :param description: A description of the experiment.
        :param tags: A dictionary of string key-value tag pairs to associate with the experiment.
        :param config: A configuration dictionary describing the experiment.
        :param summary_evaluators: A list of summary evaluator functions to evaluate the task results and evaluations
                                   to produce a single value.
                                   Must accept parameters ``inputs``, ``outputs``, ``expected_outputs``,
                                   ``evaluators_results``.
        :param runs: The number of times to run the experiment, or, run the task for every dataset record the defined
                     number of times.
        """
        if not callable(task):
            raise TypeError("task must be a callable function.")
        sig = inspect.signature(task)
        params = sig.parameters
        if "input_data" not in params or "config" not in params:
            raise TypeError("Task function must have 'input_data' and 'config' parameters.")
        if not isinstance(dataset, Dataset):
            raise TypeError("Dataset must be an LLMObs Dataset object.")
        if not evaluators or not all(callable(evaluator) for evaluator in evaluators):
            raise TypeError("Evaluators must be a list of callable functions.")
        for evaluator in evaluators:
            sig = inspect.signature(evaluator)
            params = sig.parameters
            evaluator_required_params = ("input_data", "output_data", "expected_output")
            if not all(param in params for param in evaluator_required_params):
                raise TypeError("Evaluator function must have parameters {}.".format(evaluator_required_params))

        if summary_evaluators and not all(callable(summary_evaluator) for summary_evaluator in summary_evaluators):
            raise TypeError("Summary evaluators must be a list of callable functions.")
        if summary_evaluators:
            for summary_evaluator in summary_evaluators:
                sig = inspect.signature(summary_evaluator)
                params = sig.parameters
                summary_evaluator_required_params = ("inputs", "outputs", "expected_outputs", "evaluators_results")
                if not all(param in params for param in summary_evaluator_required_params):
                    raise TypeError(
                        "Summary evaluator function must have parameters {}.".format(summary_evaluator_required_params)
                    )
        return Experiment(
            name,
            task,
            dataset,
            evaluators,
            project_name=project_name or cls._project_name,
            tags=tags,
            description=description,
            config=config,
            _llmobs_instance=cls._instance,
            summary_evaluators=summary_evaluators,
            runs=runs,
        )

    @classmethod
    def register_processor(cls, processor: Optional[Callable[[LLMObsSpan], Optional[LLMObsSpan]]] = None) -> None:
        """Register a processor to be called on each LLMObs span.

        This can be used to modify the span before it is sent to LLMObs. For example, you can modify the input/output.
        You can also return None to omit the span entirely from being sent to LLMObs.

        To deregister the processor, call `register_processor(None)`.

        :param processor: A function that takes an LLMObsSpan and returns an LLMObsSpan or None.
                         If None is returned, the span will be omitted and not sent to LLMObs.
        """
        cls._instance._user_span_processor = processor

    @classmethod
    def _integration_is_enabled(cls, integration: str) -> bool:
        if integration not in SUPPORTED_LLMOBS_INTEGRATIONS:
            return False
        return SUPPORTED_LLMOBS_INTEGRATIONS[integration] in ddtrace._monkey._get_patched_modules()

    @classmethod
    def disable(cls) -> None:
        if not cls.enabled:
            log.debug("%s not enabled", cls.__name__)
            return
        log.debug("Disabling %s", cls.__name__)
        atexit.unregister(cls.disable)

        cls._instance.stop()
        cls.enabled = False
        telemetry_writer.product_activated(TELEMETRY_APM_PRODUCT.LLMOBS, False)

        log.debug("%s disabled", cls.__name__)

    def _tag_span_links(self, span, span_links):
        if not span_links:
            return
        span_links = [
            span_link
            for span_link in span_links
            if span_link["span_id"] != LLMObs.export_span(span)["span_id"]
            and span_link["trace_id"] == LLMObs.export_span(span)["trace_id"]
        ]
        current_span_links = span._get_ctx_item(SPAN_LINKS)
        if current_span_links:
            span_links = current_span_links + span_links
        span._set_ctx_item(SPAN_LINKS, span_links)

    @classmethod
    def annotation_context(
        cls,
        tags: Optional[Dict[str, Any]] = None,
        prompt: Optional[Union[dict, Prompt]] = None,
        name: Optional[str] = None,
    ) -> AnnotationContext:
        """
        Sets specified attributes on all LLMObs spans created while the returned AnnotationContext is active.
        Annotations are applied in the order in which annotation contexts are entered.

        :param tags: Dictionary of JSON serializable key-value tag pairs to set or update on the LLMObs span
                     regarding the span's context.
        :param prompt: A dictionary that represents the prompt used for an LLM call in the following form:
                        `{
                            "id": "...",
                            "version": "...",
                            "chat_template": [{"content": "...", "role": "..."}, ...],
                            "variables": {"variable_1": "...", ...}}`.
                            "tags": {"key1": "value1", "key2": "value2"},
                        }`
                        Can also be set using the `ddtrace.llmobs.utils.Prompt` constructor class.
                        - This argument is only applicable to LLM spans.
                        - The dictionary may contain optional keys relevant to Templates and RAG applications:
                            `rag_context_variables` - a list of variable key names that contain ground
                                                        truth context information
                            `rag_query_variables` - a list of variable key names that contains query
                                                        information for an LLM call
        :param name: Set to override the span name for any spans annotated within the returned context.
        """
        # id to track an annotation for registering / de-registering
        annotation_id = rand64bits()

        def get_annotations_context_id():
            current_ctx = cls._instance.tracer.current_trace_context()
            # default the context id to the annotation id
            ctx_id = annotation_id
            if current_ctx is None:
                current_ctx = Context(is_remote=False)
                current_ctx.set_baggage_item(ANNOTATIONS_CONTEXT_ID, ctx_id)
                cls._instance.tracer.context_provider.activate(current_ctx)
            elif not current_ctx.get_baggage_item(ANNOTATIONS_CONTEXT_ID):
                current_ctx.set_baggage_item(ANNOTATIONS_CONTEXT_ID, ctx_id)
            else:
                ctx_id = current_ctx.get_baggage_item(ANNOTATIONS_CONTEXT_ID)
            return ctx_id

        def register_annotation():
            with cls._instance._annotation_context_lock:
                ctx_id = get_annotations_context_id()
                cls._instance._annotations.append(
                    (annotation_id, ctx_id, {"tags": tags, "prompt": prompt, "_name": name})
                )

        def deregister_annotation():
            with cls._instance._annotation_context_lock:
                for i, (key, _, _) in enumerate(cls._instance._annotations):
                    if key == annotation_id:
                        cls._instance._annotations.pop(i)
                        return
                else:
                    log.debug("Failed to pop annotation context")

        return AnnotationContext(register_annotation, deregister_annotation)

    @classmethod
    def flush(cls) -> None:
        """
        Flushes any remaining spans and evaluation metrics to the LLMObs backend.
        """
        if cls.enabled is False:
            log.warning("flushing when LLMObs is disabled. No spans or evaluation metrics will be sent.")
            return

        error = None
        try:
            cls._instance._evaluator_runner.periodic()
        except Exception:
            error = "evaluator_flush_error"
            log.warning("Failed to run evaluator runner.", exc_info=True)

        try:
            cls._instance._llmobs_span_writer.periodic()
            cls._instance._llmobs_eval_metric_writer.periodic()
        except Exception:
            error = "writer_flush_error"
            log.warning("Failed to flush LLMObs spans and evaluation metrics.", exc_info=True)

        telemetry.record_user_flush(error)

    @staticmethod
    def _patch_integrations() -> None:
        """
        Patch LLM integrations. Ensure that we do not ignore DD_TRACE_<MODULE>_ENABLED or DD_PATCH_MODULES settings.
        """
        integrations_to_patch: Dict[str, Union[List[str], bool]] = {
            integration: ["bedrock-runtime", "bedrock-agent-runtime"] if integration == "botocore" else True
            for integration in SUPPORTED_LLMOBS_INTEGRATIONS.values()
        }
        for module, _ in integrations_to_patch.items():
            env_var = "DD_TRACE_%s_ENABLED" % module.upper()
            if env_var in os.environ:
                integrations_to_patch[module] = asbool(os.environ[env_var])
        dd_patch_modules = os.getenv("DD_PATCH_MODULES")
        dd_patch_modules_to_str = parse_tags_str(dd_patch_modules)
        integrations_to_patch.update(
            {k: asbool(v) for k, v in dd_patch_modules_to_str.items() if k in SUPPORTED_LLMOBS_INTEGRATIONS.values()}
        )
        patch(raise_errors=True, **integrations_to_patch)
        llm_patched_modules = [k for k, v in integrations_to_patch.items() if v]
        log.debug("Patched LLM integrations: %s", llm_patched_modules)

    @classmethod
    def export_span(cls, span: Optional[Span] = None) -> Optional[ExportedLLMObsSpan]:
        """Returns a simple representation of a span to export its span and trace IDs.
        If no span is provided, the current active LLMObs-type span will be used.
        """
        if span is None:
            span = cls._instance._current_span()
            if span is None:
                telemetry.record_span_exported(span, "no_active_span")
                log.warning("No span provided and no active LLMObs-generated span found.")
                return None
        error = None
        try:
            if span.span_type != SpanTypes.LLM:
                error = "invalid_span"
                log.warning("Span must be an LLMObs-generated span.")
                return None
            return ExportedLLMObsSpan(
                span_id=str(span.span_id),
                trace_id=format_trace_id(span._get_ctx_item(LLMOBS_TRACE_ID) or span.trace_id),
            )
        except (TypeError, AttributeError):
            error = "invalid_span"
            log.warning("Failed to export span. Span must be a valid Span object.")
            return None
        finally:
            telemetry.record_span_exported(span, error)

    def _current_span(self) -> Optional[Span]:
        """Returns the currently active LLMObs-generated span.
        Note that there may be an active span represented by a context object
        (i.e. a distributed trace) which will not be returned by this method.
        """
        active = self._llmobs_context_provider.active()
        return active if isinstance(active, Span) else None

    def _current_trace_context(self) -> Optional[Context]:
        """Returns the context for the current LLMObs trace."""
        active = self._llmobs_context_provider.active()
        if isinstance(active, Context):
            return active
        elif isinstance(active, Span):
            context = active.context
            context._meta[PROPAGATED_LLMOBS_TRACE_ID_KEY] = str(active._get_ctx_item(LLMOBS_TRACE_ID)) or str(
                active.trace_id
            )
            return context
        return None

    def _activate_llmobs_span(self, span: Span) -> None:
        """Propagate the llmobs parent span's ID as the new span's parent ID and activate the new span."""
        llmobs_parent = self._llmobs_context_provider.active()
        if llmobs_parent:
            span._set_ctx_item(PARENT_ID_KEY, str(llmobs_parent.span_id))
            parent_llmobs_trace_id = (
                llmobs_parent._get_ctx_item(LLMOBS_TRACE_ID)
                if isinstance(llmobs_parent, Span)
                else llmobs_parent._meta.get(PROPAGATED_LLMOBS_TRACE_ID_KEY)
            )
            llmobs_trace_id = (
                int(parent_llmobs_trace_id) if parent_llmobs_trace_id is not None else llmobs_parent.trace_id
            )
            span._set_ctx_item(LLMOBS_TRACE_ID, llmobs_trace_id)
        else:
            span._set_ctx_item(PARENT_ID_KEY, ROOT_PARENT_ID)
            span._set_ctx_item(LLMOBS_TRACE_ID, rand128bits())
        self._llmobs_context_provider.activate(span)

    def _start_span(
        self,
        operation_kind: str,
        name: Optional[str] = None,
        session_id: Optional[str] = None,
        model_name: Optional[str] = None,
        model_provider: Optional[str] = None,
        ml_app: Optional[str] = None,
        _decorator: bool = False,
    ) -> Span:
        if name is None:
            name = operation_kind
        span = self.tracer.trace(name, resource=operation_kind, span_type=SpanTypes.LLM)

        if not self.enabled:
            return span

        span._set_ctx_item(SPAN_KIND, operation_kind)
        if model_name is not None:
            span._set_ctx_item(MODEL_NAME, model_name)
        if model_provider is not None:
            span._set_ctx_item(MODEL_PROVIDER, model_provider)
        session_id = session_id if session_id is not None else _get_session_id(span)
        if session_id is not None:
            span._set_ctx_item(SESSION_ID, session_id)

        ml_app = ml_app if ml_app is not None else _get_ml_app(span)
        if ml_app is None:
            raise ValueError(
                "ml_app is required for sending LLM Observability data. "
                "Ensure the name of your LLM application is set via `DD_LLMOBS_ML_APP` or `LLMObs.enable(ml_app='...')`"
                "before running your application."
            )
        span._set_ctx_items({DECORATOR: _decorator, SPAN_KIND: operation_kind, ML_APP: ml_app})
        log.debug("Starting LLMObs span: %s, span_kind: %s, ml_app: %s", name, operation_kind, ml_app)
        return span

    @classmethod
    def llm(
        cls,
        model_name: Optional[str] = None,
        name: Optional[str] = None,
        model_provider: Optional[str] = None,
        session_id: Optional[str] = None,
        ml_app: Optional[str] = None,
        _decorator: bool = False,
    ) -> Span:
        """
        Trace an invocation call to an LLM where inputs and outputs are represented as text.

        :param str model_name: The name of the invoked LLM. If not provided, a default value of "custom" will be set.
        :param str name: The name of the traced operation. If not provided, a default value of "llm" will be set.
        :param str model_provider: The name of the invoked LLM provider (ex: openai, bedrock).
                                   If not provided, a default value of "custom" will be set.
        :param str session_id: The ID of the underlying user session. Required for tracking sessions.
        :param str ml_app: The name of the ML application that the agent is orchestrating. If not provided, the default
                           value will be set to the value of `DD_LLMOBS_ML_APP`.

        :returns: The Span object representing the traced operation.
        """
        if cls.enabled is False:
            log.warning(SPAN_START_WHILE_DISABLED_WARNING)
        if model_name is None:
            model_name = "custom"
        if model_provider is None:
            model_provider = "custom"
        return cls._instance._start_span(
            "llm",
            name,
            model_name=model_name,
            model_provider=model_provider,
            session_id=session_id,
            ml_app=ml_app,
            _decorator=_decorator,
        )

    @classmethod
    def tool(
        cls,
        name: Optional[str] = None,
        session_id: Optional[str] = None,
        ml_app: Optional[str] = None,
        _decorator: bool = False,
    ) -> Span:
        """
        Trace a call to an external interface or API.

        :param str name: The name of the traced operation. If not provided, a default value of "tool" will be set.
        :param str session_id: The ID of the underlying user session. Required for tracking sessions.
        :param str ml_app: The name of the ML application that the agent is orchestrating. If not provided, the default
                           value will be set to the value of `DD_LLMOBS_ML_APP`.

        :returns: The Span object representing the traced operation.
        """
        if cls.enabled is False:
            log.warning(SPAN_START_WHILE_DISABLED_WARNING)
        return cls._instance._start_span("tool", name=name, session_id=session_id, ml_app=ml_app, _decorator=_decorator)

    @classmethod
    def task(
        cls,
        name: Optional[str] = None,
        session_id: Optional[str] = None,
        ml_app: Optional[str] = None,
        _decorator: bool = False,
    ) -> Span:
        """
        Trace a standalone non-LLM operation which does not involve an external request.

        :param str name: The name of the traced operation. If not provided, a default value of "task" will be set.
        :param str session_id: The ID of the underlying user session. Required for tracking sessions.
        :param str ml_app: The name of the ML application that the agent is orchestrating. If not provided, the default
                           value will be set to the value of `DD_LLMOBS_ML_APP`.

        :returns: The Span object representing the traced operation.
        """
        if cls.enabled is False:
            log.warning(SPAN_START_WHILE_DISABLED_WARNING)
        return cls._instance._start_span("task", name=name, session_id=session_id, ml_app=ml_app, _decorator=_decorator)

    @classmethod
    def agent(
        cls,
        name: Optional[str] = None,
        session_id: Optional[str] = None,
        ml_app: Optional[str] = None,
        _decorator: bool = False,
    ) -> Span:
        """
        Trace a dynamic workflow in which an embedded language model (agent) decides what sequence of actions to take.

        :param str name: The name of the traced operation. If not provided, a default value of "agent" will be set.
        :param str session_id: The ID of the underlying user session. Required for tracking sessions.
        :param str ml_app: The name of the ML application that the agent is orchestrating. If not provided, the default
                           value will be set to the value of `DD_LLMOBS_ML_APP`.

        :returns: The Span object representing the traced operation.
        """
        if cls.enabled is False:
            log.warning(SPAN_START_WHILE_DISABLED_WARNING)
        return cls._instance._start_span(
            "agent", name=name, session_id=session_id, ml_app=ml_app, _decorator=_decorator
        )

    @classmethod
    def workflow(
        cls,
        name: Optional[str] = None,
        session_id: Optional[str] = None,
        ml_app: Optional[str] = None,
        _decorator: bool = False,
    ) -> Span:
        """
        Trace a predefined or static sequence of operations.

        :param str name: The name of the traced operation. If not provided, a default value of "workflow" will be set.
        :param str session_id: The ID of the underlying user session. Required for tracking sessions.
        :param str ml_app: The name of the ML application that the agent is orchestrating. If not provided, the default
                           value will be set to the value of `DD_LLMOBS_ML_APP`.

        :returns: The Span object representing the traced operation.
        """
        if cls.enabled is False:
            log.warning(SPAN_START_WHILE_DISABLED_WARNING)
        return cls._instance._start_span(
            "workflow", name=name, session_id=session_id, ml_app=ml_app, _decorator=_decorator
        )

    @classmethod
    def embedding(
        cls,
        model_name: Optional[str] = None,
        name: Optional[str] = None,
        model_provider: Optional[str] = None,
        session_id: Optional[str] = None,
        ml_app: Optional[str] = None,
        _decorator: bool = False,
    ) -> Span:
        """
        Trace a call to an embedding model or function to create an embedding.

        :param str model_name: The name of the invoked embedding model.
                               If not provided, a default value of "custom" will be set.
        :param str name: The name of the traced operation. If not provided, a default value of "embedding" will be set.
        :param str model_provider: The name of the invoked LLM provider (ex: openai, bedrock).
                                   If not provided, a default value of "custom" will be set.
        :param str session_id: The ID of the underlying user session. Required for tracking sessions.
        :param str ml_app: The name of the ML application that the agent is orchestrating. If not provided, the default
                           value will be set to the value of `DD_LLMOBS_ML_APP`.

        :returns: The Span object representing the traced operation.
        """
        if cls.enabled is False:
            log.warning(SPAN_START_WHILE_DISABLED_WARNING)
        if model_name is None:
            model_name = "custom"
        if model_provider is None:
            model_provider = "custom"
        return cls._instance._start_span(
            "embedding",
            name,
            model_name=model_name,
            model_provider=model_provider,
            session_id=session_id,
            ml_app=ml_app,
            _decorator=_decorator,
        )

    @classmethod
    def retrieval(
        cls,
        name: Optional[str] = None,
        session_id: Optional[str] = None,
        ml_app: Optional[str] = None,
        _decorator: bool = False,
    ) -> Span:
        """
        Trace a vector search operation involving a list of documents being returned from an external knowledge base.

        :param str name: The name of the traced operation. If not provided, a default value of "workflow" will be set.
        :param str session_id: The ID of the underlying user session. Required for tracking sessions.
        :param str ml_app: The name of the ML application that the agent is orchestrating. If not provided, the default
                           value will be set to the value of `DD_LLMOBS_ML_APP`.

        :returns: The Span object representing the traced operation.
        """
        if cls.enabled is False:
            log.warning(SPAN_START_WHILE_DISABLED_WARNING)
        return cls._instance._start_span(
            "retrieval", name=name, session_id=session_id, ml_app=ml_app, _decorator=_decorator
        )

    @classmethod
    def _experiment(
        cls,
        name: Optional[str] = None,
        session_id: Optional[str] = None,
        ml_app: Optional[str] = None,
        experiment_id: Optional[str] = None,
        run_id: Optional[str] = None,
        run_iteration: Optional[int] = None,
    ) -> Span:
        """
        Trace an LLM experiment, only used internally by the experiments SDK.
        :param str name: The name of the traced operation. If not provided, a default value of "agent" will be set.
        :param str session_id: The ID of the underlying user session. Required for tracking sessions.
        :param str ml_app: The name of the ML application that the agent is orchestrating. If not provided, the default
                           value will be set to the value of `DD_LLMOBS_ML_APP`.
        :param str experiment_id: The ID of the experiment to associate with this span and its children.
        :returns: The Span object representing the traced operation.
        """
        if cls.enabled is False:
            log.warning(SPAN_START_WHILE_DISABLED_WARNING)
        span = cls._instance._start_span("experiment", name=name, session_id=session_id, ml_app=ml_app)

        # Set experiment_id in baggage if provided
        if experiment_id:
            span.context.set_baggage_item(EXPERIMENT_ID_KEY, experiment_id)

        if run_id:
            span.context.set_baggage_item(EXPERIMENT_RUN_ID_KEY, run_id)

        if run_iteration is not None:
            span.context.set_baggage_item(EXPERIMENT_RUN_ITERATION_KEY, run_iteration)

        return span

    @classmethod
    def annotate(
        cls,
        span: Optional[Span] = None,
        prompt: Optional[dict] = None,
        input_data: Optional[Any] = None,
        output_data: Optional[Any] = None,
        metadata: Optional[Dict[str, Any]] = None,
        metrics: Optional[Dict[str, Any]] = None,
        tags: Optional[Dict[str, Any]] = None,
        tool_definitions: Optional[List[Dict[str, Any]]] = None,
        _name: Optional[str] = None,
    ) -> None:
        """
        Sets metadata, inputs, outputs, tags, and metrics as provided for a given LLMObs span.
        Note that with the exception of tags, this method will override any existing values for the provided fields.

        :param Span span: Span to annotate. If no span is provided, the current active span will be used.
                          Must be an LLMObs-type span, i.e. generated by the LLMObs SDK.
        :param prompt: A dictionary that represents the prompt used for an LLM call in the following form:
                        `{
                            "id": "...",
                            "template": "...",
                            "chat_template": [{"content": "...", "role": "..."}, ...])
                            "version": "...",
                            "variables": {"variable_1": "...", ...},
                            tags": {"tag_1": "...", ...},
                        }`.
                        Can also be set using the `ddtrace.llmobs.utils.Prompt` constructor class.
                        - This argument is only applicable to LLM spans.
                        - The dictionary may contain two optional keys relevant to RAG applications:
                            `rag_context_variables` - a list of variable key names that contain ground
                                                        truth context information
                            `rag_query_variables` - a list of variable key names that contains query
                                                        information for an LLM call
        :param input_data: A single input string, dictionary, or a list of dictionaries based on the span kind:
                           - llm spans: accepts a string, or a dictionary of form {"content": "...", "role": "...",
                                        "tool_calls": ..., "tool_results": ...}, where "tool_calls" are an optional
                                        list of tool call dictionaries with required keys: "name", "arguments", and
                                        optional keys: "tool_id", "type", and "tool_results" are an optional list of
                                        tool result dictionaries with required key: "result", and optional keys:
                                        "name", "tool_id", "type" for function calling scenarios.
                           - embedding spans: accepts a string, list of strings, or a dictionary of form
                                              {"text": "...", ...} or a list of dictionaries with the same signature.
                           - other: any JSON serializable type.
        :param output_data: A single output string, dictionary, or a list of dictionaries based on the span kind:
                           - llm spans: accepts a string, or a dictionary of form {"content": "...", "role": "...",
                                        "tool_calls": ...}, where "tool_calls" are an optional list of tool call
                                        dictionaries with required keys: "name", "arguments", and optional keys:
                                        "tool_id", "type" for function calling scenarios.
                           - retrieval spans: a dictionary containing any of the key value pairs
                                              {"name": str, "id": str, "text": str, "score": float},
                                              or a list of dictionaries with the same signature.
                           - other: any JSON serializable type.
        :param metadata: Dictionary of JSON serializable key-value metadata pairs relevant to the input/output operation
                         described by the LLMObs span.
        :param tags: Dictionary of JSON serializable key-value tag pairs to set or update on the LLMObs span
                     regarding the span's context.
        :param tool_definitions: List of tool definition dictionaries for tool calling scenarios.
                            - This argument is only applicable to LLM spans.
                            - Each tool definition is a dictionary containing a required "name" (string),
                                   and optional "description" (string) and "schema" (JSON serializable dictionary) keys.
        :param metrics: Dictionary of JSON serializable key-value metric pairs,
                        such as `{prompt,completion,total}_tokens`.
        """
        error = None
        try:
            if span is None:
                span = cls._instance._current_span()
                if span is None:
                    error = "invalid_span_no_active_spans"
                    log.warning("No span provided and no active LLMObs-generated span found.")
                    return
            if span.span_type != SpanTypes.LLM:
                error = "invalid_span_type"
                log.warning("Span must be an LLMObs-generated span.")
                return
            if span.finished:
                error = "invalid_finished_span"
                log.warning("Cannot annotate a finished span.")
                return
            if metadata is not None:
                if not isinstance(metadata, dict):
                    error = "invalid_metadata"
                    log.warning("metadata must be a dictionary")
                else:
                    cls._set_dict_attribute(span, METADATA, metadata)
            if metrics is not None:
                if not isinstance(metrics, dict) or not all(isinstance(v, (int, float)) for v in metrics.values()):
                    error = "invalid_metrics"
                    log.warning("metrics must be a dictionary of string key - numeric value pairs.")
                else:
                    cls._set_dict_attribute(span, METRICS, metrics)
            if tags is not None:
                if not isinstance(tags, dict):
                    error = "invalid_tags"
                    log.warning("span tags must be a dictionary of string key - primitive value pairs.")
                else:
                    session_id = tags.get("session_id")
                    if session_id:
                        span._set_ctx_item(SESSION_ID, str(session_id))
                    cls._set_dict_attribute(span, TAGS, tags)
            if tool_definitions is not None:
                validated_tool_definitions = extract_tool_definitions(tool_definitions)
                if validated_tool_definitions:
                    span._set_ctx_item(TOOL_DEFINITIONS, validated_tool_definitions)
            span_kind = span._get_ctx_item(SPAN_KIND)
            if _name is not None:
                span.name = _name
            if prompt is not None:
                try:
                    validated_prompt = _validate_prompt(prompt, strict_validation=False)
                    cls._set_dict_attribute(span, INPUT_PROMPT, validated_prompt)
                except (ValueError, TypeError) as e:
                    error = "invalid_prompt"
                    log.warning("Failed to validate prompt with error:", str(e), exc_info=True)
            if not span_kind:
                log.debug("Span kind not specified, skipping annotation for input/output data")
                return
            if input_data is not None or output_data is not None:
                if span_kind == "llm":
                    error = cls._tag_llm_io(span, input_messages=input_data, output_messages=output_data)
                elif span_kind == "embedding":
                    error = cls._tag_embedding_io(span, input_documents=input_data, output_text=output_data)
                elif span_kind == "retrieval":
                    error = cls._tag_retrieval_io(span, input_text=input_data, output_documents=output_data)
                elif span_kind == "experiment":
                    cls._tag_freeform_io(span, input_value=input_data, output_value=output_data)
                else:
                    cls._tag_text_io(span, input_value=input_data, output_value=output_data)
        finally:
            telemetry.record_llmobs_annotate(span, error)

    @classmethod
    def _tag_llm_io(cls, span, input_messages=None, output_messages=None) -> Optional[str]:
        """Tags input/output messages for LLM-kind spans.
        Will be mapped to span's `meta.{input,output}.messages` fields.
        """
        if input_messages is not None:
            try:
                if not isinstance(input_messages, Messages):
                    input_messages = Messages(input_messages)
                if input_messages.messages:
                    span._set_ctx_item(INPUT_MESSAGES, input_messages.messages)
            except TypeError:
                log.warning("Failed to parse input messages.", exc_info=True)
                return "invalid_io_messages"
        if output_messages is None:
            return None
        try:
            if not isinstance(output_messages, Messages):
                output_messages = Messages(output_messages)
            if not output_messages.messages:
                return None
            span._set_ctx_item(OUTPUT_MESSAGES, output_messages.messages)
        except TypeError:
            log.warning("Failed to parse output messages.", exc_info=True)
            return "invalid_io_messages"
        return None

    @classmethod
    def _tag_embedding_io(cls, span, input_documents=None, output_text=None) -> Optional[str]:
        """Tags input documents and output text for embedding-kind spans.
        Will be mapped to span's `meta.{input,output}.text` fields.
        """
        if input_documents is not None:
            try:
                if not isinstance(input_documents, Documents):
                    input_documents = Documents(input_documents)
                if input_documents.documents:
                    span._set_ctx_item(INPUT_DOCUMENTS, input_documents.documents)
            except TypeError:
                log.warning("Failed to parse input documents.", exc_info=True)
                return "invalid_embedding_io"
        if output_text is None:
            return None
        span._set_ctx_item(OUTPUT_VALUE, str(output_text))
        return None

    @classmethod
    def _tag_retrieval_io(cls, span, input_text=None, output_documents=None) -> Optional[str]:
        """Tags input text and output documents for retrieval-kind spans.
        Will be mapped to span's `meta.{input,output}.text` fields.
        """
        if input_text is not None:
            span._set_ctx_item(INPUT_VALUE, safe_json(input_text))
        if output_documents is None:
            return None
        try:
            if not isinstance(output_documents, Documents):
                output_documents = Documents(output_documents)
            if not output_documents.documents:
                return None
            span._set_ctx_item(OUTPUT_DOCUMENTS, output_documents.documents)
        except TypeError:
            log.warning("Failed to parse output documents.", exc_info=True)
            return "invalid_retrieval_io"
        return None

    @classmethod
    def _tag_text_io(cls, span, input_value=None, output_value=None):
        """Tags input/output values for non-LLM kind spans.
        Will be mapped to span's `meta.{input,output}.values` fields.
        """
        if input_value is not None:
            span._set_ctx_item(INPUT_VALUE, safe_json(input_value))
        if output_value is not None:
            span._set_ctx_item(OUTPUT_VALUE, safe_json(output_value))

    @classmethod
    def _tag_freeform_io(cls, span, input_value=None, output_value=None):
        """Tags input/output values for experient spans.
        Will be mapped to span's `meta.{input,output}` fields.
        this is meant to be non restrictive on user's data, experiments allow
        arbitrary structured or non structured IO values in its spans
        """
        if input_value is not None:
            span._set_ctx_item(EXPERIMENTS_INPUT, safe_json(input_value))
        if output_value is not None:
            span._set_ctx_item(EXPERIMENTS_OUTPUT, safe_json(output_value))

    @staticmethod
    def _set_dict_attribute(span: Span, key, value: Dict[str, Any]) -> None:
        """Sets a given LLM Obs span attribute with a dictionary key/values.
        If the attribute is already set on the span, the new dict with be merged with the existing
        dict.
        """
        existing_value = span._get_ctx_item(key) or {}
        existing_value.update(value)
        span._set_ctx_item(key, existing_value)

    @classmethod
    def submit_evaluation_for(
        cls,
        label: str,
        metric_type: str,
        value: Union[str, int, float, bool],
        span: Optional[dict] = None,
        span_with_tag_value: Optional[Dict[str, str]] = None,
        tags: Optional[Dict[str, str]] = None,
        ml_app: Optional[str] = None,
        timestamp_ms: Optional[int] = None,
        metadata: Optional[Dict[str, object]] = None,
        assessment: Optional[str] = None,
        reasoning: Optional[str] = None,
    ) -> None:
        """
        Submits a custom evaluation metric for a given span. This method is deprecated and will be
        removed in the next major version of ddtrace (4.0). Please use `LLMObs.submit_evaluation()` instead.
        """
        log.warning(
            "LLMObs.submit_evaluation_for() is deprecated and will be removed in the next major "
            "version of ddtrace (4.0). Please use LLMObs.submit_evaluation() instead."
        )
        return cls.submit_evaluation(
            label=label,
            metric_type=metric_type,
            value=value,
            span=span,
            span_with_tag_value=span_with_tag_value,
            tags=tags,
            ml_app=ml_app,
            timestamp_ms=timestamp_ms,
            metadata=metadata,
            assessment=assessment,
            reasoning=reasoning,
        )

    @classmethod
    def submit_evaluation(
        cls,
        label: str,
        metric_type: str,
        value: Union[str, int, float, bool],
        span_context: Optional[Dict[str, str]] = None,
        span: Optional[dict] = None,
        span_with_tag_value: Optional[Dict[str, str]] = None,
        tags: Optional[Dict[str, str]] = None,
        ml_app: Optional[str] = None,
        timestamp_ms: Optional[int] = None,
        metadata: Optional[Dict[str, object]] = None,
        assessment: Optional[str] = None,
        reasoning: Optional[str] = None,
    ) -> None:
        """
        Submits a custom evaluation metric for a given span.

        :param str label: The name of the evaluation metric.
        :param str metric_type: The type of the evaluation metric. One of "categorical", "score", "boolean".
        :param value: The value of the evaluation metric.
                      Must be a string (categorical), integer (score), float (score), or boolean (boolean).
        :param dict span_context: A dictionary containing the span_id and trace_id of interest. This is a
                            deprecated parameter and will be removed in the next major version of
                            ddtrace (4.0). Please use `span` or `span_with_tag_value` instead.
        :param dict span: A dictionary of shape {'span_id': str, 'trace_id': str} uniquely identifying
                            the span associated with this evaluation.
        :param dict span_with_tag_value: A dictionary with the format {'tag_key': str, 'tag_value': str}
                            uniquely identifying the span associated with this evaluation.
        :param tags: A dictionary of string key-value pairs to tag the evaluation metric with.
        :param str ml_app: The name of the ML application
        :param int timestamp_ms: The unix timestamp in milliseconds when the evaluation metric result was generated.
                                    If not set, the current time will be used.
        :param dict metadata: A JSON serializable dictionary of key-value metadata pairs relevant to the
                                evaluation metric.
        :param str assessment: An assessment of this evaluation. Must be either "pass" or "fail".
        :param str reasoning: An explanation of the evaluation result.
        """
        if span_context is not None:
            log.warning(
                "The `span_context` parameter is deprecated and will be removed in the next major version of "
                "ddtrace (4.0). Please use `span` or `span_with_tag_value` instead."
            )
            span = span or span_context

        if cls.enabled is False:
            log.debug(
                "LLMObs.submit_evaluation() called when LLMObs is not enabled. ",
                "Evaluation metric data will not be sent.",
            )
            return

        error = None
        join_on = {}
        try:
            has_exactly_one_joining_key = (span is not None) ^ (span_with_tag_value is not None)

            if not has_exactly_one_joining_key:
                error = "provided_both_span_and_tag_joining_key"
                raise ValueError(
                    "Exactly one of `span` or `span_with_tag_value` must be specified to submit an evaluation metric."
                )

            if span is not None:
                if (
                    not isinstance(span, dict)
                    or not isinstance(span.get("span_id"), str)
                    or not isinstance(span.get("trace_id"), str)
                ):
                    error = "invalid_span"
                    raise TypeError(
                        "`span` must be a dictionary containing both span_id and trace_id keys. "
                        "LLMObs.export_span() can be used to generate this dictionary from a given span."
                    )
                join_on["span"] = span
            elif span_with_tag_value is not None:
                if (
                    not isinstance(span_with_tag_value, dict)
                    or not isinstance(span_with_tag_value.get("tag_key"), str)
                    or not isinstance(span_with_tag_value.get("tag_value"), str)
                ):
                    error = "invalid_joining_key"
                    raise TypeError(
                        "`span_with_tag_value` must be a dict with keys 'tag_key' and 'tag_value' "
                        "containing string values"
                    )
                join_on["tag"] = {
                    "key": span_with_tag_value.get("tag_key"),
                    "value": span_with_tag_value.get("tag_value"),
                }

            timestamp_ms = timestamp_ms if timestamp_ms else int(time.time() * 1000)

            if not isinstance(timestamp_ms, int) or timestamp_ms < 0:
                error = "invalid_timestamp"
                raise ValueError("timestamp_ms must be a non-negative integer. Evaluation metric data will not be sent")

            if not label:
                error = "invalid_metric_label"
                raise ValueError("label must be the specified name of the evaluation metric.")

            metric_type = metric_type.lower()
            if metric_type not in ("categorical", "score", "boolean"):
                error = "invalid_metric_type"
                raise ValueError("metric_type must be one of 'categorical', 'score', or 'boolean'.")

            if metric_type == "categorical" and not isinstance(value, str):
                error = "invalid_metric_value"
                raise TypeError("value must be a string for a categorical metric.")
            if metric_type == "score" and not isinstance(value, (int, float)):
                error = "invalid_metric_value"
                raise TypeError("value must be an integer or float for a score metric.")
            if metric_type == "boolean" and not isinstance(value, bool):
                error = "invalid_metric_value"
                raise TypeError("value must be a boolean for a boolean metric.")

            if tags is not None and not isinstance(tags, dict):
                log.warning("tags must be a dictionary of string key-value pairs.")
                tags = {}

            ml_app = ml_app if ml_app else config._llmobs_ml_app
            if not ml_app:
                error = "missing_ml_app"
                log.warning(
                    "ML App name is required for sending evaluation metrics. Evaluation metric data will not be sent. "
                    "Ensure this configuration is set before running your application."
                )
                return

            evaluation_tags = {
                "ddtrace.version": ddtrace.__version__,
                "ml_app": ml_app,
            }

            if tags:
                for k, v in tags.items():
                    try:
                        evaluation_tags[ensure_text(k)] = ensure_text(v)
                    except TypeError:
                        error = "invalid_tags"
                        log.warning("Failed to parse tags. Tags for evaluation metrics must be strings.")

            evaluation_metric: LLMObsEvaluationMetricEvent = {
                "join_on": join_on,
                "label": str(label),
                "metric_type": metric_type,
                "timestamp_ms": timestamp_ms,
                "{}_value".format(metric_type): value,  # type: ignore
                "ml_app": ml_app,
                "tags": ["{}:{}".format(k, v) for k, v in evaluation_tags.items()],
            }

            if assessment:
                if not isinstance(assessment, str) or assessment not in ("pass", "fail"):
                    error = "invalid_assessment"
                    log.warning("Failed to parse assessment. assessment must be either 'pass' or 'fail'.")
                else:
                    evaluation_metric["assessment"] = assessment
            if reasoning:
                if not isinstance(reasoning, str):
                    error = "invalid_reasoning"
                    log.warning("Failed to parse reasoning. reasoning must be a string.")
                else:
                    evaluation_metric["reasoning"] = reasoning

            if metadata:
                if not isinstance(metadata, dict):
                    error = "invalid_metadata"
                    log.warning("metadata must be json serializable dictionary.")
                else:
                    metadata = safe_json(metadata)
                    if metadata and isinstance(metadata, str):
                        evaluation_metric["metadata"] = json.loads(metadata)

            cls._instance._llmobs_eval_metric_writer.enqueue(evaluation_metric)
        finally:
            telemetry.record_llmobs_submit_evaluation(join_on, metric_type, error)

    @classmethod
    def _inject_llmobs_context(cls, span_context: Context, request_headers: Dict[str, str]) -> None:
        if cls.enabled is False:
            return

        active_span = cls._instance._llmobs_context_provider.active()
        active_context = active_span.context if isinstance(active_span, Span) else active_span

        parent_id = str(active_context.span_id) if active_context is not None else ROOT_PARENT_ID

        ml_app = None
        if isinstance(active_span, Span):
            ml_app = active_span._get_ctx_item(ML_APP)
            llmobs_trace_id = active_span._get_ctx_item(LLMOBS_TRACE_ID)
        elif active_context is not None:
            ml_app = active_context._meta.get(PROPAGATED_ML_APP_KEY) or config._llmobs_ml_app
            llmobs_trace_id = active_context._meta.get(PROPAGATED_LLMOBS_TRACE_ID_KEY) or rand128bits()
        else:
            ml_app = config._llmobs_ml_app
            llmobs_trace_id = rand128bits()

        span_context._meta[PROPAGATED_PARENT_ID_KEY] = parent_id
        span_context._meta[PROPAGATED_LLMOBS_TRACE_ID_KEY] = str(llmobs_trace_id)

        if ml_app is not None:
            span_context._meta[PROPAGATED_ML_APP_KEY] = ml_app

    @classmethod
    def inject_distributed_headers(cls, request_headers: Dict[str, str], span: Optional[Span] = None) -> Dict[str, str]:
        """Injects the span's distributed context into the given request headers."""
        if cls.enabled is False:
            log.warning(
                "LLMObs.inject_distributed_headers() called when LLMObs is not enabled. "
                "Distributed context will not be injected."
            )
            return request_headers
        error = None
        try:
            if not isinstance(request_headers, dict):
                error = "invalid_request_headers"
                log.warning("request_headers must be a dictionary of string key-value pairs.")
                return request_headers
            if span is None:
                span = cls._instance.tracer.current_span()
            if span is None:
                error = "no_active_span"
                log.warning("No span provided and no currently active span found.")
                return request_headers
            if not isinstance(span, Span):
                error = "invalid_span"
                log.warning("span must be a valid Span object. Distributed context will not be injected.")
                return request_headers
            HTTPPropagator.inject(span.context, request_headers)
            return request_headers
        finally:
            telemetry.record_inject_distributed_headers(error)

    @classmethod
    def _activate_llmobs_distributed_context(cls, request_headers: Dict[str, str], context: Context) -> Optional[str]:
        if cls.enabled is False:
            return None
        if not context.trace_id or not context.span_id:
            log.warning("Failed to extract trace/span ID from request headers.")
            return "missing_context"
        _parent_id = context._meta.get(PROPAGATED_PARENT_ID_KEY)
        if _parent_id is None:
            log.debug("Failed to extract LLMObs parent ID from request headers.")
            return "missing_parent_id"
        try:
            parent_id = int(_parent_id)
        except ValueError:
            log.warning("Failed to parse LLMObs parent ID from request headers.")
            return "invalid_parent_id"
        parent_llmobs_trace_id = context._meta.get(PROPAGATED_LLMOBS_TRACE_ID_KEY)
        if parent_llmobs_trace_id is None:
            log.debug("Failed to extract LLMObs trace ID from request headers. Expected string, got None.")
            llmobs_context = Context(trace_id=context.trace_id, span_id=parent_id)
            llmobs_context._meta[PROPAGATED_LLMOBS_TRACE_ID_KEY] = str(context.trace_id)
            cls._instance._llmobs_context_provider.activate(llmobs_context)
            return "missing_parent_llmobs_trace_id"
        llmobs_context = Context(trace_id=context.trace_id, span_id=parent_id)
        llmobs_context._meta[PROPAGATED_LLMOBS_TRACE_ID_KEY] = str(parent_llmobs_trace_id)
        cls._instance._llmobs_context_provider.activate(llmobs_context)
        return None

    @classmethod
    def activate_distributed_headers(cls, request_headers: Dict[str, str]) -> None:
        """
        Activates distributed tracing headers for the current request.

        :param request_headers: A dictionary containing the headers for the current request.
        """
        if cls.enabled is False:
            log.warning(
                "LLMObs.activate_distributed_headers() called when LLMObs is not enabled. "
                "Distributed context will not be activated."
            )
            return
        context = HTTPPropagator.extract(request_headers)
        cls._instance.tracer.context_provider.activate(context)
        error = cls._instance._activate_llmobs_distributed_context(request_headers, context)
        telemetry.record_activate_distributed_headers(error)


# initialize the default llmobs instance
LLMObs._instance = LLMObs()
