diff --git a/.agents/skills/llmobs-integrations/SKILL.md b/.agents/skills/llmobs-integrations/SKILL.md index de3537c5ea7..c4817d3cf73 100644 --- a/.agents/skills/llmobs-integrations/SKILL.md +++ b/.agents/skills/llmobs-integrations/SKILL.md @@ -12,13 +12,16 @@ Use `apm-integrations` as well when changing contrib patching, registration, or Standard request spans use `LlmRequestEvent` and `core.context_with_event()`; the tracing subscriber owns span creation, errors, and finishing. For new integrations, leave the event's -`llmobs_integration` unset and connect extraction through `LlmEvents` subscribers under -`ddtrace/llmobs/_contrib//`. Specialized agent/tool/duplex integrations may use direct spans. +`llmobs_integration` unset and register the integration class under +`ddtrace/llmobs/_contrib//`. Shared `LlmEvents` subscribers in +`ddtrace/llmobs/_contrib/_subscribers.py` select the integration using `ctx.event.component`. +Specialized agent/tool/duplex integrations may use direct spans. Preserve the existing lifecycle when extending an integration. - Standard requests: `ddtrace/contrib/internal/anthropic/patch.py` and `ddtrace/llmobs/_integrations/anthropic.py`. -- Product subscribers: `ddtrace/llmobs/_contrib/anthropic/`; patch hooks: +- Product registration: `ddtrace/llmobs/_contrib/anthropic/`; shared subscribers: + `ddtrace/llmobs/_contrib/_subscribers.py`; patch hooks: `ddtrace/internal/integration_hooks.py` and `pyproject.toml`'s `ddtrace.integration_hooks` group. - Agent/tool spans: `ddtrace/llmobs/_integrations/claude_agent_sdk.py`. - Base contract: `ddtrace/llmobs/_integrations/base.py`. @@ -39,8 +42,9 @@ Preserve the existing lifecycle when extending an integration. - Reuse extraction and metadata sanitizers; retain provider-specific token accounting and media limits. - Build subscriber-owned integration instances lazily. Register through the contrib's `integration_hooks.on_patch()` / `on_unpatch()` calls so direct patching also works. - Keep subscribers available when LLMObs is disabled for APM shadow tags, and retain finishing - handlers for requests/streams in flight at unpatch. See the implementation guide for lifecycle details. + Keep shared subscribers available when LLMObs is disabled for APM shadow tags. Retain each + started request's integration on its execution context so requests and streams still finalize + after unpatch or repatch. See the implementation guide for lifecycle details. Read only the reference relevant to the change: [implementation](references/implementation-guide.md), [failure modes](references/failure-modes.md), diff --git a/.agents/skills/llmobs-integrations/references/failure-modes.md b/.agents/skills/llmobs-integrations/references/failure-modes.md index d43d658f908..7d3352ba931 100644 --- a/.agents/skills/llmobs-integrations/references/failure-modes.md +++ b/.agents/skills/llmobs-integrations/references/failure-modes.md @@ -21,7 +21,7 @@ Comprehensive debugging guide for all known LLMObs integration failure modes. Ea - Verify `LlmRequestEvent(..., submit_to_llmobs=True, request_kwargs=kwargs, ...)` in event-based patch wrappers, plus either registered LLMObs subscribers or `llmobs_integration=integration` for older integrations - Verify success and error paths call `ctx.dispatch_ended_event(...)` - Legacy integrations only: verify `patch()` stores integration: `module._datadog_integration = MyIntegration(integration_config=config.mylib)` -- Subscriber-based integrations: do not store the integration on the module (that re-couples the patch layer to LLMObs). Instead verify the subscriber's lazy `integration()` classmethod builds and caches the instance (e.g. `LLMObsAnthropicSubscriber.integration()` in `ddtrace/llmobs/_contrib/anthropic/subscribers.py`) +- Subscriber-based integrations: do not store the integration on the module (that re-couples the patch layer to LLMObs). Instead verify that the component is registered by its entry point and the shared starting subscriber in `ddtrace/llmobs/_contrib/_subscribers.py` builds and stores the integration on the request context - Check `DD_LLMOBS_ENABLED=1` is set --- diff --git a/.agents/skills/llmobs-integrations/references/implementation-guide.md b/.agents/skills/llmobs-integrations/references/implementation-guide.md index 139677a3192..cf2e42dea0e 100644 --- a/.agents/skills/llmobs-integrations/references/implementation-guide.md +++ b/.agents/skills/llmobs-integrations/references/implementation-guide.md @@ -21,7 +21,7 @@ the provider's `_integration_name`. Supply the two provider-specific behaviors: Use the actual base/reference signatures and implement or inherit the behavior needed by the integration. For standard request integrations, construct it lazily in the LLMObs subscriber -with `integration_config=config.`; see `ddtrace/llmobs/_contrib/anthropic/subscribers.py`. +with `integration_config=config.`; see `ddtrace/llmobs/_contrib/_subscribers.py`. Keep contrib free of LLMObs imports. When extending an integration that stores its instance on `module._datadog_integration`, preserve that ownership and cleanup unless migrating its lifecycle. @@ -62,8 +62,9 @@ standard request integrations so contrib does not import LLMObs. ## LLMObs subscribers and registration -Follow `ddtrace/llmobs/_contrib/anthropic/`. Each handler receives an `ExecutionContext` and -must filter on `ctx.event.component`, because the events are shared across integrations: +Follow `ddtrace/llmobs/_contrib/anthropic/` for integration registration. The shared subscribers +in `ddtrace/llmobs/_contrib/_subscribers.py` receive an `ExecutionContext` and select the +registered integration using `ctx.event.component`: - `LlmEvents.SPAN_STARTING`: before span creation, set `event.span_type = SpanTypes.LLM` when `event.submit_to_llmobs` and the integration's LLMObs enablement require it. @@ -71,11 +72,14 @@ must filter on `ctx.event.component`, because the events are shared across integ before children start. Preserve APM shadow tags and proxy handling from the reference. - `LlmEvents.SPAN_FINISHING`: extract data with `integration.llmobs_set_tags()` before finishing. -Set `auto_register = False`. Provide `listen()` / `unlisten()` in -`ddtrace/llmobs/_contrib//__init__.py`; register all three handlers in `listen()`. -During `unlisten()`, retain the finishing handler so in-flight requests and deferred streams -still receive output and usage data. Build the integration lazily after its config is registered. -Subscribers also set APM shadow tags, so registration must work with LLMObs disabled. +Provide `listen()` / `unlisten()` in `ddtrace/llmobs/_contrib//__init__.py`. +Call the shared `register(component, integration_cls)` from `listen()` and +`unregister(component)` from `unlisten()`; do not create per-integration subscriber classes. +Shared listeners remain registered; unregistering removes only that component from the registry +and cache. Each started request retains its integration on its execution context, so in-flight +requests and deferred streams still finalize after unpatch or repatch. Build the integration +lazily after its config is registered. Subscribers also set APM shadow tags, so registration +must work with LLMObs disabled. Declare the subscriber module in `pyproject.toml`, substituting the integration's name: diff --git a/ddtrace/contrib/internal/llama_index/_constants.py b/ddtrace/contrib/internal/llama_index/_constants.py new file mode 100644 index 00000000000..e78d34970a3 --- /dev/null +++ b/ddtrace/contrib/internal/llama_index/_constants.py @@ -0,0 +1,2 @@ +TAG_MODEL = "llama_index.request.model" +TAG_PROVIDER = "llama_index.request.provider" diff --git a/ddtrace/contrib/internal/llama_index/_streaming.py b/ddtrace/contrib/internal/llama_index/_streaming.py index 579f22973de..637b0d0b65a 100644 --- a/ddtrace/contrib/internal/llama_index/_streaming.py +++ b/ddtrace/contrib/internal/llama_index/_streaming.py @@ -6,13 +6,12 @@ from typing import Optional from typing import Union -from ddtrace.contrib.internal.stream_handler import AsyncStreamHandler -from ddtrace.contrib.internal.stream_handler import StreamHandler -from ddtrace.contrib.internal.stream_handler import make_traced_stream from ddtrace.internal import core from ddtrace.internal.logger import get_logger from ddtrace.internal.span_bus import span_from_context -from ddtrace.llmobs._integrations import LlamaIndexIntegration +from ddtrace.internal.utils.streaming import AsyncStreamHandler +from ddtrace.internal.utils.streaming import StreamHandler +from ddtrace.internal.utils.streaming import make_traced_stream if TYPE_CHECKING: @@ -28,7 +27,6 @@ class _BaseLlamaIndexStreamHandler: which provide the attributes below via BaseStreamHandler.__init__. """ - integration: LlamaIndexIntegration primary_span: Span request_args: tuple request_kwargs: dict[str, Any] @@ -63,16 +61,19 @@ async def process_chunk(self, chunk: Any, iterator: Any = None) -> None: def handle_streamed_response( - integration: LlamaIndexIntegration, resp: Any, args: tuple, kwargs: dict[str, Any], ctx: core.ExecutionContext, ) -> Any: - """Wrap a sync or async LlamaIndex stream for tracing.""" + """Wrap a sync or async LlamaIndex stream for tracing. + + The handlers only need the context to dispatch the deferred ended event, so the + integration slot BaseStreamHandler keeps for other integrations goes unused. + """ handler: Union[LlamaIndexStreamHandler, LlamaIndexAsyncStreamHandler] if inspect.isasyncgen(resp): - handler = LlamaIndexAsyncStreamHandler(integration, span_from_context(ctx), args, kwargs, ctx=ctx) + handler = LlamaIndexAsyncStreamHandler(None, span_from_context(ctx), args, kwargs, ctx=ctx) else: - handler = LlamaIndexStreamHandler(integration, span_from_context(ctx), args, kwargs, ctx=ctx) + handler = LlamaIndexStreamHandler(None, span_from_context(ctx), args, kwargs, ctx=ctx) return make_traced_stream(resp, handler) diff --git a/ddtrace/contrib/internal/llama_index/patch.py b/ddtrace/contrib/internal/llama_index/patch.py index c0a3b3e146a..0815598c8ee 100644 --- a/ddtrace/contrib/internal/llama_index/patch.py +++ b/ddtrace/contrib/internal/llama_index/patch.py @@ -10,6 +10,8 @@ from ddtrace import config from ddtrace.contrib._events.llm import LlmRequestEvent +from ddtrace.contrib.internal.llama_index._constants import TAG_MODEL +from ddtrace.contrib.internal.llama_index._constants import TAG_PROVIDER from ddtrace.contrib.internal.llama_index._streaming import handle_streamed_response from ddtrace.contrib.internal.llama_index._utils import build_agent_call_tool_request_kwargs from ddtrace.contrib.internal.llama_index._utils import build_agent_run_request_kwargs @@ -21,8 +23,8 @@ from ddtrace.contrib.internal.llama_index._utils import get_model_provider from ddtrace.contrib.internal.trace_utils import int_service from ddtrace.internal import core +from ddtrace.internal import integration_hooks from ddtrace.internal.logger import get_logger -from ddtrace.llmobs._integrations import LlamaIndexIntegration log = get_logger(__name__) @@ -43,11 +45,6 @@ def _supported_versions() -> dict[str, str]: _DD_WRAPPED = "__dd_wrapped__" -def _get_integration() -> LlamaIndexIntegration: - """Retrieve the integration instance stored on the ``llama_index.core`` module by ``patch()``.""" - return llama_core._datadog_integration - - # LlamaIndex LLM methods (chat, complete, etc.) are overridden on concrete subclasses, # so wrapping BaseLLM alone has no effect — we must setattr on each subclass individually. @@ -82,17 +79,23 @@ def _create_event( cardinality low. For non-LLM operations (query, retrieval, embedding, agent) the resource is the class name. """ - integration = _get_integration() resource = model if (model and not operation) else instance.__class__.__name__ provider = get_model_provider(instance) + # Operations without a model omit the model tag. Provider detection always + # returns a string, falling back to "unknown" when no provider is found. + tags = {} + if model is not None: + tags[TAG_MODEL] = model + if provider is not None: + tags[TAG_PROVIDER] = provider return LlmRequestEvent( - component="llama_index", - integration_config=integration.integration_config, - service=int_service(None, integration.integration_config), + component=config.llama_index.integration_name, + integration_config=config.llama_index, + service=int_service(None, config.llama_index), resource=resource, provider=provider, model=model, - llmobs_integration=integration, + tags=tags, submit_to_llmobs=True, request_kwargs=request_kwargs, instance=instance, @@ -123,7 +126,7 @@ def wrapper(func, instance, args, kwargs): ctx.dispatch_ended_event(*sys.exc_info()) raise if always_stream or isinstance(resp, Generator): - return handle_streamed_response(_get_integration(), resp, args, request_kwargs, ctx) + return handle_streamed_response(resp, args, request_kwargs, ctx) event.response = resp ctx.dispatch_ended_event() return resp @@ -147,7 +150,7 @@ async def wrapper(func, instance, args, kwargs): ctx.dispatch_ended_event(*sys.exc_info()) raise if always_stream or inspect.isasyncgen(resp): - return handle_streamed_response(_get_integration(), resp, args, request_kwargs, ctx) + return handle_streamed_response(resp, args, request_kwargs, ctx) event.response = resp ctx.dispatch_ended_event() return resp @@ -289,9 +292,6 @@ def patch(): return llama_core._datadog_patch = True - integration = LlamaIndexIntegration(integration_config=config.llama_index) - llama_core._datadog_integration = integration - # LlamaIndex LLM methods (chat, complete, etc.) are abstract on BaseLLM — # concrete subclasses override them entirely, so we must wrap each subclass. base = llama_core.base @@ -319,12 +319,18 @@ def patch(): _originals[(BaseLLM, "__init__")] = BaseLLM.__init__ BaseLLM.__init__ = _patched_init(BaseLLM.__init__) + # Let products (LLMObs) attach their own LlmEvents subscribers without this + # module importing them. + integration_hooks.on_patch(config.llama_index.integration_name) + def unpatch(): if not getattr(llama_core, "_datadog_patch", False): return llama_core._datadog_patch = False + integration_hooks.on_unpatch(config.llama_index.integration_name) + for (cls, method_name), original in _originals.items(): try: setattr(cls, method_name, original) @@ -332,6 +338,3 @@ def unpatch(): log.debug("Failed to unwrap %s.%s", cls.__name__, method_name, exc_info=True) _originals.clear() _wrapped_classes.clear() - - if hasattr(llama_core, "_datadog_integration"): - delattr(llama_core, "_datadog_integration") diff --git a/ddtrace/llmobs/_contrib/_subscribers.py b/ddtrace/llmobs/_contrib/_subscribers.py new file mode 100644 index 00000000000..d826cfc6116 --- /dev/null +++ b/ddtrace/llmobs/_contrib/_subscribers.py @@ -0,0 +1,97 @@ +"""Shared LLMObs lifecycle subscribers selected by the request's component.""" + +from typing import Optional + +from ddtrace import config +from ddtrace.contrib._events.llm import LlmEvents +from ddtrace.contrib._events.llm import LlmRequestEvent +from ddtrace.ext import SpanTypes +from ddtrace.internal import core +from ddtrace.internal.core.subscriber import Subscriber +from ddtrace.internal.span_bus import span_from_context +from ddtrace.llmobs._constants import LLMOBS_APM_SHADOW_ENABLED_METRIC_KEY +from ddtrace.llmobs._constants import PROXY_REQUEST +from ddtrace.llmobs._integrations.base import BaseLLMIntegration + + +_integration_classes: dict[str, type[BaseLLMIntegration]] = {} +_integrations: dict[str, BaseLLMIntegration] = {} +_CONTEXT_KEY = "llmobs.integration" + + +def register(component: str, integration_cls: type[BaseLLMIntegration]) -> None: + _integration_classes[component] = integration_cls + LLMObsSpanStartingSubscriber.register() + LLMObsSpanStartedSubscriber.register() + LLMObsSpanFinishingSubscriber.register() + + +def unregister(component: str) -> None: + _integration_classes.pop(component, None) + _integrations.pop(component, None) + # Shared listeners remain registered for other components and in-flight requests. + + +class LLMObsSpanStartingSubscriber(Subscriber): + auto_register = False + event_names = (LlmEvents.SPAN_STARTING.value,) + + @classmethod + def on_event(cls, ctx: core.ExecutionContext[LlmRequestEvent]) -> None: + component = ctx.event.component + integration_cls = _integration_classes.get(component) + if integration_cls is None: + return + integration = _integrations.get(component) + if integration is None: + integration = integration_cls(integration_config=getattr(config, component)) + _integrations[component] = integration + # Keep the request's integration alive across unpatch and repatch. + ctx.set_item(_CONTEXT_KEY, integration) + if ctx.event.submit_to_llmobs and integration.llmobs_enabled: + ctx.event.span_type = SpanTypes.LLM + + +class LLMObsSpanStartedSubscriber(Subscriber): + auto_register = False + event_names = (LlmEvents.SPAN_STARTED.value,) + + @classmethod + def on_event(cls, ctx: core.ExecutionContext[LlmRequestEvent]) -> None: + integration: Optional[BaseLLMIntegration] = ctx.get_item(_CONTEXT_KEY) + if integration is None: + return + event = ctx.event + span = span_from_context(ctx) + if event.submit_to_llmobs: + span._set_attribute(LLMOBS_APM_SHADOW_ENABLED_METRIC_KEY, 1 if integration.llmobs_enabled else 0) + integration._set_base_span_tags( + span, + model=event.model, + provider=event.provider, + instance=event.instance, + ) + base_url = integration._get_base_url(instance=event.instance) # type: ignore[arg-type] + if integration._is_instrumented_proxy_url(base_url): + span._set_ctx_item(PROXY_REQUEST, True) + integration._annotate_integration_tag(span) + integration._stamp_llmobs_span_kind_at_start(span, event.operation, operation=event.operation) + + +class LLMObsSpanFinishingSubscriber(Subscriber): + auto_register = False + event_names = (LlmEvents.SPAN_FINISHING.value,) + + @classmethod + def on_event(cls, ctx: core.ExecutionContext[LlmRequestEvent]) -> None: + integration: Optional[BaseLLMIntegration] = ctx.get_item(_CONTEXT_KEY) + if integration is None: + return + event = ctx.event + integration.llmobs_set_tags( + span_from_context(ctx), + args=[], + kwargs=event.request_kwargs, + response=event.response, + operation=event.operation, + ) diff --git a/ddtrace/llmobs/_contrib/anthropic/__init__.py b/ddtrace/llmobs/_contrib/anthropic/__init__.py index 6a0cf4d059a..8ab07ce5e44 100644 --- a/ddtrace/llmobs/_contrib/anthropic/__init__.py +++ b/ddtrace/llmobs/_contrib/anthropic/__init__.py @@ -1,17 +1,11 @@ -from ddtrace.llmobs._contrib.anthropic.subscribers import LLMObsAnthropicSpanFinishingSubscriber -from ddtrace.llmobs._contrib.anthropic.subscribers import LLMObsAnthropicSpanStartedSubscriber -from ddtrace.llmobs._contrib.anthropic.subscribers import LLMObsAnthropicSpanStartingSubscriber +from ddtrace.llmobs._contrib._subscribers import register +from ddtrace.llmobs._contrib._subscribers import unregister +from ddtrace.llmobs._integrations.anthropic import AnthropicIntegration def listen() -> None: - LLMObsAnthropicSpanStartingSubscriber.register() - LLMObsAnthropicSpanStartedSubscriber.register() - LLMObsAnthropicSpanFinishingSubscriber.register() + register("anthropic", AnthropicIntegration) def unlisten() -> None: - LLMObsAnthropicSpanStartingSubscriber.unregister() - LLMObsAnthropicSpanStartedSubscriber.unregister() - # The finishing subscriber stays registered: requests and deferred streams that started - # before unpatch() still need their output, token metrics, and shadow tags when they end. - # It only acts on Anthropic contexts, and an unpatched client creates no new ones. + unregister("anthropic") diff --git a/ddtrace/llmobs/_contrib/anthropic/subscribers.py b/ddtrace/llmobs/_contrib/anthropic/subscribers.py deleted file mode 100644 index 40b78488a04..00000000000 --- a/ddtrace/llmobs/_contrib/anthropic/subscribers.py +++ /dev/null @@ -1,90 +0,0 @@ -from typing import ClassVar -from typing import Optional - -from ddtrace import config -from ddtrace.contrib._events.llm import LlmEvents -from ddtrace.contrib._events.llm import LlmRequestEvent -from ddtrace.contrib.internal.anthropic._constants import COMPONENT -from ddtrace.ext import SpanTypes -from ddtrace.internal import core -from ddtrace.internal.core.subscriber import Subscriber -from ddtrace.internal.span_bus import span_from_context -from ddtrace.llmobs._constants import LLMOBS_APM_SHADOW_ENABLED_METRIC_KEY -from ddtrace.llmobs._constants import PROXY_REQUEST -from ddtrace.llmobs._integrations.anthropic import AnthropicIntegration - - -class LLMObsAnthropicSubscriber(Subscriber): - """Base for LLMObs subscribers to the Anthropic contrib's LlmEvents. - - The events are shared by every LLM integration, so handlers ignore contexts - whose event belongs to another component. - """ - - auto_register = False - _integration: ClassVar[Optional[AnthropicIntegration]] = None - - @classmethod - def integration(cls) -> AnthropicIntegration: - # config.anthropic is only defined once the contrib is imported, so build lazily. - if LLMObsAnthropicSubscriber._integration is None: - LLMObsAnthropicSubscriber._integration = AnthropicIntegration(integration_config=config.anthropic) - return LLMObsAnthropicSubscriber._integration - - @staticmethod - def _is_anthropic(ctx: core.ExecutionContext[LlmRequestEvent]) -> bool: - return ctx.event.component == COMPONENT - - -class LLMObsAnthropicSpanStartingSubscriber(LLMObsAnthropicSubscriber): - event_names = (LlmEvents.SPAN_STARTING.value,) - - @classmethod - def on_event(cls, event_instance: core.ExecutionContext[LlmRequestEvent]) -> None: - # LlmTracingSubscriber dispatches the context rather than the event. - ctx = event_instance - if not cls._is_anthropic(ctx): - return - event = ctx.event - if event.submit_to_llmobs and cls.integration().llmobs_enabled: - event.span_type = SpanTypes.LLM - - -class LLMObsAnthropicSpanStartedSubscriber(LLMObsAnthropicSubscriber): - event_names = (LlmEvents.SPAN_STARTED.value,) - - @classmethod - def on_event(cls, event_instance: core.ExecutionContext[LlmRequestEvent]) -> None: - # LlmTracingSubscriber dispatches the context rather than the event. - ctx = event_instance - if not cls._is_anthropic(ctx): - return - event = ctx.event - span = span_from_context(ctx) - integration = cls.integration() - if event.submit_to_llmobs: - span._set_attribute(LLMOBS_APM_SHADOW_ENABLED_METRIC_KEY, 1 if integration.llmobs_enabled else 0) - integration._set_base_span_tags(span, instance=event.instance) - if integration._is_instrumented_proxy_url(integration._get_base_url(instance=event.instance)): # type: ignore[arg-type] - span._set_ctx_item(PROXY_REQUEST, True) - integration._annotate_integration_tag(span) - integration._stamp_llmobs_span_kind_at_start(span, event.operation, operation=event.operation) - - -class LLMObsAnthropicSpanFinishingSubscriber(LLMObsAnthropicSubscriber): - event_names = (LlmEvents.SPAN_FINISHING.value,) - - @classmethod - def on_event(cls, event_instance: core.ExecutionContext[LlmRequestEvent]) -> None: - # LlmTracingSubscriber dispatches the context rather than the event. - ctx = event_instance - if not cls._is_anthropic(ctx): - return - event = ctx.event - cls.integration().llmobs_set_tags( - span_from_context(ctx), - args=[], - kwargs=event.request_kwargs, - response=event.response, - operation=event.operation, - ) diff --git a/ddtrace/llmobs/_contrib/llama_index/__init__.py b/ddtrace/llmobs/_contrib/llama_index/__init__.py new file mode 100644 index 00000000000..47fd813c254 --- /dev/null +++ b/ddtrace/llmobs/_contrib/llama_index/__init__.py @@ -0,0 +1,11 @@ +from ddtrace.llmobs._contrib._subscribers import register +from ddtrace.llmobs._contrib._subscribers import unregister +from ddtrace.llmobs._integrations.llama_index import LlamaIndexIntegration + + +def listen() -> None: + register("llama_index", LlamaIndexIntegration) + + +def unlisten() -> None: + unregister("llama_index") diff --git a/ddtrace/llmobs/_integrations/llama_index.py b/ddtrace/llmobs/_integrations/llama_index.py index 1c41618b119..7f2d38eab4c 100644 --- a/ddtrace/llmobs/_integrations/llama_index.py +++ b/ddtrace/llmobs/_integrations/llama_index.py @@ -1,6 +1,8 @@ from typing import Any from typing import Optional +from ddtrace.contrib.internal.llama_index._constants import TAG_MODEL +from ddtrace.contrib.internal.llama_index._constants import TAG_PROVIDER from ddtrace.internal.logger import get_logger from ddtrace.internal.utils.formats import _get_attr from ddtrace.llmobs._constants import INPUT_TOKENS_METRIC_KEY @@ -16,24 +18,15 @@ log = get_logger(__name__) -MODEL = "llama_index.request.model" -PROVIDER = "llama_index.request.provider" - - class LlamaIndexIntegration(BaseLLMIntegration): _integration_name = "llama_index" - def _set_base_span_tags( - self, - span: Span, - model: Optional[str] = None, - provider: Optional[str] = None, - **kwargs: dict[str, Any], - ) -> None: - if model is not None: - span._set_attribute(MODEL, model) - if provider is not None: - span._set_attribute(PROVIDER, provider) + def _set_base_span_tags(self, span: Span, **kwargs: Any) -> None: + """Nothing to do: the contrib sets the APM model/provider tags itself. + + LlamaIndex has no per-request base_url to record, so unlike other integrations + there is no LLMObs-only span state to set at start. + """ def _llmobs_span_kind(self, operation_id: str, span: Span, **kwargs: Any) -> Optional[str]: # Event-based: on_started calls this with operation=event.operation (no trace() path). @@ -155,8 +148,8 @@ def _llmobs_set_embedding_tags( _annotate_llmobs_span_data( span, kind="embedding", - model_name=span.get_tag(MODEL) or "", - model_provider=span.get_tag(PROVIDER) or "", + model_name=span.get_tag(TAG_MODEL) or "", + model_provider=span.get_tag(TAG_PROVIDER) or "", input_documents=input_documents, output_value=output_value, ) @@ -230,8 +223,8 @@ def _llmobs_set_llm_tags( _annotate_llmobs_span_data( span, kind="llm", - model_name=span.get_tag(MODEL) or "", - model_provider=span.get_tag(PROVIDER) or "", + model_name=span.get_tag(TAG_MODEL) or "", + model_provider=span.get_tag(TAG_PROVIDER) or "", input_messages=input_messages, metadata=metadata, output_messages=output_messages, @@ -256,8 +249,8 @@ def _set_apm_shadow_tags(self, span, args, kwargs, response=None, operation=""): span, metrics, span_kind, - model_name=span.get_tag(MODEL), - model_provider=span.get_tag(PROVIDER), + model_name=span.get_tag(TAG_MODEL), + model_provider=span.get_tag(TAG_PROVIDER), ) def _extract_input_messages(self, kwargs: dict[str, Any], is_chat: bool) -> list[Message]: diff --git a/pyproject.toml b/pyproject.toml index 0e42beef600..46d73ff435d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -79,6 +79,7 @@ datadog = "ddtrace.contrib.internal.mlflow.auth_plugin:DatadogHeaderProvider" [project.entry-points.'ddtrace.integration_hooks'] anthropic = "ddtrace.llmobs._contrib.anthropic" +llama_index = "ddtrace.llmobs._contrib.llama_index" [project.entry-points.'ddtrace.products'] "apm-tracing-rc" = "ddtrace.internal.remoteconfig.products.apm_tracing" diff --git a/tests/contrib/llama_index/test_llama_index_llmobs.py b/tests/contrib/llama_index/test_llama_index_llmobs.py index e376e2a3053..82795fdf7cc 100644 --- a/tests/contrib/llama_index/test_llama_index_llmobs.py +++ b/tests/contrib/llama_index/test_llama_index_llmobs.py @@ -1,6 +1,7 @@ from unittest import mock from llama_index.core.llms import ChatMessage +from llama_index.core.llms import ChatResponse from llama_index.llms.openai import OpenAI import pytest @@ -416,3 +417,81 @@ def test_shadow_tags_llm_when_llmobs_disabled(tracer): assert span.get_metric("_dd.llmobs.input_tokens") == 14 assert span.get_metric("_dd.llmobs.output_tokens") == 15 assert span.get_metric("_dd.llmobs.total_tokens") == 29 + + +@pytest.mark.parametrize("repatch", [False, True]) +def test_stream_finalizes_after_unpatch(llama_index, llama_index_llmobs, test_spans, repatch): + from ddtrace.contrib.internal.llama_index.patch import patch + from ddtrace.contrib.internal.llama_index.patch import unpatch + + llm = OpenAI(model="gpt-4o-mini", max_tokens=100) + response = ChatResponse( + message=ChatMessage(role="assistant", content="Hello"), + delta="Hello", + raw={"usage": {"prompt_tokens": 8, "completion_tokens": 1, "total_tokens": 9}}, + ) + with mock.patch.object(OpenAI, "_stream_chat", return_value=(chunk for chunk in [response])): + stream = llm.stream_chat(messages=[ChatMessage(role="user", content="Hi")]) + unpatch() + if repatch: + patch() + list(stream) + + spans = [span for trace in test_spans.pop_traces() for span in trace] + assert len(spans) == 1 + assert_llmobs_span_data( + _get_llmobs_data_metastruct(spans[0]), + span_kind="llm", + model_name="gpt-4o-mini", + model_provider="openai", + input_messages=[{"content": "Hi", "role": "user"}], + output_messages=[{"content": "Hello", "role": "assistant"}], + metadata={"max_tokens": 100}, + metrics={"input_tokens": 8, "output_tokens": 1, "total_tokens": 9}, + tags=EXPECTED_TAGS, + ) + + +def test_shared_subscribers_isolate_components_and_repatch(monkeypatch): + from ddtrace import config + from ddtrace.contrib._events.llm import LlmRequestEvent + from ddtrace.internal import core + from ddtrace.llmobs._contrib import _subscribers + + monkeypatch.setattr(_subscribers, "_integration_classes", {}) + monkeypatch.setattr(_subscribers, "_integrations", {}) + first = mock.Mock(llmobs_enabled=False) + replacement = mock.Mock(llmobs_enabled=False) + other = mock.Mock(llmobs_enabled=False) + for integration in (first, replacement, other): + integration._is_instrumented_proxy_url.return_value = False + factory = mock.Mock(side_effect=[first, replacement]) + other_factory = mock.Mock(return_value=other) + _subscribers.register("llama_index", factory) + _subscribers.register("anthropic", other_factory) + + def request(component): + return core.context_with_event( + LlmRequestEvent( + component=component, + integration_config=getattr(config, component), + provider="test", + ) + ) + + with request("llama_index"): + _subscribers.unregister("llama_index") + with request("llama_index"): + pass + first.llmobs_set_tags.assert_not_called() + with request("anthropic"): + pass + other.llmobs_set_tags.assert_called_once() + _subscribers.register("llama_index", factory) + _subscribers.register("llama_index", factory) + with request("llama_index"): + pass + replacement.llmobs_set_tags.assert_called_once() + first.llmobs_set_tags.assert_not_called() + first.llmobs_set_tags.assert_called_once() + assert factory.call_count == 2 diff --git a/tests/contrib/llama_index/test_llama_index_patch.py b/tests/contrib/llama_index/test_llama_index_patch.py index 9a47ffa627b..2d687c4451b 100644 --- a/tests/contrib/llama_index/test_llama_index_patch.py +++ b/tests/contrib/llama_index/test_llama_index_patch.py @@ -1,3 +1,5 @@ +import pytest + from ddtrace.contrib.internal.llama_index.patch import _DD_WRAPPED from ddtrace.contrib.internal.llama_index.patch import get_version from ddtrace.contrib.internal.llama_index.patch import patch @@ -109,3 +111,38 @@ def assert_not_module_double_patched(self, llama_index): if BaseWorkflowAgent is not None: self.assert_not_double_wrapped(BaseWorkflowAgent.run) self.assert_not_double_wrapped(BaseWorkflowAgent.call_tool) + + +@pytest.mark.subprocess(env=dict(DD_LLMOBS_ENABLED="false"), err=None) +def test_manual_patch_emits_shadow_tags(): + from unittest.mock import patch as mock_patch + + from llama_index.core.llms import ChatMessage + from llama_index.core.llms import ChatResponse + from llama_index.llms.openai import OpenAI + + import ddtrace + from tests.utils import DummyWriter + + writer = DummyWriter() + ddtrace.tracer._span_aggregator.writer = writer + ddtrace.patch(llama_index=True) + + response = ChatResponse( + message=ChatMessage(role="assistant", content="Hello"), + raw={"usage": {"prompt_tokens": 14, "completion_tokens": 15, "total_tokens": 29}}, + ) + llm = OpenAI(model="gpt-4o-mini", api_key="not-a-real-key") + with mock_patch.object(OpenAI, "_chat", return_value=response): + llm.chat([ChatMessage(role="user", content="Hello")]) + + spans = [span for span in writer.pop() if span.name == "llama_index.request"] + assert len(spans) == 1 + span = spans[0] + assert span.get_tag("_dd.llmobs.span_kind") == "llm" + assert span.get_tag("_dd.llmobs.model_name") == "gpt-4o-mini" + assert span.get_tag("_dd.llmobs.model_provider") == "openai" + assert span.get_metric("_dd.llmobs.enabled") == 0 + assert span.get_metric("_dd.llmobs.input_tokens") == 14 + assert span.get_metric("_dd.llmobs.output_tokens") == 15 + assert span.get_metric("_dd.llmobs.total_tokens") == 29