feat(llmobs): llmobs-specific context manager (#12236)

[MLOB-1342] This is a manual copy of #10767 to 3.x-staging due to merge
conflicts between 3.x-staging and main.

## Summary
Public facing changes:
- Any LLMObs method (`export_span(), annotate()`) that allow an optional
span argument will now default to finding the current active LLMObs span
rather than the current active APM span.
- ~Adds multithreading (futures.multithreading) support for LLMObs.
Previously multithreaded apps would result in broken traces.~ UPDATE:
will be moving multithreading support to a separate PR.

Private changes:
- LLMObs has its own context provider which keeps track of the active
LLM-type span (generated by both LLMObs._start_span() and LLM
integrations)
- HTTPPropagation now adds LLMObs parent ID as a field on the request
headers directly, rather than through the context object.
- Adds private helper method `LLMObs._instance.current_span()`, returns
the current active LLMObs-generated (integration, SDK) span.
- Adds private helper method
`LLMObs._instance._current_trace_context()`, returns current LLMObs
context (which can represent both a span or a distributed span)
- Adds a new field to the LLMObs span event struct, `_dd` which is a
str/str dictionary containing the span/trace IDs of the APM span to
correlate with. Currently these are the same span/trace IDs as the
LLMObs span/trace ID, but this unlocks future steps of using independent
span/trace IDs.
 
## Previous behavior
LLMObs spans are based on APM spans, except LLMObs spans' parenting
involves only other LLMObs spans. So with a potential trace structure
containing a mixture of APM-specific and LLMObs spans, like:
```
Span A (LLMObs span) --> Span B (Apm-specific span) --> Span C (LLMObs span)
```
LLMObs only cares about the LLMObs spans, where span C's parent is the
root span, even though in APM it would be span B. Combined with
distributed tracing and multithreading, this makes it not so easy to
determine that "correct" (read LLMObs) parenting tree for traces
submitted to LLM Observability.

### Problems with previous approach

Previously we worked around this by traversing the span's local parent
tree and finding the next LLM-type span on both span start and finish
for non-distributed cases, and for distributed cases we would attach the
parent ID on the span context's meta field to be propagated in
distributed request headers. However attaching things to the span
context meta was not suitable long-term due to a couple factors:
1. Context objects are not thread-safe: in a multithreading case with
n>1 child threads creating their own spans, the parent ID stored in the
context object could be overwritten during thread execution, therefore
incorrectly propagating parent IDs.
2. Context objects store trace-specific information, and are not
designed for our use case where we skip spans here and there in the
trace. This also leads to edge cases that were handled with [ugly
workaround
code](https://github.com/DataDog/dd-trace-py/blob/3bffd02a071db91a6fdc56ae12b496dc84ff8abf/ddtrace/llmobs/_llmobs.py#L312-L317):
<details>
  <summary><i>Example ugly workaround</i></summary>
<b>Any meta fields set on the context object gets propagated as span
tags on all subsequent spans in the trace on span start time, except for
the spans in the first service of a trace which get propagated at span
finish time. Fixing this resulted in overriding these span tags on span
start and more checks on span finish.</b>
</details>

## Current approach
Instead of being dependent on a Context object that doesn't quite fit
our use case and trying to make it fit our use case, we simply keep
track of our own active LLMObs span/context:
- `LLMObsContextProvider` handles keeping track of the current active
LLMObs span via `active()` and `activate()`
- Instead of traversing a span's local ancestor tree to solve for a
span's llmobs parent ID, we just use
`LLMObsContextProvider._activate_llmobs_span()` and set the llmobs
parent ID as a tag at span start time.
(called by `LLMObs._start_span()` and
`BaseLLMIntegration.trace(submit_to_llmobs=True)` and the bedrock
integration).
- `LLMObs.inject_distributed_headers` now uses the LLMObsContextProvider
to inject the active llmobs span's ID into span context and request
headers
- `LLMObs.activate_distributed_headers()` now uses the
LLMObsContextProvider to activate the extracted llmobs context to
continue the trace in a distributed case.
- `trace_utils.activate_distributed_headers()` now includes automatic
llmobs context activation if llmobs is enabled. I've used signal
handling to avoid importing LLMObs entirely for non-LLMObs users (same
for `HTTPPropagator.inject()`).

By keeping track of our own active LLMObs spans, spans submitted to LLM
Observability have an independent set of span and parent IDs, even if
the span and trace IDs are shared with APM spans for now. This is the
first step to decoupling from tracer internals.

## Next steps
We can go further by generating LLMObs-specific span/trace IDs which are
separate from APM. This will solve some edge cases with traces involving
mixed APM/LLMObs spans.


## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met 
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)


[MLOB-1342]:
https://datadoghq.atlassian.net/browse/MLOB-1342?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ

Author: Yun-Kim

Diff for: ddtrace/contrib/internal/trace_utils.py

@@ -596,6 +596,7 @@ def activate_distributed_headers(tracer, int_config=None, request_headers=None,
         # We have parsed a trace id from headers, and we do not already
         # have a context with the same trace id active
         tracer.context_provider.activate(context)
+        core.dispatch("http.activate_distributed_headers", (request_headers, context))
 
         dispatch("distributed_context.activated", (context,))
 

Diff for: ddtrace/llmobs/_constants.py

@@ -45,6 +45,8 @@
 DROPPED_IO_COLLECTION_ERROR = "dropped_io"
 DROPPED_VALUE_TEXT = "[This value has been dropped because this span's size exceeds the 1MB size limit.]"
 
+ROOT_PARENT_ID = "undefined"
+
 # Set for traces of evaluator integrations e.g. `runner.integration:ragas`.
 # Used to differentiate traces of Datadog-run operations vs user-application operations.
 RUNNER_IS_INTEGRATION_SPAN_TAG = "runner.integration"

Diff for: ddtrace/llmobs/_context.py

@@ -0,0 +1,59 @@
+import contextvars
+from typing import Optional
+from typing import Union
+
+from ddtrace._trace.context import Context
+from ddtrace._trace.provider import DefaultContextProvider
+from ddtrace._trace.span import Span
+from ddtrace.ext import SpanTypes
+
+
+ContextTypeValue = Optional[Union[Context, Span]]
+
+
+_DD_LLMOBS_CONTEXTVAR: contextvars.ContextVar[ContextTypeValue] = contextvars.ContextVar(
+    "datadog_llmobs_contextvar",
+    default=None,
+)
+
+
+class LLMObsContextProvider(DefaultContextProvider):
+    """Context provider that retrieves contexts from a context variable.
+    It is suitable for synchronous programming and for asynchronous executors
+    that support contextvars.
+    """
+
+    def __init__(self) -> None:
+        super(DefaultContextProvider, self).__init__()
+        _DD_LLMOBS_CONTEXTVAR.set(None)
+
+    def _has_active_context(self) -> bool:
+        """Returns whether there is an active context in the current execution."""
+        ctx = _DD_LLMOBS_CONTEXTVAR.get()
+        return ctx is not None
+
+    def _update_active(self, span: Span) -> Optional[Span]:
+        """Updates the active LLMObs span.
+        The active span is updated to be the span's closest unfinished LLMObs ancestor span.
+        """
+        if not span.finished:
+            return span
+        new_active: Optional[Span] = span
+        while new_active and new_active.finished:
+            new_active = new_active._parent
+            if new_active and not new_active.finished and new_active.span_type == SpanTypes.LLM:
+                break
+        self.activate(new_active)
+        return new_active
+
+    def activate(self, ctx: ContextTypeValue) -> None:
+        """Makes the given context active in the current execution."""
+        _DD_LLMOBS_CONTEXTVAR.set(ctx)
+        super(DefaultContextProvider, self).activate(ctx)
+
+    def active(self) -> ContextTypeValue:
+        """Returns the active span or context for the current execution."""
+        item = _DD_LLMOBS_CONTEXTVAR.get()
+        if isinstance(item, Span):
+            return self._update_active(item)
+        return item

Diff for: ddtrace/llmobs/_integrations/base.py

@@ -18,11 +18,8 @@
 from ddtrace.internal.telemetry import telemetry_writer
 from ddtrace.internal.telemetry.constants import TELEMETRY_NAMESPACE
 from ddtrace.internal.utils.formats import asbool
-from ddtrace.llmobs._constants import PARENT_ID_KEY
-from ddtrace.llmobs._constants import PROPAGATED_PARENT_ID_KEY
 from ddtrace.llmobs._llmobs import LLMObs
 from ddtrace.llmobs._log_writer import V2LogWriter
-from ddtrace.llmobs._utils import _get_llmobs_parent_id
 from ddtrace.settings import IntegrationConfig
 from ddtrace.trace import Pin
 from ddtrace.trace import Span
@@ -138,21 +135,16 @@ def trace(self, pin: Pin, operation_id: str, submit_to_llmobs: bool = False, **k
         span.set_tag(_SPAN_MEASURED_KEY)
         self._set_base_span_tags(span, **kwargs)
         if submit_to_llmobs and self.llmobs_enabled:
-            if span.get_tag(PROPAGATED_PARENT_ID_KEY) is None:
-                # For non-distributed traces or spans in the first service of a distributed trace,
-                # The LLMObs parent ID tag is not set at span start time. We need to manually set the parent ID tag now
-                # in these cases to avoid conflicting with the later propagated tags.
-                parent_id = _get_llmobs_parent_id(span) or "undefined"
-                span._set_ctx_item(PARENT_ID_KEY, str(parent_id))
-        telemetry_writer.add_count_metric(
-            namespace=TELEMETRY_NAMESPACE.MLOBS,
-            name="span.start",
-            value=1,
-            tags=(
-                ("integration", self._integration_name),
-                ("autoinstrumented", "true"),
-            ),
-        )
+            LLMObs._instance._activate_llmobs_span(span)
+            telemetry_writer.add_count_metric(
+                namespace=TELEMETRY_NAMESPACE.MLOBS,
+                name="span.start",
+                value=1,
+                tags=(
+                    ("integration", self._integration_name),
+                    ("autoinstrumented", "true"),
+                ),
+            )
         return span
 
     @classmethod

Diff for: ddtrace/llmobs/_integrations/bedrock.py

@@ -4,18 +4,16 @@
 from typing import Optional
 
 from ddtrace.internal.logger import get_logger
+from ddtrace.llmobs import LLMObs
 from ddtrace.llmobs._constants import INPUT_MESSAGES
 from ddtrace.llmobs._constants import METADATA
 from ddtrace.llmobs._constants import METRICS
 from ddtrace.llmobs._constants import MODEL_NAME
 from ddtrace.llmobs._constants import MODEL_PROVIDER
 from ddtrace.llmobs._constants import OUTPUT_MESSAGES
-from ddtrace.llmobs._constants import PARENT_ID_KEY
-from ddtrace.llmobs._constants import PROPAGATED_PARENT_ID_KEY
 from ddtrace.llmobs._constants import SPAN_KIND
 from ddtrace.llmobs._integrations import BaseLLMIntegration
 from ddtrace.llmobs._integrations.utils import get_llmobs_metrics_tags
-from ddtrace.llmobs._utils import _get_llmobs_parent_id
 from ddtrace.trace import Span
 
 
@@ -34,9 +32,7 @@ def _llmobs_set_tags(
         operation: str = "",
     ) -> None:
         """Extract prompt/response tags from a completion and set them as temporary "_ml_obs.*" tags."""
-        if span.get_tag(PROPAGATED_PARENT_ID_KEY) is None:
-            parent_id = _get_llmobs_parent_id(span) or "undefined"
-            span._set_ctx_item(PARENT_ID_KEY, parent_id)
+        LLMObs._instance._activate_llmobs_span(span)
         parameters = {}
         if span.get_tag("bedrock.request.temperature"):
             parameters["temperature"] = float(span.get_tag("bedrock.request.temperature") or 0.0)

Diff for: ddtrace/llmobs/_integrations/langgraph.py

@@ -8,12 +8,13 @@
 from ddtrace.llmobs._constants import INPUT_VALUE
 from ddtrace.llmobs._constants import NAME
 from ddtrace.llmobs._constants import OUTPUT_VALUE
+from ddtrace.llmobs._constants import PARENT_ID_KEY
+from ddtrace.llmobs._constants import ROOT_PARENT_ID
 from ddtrace.llmobs._constants import SPAN_KIND
 from ddtrace.llmobs._constants import SPAN_LINKS
 from ddtrace.llmobs._integrations.base import BaseLLMIntegration
 from ddtrace.llmobs._integrations.utils import format_langchain_io
 from ddtrace.llmobs._utils import _get_attr
-from ddtrace.llmobs._utils import _get_llmobs_parent_id
 from ddtrace.llmobs._utils import _get_nearest_llmobs_ancestor
 from ddtrace.trace import Span
 from ddtrace.trace import tracer
@@ -175,7 +176,7 @@ def _default_span_link(span: Span):
     the span is linked to its parent's input.
     """
     return {
-        "span_id": str(_get_llmobs_parent_id(span)) or "undefined",
+        "span_id": span._get_ctx_item(PARENT_ID_KEY) or ROOT_PARENT_ID,
         "trace_id": "{:x}".format(span.trace_id),
         "attributes": {"from": "input", "to": "input"},
     }